From e7b4988d82172dfb0f909bdb025af218fb346eca Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sun, 28 Jun 2026 10:13:03 +0300 Subject: [PATCH] Compute document centroids during the vector scan --- haiku_rag_slim/haiku/rag/doctor.py | 117 ++++++++++++++++------------- tests/test_doctor.py | 58 ++++++++------ 2 files changed, 101 insertions(+), 74 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/doctor.py b/haiku_rag_slim/haiku/rag/doctor.py index 7b2d7542..732cac63 100644 --- a/haiku_rag_slim/haiku/rag/doctor.py +++ b/haiku_rag_slim/haiku/rag/doctor.py @@ -243,43 +243,36 @@ class _DuplicateFamily(BaseModel): def _duplicate_families( - doc_vectors: dict[str, np.ndarray], cfg: DuplicateDetectionConfig + doc_ids: list[str], + centroids: np.ndarray, + counts: np.ndarray, + cfg: DuplicateDetectionConfig, ) -> list[_DuplicateFamily]: """Cluster documents whose embedding centroids are nearly identical. - One unit centroid per document, pairwise cosine, then connected components - of pairs above ``similarity_threshold``. One family per component, each - carrying every member's highest cosine to another member of the family. + ``centroids`` holds one summed (unnormalized) centroid per document and + ``counts`` its embedded-chunk count. Documents below the small-document + floor are dropped; the rest are normalized and clustered by union-find over + pairwise cosine above ``similarity_threshold``. One family per component, + each carrying every member's highest cosine to another member. """ - # Drop unembedded (zero) vectors and documents below the small-document - # floor; reduce each remaining document to a unit centroid. - centroids_by_doc: dict[str, np.ndarray] = {} - sizes: dict[str, int] = {} - for doc_id, matrix in doc_vectors.items(): - m = np.asarray(matrix, dtype=np.float32) - if m.ndim != 2 or m.shape[0] == 0: - continue - m = m[np.linalg.norm(m, axis=1) > 0] - if m.shape[0] < cfg.min_chunks: - continue - centroid = m.mean(axis=0) - norm = np.linalg.norm(centroid) - if norm == 0: - continue - centroids_by_doc[doc_id] = centroid / norm - sizes[doc_id] = int(m.shape[0]) - if len(centroids_by_doc) < 2: + centroids = np.asarray(centroids, dtype=np.float32) + counts = np.asarray(counts) + norms = np.linalg.norm(centroids, axis=1) + eligible = np.nonzero((counts >= cfg.min_chunks) & (norms > 0))[0] + if eligible.size < 2: return [] - - order = sorted(centroids_by_doc) - centroids = np.array([centroids_by_doc[d] for d in order]) + unit = centroids[eligible] / norms[eligible][:, None] + ids = [doc_ids[i] for i in eligible] + sizes = {doc_ids[i]: int(counts[i]) for i in eligible} + n = len(ids) # Pairwise cosine, block-wise to avoid a full D×D matrix at once. Each row # only compares against higher-indexed documents (upper triangle). Cluster # with union-find and keep only each document's best similarity to a twin — # a self-similar corpus forms one clique, so storing every pair would be # O(D²) objects. - parent = list(range(len(order))) + parent = list(range(n)) def find(x: int) -> int: while parent[x] != x: @@ -287,11 +280,11 @@ def _duplicate_families( x = parent[x] return x - best = np.zeros(len(order), dtype=np.float32) + best = np.zeros(n, dtype=np.float32) linked = False block = 512 - for start in range(0, len(order), block): - sims = centroids[start : start + block] @ centroids.T + for start in range(0, n, block): + sims = unit[start : start + block] @ unit.T for row in range(sims.shape[0]): gi = start + row cols = ( @@ -310,21 +303,21 @@ def _duplicate_families( return [] components: dict[int, list[int]] = {} - for idx in range(len(order)): + for idx in range(n): components.setdefault(find(idx), []).append(idx) families: list[_DuplicateFamily] = [] for indices in components.values(): if len(indices) < 2: continue - members = sorted(order[i] for i in indices) + members = sorted(ids[i] for i in indices) # Largest document (most chunks) is the one to keep; smallest id on a tie. keep = min(members, key=lambda d: (-sizes[d], d)) families.append( _DuplicateFamily( members=members, keep=keep, - similarity={order[i]: round(float(best[i]), 3) for i in indices}, + similarity={ids[i]: round(float(best[i]), 3) for i in indices}, sizes={d: sizes[d] for d in members}, ) ) @@ -375,13 +368,15 @@ def _write_duplicates_out( def _check_duplicate_documents( - doc_vectors: dict[str, np.ndarray], + doc_ids: list[str], + centroids: np.ndarray, + counts: np.ndarray, uri_by_doc: Mapping[str, str | None], title_by_doc: Mapping[str, str | None], cfg: DuplicateDetectionConfig, yaml_path: Path | None = None, ) -> CheckResult: - families = _duplicate_families(doc_vectors, cfg) + families = _duplicate_families(doc_ids, centroids, counts, cfg) def label(doc_id: str) -> str: return uri_by_doc.get(doc_id) or title_by_doc.get(doc_id) or doc_id @@ -618,43 +613,59 @@ async def run_db_checks( ) ) - ids = arrow.column("id").to_pylist() # Reshape the Arrow fixed-size-list child buffer directly into an (N, dim) # float32 matrix. Going through to_pylist() would box N*dim Python floats # (tens of GB and most of the wall-clock on large corpora); the stored # vectors are already float32, so this keeps the layout and the dtype. vec_col = arrow.column("vector").combine_chunks() vectors = vec_col.values.to_numpy(zero_copy_only=False).reshape(-1, actual_dim) - zero_ids: list[str] = [] - if vectors.size: - zero_ids = [ids[i] for i in np.nonzero(~vectors.any(axis=1))[0]] + embedded = vectors.any(axis=1) if vectors.size else np.zeros(0, dtype=bool) + + # Unembedded (all-zero) chunks: report a count and a few sampled ids without + # materializing every chunk id. + zero_rows = np.nonzero(~embedded)[0] + zero_count = int(zero_rows.size) + id_col = arrow.column("id") + zero_sample = [id_col[int(i)].as_py() for i in zero_rows[:_SAMPLE_LIMIT]] + if zero_count > _SAMPLE_LIMIT: + zero_sample.append(f"... (+{zero_count - _SAMPLE_LIMIT} more)") results.append( CheckResult( name="unembedded_chunks", - severity=Severity.WARN if zero_ids else Severity.OK, + severity=Severity.WARN if zero_count else Severity.OK, message=( - f"{len(zero_ids)} chunk(s) have an all-zero (unembedded) vector." - if zero_ids + f"{zero_count} chunk(s) have an all-zero (unembedded) vector." + if zero_count else "All chunks are embedded." ), - remediation="haiku-rag rebuild --embed-only" if zero_ids else None, - details=_sample(zero_ids), + remediation="haiku-rag rebuild --embed-only" if zero_count else None, + details=zero_sample, ) ) - # Near-identical documents (centroid cosine), grouped from the same vector - # scan rather than a second pass. - chunk_doc_ids_ordered = arrow.column("document_id").to_pylist() - indices_by_doc: dict[str, list[int]] = {} - for index, doc_id in enumerate(chunk_doc_ids_ordered): - indices_by_doc.setdefault(doc_id, []).append(index) - doc_vectors = {doc_id: vectors[idx] for doc_id, idx in indices_by_doc.items()} - # Per-doc fancy indexing has copied every vector; release the full matrix so - # both copies are not resident during duplicate detection. + # Near-identical documents (centroid cosine). Reduce each document's chunk + # vectors to one summed centroid during the scan: dictionary-encode the + # document ids into integer codes, then sum each document's embedded rows in + # a single pass per document — no second full copy of the vector matrix. + encoded = arrow.column("document_id").combine_chunks().dictionary_encode() + doc_ids = encoded.dictionary.to_pylist() + codes = encoded.indices.to_numpy(zero_copy_only=False) + centroids = np.zeros((len(doc_ids), actual_dim), dtype=np.float32) + counts = np.zeros(len(doc_ids), dtype=np.int64) + order = np.argsort(codes, kind="stable") + bounds = np.searchsorted(codes, np.arange(len(doc_ids) + 1), sorter=order) + for d in range(len(doc_ids)): + rows = order[bounds[d] : bounds[d + 1]] + rows = rows[embedded[rows]] + counts[d] = rows.size + if rows.size: + centroids[d] = vectors[rows].sum(axis=0) del vectors results.append( _check_duplicate_documents( - doc_vectors, + doc_ids, + centroids, + counts, uri_by_doc, title_by_doc, config.doctor.duplicates, diff --git a/tests/test_doctor.py b/tests/test_doctor.py index c888a18b..619a0b6d 100644 --- a/tests/test_doctor.py +++ b/tests/test_doctor.py @@ -944,22 +944,28 @@ async def test_probe_endpoint_connection_error(): # --- Duplicate-document detection ---------------------------------------- -def _docs(spec: dict[str, list[int]], dim: int = 8) -> dict[str, np.ndarray]: - """Build per-document chunk matrices from one-hot indices. +def _centroids( + spec: dict[str, list[int]], dim: int = 8 +) -> tuple[list[str], np.ndarray, np.ndarray]: + """Summed one-hot centroids + chunk counts per document, as + ``_duplicate_families`` consumes them. Orthogonal one-hot chunks make the centroid cosine of two documents equal to ``shared / sqrt(len(a) * len(b))``: identical documents score 1.0, fully distinct documents score 0.0. """ eye = np.eye(dim) - return { - doc: np.array([eye[i] for i in idxs], dtype=float) for doc, idxs in spec.items() - } + doc_ids = list(spec) + centroids = np.array( + [eye[idxs].sum(axis=0) for idxs in spec.values()], dtype=np.float32 + ) + counts = np.array([len(idxs) for idxs in spec.values()], dtype=np.int64) + return doc_ids, centroids, counts def test_duplicate_families_identical_docs_flagged(): families = _duplicate_families( - _docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3]}), DuplicateDetectionConfig() + *_centroids({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3]}), DuplicateDetectionConfig() ) assert len(families) == 1 assert set(families[0].members) == {"a", "b"} @@ -970,21 +976,22 @@ def test_duplicate_families_append_only_not_flagged(): # A is fully contained in the larger B, but their centroids diverge # (cosine sqrt(3/6) ~= 0.71), so it stays below the similarity cutoff. families = _duplicate_families( - _docs({"a": [0, 1, 2], "b": [0, 1, 2, 3, 4, 5]}), DuplicateDetectionConfig() + *_centroids({"a": [0, 1, 2], "b": [0, 1, 2, 3, 4, 5]}), + DuplicateDetectionConfig(), ) assert families == [] def test_duplicate_families_distinct_docs_none(): families = _duplicate_families( - _docs({"a": [0, 1, 2], "b": [3, 4, 5]}), DuplicateDetectionConfig() + *_centroids({"a": [0, 1, 2], "b": [3, 4, 5]}), DuplicateDetectionConfig() ) assert families == [] def test_duplicate_families_three_way_one_family(): families = _duplicate_families( - _docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [0, 1, 2, 3]}), + *_centroids({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [0, 1, 2, 3]}), DuplicateDetectionConfig(), ) assert len(families) == 1 @@ -997,7 +1004,7 @@ def test_duplicate_families_clique_single_family(): # A self-similar corpus (all identical) is one clique. Union-find collapses # it to a single family without materializing every pair. spec = {chr(ord("a") + k): [0, 1, 2, 3] for k in range(8)} - families = _duplicate_families(_docs(spec), DuplicateDetectionConfig()) + families = _duplicate_families(*_centroids(spec), DuplicateDetectionConfig()) assert len(families) == 1 assert set(families[0].members) == set(spec) assert all(s == pytest.approx(1.0) for s in families[0].similarity.values()) @@ -1006,7 +1013,7 @@ def test_duplicate_families_clique_single_family(): def test_duplicate_families_tiny_docs_ignored(): # min_chunks = 3 excludes the one-chunk documents. families = _duplicate_families( - _docs({"a": [0], "b": [0]}), DuplicateDetectionConfig() + *_centroids({"a": [0], "b": [0]}), DuplicateDetectionConfig() ) assert families == [] @@ -1014,9 +1021,9 @@ def test_duplicate_families_tiny_docs_ignored(): def test_duplicate_families_threshold_is_configurable(): # Share 3 of 4 chunks each -> centroid cosine 0.75. spec = {"a": [0, 1, 2, 3], "b": [0, 1, 2, 4]} - assert _duplicate_families(_docs(spec), DuplicateDetectionConfig()) == [] + assert _duplicate_families(*_centroids(spec), DuplicateDetectionConfig()) == [] flagged = _duplicate_families( - _docs(spec), DuplicateDetectionConfig(similarity_threshold=0.7) + *_centroids(spec), DuplicateDetectionConfig(similarity_threshold=0.7) ) assert len(flagged) == 1 assert set(flagged[0].members) == {"a", "b"} @@ -1029,9 +1036,10 @@ def test_duplicate_documents_report_truncates_summary(): idx = [3 * k, 3 * k + 1, 3 * k + 2] spec[f"a{k}"] = idx spec[f"b{k}"] = list(idx) - docs = _docs(spec, dim=3 * pairs) uris = {d: f"file:///srv/shared/library/docs/{d}.pdf" for d in spec} - result = _check_duplicate_documents(docs, uris, {}, DuplicateDetectionConfig()) + result = _check_duplicate_documents( + *_centroids(spec, dim=3 * pairs), uris, {}, DuplicateDetectionConfig() + ) assert result.severity is Severity.WARN # The summary message still reports the full total. assert f"{pairs} group(s)" in result.message @@ -1042,10 +1050,14 @@ def test_duplicate_documents_report_truncates_summary(): def test_duplicate_documents_report_factors_common_path(): - docs = _docs({"a": [0, 1, 2], "b": [0, 1, 2]}, dim=3) base = "file:///srv/shared/library/docs/" uris = {"a": base + "alpha.pdf", "b": base + "beta.pdf"} - result = _check_duplicate_documents(docs, uris, {}, DuplicateDetectionConfig()) + result = _check_duplicate_documents( + *_centroids({"a": [0, 1, 2], "b": [0, 1, 2]}, dim=3), + uris, + {}, + DuplicateDetectionConfig(), + ) assert f"common path: {base}" in result.details member_lines = [d for d in result.details if d.lstrip().startswith("#")] assert {d.strip() for d in member_lines} == {"#1 alpha.pdf", "#2 beta.pdf"} @@ -1054,11 +1066,11 @@ def test_duplicate_documents_report_factors_common_path(): def test_duplicate_documents_writes_yaml(tmp_path): # a,b identical (a 4-chunk duplicate); c distinct and excluded. - docs = _docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [4, 5, 6]}, dim=8) + spec = {"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [4, 5, 6]} uris = {"a": "file:///x/a.pdf", "b": "file:///x/b.pdf", "c": "file:///x/c.pdf"} out = tmp_path / "dups.yaml" _check_duplicate_documents( - docs, uris, {}, DuplicateDetectionConfig(), yaml_path=out + *_centroids(spec, dim=8), uris, {}, DuplicateDetectionConfig(), yaml_path=out ) data = yaml.safe_load(out.read_text()) assert len(data["groups"]) == 1 @@ -1076,10 +1088,14 @@ def test_duplicate_documents_writes_yaml(tmp_path): def test_duplicate_documents_writes_empty_yaml_when_none(tmp_path): - docs = _docs({"a": [0, 1, 2], "b": [3, 4, 5]}, dim=6) # distinct + spec = {"a": [0, 1, 2], "b": [3, 4, 5]} # distinct out = tmp_path / "dups.yaml" _check_duplicate_documents( - docs, {"a": "u", "b": "v"}, {}, DuplicateDetectionConfig(), yaml_path=out + *_centroids(spec, dim=6), + {"a": "u", "b": "v"}, + {}, + DuplicateDetectionConfig(), + yaml_path=out, ) assert yaml.safe_load(out.read_text()) == {"groups": []}