danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

performance updates to duplicate checking

Browse source
This commit is contained in:
bryan davis 2026-06-26 12:06:31 -05:00
parent cd1a73b389
commit d1a873b311
No known key found for this signature in database
GPG key ID: D11B4A4C0C731E5E
1 changed files with 23 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
haiku_rag_slim/haiku/rag/doctor.py
View file

@ -252,9 +252,18 @@ def _unit(vector: np.ndarray) -> np.ndarray:
return vector / norm if norm else vector
def _containment(source: np.ndarray, target: np.ndarray, twin: float) -> float:
"""Fraction of ``source`` chunks with a near-identical chunk in ``target``."""
return float(((source @ target.T).max(axis=1) >= twin).mean())
def _containment(a: np.ndarray, b: np.ndarray, twin: float) -> tuple[float, float]:
"""Directed chunk-overlap containment in both directions from one product.
Returns ``(a_in_b, b_in_a)`` where each value is the fraction of one
document's chunks that have a near-identical chunk in the other. Both
directions are the row-max and column-max of the same similarity matrix, so
the matmul is computed once rather than once per direction.
"""
sims = a @ b.T
a_in_b = float((sims.max(axis=1) >= twin).mean())
b_in_a = float((sims.max(axis=0) >= twin).mean())
return a_in_b, b_in_a
def _duplicate_families(
@ -270,7 +279,7 @@ def _duplicate_families(
# floor; normalize the rest to unit length.
normalized: dict[str, np.ndarray] = {}
for doc_id, matrix in doc_vectors.items():
m = np.asarray(matrix, dtype=float)
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]
@ -321,12 +330,9 @@ def _duplicate_families(
adjacency: dict[int, set[int]] = {}
edges: dict[tuple[int, int], tuple[float, float]] = {}
for i, j in candidates:
a_to_b = _containment(
a_to_b, b_to_a = _containment(
normalized[order[i]], normalized[order[j]], cfg.twin_similarity
)
b_to_a = _containment(
normalized[order[j]], normalized[order[i]], cfg.twin_similarity
)
if max(a_to_b, b_to_a) >= cfg.containment_threshold:
adjacency.setdefault(i, set()).add(j)
adjacency.setdefault(j, set()).add(i)
@ -661,7 +667,12 @@ async def run_db_checks(
)
ids = arrow.column("id").to_pylist()
vectors = np.asarray(arrow.column("vector").to_pylist(), dtype=float)
# 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 the bulk 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]]
@ -686,6 +697,9 @@ async def run_db_checks(
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.
del vectors
results.append(
_check_duplicate_documents(
doc_vectors,