93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
"""Candidate pair generation for the Lint pass.
|
|
|
|
For each source claim atom, we pull its nearest neighbors via the Atomic
|
|
/api/atoms/{id}/similar endpoint (which uses pgvector kNN under the hood).
|
|
Only Type/Claim neighbors are kept; Type/Document neighbors are dropped so
|
|
we only compare claim-to-claim.
|
|
|
|
Pairs are canonicalized as (lower_id, higher_id) so (A, B) and (B, A)
|
|
produce the same CandidatePair and don't get evaluated twice.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from shared.atomic_api import AtomicClient
|
|
from shared.logging import get_logger
|
|
|
|
log = get_logger(__name__)
|
|
|
|
# Atoms with similarity below this are considered too unrelated to be worth
|
|
# an NLI call. Between 0.55 and 0.95 is the interesting band — below is
|
|
# "probably different topic", above is "almost certainly the same text".
|
|
MIN_PAIR_SIMILARITY = 0.55
|
|
NEIGHBORS_PER_CLAIM = 15
|
|
|
|
|
|
@dataclass(slots=True, frozen=True)
|
|
class CandidatePair:
|
|
"""A pair of claim atom ids to evaluate, canonicalized lo<hi."""
|
|
|
|
atom_a_id: str
|
|
atom_b_id: str
|
|
similarity: float
|
|
|
|
|
|
def _is_claim_hit(hit_tags: list[dict]) -> bool:
|
|
return any(t.get("name") == "Claim" for t in hit_tags)
|
|
|
|
|
|
async def generate_candidates(
|
|
atomic: AtomicClient,
|
|
*,
|
|
source_atom_ids: list[str],
|
|
min_similarity: float = MIN_PAIR_SIMILARITY,
|
|
neighbors_per: int = NEIGHBORS_PER_CLAIM,
|
|
) -> list[CandidatePair]:
|
|
"""Walk each source atom and collect canonical candidate pairs.
|
|
|
|
Limitation: Atomic's /similar endpoint returns Document and Claim atoms
|
|
mixed together. We filter client-side for Type/Claim. It might be worth
|
|
upstream adding a tag filter later, but for this corpus size the current
|
|
approach is fine.
|
|
"""
|
|
seen: set[tuple[str, str]] = set()
|
|
out: list[CandidatePair] = []
|
|
|
|
for idx, src_id in enumerate(source_atom_ids):
|
|
try:
|
|
hits = await atomic.find_similar(
|
|
src_id, threshold=min_similarity, limit=neighbors_per
|
|
)
|
|
except Exception as e: # noqa: BLE001
|
|
log.warning(
|
|
"find_similar_failed", atom_id=src_id, error=f"{type(e).__name__}:{e}"
|
|
)
|
|
continue
|
|
|
|
for h in hits:
|
|
if h.atom_id == src_id:
|
|
continue
|
|
if not _is_claim_hit(h.tags):
|
|
continue
|
|
lo, hi = sorted((src_id, h.atom_id))
|
|
key = (lo, hi)
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
out.append(
|
|
CandidatePair(
|
|
atom_a_id=lo, atom_b_id=hi, similarity=round(h.similarity, 4)
|
|
)
|
|
)
|
|
|
|
if (idx + 1) % 50 == 0:
|
|
log.info(
|
|
"candidates_progress",
|
|
processed=idx + 1,
|
|
total=len(source_atom_ids),
|
|
pairs=len(out),
|
|
)
|
|
|
|
return out
|