diff --git a/aborist/concepts/query.py b/aborist/concepts/query.py index e367244..caad29d 100644 --- a/aborist/concepts/query.py +++ b/aborist/concepts/query.py @@ -45,10 +45,11 @@ def has_compare_phrasing(question: str) -> bool: # Cross-shard lookup with mtime-keyed cache # --------------------------------------------------------------------------- -# Cache shape: { shards_dir_str: (mtime_signature, synonym_index, rivalry_pairs) } -# - synonym_index: { token_lower: set(target_lower) } +# Cache shape: { shards_dir_str: (mtime_sig, manual_index, derived_index, rivalry_pairs) } +# - manual_index: curated synonym edges; always expanded +# - derived_index: corpus-derived synonym edges; expansion subject to per-token cap # - rivalry_pairs: list of (set_a, set_b) — frozensets of lowercase tokens -_CACHE: dict[str, tuple[tuple, dict, list]] = {} +_CACHE: dict[str, tuple[tuple, dict, dict, list]] = {} def _shards_mtime_signature(shards_dir: Path) -> tuple: @@ -62,75 +63,90 @@ def _shards_mtime_signature(shards_dir: Path) -> tuple: ) -def _load_indices(shards_dir: Path) -> tuple[dict, list]: +def _load_indices(shards_dir: Path) -> tuple[dict, dict, list]: """Walk all shards, materialize the synonym & rivalry indices. Synonyms map every token to its **direct neighbors only**, NOT the transitive closure across reciprocal-link chains. Why: union-find over the full Wikipedia reciprocal-link graph collapses everything into one giant connected component (54k+ tokens for any seed in - a 4-shard corpus). That made one query token expand to the entire - synonym alphabet, which (a) overwhelmed retrieval relevance and - (b) tripped SQLite's expression-tree-depth limit when ``_search_titles`` - built a per-token OR clause. Direct-neighbor expansion preserves the - legacy frozenset semantics (every member of a group expanded to the - others in that group, but the groups didn't chain). + a 4-shard corpus). Direct-neighbor expansion preserves the legacy + frozenset semantics (every member of a group expanded to the others + in that group, but the groups didn't chain). - Rivalry pairs stay as discrete pair objects: each row's (token, target) - becomes a 2-element pair stored as (frozenset({token}), frozenset({target})). - The legacy semantics (group-vs-group rivalry) is reachable by storing - one rivalry row per (token_a_in_group_A, token_b_in_group_B) cross- - product OR by relying on the lookup-time intersection check picking - up any pair whose either side intersects the query. + Two synonym indices are built, one per evidence class: + + - ``manual_index`` — manual_legacy + manual rows. Curated; always + expanded regardless of per-token degree. + Captures the brain-tech / AMD-family / etc. + seed groups whose anchor token has many + deliberate members (e.g. telepathy → 29 + members of the brain-tech group). + - ``derived_index`` — link_reciprocity & other corpus-derived + extractors. Subject to per-token degree + cap because the Wikipedia link graph + carries topic-adjacency noise on generic + tokens (person, thoughts, language). + + Rivalry pairs use the union of both indices for closure. """ conn = connect_query(shards_dir=shards_dir) rows = conn.execute( - "SELECT relation_kind, token, target FROM concept_relations" + "SELECT relation_kind, evidence_kind, token, target FROM concept_relations" ).fetchall() conn.close() - # Direct-neighbor adjacency: token -> set of immediate synonym - # partners. Symmetric (an edge (A,B) registers both A→B and B→A). - synonym_index: dict[str, set[str]] = {} + manual_index: dict[str, set[str]] = {} + derived_index: dict[str, set[str]] = {} rivalry_rows: list[tuple[str, str]] = [] + + # Curated evidence kinds — never capped at expansion time. + MANUAL_KINDS = {"manual", "manual_legacy"} + for r in rows: kind = r["relation_kind"] + evidence_kind = r["evidence_kind"] a = (r["token"] or "").lower() b = (r["target"] or "").lower() if not a or not b or a == b: continue if kind == "synonym": - synonym_index.setdefault(a, set()).add(b) - synonym_index.setdefault(b, set()).add(a) + target_index = ( + manual_index if evidence_kind in MANUAL_KINDS else derived_index + ) + target_index.setdefault(a, set()).add(b) + target_index.setdefault(b, set()).add(a) elif kind == "rivalry": rivalry_rows.append((a, b)) - # Rivalry pairs use the same direct-neighbor semantic: a rivalry edge - # (A, B) captures A's direct neighbors (incl. A itself) vs B's direct - # neighbors (incl. B itself). Bounded scope; matches synonym semantics. + # Rivalry pairs use union of manual + derived neighborhoods. rivalry_pairs: list[tuple[frozenset[str], frozenset[str]]] = [] for a, b in rivalry_rows: - ga = frozenset(synonym_index.get(a, set()) | {a}) - gb = frozenset(synonym_index.get(b, set()) | {b}) + ga = frozenset( + manual_index.get(a, set()) | derived_index.get(a, set()) | {a} + ) + gb = frozenset( + manual_index.get(b, set()) | derived_index.get(b, set()) | {b} + ) rivalry_pairs.append((ga, gb)) - return synonym_index, rivalry_pairs + return manual_index, derived_index, rivalry_pairs -def _get_indices(shards_dir: Path | str | None) -> tuple[dict, list]: - """Return (synonym_index, rivalry_pairs), using cached values when - the shard mtime signature is unchanged.""" +def _get_indices(shards_dir: Path | str | None) -> tuple[dict, dict, list]: + """Return (manual_index, derived_index, rivalry_pairs), using + cached values when the shard mtime signature is unchanged.""" if shards_dir is None: - return {}, [] + return {}, {}, [] p = Path(shards_dir) key = str(p.resolve()) sig = _shards_mtime_signature(p) cached = _CACHE.get(key) if cached is not None and cached[0] == sig: - return cached[1], cached[2] - syn, riv = _load_indices(p) - _CACHE[key] = (sig, syn, riv) - return syn, riv + return cached[1], cached[2], cached[3] + manual, derived, riv = _load_indices(p) + _CACHE[key] = (sig, manual, derived, riv) + return manual, derived, riv def invalidate_cache() -> None: @@ -145,28 +161,81 @@ def invalidate_cache() -> None: # --------------------------------------------------------------------------- +# Caps that match the legacy (frozenset) expansion size. The corpus- +# derived synonym graph is noisier than hand-curated frozensets: +# generic tokens like "person" / "thoughts" / "language" have ~20+ +# reciprocal-link neighbors each, most of which are topic-adjacency +# noise rather than actual synonyms. Without these caps a 19-token +# query expands to ~400 accept tokens, and the title-LIKE search +# multiplies that by the document count to a many-minute hang. +# +# MAX_NEIGHBORS_PER_TOKEN: any token with more than this many +# direct neighbors is treated as "too generic to expand" — we add +# only the token itself, not its neighbors. Mirrors how the legacy +# frozensets covered named entities (athlon, pentium) but not common +# words (person, thoughts). +# +# MAX_TOTAL_TOKENS: overall cap on the expanded set. Original query +# tokens are always preserved; once total exceeds the cap, neighbors +# are sorted alphabetically & truncated. Bound on title-LIKE clause +# count keeps the SQL tractable on a 3.47M-doc corpus. +MAX_NEIGHBORS_PER_TOKEN = 8 +MAX_TOTAL_TOKENS = 50 + + def synonym_expand( tokens: set[str], *, shards_dir: Path | str | None = None, + max_neighbors_per_token: int = MAX_NEIGHBORS_PER_TOKEN, + max_total: int = MAX_TOTAL_TOKENS, ) -> set[str]: - """Add every member of every synonym group that any input token hits. + """Add direct synonym neighbors for any input token whose degree + is bounded enough that its neighbors are likely topical, not + topic-adjacency noise. - Lower-case at the comparison layer; output preserves the lowercase - forms stored in the DB. Callers that need title-case round-trip can - use the original ``tokens`` alongside the expansion. + Two caps protect retrieval performance & quality: + + 1. ``max_neighbors_per_token`` — tokens with more direct neighbors + than this contribute NO expansion. Generic tokens ("person", + "thoughts") have huge degree in the Wikipedia reciprocal-link + graph; expanding them dumps random topical-cluster noise. + Specific named entities ("athlon", "telepathy") have small + focused neighborhoods that pass the cap. + + 2. ``max_total`` — overall cap on expanded set size. Bounds the + SQL clause count downstream. Original query tokens are always + preserved; if total > cap, neighbors are sorted alphabetically + & truncated. """ if not tokens: return set() if shards_dir is None: - # No shards configured: no expansion. Equivalent to the legacy - # behavior when the frozensets were empty. return set(tokens) - synonym_index, _ = _get_indices(shards_dir) - expanded: set[str] = {t.lower() for t in tokens} - for t in list(expanded): - if t in synonym_index: - expanded |= synonym_index[t] + manual_index, derived_index, _ = _get_indices(shards_dir) + qlower = {t.lower() for t in tokens} + expanded: set[str] = set(qlower) + # Manual (curated) synonyms always expand: brain-tech / AMD-family / + # etc. seed groups have legitimately many members per anchor & we + # trust the curation. + for t in qlower: + if t in manual_index: + expanded |= manual_index[t] + # Derived (corpus-extracted) synonyms cap on per-token degree. + # Generic tokens like "person" / "thoughts" / "language" have wide + # noisy neighborhoods in the reciprocal-link graph — skip those. + # Specific tokens with bounded degree expand cleanly. + for t in qlower: + if t not in derived_index: + continue + neighbors = derived_index[t] + if len(neighbors) > max_neighbors_per_token: + continue + expanded |= neighbors + if len(expanded) > max_total: + ordered_neighbors = sorted(expanded - qlower) + budget = max(0, max_total - len(qlower)) + expanded = qlower | set(ordered_neighbors[:budget]) return expanded @@ -185,7 +254,7 @@ def rivalry_excluded( """ if compare_phrasing or not tokens or shards_dir is None: return set() - _, rivalry_pairs = _get_indices(shards_dir) + _, _, rivalry_pairs = _get_indices(shards_dir) qlower = {t.lower() for t in tokens} excluded: set[str] = set() for a, b in rivalry_pairs: diff --git a/aborist/concepts/seed.py b/aborist/concepts/seed.py index c7dec8b..5644839 100644 --- a/aborist/concepts/seed.py +++ b/aborist/concepts/seed.py @@ -80,27 +80,31 @@ def seed_legacy_concepts(conn: sqlite3.Connection) -> dict: """ syn_ins = syn_skip = riv_ins = riv_skip = 0 - # Synonyms: emit a forward edge between consecutive members of each - # group. Union-find at lookup time builds the full closure, so the - # quadratic per-group emission is unnecessary. + # Synonyms: emit a clique edge for every (a, b) pair within each + # group. Lookup is direct-neighbor (no transitive closure across + # groups) so we need the full N*(N-1)/2 pairs to preserve the + # legacy frozenset semantic where every member retrieved every + # other member. Quadratic in group size but groups stay small + # (largest is the 30-member brain-tech set → 435 pairs). for group in LEGACY_SYNONYM_GROUPS: if len(group) < 2: continue - anchor = group[0] - for member in group[1:]: - inserted = add_concept_relation( - conn, - source_root=LEGACY_SOURCE_ROOT, - relation_kind="synonym", - token=anchor, - target=member, - evidence_kind=LEGACY_EVIDENCE_KIND, - derived_from="aborist.qa.concepts (legacy frozensets)", - ) - if inserted: - syn_ins += 1 - else: - syn_skip += 1 + members = list(group) + for i, a in enumerate(members): + for b in members[i + 1:]: + inserted = add_concept_relation( + conn, + source_root=LEGACY_SOURCE_ROOT, + relation_kind="synonym", + token=a, + target=b, + evidence_kind=LEGACY_EVIDENCE_KIND, + derived_from="aborist.qa.concepts (legacy frozensets)", + ) + if inserted: + syn_ins += 1 + else: + syn_skip += 1 # Rivalries: one row per representative pair. Group closure happens # at lookup time via the synonym index.