diff --git a/haystack/components/preprocessors/embedding_based_document_splitter.py b/haystack/components/preprocessors/embedding_based_document_splitter.py index 17acefd3b7..ba0df72c68 100644 --- a/haystack/components/preprocessors/embedding_based_document_splitter.py +++ b/haystack/components/preprocessors/embedding_based_document_splitter.py @@ -79,6 +79,7 @@ def __init__( language: Language = "en", use_split_rules: bool = True, extend_abbreviations: bool = True, + min_distance_for_two_groups: float | None = None, ) -> None: """ Initialize EmbeddingBasedDocumentSplitter. @@ -86,7 +87,9 @@ def __init__( :param document_embedder: The DocumentEmbedder to use for calculating embeddings. :param sentences_per_group: Number of sentences to group together before embedding. :param percentile: Percentile threshold for cosine distance. Distances above this percentile - are treated as break points. + are treated as break points. Does not apply to documents that tokenize into exactly two sentence + groups: with only one distance to evaluate there is no distribution to compute a percentile against, + so `min_distance_for_two_groups` is used instead for that case. :param min_length: Minimum length of splits in characters. Splits below this length will be merged. :param max_length: Maximum length of splits in characters. Splits above this length will be recursively split. :param language: Language for sentence tokenization. @@ -95,6 +98,17 @@ def __init__( :param extend_abbreviations: If True, the abbreviations used by NLTK's PunktTokenizer are extended by a list of curated abbreviations. Currently supported languages are: en, de. If False, the default abbreviations are used. + :param min_distance_for_two_groups: Cosine distance threshold used only for documents that tokenize into + exactly two sentence groups. With a single distance to evaluate there is no distribution to compute + `percentile` against, so this fixed, absolute threshold is used instead: the lone gap becomes a split + point if its distance exceeds this value. There is no universally correct default: what counts as + "the same topic, reworded" versus "a genuine topic change" is specific to the embedding model in use, + and the two can differ in cosine distance by an order of magnitude or more between models (some + embedders compress all pairwise distances into a narrow band, others spread them out). Defaults to + `None`, meaning documents with exactly two sentence groups are never split, matching the behavior of + documents with only one group. To enable splitting for this case, measure cosine distances from your + own embedder on a few known same-topic and known different-topic sentence pairs, and set this to a + value between the two. """ self.document_embedder = document_embedder @@ -114,6 +128,10 @@ def __init__( raise ValueError("max_length must be greater than min_length.") self.max_length = max_length + if min_distance_for_two_groups is not None and not 0.0 <= min_distance_for_two_groups <= 2.0: + raise ValueError("min_distance_for_two_groups must be between 0.0 and 2.0.") + self.min_distance_for_two_groups = min_distance_for_two_groups + self.language = language self.use_split_rules = use_split_rules self.extend_abbreviations = extend_abbreviations @@ -367,6 +385,17 @@ def _find_split_points(self, embeddings: list[list[float]]) -> list[int]: ) distances.append(distance) + # With a single distance (i.e. exactly two sentence groups) there is no distribution to compute a + # percentile against: np.percentile of a one-element list always returns that same element, so + # `distance > threshold` below would always compare the value to itself and could never be True. + # `percentile` cannot apply here, so we fall back to `min_distance_for_two_groups` instead, which + # defaults to `None` (never split) since there's no universally correct absolute distance threshold + # across embedding models -- see its docstring in `__init__`. + if len(distances) == 1: + if self.min_distance_for_two_groups is None: + return [] + return [1] if distances[0] > self.min_distance_for_two_groups else [] + # Calculate threshold based on percentile threshold = np.percentile(distances, self.percentile * 100) @@ -572,6 +601,7 @@ def to_dict(self) -> dict[str, Any]: language=self.language, use_split_rules=self.use_split_rules, extend_abbreviations=self.extend_abbreviations, + min_distance_for_two_groups=self.min_distance_for_two_groups, ) @classmethod diff --git a/releasenotes/notes/fix-embedding-splitter-two-groups-48201a43ca06f07b.yaml b/releasenotes/notes/fix-embedding-splitter-two-groups-48201a43ca06f07b.yaml new file mode 100644 index 0000000000..754faae6f0 --- /dev/null +++ b/releasenotes/notes/fix-embedding-splitter-two-groups-48201a43ca06f07b.yaml @@ -0,0 +1,18 @@ +--- +fixes: + - | + `EmbeddingBasedDocumentSplitter` could never split a document that tokenized into exactly two sentence + groups, no matter how dissimilar the groups were or how low `percentile` was set. With a single distance + to compare, `numpy.percentile` always returns that same value, so the `distance > threshold` check was + comparing the value to itself and could never be `True`. `percentile` cannot decide this case, since it + requires a distribution of distances to be meaningful and a single gap doesn't provide one. Splitting a + two-group document is now available via the new `min_distance_for_two_groups` parameter (see below); + without it, two-group documents keep the previous, safe behavior of never being split. +enhancements: + - | + Added `min_distance_for_two_groups` to `EmbeddingBasedDocumentSplitter`, a cosine-distance threshold used + only for documents that tokenize into exactly two sentence groups, where `percentile` cannot apply. + Defaults to `None` (never split this case). There is no universally correct default value: how far apart + "the same topic, reworded" and "a genuine topic change" sit in cosine distance is specific to the + embedding model in use and can vary by an order of magnitude or more between models, so this is left for + users to calibrate against their own embedder. diff --git a/test/components/preprocessors/test_embedding_based_document_splitter.py b/test/components/preprocessors/test_embedding_based_document_splitter.py index 50db2a8aa8..af80692e16 100644 --- a/test/components/preprocessors/test_embedding_based_document_splitter.py +++ b/test/components/preprocessors/test_embedding_based_document_splitter.py @@ -32,6 +32,7 @@ def test_init(self): assert splitter.percentile == 0.9 assert splitter.min_length == 50 assert splitter.max_length == 1000 + assert splitter.min_distance_for_two_groups is None def test_init_invalid_sentences_per_group(self): mock_embedder = Mock() @@ -43,6 +44,11 @@ def test_init_invalid_percentile(self): with pytest.raises(ValueError, match="percentile must be between 0.0 and 1.0"): EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, percentile=1.5) + def test_init_invalid_min_distance_for_two_groups(self): + mock_embedder = Mock() + with pytest.raises(ValueError, match="min_distance_for_two_groups must be between 0.0 and 2.0"): + EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, min_distance_for_two_groups=2.5) + def test_init_invalid_min_length(self): mock_embedder = Mock() with pytest.raises(ValueError, match="min_length must be greater than or equal to 0"): @@ -170,6 +176,95 @@ def test_find_split_points(self): # Should find a split point after the second embedding (index 2) assert 2 in split_points + def test_find_split_points_two_groups_default_never_splits(self): + """ + With exactly two sentence groups there is only one distance to compare, so np.percentile of that + single-element list always returns the value itself: `distance > threshold` is comparing the value + to itself and can never be True. `percentile` therefore cannot decide this case, and there is no + universally correct absolute distance threshold across embedding models to fall back on either (see + `min_distance_for_two_groups`'s docstring). The safe default is to never split here, matching the + pre-existing behavior for documents with a single sentence group -- even for maximally dissimilar + (orthogonal) embeddings, at the most aggressive percentile=0.0 setting. + """ + mock_embedder = Mock() + splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder, percentile=0.0) + assert splitter.min_distance_for_two_groups is None + + embeddings = [[1.0, 0.0], [0.0, 1.0]] # orthogonal -> maximum possible cosine distance + assert splitter._find_split_points(embeddings) == [] + + def test_find_split_points_two_groups_splits_once_min_distance_for_two_groups_is_configured(self): + """ + Splitting a two-group document is opt-in: once `min_distance_for_two_groups` is set, the lone + distance is compared against it directly, independent of `percentile` (which cannot apply here). + """ + mock_embedder = Mock() + embeddings = [[1.0, 0.0], [0.0, 1.0]] # orthogonal -> cosine distance of 1.0 + + for percentile in (0.0, 0.5, 0.95, 1.0): + splitter = EmbeddingBasedDocumentSplitter( + document_embedder=mock_embedder, percentile=percentile, min_distance_for_two_groups=0.5 + ) + assert splitter._find_split_points(embeddings) == [1], f"expected split at percentile={percentile}" + + def test_find_split_points_two_identical_groups_do_not_split(self): + """Two identical groups (distance == 0) must not produce a split point even at the threshold's edge.""" + mock_embedder = Mock() + splitter = EmbeddingBasedDocumentSplitter( + document_embedder=mock_embedder, percentile=0.0, min_distance_for_two_groups=0.0 + ) + + embeddings = [[1.0, 0.0], [1.0, 0.0]] + assert splitter._find_split_points(embeddings) == [] + + def test_find_split_points_two_near_duplicate_groups_do_not_split_regardless_of_percentile(self): + """ + Regression test: comparing the lone distance to itself (the original bug) or splitting on any + nonzero distance (an earlier, overly aggressive fix) both fail on near-duplicate groups. Two groups + that are almost, but not exactly, identical (a tiny cosine distance well under a realistic + `min_distance_for_two_groups`) should not split, and since `percentile` does not apply to the + two-group case, this must hold at every percentile setting, including the most aggressive (0.0). + """ + mock_embedder = Mock() + embeddings = [[1.0, 0.0], [1.0, 0.001]] # cosine distance ~5e-7: near-duplicate, not identical + + for percentile in (0.0, 0.5, 0.95, 1.0): + splitter = EmbeddingBasedDocumentSplitter( + document_embedder=mock_embedder, percentile=percentile, min_distance_for_two_groups=0.3 + ) + assert splitter._find_split_points(embeddings) == [], f"unexpected split at percentile={percentile}" + + def test_run_splits_document_with_exactly_two_sentence_groups(self): + """ + End-to-end regression test: a document that tokenizes into exactly two sentence groups with + maximally dissimilar content is split by run() once `min_distance_for_two_groups` is configured, + even at the most aggressive percentile=0.0 setting. + """ + text = ( + "Astronomy studies the stars and galaxies of outer space. " + "Sourdough baking relies on wild yeast fermentation." + ) + + def mock_run(documents): + embeddings = [[1.0, 0.0] if i == 0 else [0.0, 1.0] for i in range(len(documents))] + return {"documents": [replace(doc, embedding=emb) for doc, emb in zip(documents, embeddings, strict=True)]} + + mock_embedder = Mock() + mock_embedder.run = Mock(side_effect=mock_run) + + splitter = EmbeddingBasedDocumentSplitter( + document_embedder=mock_embedder, + sentences_per_group=1, + percentile=0.0, + min_length=0, + max_length=10000, + min_distance_for_two_groups=0.5, + ) + splitter.warm_up() + + result = splitter.run(documents=[Document(content=text)]) + assert len(result["documents"]) == 2 + def test_create_splits_from_points(self): mock_embedder = Mock() splitter = EmbeddingBasedDocumentSplitter(document_embedder=mock_embedder) @@ -379,6 +474,7 @@ def test_to_dict(self): assert result["init_parameters"]["percentile"] == 0.9 assert result["init_parameters"]["min_length"] == 50 assert result["init_parameters"]["max_length"] == 1000 + assert result["init_parameters"]["min_distance_for_two_groups"] is None assert "document_embedder" in result["init_parameters"] @pytest.mark.skipif(os.environ.get("OPENAI_API_KEY", "") == "", reason="OPENAI_API_KEY is not set")