diff --git a/haiku_rag_slim/haiku/rag/doctor.py b/haiku_rag_slim/haiku/rag/doctor.py index 2e3a9058..427dd53e 100644 --- a/haiku_rag_slim/haiku/rag/doctor.py +++ b/haiku_rag_slim/haiku/rag/doctor.py @@ -282,19 +282,40 @@ def _duplicate_families( order = sorted(normalized) centroids = np.array([_unit(normalized[d].mean(axis=0)) for d in order]) + sizes = np.array([normalized[d].shape[0] for d in order], dtype=float) # Stage 1: centroid candidate pairs, block-wise to avoid a full D×D matrix. + # The cap is enforced per row (truncating each row's matches) so a + # self-similar corpus can never allocate beyond MAX_CANDIDATE_PAIRS. candidates: list[tuple[int, int]] = [] block = 512 + capped = False for start in range(0, len(order), block): + if capped: + break sims = centroids[start : start + block] @ centroids.T for row in range(sims.shape[0]): gi = start + row - above = np.nonzero(sims[row, gi + 1 :] >= cfg.candidate_threshold)[0] + targets = np.arange(gi + 1, len(order)) + if targets.size == 0: + continue + # A smaller document can be fully contained in a larger append-only + # revision even when the fixed centroid threshold would fail: + # with orthogonal chunks, cosine falls to sqrt(small / large). + # Scale the candidate gate by that size ratio, then let directed + # containment make the actual duplicate decision. + ratios = np.minimum(sizes[gi], sizes[targets]) / np.maximum( + sizes[gi], sizes[targets] + ) + thresholds = cfg.candidate_threshold * np.sqrt(ratios) + above = np.nonzero(sims[row, gi + 1 :] >= thresholds)[0] + remaining = MAX_CANDIDATE_PAIRS - len(candidates) + if len(above) >= remaining: + above = above[:remaining] + capped = True candidates.extend((gi, gi + 1 + int(j)) for j in above) - if len(candidates) >= MAX_CANDIDATE_PAIRS: - candidates = candidates[:MAX_CANDIDATE_PAIRS] - break + if capped: + break # Stage 2: confirm candidates with directed chunk-overlap containment. adjacency: dict[int, set[int]] = {} diff --git a/tests/test_doctor.py b/tests/test_doctor.py index de93e2fd..343b5c19 100644 --- a/tests/test_doctor.py +++ b/tests/test_doctor.py @@ -985,6 +985,17 @@ def test_duplicate_families_append_only_is_asymmetric(): assert a_to_b[3] == pytest.approx(0.5) +def test_duplicate_families_append_only_passes_default_centroid_gate(): + # Fixed 0.85 centroid gating misses this: centroid cosine is sqrt(3 / 6), + # but A is fully contained in B and should reach the containment verifier. + families = _duplicate_families( + _docs({"a": [0, 1, 2], "b": [0, 1, 2, 3, 4, 5]}), + DuplicateDetectionConfig(), + ) + assert len(families) == 1 + assert set(families[0].members) == {"a", "b"} + + def test_duplicate_families_distinct_docs_none(): families = _duplicate_families( _docs({"a": [0, 1, 2], "b": [3, 4, 5]}), _stage2_cfg() @@ -1014,6 +1025,16 @@ def test_duplicate_families_tiny_docs_ignored(): assert families == [] +def test_duplicate_families_caps_candidates_during_collection(monkeypatch): + # Three mutually-identical docs would yield 3 candidate pairs, but a cap of 1 + # must stop collection after the first (a,b), leaving c unconfirmed. + monkeypatch.setattr("haiku.rag.doctor.MAX_CANDIDATE_PAIRS", 1) + docs = _docs({"a": [0, 1, 2], "b": [0, 1, 2], "c": [0, 1, 2]}, dim=3) + families = _duplicate_families(docs, _stage2_cfg()) + assert len(families) == 1 + assert set(families[0].members) == {"a", "b"} + + def test_duplicate_families_threshold_is_configurable(): # Share 3 of 5 chunks each -> containment 0.6 both ways. spec = {"a": [0, 1, 2, 3, 4], "b": [0, 1, 2, 5, 6]}