Implementation:Datajuicer Data juicer TextEmbdSimilarityFilter

Knowledge Sources	Datajuicer_Data_juicer
Domains	Data_Quality, Filtering
Last Updated	2026-02-14 16:00 GMT

Overview

Concrete tool for filtering data samples based on text embedding similarity provided by Data-Juicer.

Description

TextEmbdSimilarityFilter is a filter operator that keeps samples whose average embedding similarity to a set of validation texts falls within a specific range. It extends Filter and uses the two-phase compute_stats/process pattern. It computes text embeddings using either an API-based model (default: text-embedding-v4) or a HuggingFace model, then calculates cosine similarity between each sample's embedding and precomputed validation embeddings. Supports multiple pooling strategies (last token, mean, weighted mean). Validation features must be prepared via prepare_valid_feature(). Registered as an ATTRIBUTION_FILTER. Supports CUDA acceleration.

Usage

Import when filtering based on text embedding similarity. Configure in YAML or Python.

Code Reference

Source Location

• Repository: Datajuicer_Data_juicer
• File: data_juicer/ops/filter/text_embd_similarity_filter.py

Signature

@OPERATORS.register_module("text_embd_similarity_filter")
class TextEmbdSimilarityFilter(Filter):
    def __init__(self, api_or_hf_model: str = "text-embedding-v4", is_hf_model: bool = False, api_endpoint: str = "embeddings", response_path: str = "data.0.embedding", model_params: Optional[Dict] = None, min_score: ClosedUnitInterval = 0.1, max_score: ClosedUnitInterval = 1.0, valid_dataset: Optional[List[Dict]] = None, ebd_dim: int = 4096, pooling: Optional[str] = None, input_template: Optional[str] = None, *args, **kwargs):

Import

from data_juicer.ops.filter.text_embd_similarity_filter import TextEmbdSimilarityFilter

I/O Contract

Inputs

Name	Type	Required	Description
api_or_hf_model	str	No	API or HuggingFace embedding model name (default: "text-embedding-v4")
is_hf_model	bool	No	Whether the model is from HuggingFace (default: False)
min_score	ClosedUnitInterval	No	Minimum average similarity to keep samples (default: 0.1)
max_score	ClosedUnitInterval	No	Maximum average similarity to keep samples (default: 1.0)
valid_dataset	Optional[List[Dict]]	No	Validation dataset for computing reference embeddings
ebd_dim	int	No	Embedding dimension for API models (default: 4096)
pooling	Optional[str]	No	Pooling strategy: None, "mean", or "weighted_mean"

Outputs

Name	Type	Description
samples	Dict	Filtered samples with text_embd_similarity stat computed

Usage Examples

YAML Configuration

process:
  - text_embd_similarity_filter:
      min_score: 0.1
      max_score: 1.0

Python API

from data_juicer.ops.filter.text_embd_similarity_filter import TextEmbdSimilarityFilter
op = TextEmbdSimilarityFilter(min_score=0.1, max_score=1.0)

Related Pages

• Environment:Datajuicer_Data_juicer_Python_Runtime_Environment