Simplify duplicate detection to whole-document centroid similarity
This commit is contained in:
parent
c7c3f75508
commit
15aae0f242
3 changed files with 133 additions and 194 deletions
|
|
@ -106,15 +106,13 @@ class AnalysisConfig(BaseModel):
|
||||||
class DuplicateDetectionConfig(BaseModel):
|
class DuplicateDetectionConfig(BaseModel):
|
||||||
"""Thresholds for doctor's near-duplicate document detection.
|
"""Thresholds for doctor's near-duplicate document detection.
|
||||||
|
|
||||||
Detection clusters documents that share most of their chunks (revisions of
|
Detection clusters whole documents whose embedding centroids are nearly
|
||||||
one another). ``containment_threshold`` is the decision knob; the others
|
identical — the same document ingested twice, or a light revision.
|
||||||
tune the cheap centroid pre-filter, what counts as a shared chunk, and which
|
``similarity_threshold`` is the cosine cutoff; ``min_chunks`` skips
|
||||||
tiny documents to skip.
|
documents too small to compare meaningfully.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
containment_threshold: float = 0.75
|
similarity_threshold: float = 0.97
|
||||||
candidate_threshold: float = 0.85
|
|
||||||
twin_similarity: float = 0.95
|
|
||||||
min_chunks: int = 3
|
min_chunks: int = 3
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -235,134 +235,97 @@ async def _column_values(table, column: str) -> list:
|
||||||
return [row[column] for row in rows]
|
return [row[column] for row in rows]
|
||||||
|
|
||||||
|
|
||||||
# Backstop on how many centroid candidate pairs we verify, bounding memory and
|
|
||||||
# runtime on a pathologically self-similar corpus.
|
|
||||||
MAX_CANDIDATE_PAIRS = 200_000
|
|
||||||
|
|
||||||
|
|
||||||
class _DuplicateFamily(BaseModel):
|
class _DuplicateFamily(BaseModel):
|
||||||
members: list[str]
|
members: list[str]
|
||||||
superset: str
|
keep: str
|
||||||
pairs: list[tuple[str, str, float, float]]
|
similarity: dict[str, float]
|
||||||
sizes: dict[str, int]
|
sizes: dict[str, int]
|
||||||
|
|
||||||
|
|
||||||
def _unit(vector: np.ndarray) -> np.ndarray:
|
|
||||||
norm = np.linalg.norm(vector)
|
|
||||||
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 _duplicate_families(
|
def _duplicate_families(
|
||||||
doc_vectors: dict[str, np.ndarray], cfg: DuplicateDetectionConfig
|
doc_vectors: dict[str, np.ndarray], cfg: DuplicateDetectionConfig
|
||||||
) -> list[_DuplicateFamily]:
|
) -> list[_DuplicateFamily]:
|
||||||
"""Cluster documents that share most of their chunks (revisions of one another).
|
"""Cluster documents whose embedding centroids are nearly identical.
|
||||||
|
|
||||||
Block-then-verify: cheap centroid similarity proposes candidate document
|
One unit centroid per document, pairwise cosine, then connected components
|
||||||
pairs, then directed chunk-overlap containment confirms them. Returns one
|
of pairs above ``similarity_threshold``. One family per component, each
|
||||||
entry per connected component of confirmed pairs.
|
carrying every member's highest cosine to another member of the family.
|
||||||
"""
|
"""
|
||||||
# Drop unembedded (zero) vectors and documents below the small-document
|
# Drop unembedded (zero) vectors and documents below the small-document
|
||||||
# floor; normalize the rest to unit length.
|
# floor; reduce each remaining document to a unit centroid.
|
||||||
normalized: dict[str, np.ndarray] = {}
|
centroids_by_doc: dict[str, np.ndarray] = {}
|
||||||
|
sizes: dict[str, int] = {}
|
||||||
for doc_id, matrix in doc_vectors.items():
|
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:
|
if m.ndim != 2 or m.shape[0] == 0:
|
||||||
continue
|
continue
|
||||||
m = m[np.linalg.norm(m, axis=1) > 0]
|
m = m[np.linalg.norm(m, axis=1) > 0]
|
||||||
if m.shape[0] < cfg.min_chunks:
|
if m.shape[0] < cfg.min_chunks:
|
||||||
continue
|
continue
|
||||||
normalized[doc_id] = m / np.linalg.norm(m, axis=1)[:, None]
|
centroid = m.mean(axis=0)
|
||||||
if len(normalized) < 2:
|
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:
|
||||||
return []
|
return []
|
||||||
|
|
||||||
order = sorted(normalized)
|
order = sorted(centroids_by_doc)
|
||||||
centroids = np.array([_unit(normalized[d].mean(axis=0)) for d in order])
|
centroids = np.array([centroids_by_doc[d] 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.
|
# Pairwise cosine, block-wise to avoid a full D×D matrix at once. Each row
|
||||||
# The cap is enforced per row (truncating each row's matches) so a
|
# only compares against higher-indexed documents (upper triangle). Cluster
|
||||||
# self-similar corpus can never allocate beyond MAX_CANDIDATE_PAIRS.
|
# with union-find and keep only each document's best similarity to a twin —
|
||||||
candidates: list[tuple[int, int]] = []
|
# a self-similar corpus forms one clique, so storing every pair would be
|
||||||
|
# O(D²) objects.
|
||||||
|
parent = list(range(len(order)))
|
||||||
|
|
||||||
|
def find(x: int) -> int:
|
||||||
|
while parent[x] != x:
|
||||||
|
parent[x] = parent[parent[x]]
|
||||||
|
x = parent[x]
|
||||||
|
return x
|
||||||
|
|
||||||
|
best = np.zeros(len(order), dtype=np.float32)
|
||||||
|
linked = False
|
||||||
block = 512
|
block = 512
|
||||||
capped = False
|
|
||||||
for start in range(0, len(order), block):
|
for start in range(0, len(order), block):
|
||||||
if capped:
|
|
||||||
break
|
|
||||||
sims = centroids[start : start + block] @ centroids.T
|
sims = centroids[start : start + block] @ centroids.T
|
||||||
for row in range(sims.shape[0]):
|
for row in range(sims.shape[0]):
|
||||||
gi = start + row
|
gi = start + row
|
||||||
targets = np.arange(gi + 1, len(order))
|
cols = (
|
||||||
if targets.size == 0:
|
gi + 1 + np.nonzero(sims[row, gi + 1 :] >= cfg.similarity_threshold)[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)
|
if cols.size == 0:
|
||||||
above = np.nonzero(sims[row, gi + 1 :] >= thresholds)[0]
|
continue
|
||||||
remaining = MAX_CANDIDATE_PAIRS - len(candidates)
|
linked = True
|
||||||
if len(above) >= remaining:
|
row_best = sims[row, cols]
|
||||||
above = above[:remaining]
|
best[gi] = max(best[gi], float(row_best.max()))
|
||||||
capped = True
|
best[cols] = np.maximum(best[cols], row_best)
|
||||||
candidates.extend((gi, gi + 1 + int(j)) for j in above)
|
ri = find(gi)
|
||||||
if capped:
|
for gj in cols.tolist():
|
||||||
break
|
parent[find(gj)] = ri
|
||||||
|
if not linked:
|
||||||
# Stage 2: confirm candidates with directed chunk-overlap containment.
|
|
||||||
adjacency: dict[int, set[int]] = {}
|
|
||||||
edges: dict[tuple[int, int], tuple[float, float]] = {}
|
|
||||||
for i, j in candidates:
|
|
||||||
a_to_b = _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)
|
|
||||||
edges[(i, j)] = (a_to_b, b_to_a)
|
|
||||||
if not edges:
|
|
||||||
return []
|
return []
|
||||||
|
|
||||||
# Cluster confirmed pairs into families (connected components).
|
components: dict[int, list[int]] = {}
|
||||||
|
for idx in range(len(order)):
|
||||||
|
components.setdefault(find(idx), []).append(idx)
|
||||||
|
|
||||||
families: list[_DuplicateFamily] = []
|
families: list[_DuplicateFamily] = []
|
||||||
seen: set[int] = set()
|
for indices in components.values():
|
||||||
for node in adjacency:
|
if len(indices) < 2:
|
||||||
if node in seen:
|
|
||||||
continue
|
continue
|
||||||
component: set[int] = set()
|
members = sorted(order[i] for i in indices)
|
||||||
stack = [node]
|
# Largest document (most chunks) is the one to keep; smallest id on a tie.
|
||||||
while stack:
|
keep = min(members, key=lambda d: (-sizes[d], d))
|
||||||
cur = stack.pop()
|
|
||||||
if cur in seen:
|
|
||||||
continue
|
|
||||||
seen.add(cur)
|
|
||||||
component.add(cur)
|
|
||||||
stack.extend(adjacency[cur] - seen)
|
|
||||||
members = sorted(order[i] for i in component)
|
|
||||||
# Largest document (most chunks) is the likely superset; smallest id on a tie.
|
|
||||||
superset = min(members, key=lambda d: (-normalized[d].shape[0], d))
|
|
||||||
pairs = sorted(
|
|
||||||
(order[i], order[j], round(ab, 3), round(ba, 3))
|
|
||||||
for (i, j), (ab, ba) in edges.items()
|
|
||||||
if i in component and j in component
|
|
||||||
)
|
|
||||||
families.append(
|
families.append(
|
||||||
_DuplicateFamily(
|
_DuplicateFamily(
|
||||||
members=members,
|
members=members,
|
||||||
superset=superset,
|
keep=keep,
|
||||||
pairs=pairs,
|
similarity={order[i]: round(float(best[i]), 3) for i in indices},
|
||||||
sizes={d: int(normalized[d].shape[0]) for d in members},
|
sizes={d: sizes[d] for d in members},
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
return sorted(families, key=lambda f: f.members)
|
return sorted(families, key=lambda f: f.members)
|
||||||
|
|
@ -387,25 +350,21 @@ def _common_path_prefix(labels: list[str]) -> str:
|
||||||
def _write_duplicates_out(
|
def _write_duplicates_out(
|
||||||
path: Path, families: list[_DuplicateFamily], label: Callable[[str], str]
|
path: Path, families: list[_DuplicateFamily], label: Callable[[str], str]
|
||||||
) -> None:
|
) -> None:
|
||||||
"""One block per group; ``keep_suggested`` marks the superset and
|
"""One block per group; ``keep_suggested`` marks the document to keep and
|
||||||
``contained_fraction`` is how much of the document is covered by the rest."""
|
``similarity`` is the highest centroid cosine to another group member."""
|
||||||
groups = []
|
groups = []
|
||||||
for n, family in enumerate(families, start=1):
|
for n, family in enumerate(families, start=1):
|
||||||
contained = dict.fromkeys(family.members, 0.0)
|
|
||||||
for a, b, a_to_b, b_to_a in family.pairs:
|
|
||||||
contained[a] = max(contained[a], a_to_b)
|
|
||||||
contained[b] = max(contained[b], b_to_a)
|
|
||||||
groups.append(
|
groups.append(
|
||||||
{
|
{
|
||||||
"group": n,
|
"group": n,
|
||||||
"keep": family.superset,
|
"keep": family.keep,
|
||||||
"documents": [
|
"documents": [
|
||||||
{
|
{
|
||||||
"document_id": member,
|
"document_id": member,
|
||||||
"document": label(member),
|
"document": label(member),
|
||||||
"chunks": family.sizes[member],
|
"chunks": family.sizes[member],
|
||||||
"contained_fraction": round(contained[member], 3),
|
"similarity": family.similarity[member],
|
||||||
"keep_suggested": member == family.superset,
|
"keep_suggested": member == family.keep,
|
||||||
}
|
}
|
||||||
for member in family.members
|
for member in family.members
|
||||||
],
|
],
|
||||||
|
|
@ -439,7 +398,7 @@ def _check_duplicate_documents(
|
||||||
|
|
||||||
# The terminal report is a summary: show the first few groups whole and
|
# The terminal report is a summary: show the first few groups whole and
|
||||||
# point at the YAML export for the rest. One block per shown group — a
|
# point at the YAML export for the rest. One block per shown group — a
|
||||||
# header, each member on its own numbered line, then a compact overlap line.
|
# header, each member on its own numbered line, then a compact similarity line.
|
||||||
shown = families[:_SAMPLE_LIMIT]
|
shown = families[:_SAMPLE_LIMIT]
|
||||||
prefix = _common_path_prefix([label(m) for f in shown for m in f.members])
|
prefix = _common_path_prefix([label(m) for f in shown for m in f.members])
|
||||||
|
|
||||||
|
|
@ -453,15 +412,14 @@ def _check_duplicate_documents(
|
||||||
for n, family in enumerate(shown, start=1):
|
for n, family in enumerate(shown, start=1):
|
||||||
number = {member: i for i, member in enumerate(family.members, start=1)}
|
number = {member: i for i, member in enumerate(family.members, start=1)}
|
||||||
details.append(
|
details.append(
|
||||||
f"group {n} — {len(family.members)} docs, keep #{number[family.superset]}:"
|
f"group {n} — {len(family.members)} docs, keep #{number[family.keep]}:"
|
||||||
)
|
)
|
||||||
for member in family.members:
|
for member in family.members:
|
||||||
details.append(f" #{number[member]} {short(member)}")
|
details.append(f" #{number[member]} {short(member)}")
|
||||||
overlaps = ", ".join(
|
sims = ", ".join(
|
||||||
f"#{number[a]}→#{number[b]} {ab:.0%}, #{number[b]}→#{number[a]} {ba:.0%}"
|
f"#{number[m]} {family.similarity[m]:.0%}" for m in family.members
|
||||||
for a, b, ab, ba in family.pairs
|
|
||||||
)
|
)
|
||||||
details.append(f" overlap: {overlaps}")
|
details.append(f" similarity: {sims}")
|
||||||
if len(families) > len(shown):
|
if len(families) > len(shown):
|
||||||
details.append(
|
details.append(
|
||||||
f"... (+{len(families) - len(shown)} more groups; "
|
f"... (+{len(families) - len(shown)} more groups; "
|
||||||
|
|
@ -473,11 +431,11 @@ def _check_duplicate_documents(
|
||||||
name="duplicate_documents",
|
name="duplicate_documents",
|
||||||
severity=Severity.WARN,
|
severity=Severity.WARN,
|
||||||
message=(
|
message=(
|
||||||
f"{len(families)} group(s) of documents with substantial chunk overlap "
|
f"{len(families)} group(s) of near-identical documents "
|
||||||
f"(potential duplicates/revisions), {total_docs} documents."
|
f"(potential duplicates), {total_docs} documents."
|
||||||
),
|
),
|
||||||
remediation=(
|
remediation=(
|
||||||
"Review each group and remove redundant revisions; overlap may be intentional."
|
"Review each group and remove redundant copies; duplication may be intentional."
|
||||||
),
|
),
|
||||||
details=details,
|
details=details,
|
||||||
)
|
)
|
||||||
|
|
@ -661,7 +619,12 @@ async def run_db_checks(
|
||||||
)
|
)
|
||||||
|
|
||||||
ids = arrow.column("id").to_pylist()
|
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 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] = []
|
zero_ids: list[str] = []
|
||||||
if vectors.size:
|
if vectors.size:
|
||||||
zero_ids = [ids[i] for i in np.nonzero(~vectors.any(axis=1))[0]]
|
zero_ids = [ids[i] for i in np.nonzero(~vectors.any(axis=1))[0]]
|
||||||
|
|
@ -679,13 +642,16 @@ async def run_db_checks(
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
# Near-duplicate documents (revisions sharing most chunks), grouped from the
|
# Near-identical documents (centroid cosine), grouped from the same vector
|
||||||
# same vector scan rather than a second pass.
|
# scan rather than a second pass.
|
||||||
chunk_doc_ids_ordered = arrow.column("document_id").to_pylist()
|
chunk_doc_ids_ordered = arrow.column("document_id").to_pylist()
|
||||||
indices_by_doc: dict[str, list[int]] = {}
|
indices_by_doc: dict[str, list[int]] = {}
|
||||||
for index, doc_id in enumerate(chunk_doc_ids_ordered):
|
for index, doc_id in enumerate(chunk_doc_ids_ordered):
|
||||||
indices_by_doc.setdefault(doc_id, []).append(index)
|
indices_by_doc.setdefault(doc_id, []).append(index)
|
||||||
doc_vectors = {doc_id: vectors[idx] for doc_id, idx in indices_by_doc.items()}
|
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(
|
results.append(
|
||||||
_check_duplicate_documents(
|
_check_duplicate_documents(
|
||||||
doc_vectors,
|
doc_vectors,
|
||||||
|
|
|
||||||
|
|
@ -947,8 +947,9 @@ async def test_probe_endpoint_connection_error():
|
||||||
def _docs(spec: dict[str, list[int]], dim: int = 8) -> dict[str, np.ndarray]:
|
def _docs(spec: dict[str, list[int]], dim: int = 8) -> dict[str, np.ndarray]:
|
||||||
"""Build per-document chunk matrices from one-hot indices.
|
"""Build per-document chunk matrices from one-hot indices.
|
||||||
|
|
||||||
A shared index across documents is a shared (identical) chunk; distinct
|
Orthogonal one-hot chunks make the centroid cosine of two documents equal to
|
||||||
indices are orthogonal, so they never count as twins.
|
``shared / sqrt(len(a) * len(b))``: identical documents score 1.0, fully
|
||||||
|
distinct documents score 0.0.
|
||||||
"""
|
"""
|
||||||
eye = np.eye(dim)
|
eye = np.eye(dim)
|
||||||
return {
|
return {
|
||||||
|
|
@ -956,90 +957,67 @@ def _docs(spec: dict[str, list[int]], dim: int = 8) -> dict[str, np.ndarray]:
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
# Stage-2 (containment/clustering) unit tests disable the centroid gate
|
def test_duplicate_families_identical_docs_flagged():
|
||||||
# (candidate_threshold=0.0) so every pair is verified; orthogonal one-hot chunks
|
|
||||||
# would otherwise drop centroids below the default gate. The gate itself is
|
|
||||||
# exercised by the end-to-end tests below.
|
|
||||||
def _stage2_cfg(**kw) -> DuplicateDetectionConfig:
|
|
||||||
return DuplicateDetectionConfig(candidate_threshold=0.0, **kw)
|
|
||||||
|
|
||||||
|
|
||||||
def test_duplicate_families_revision_pair():
|
|
||||||
families = _duplicate_families(
|
families = _duplicate_families(
|
||||||
_docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 4]}), _stage2_cfg()
|
_docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3]}), DuplicateDetectionConfig()
|
||||||
)
|
)
|
||||||
assert len(families) == 1
|
assert len(families) == 1
|
||||||
assert set(families[0].members) == {"a", "b"}
|
assert set(families[0].members) == {"a", "b"}
|
||||||
|
assert families[0].similarity == {"a": pytest.approx(1.0), "b": pytest.approx(1.0)}
|
||||||
|
|
||||||
|
|
||||||
def test_duplicate_families_append_only_is_asymmetric():
|
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(
|
families = _duplicate_families(
|
||||||
_docs({"a": [0, 1, 2], "b": [0, 1, 2, 3, 4, 5]}), _stage2_cfg()
|
_docs({"a": [0, 1, 2], "b": [0, 1, 2, 3, 4, 5]}), DuplicateDetectionConfig()
|
||||||
)
|
)
|
||||||
assert len(families) == 1
|
assert families == []
|
||||||
fam = families[0]
|
|
||||||
assert fam.superset == "b" # the larger document
|
|
||||||
# directed containment: all of A is in B (1.0); only half of B is in A.
|
|
||||||
a_to_b = next(p for p in fam.pairs if p[:2] == ("a", "b"))
|
|
||||||
assert a_to_b[2] == pytest.approx(1.0)
|
|
||||||
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():
|
def test_duplicate_families_distinct_docs_none():
|
||||||
families = _duplicate_families(
|
families = _duplicate_families(
|
||||||
_docs({"a": [0, 1, 2], "b": [3, 4, 5]}), _stage2_cfg()
|
_docs({"a": [0, 1, 2], "b": [3, 4, 5]}), DuplicateDetectionConfig()
|
||||||
)
|
)
|
||||||
assert families == []
|
assert families == []
|
||||||
|
|
||||||
|
|
||||||
def test_duplicate_families_three_way_chain_one_family():
|
def test_duplicate_families_three_way_one_family():
|
||||||
families = _duplicate_families(
|
families = _duplicate_families(
|
||||||
_docs(
|
_docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [0, 1, 2, 3]}),
|
||||||
{
|
DuplicateDetectionConfig(),
|
||||||
"a": [0, 1, 2, 3],
|
|
||||||
"b": [0, 1, 2, 3, 4],
|
|
||||||
"c": [0, 1, 2, 3, 4, 5],
|
|
||||||
}
|
|
||||||
),
|
|
||||||
_stage2_cfg(),
|
|
||||||
)
|
)
|
||||||
assert len(families) == 1
|
assert len(families) == 1
|
||||||
assert set(families[0].members) == {"a", "b", "c"}
|
assert set(families[0].members) == {"a", "b", "c"}
|
||||||
assert families[0].superset == "c"
|
# Equal sizes -> smallest id is kept.
|
||||||
|
assert families[0].keep == "a"
|
||||||
|
|
||||||
|
|
||||||
|
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())
|
||||||
|
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())
|
||||||
|
|
||||||
|
|
||||||
def test_duplicate_families_tiny_docs_ignored():
|
def test_duplicate_families_tiny_docs_ignored():
|
||||||
# min_chunks = 3 excludes the one-chunk documents.
|
# min_chunks = 3 excludes the one-chunk documents.
|
||||||
families = _duplicate_families(_docs({"a": [0], "b": [0]}), _stage2_cfg())
|
families = _duplicate_families(
|
||||||
|
_docs({"a": [0], "b": [0]}), DuplicateDetectionConfig()
|
||||||
|
)
|
||||||
assert families == []
|
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():
|
def test_duplicate_families_threshold_is_configurable():
|
||||||
# Share 3 of 5 chunks each -> containment 0.6 both ways.
|
# Share 3 of 4 chunks each -> centroid cosine 0.75.
|
||||||
spec = {"a": [0, 1, 2, 3, 4], "b": [0, 1, 2, 5, 6]}
|
spec = {"a": [0, 1, 2, 3], "b": [0, 1, 2, 4]}
|
||||||
assert _duplicate_families(_docs(spec), _stage2_cfg()) == []
|
assert _duplicate_families(_docs(spec), DuplicateDetectionConfig()) == []
|
||||||
flagged = _duplicate_families(_docs(spec), _stage2_cfg(containment_threshold=0.6))
|
flagged = _duplicate_families(
|
||||||
|
_docs(spec), DuplicateDetectionConfig(similarity_threshold=0.7)
|
||||||
|
)
|
||||||
assert len(flagged) == 1
|
assert len(flagged) == 1
|
||||||
assert set(flagged[0].members) == {"a", "b"}
|
assert set(flagged[0].members) == {"a", "b"}
|
||||||
|
|
||||||
|
|
@ -1053,7 +1031,7 @@ def test_duplicate_documents_report_truncates_summary():
|
||||||
spec[f"b{k}"] = list(idx)
|
spec[f"b{k}"] = list(idx)
|
||||||
docs = _docs(spec, dim=3 * pairs)
|
docs = _docs(spec, dim=3 * pairs)
|
||||||
uris = {d: f"file:///srv/shared/library/docs/{d}.pdf" for d in spec}
|
uris = {d: f"file:///srv/shared/library/docs/{d}.pdf" for d in spec}
|
||||||
result = _check_duplicate_documents(docs, uris, {}, _stage2_cfg())
|
result = _check_duplicate_documents(docs, uris, {}, DuplicateDetectionConfig())
|
||||||
assert result.severity is Severity.WARN
|
assert result.severity is Severity.WARN
|
||||||
# The summary message still reports the full total.
|
# The summary message still reports the full total.
|
||||||
assert f"{pairs} group(s)" in result.message
|
assert f"{pairs} group(s)" in result.message
|
||||||
|
|
@ -1067,7 +1045,7 @@ def test_duplicate_documents_report_factors_common_path():
|
||||||
docs = _docs({"a": [0, 1, 2], "b": [0, 1, 2]}, dim=3)
|
docs = _docs({"a": [0, 1, 2], "b": [0, 1, 2]}, dim=3)
|
||||||
base = "file:///srv/shared/library/docs/"
|
base = "file:///srv/shared/library/docs/"
|
||||||
uris = {"a": base + "alpha.pdf", "b": base + "beta.pdf"}
|
uris = {"a": base + "alpha.pdf", "b": base + "beta.pdf"}
|
||||||
result = _check_duplicate_documents(docs, uris, {}, _stage2_cfg())
|
result = _check_duplicate_documents(docs, uris, {}, DuplicateDetectionConfig())
|
||||||
assert f"common path: {base}" in result.details
|
assert f"common path: {base}" in result.details
|
||||||
member_lines = [d for d in result.details if d.lstrip().startswith("#")]
|
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"}
|
assert {d.strip() for d in member_lines} == {"#1 alpha.pdf", "#2 beta.pdf"}
|
||||||
|
|
@ -1079,7 +1057,9 @@ def test_duplicate_documents_writes_yaml(tmp_path):
|
||||||
docs = _docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [4, 5, 6]}, dim=8)
|
docs = _docs({"a": [0, 1, 2, 3], "b": [0, 1, 2, 3], "c": [4, 5, 6]}, dim=8)
|
||||||
uris = {"a": "file:///x/a.pdf", "b": "file:///x/b.pdf", "c": "file:///x/c.pdf"}
|
uris = {"a": "file:///x/a.pdf", "b": "file:///x/b.pdf", "c": "file:///x/c.pdf"}
|
||||||
out = tmp_path / "dups.yaml"
|
out = tmp_path / "dups.yaml"
|
||||||
_check_duplicate_documents(docs, uris, {}, _stage2_cfg(), yaml_path=out)
|
_check_duplicate_documents(
|
||||||
|
docs, uris, {}, DuplicateDetectionConfig(), yaml_path=out
|
||||||
|
)
|
||||||
data = yaml.safe_load(out.read_text())
|
data = yaml.safe_load(out.read_text())
|
||||||
assert len(data["groups"]) == 1
|
assert len(data["groups"]) == 1
|
||||||
group = data["groups"][0]
|
group = data["groups"][0]
|
||||||
|
|
@ -1088,6 +1068,7 @@ def test_duplicate_documents_writes_yaml(tmp_path):
|
||||||
assert [d["document_id"] for d in docs_out] == ["a", "b"]
|
assert [d["document_id"] for d in docs_out] == ["a", "b"]
|
||||||
assert [d["document"] for d in docs_out] == ["file:///x/a.pdf", "file:///x/b.pdf"]
|
assert [d["document"] for d in docs_out] == ["file:///x/a.pdf", "file:///x/b.pdf"]
|
||||||
assert all(d["chunks"] == 4 for d in docs_out)
|
assert all(d["chunks"] == 4 for d in docs_out)
|
||||||
|
assert all(d["similarity"] == pytest.approx(1.0) for d in docs_out)
|
||||||
assert {d["document_id"]: d["keep_suggested"] for d in docs_out} == {
|
assert {d["document_id"]: d["keep_suggested"] for d in docs_out} == {
|
||||||
"a": True,
|
"a": True,
|
||||||
"b": False,
|
"b": False,
|
||||||
|
|
@ -1098,7 +1079,7 @@ def test_duplicate_documents_writes_empty_yaml_when_none(tmp_path):
|
||||||
docs = _docs({"a": [0, 1, 2], "b": [3, 4, 5]}, dim=6) # distinct
|
docs = _docs({"a": [0, 1, 2], "b": [3, 4, 5]}, dim=6) # distinct
|
||||||
out = tmp_path / "dups.yaml"
|
out = tmp_path / "dups.yaml"
|
||||||
_check_duplicate_documents(
|
_check_duplicate_documents(
|
||||||
docs, {"a": "u", "b": "v"}, {}, _stage2_cfg(), yaml_path=out
|
docs, {"a": "u", "b": "v"}, {}, DuplicateDetectionConfig(), yaml_path=out
|
||||||
)
|
)
|
||||||
assert yaml.safe_load(out.read_text()) == {"groups": []}
|
assert yaml.safe_load(out.read_text()) == {"groups": []}
|
||||||
|
|
||||||
|
|
@ -1162,7 +1143,7 @@ async def _build_dup_db(path, docs: dict[str, list[int]], *, vector_dim: int = 8
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_duplicate_documents_check_warns_end_to_end(temp_db_path):
|
async def test_duplicate_documents_check_warns_end_to_end(temp_db_path):
|
||||||
await _build_dup_db(temp_db_path, {"a": [0, 1, 2, 3], "b": [0, 1, 2, 3, 4]})
|
await _build_dup_db(temp_db_path, {"a": [0, 1, 2, 3], "b": [0, 1, 2, 3]})
|
||||||
report = await run_doctor(_config(vector_dim=8), temp_db_path, {})
|
report = await run_doctor(_config(vector_dim=8), temp_db_path, {})
|
||||||
result = _result(report, "duplicate_documents")
|
result = _result(report, "duplicate_documents")
|
||||||
assert result.severity is Severity.WARN
|
assert result.severity is Severity.WARN
|
||||||
|
|
@ -1179,14 +1160,10 @@ async def test_duplicate_documents_check_ok_when_distinct(temp_db_path):
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_duplicate_documents_check_reads_config(temp_db_path):
|
async def test_duplicate_documents_check_reads_config(temp_db_path):
|
||||||
# Share 3 of 5 -> containment 0.6. Disable the centroid gate on both runs so
|
# Share 3 of 5 -> centroid cosine 0.6, below the default 0.97 cutoff.
|
||||||
# only containment_threshold decides the outcome.
|
|
||||||
await _build_dup_db(temp_db_path, {"a": [0, 1, 2, 3, 4], "b": [0, 1, 2, 5, 6]})
|
await _build_dup_db(temp_db_path, {"a": [0, 1, 2, 3, 4], "b": [0, 1, 2, 5, 6]})
|
||||||
|
|
||||||
base = _config(vector_dim=8)
|
base = _config(vector_dim=8)
|
||||||
base.doctor = DoctorConfig(
|
|
||||||
duplicates=DuplicateDetectionConfig(candidate_threshold=0.0)
|
|
||||||
)
|
|
||||||
assert (
|
assert (
|
||||||
_result(
|
_result(
|
||||||
await run_doctor(base, temp_db_path, {}), "duplicate_documents"
|
await run_doctor(base, temp_db_path, {}), "duplicate_documents"
|
||||||
|
|
@ -1196,9 +1173,7 @@ async def test_duplicate_documents_check_reads_config(temp_db_path):
|
||||||
|
|
||||||
tuned = _config(vector_dim=8)
|
tuned = _config(vector_dim=8)
|
||||||
tuned.doctor = DoctorConfig(
|
tuned.doctor = DoctorConfig(
|
||||||
duplicates=DuplicateDetectionConfig(
|
duplicates=DuplicateDetectionConfig(similarity_threshold=0.5)
|
||||||
candidate_threshold=0.0, containment_threshold=0.6
|
|
||||||
)
|
|
||||||
)
|
)
|
||||||
assert (
|
assert (
|
||||||
_result(
|
_result(
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue