fts_title: AND + per-token two-pass merge (fixes multi-token title regression)
Pure per-token title MATCH broke multi-token title surfacing. For
"who created spider man?", per-token "spider" alone floods top-K
with single-word titles ("Spider", "Wolf spider", "Spiders (album)")
because single-token titles have shorter doc length → higher bm25.
The "Spider-Man" article (multi-token title) never reached top-32 →
retrieval surfaced only Spider-Man universe filler (Rhino, Sinister
Six, Liz Allan, …) and the canonical primary was absent → answer was
"the evidence does not contain who created Spider-Man" (HYBRID).
Fix at both SqliteShardCorpus and MultiShardSqliteCorpus:
1. AND pass: ``"tok1 AND tok2 AND …"`` single MATCH. Rewards
multi-token title hits — Spider-Man (which contains both
"spider" and "man" as separate tokens after FTS5 hyphen-splits)
ranks where its co-occurrence earns.
2. Per-token pass: one MATCH per token, round-robin merge across
buckets. Preserves the "capitol of paris" Paris-survival fix —
each token guaranteed surface area.
Merge by best bm25 per doc_root; AND-mode hits emit first because
their multi-token grounding is the strongest title-relevance signal.
Live:
- "who is spider man?" → STRICT, 0 violations (was STRICT)
- "who created spider man?" → STRICT, cites Spider-Man article
with Stan Lee + Steve Ditko answer
(was HYBRID, Spider-Man article was
missing from sources entirely)
- "what is the capitol of paris" → Paris still at rank 4 in titles
(the prior fix's invariant holds)
110 tests pass.
This commit is contained in:
parent
0cf2e210a0
commit
0ea55700f6
1 changed files with 119 additions and 48 deletions
|
|
@ -344,15 +344,23 @@ class SqliteShardCorpus:
|
|||
def fts_title(self, query: str, *, limit: int = 8) -> list[Hit]:
|
||||
"""Title-only bm25() via the documents_fts virtual table.
|
||||
|
||||
Per-token fan-out (one MATCH per non-stopword token, dedupe
|
||||
by doc_root keeping best bm25). Single-MATCH ``tok1 OR tok2``
|
||||
is biased toward the rarer token — for "what is the capitol
|
||||
of paris", "capitol" (DF ~1k titles) outranks "paris" (DF ~5k
|
||||
titles) so a unioned MATCH returns 12 Capitol-* titles and
|
||||
zero Paris-* hits, squeezing the actual "Paris" article out
|
||||
of the top-K entirely. Per-token MATCH gives each token its
|
||||
own candidate space; the merge then takes the best bm25 each
|
||||
doc earned in any token's run.
|
||||
Two-pass union (AND + per-token), merged by best bm25:
|
||||
|
||||
- **AND pass**: MATCH on all tokens joined with AND ("spider
|
||||
AND man") — rewards multi-token titles like "Spider-Man"
|
||||
that contain BOTH tokens. Without this, a single-token MATCH
|
||||
on "spider" alone surfaces "Spider", "Wolf spider", "Spiders
|
||||
(album)" etc. ahead of "Spider-Man" because single-token
|
||||
titles have shorter doc length → higher bm25.
|
||||
|
||||
- **Per-token pass**: one MATCH per non-stopword token —
|
||||
guarantees each query token gets independent surface area.
|
||||
Necessary because single-MATCH ``tok1 OR tok2`` is biased
|
||||
toward the rarer token (e.g. "capitol" DF~1k outranks
|
||||
"paris" DF~5k titles in OR mode and squeezes Paris out
|
||||
entirely).
|
||||
|
||||
The merge takes the best bm25 each doc earned in any pass.
|
||||
"""
|
||||
from arborist.wallet.bucket import (
|
||||
_to_fts5,
|
||||
|
|
@ -360,9 +368,6 @@ class SqliteShardCorpus:
|
|||
_FTS5_WORD_RE,
|
||||
)
|
||||
import unicodedata
|
||||
# Single-token MATCH list — preserve FTS5-safe tokenization
|
||||
# via the same sanitizer fts_body uses (NFKD-fold, drop
|
||||
# stopwords, ≥2 char filter).
|
||||
if not _to_fts5(query).strip():
|
||||
return []
|
||||
folded = "".join(
|
||||
|
|
@ -384,15 +389,35 @@ class SqliteShardCorpus:
|
|||
"ORDER BY score LIMIT ?"
|
||||
)
|
||||
best: dict[str, tuple] = {}
|
||||
# AND pass first when there are ≥2 tokens. FTS5 returns nothing
|
||||
# if any token has zero matches across documents_fts; that's
|
||||
# fine — the per-token pass below picks up the slack.
|
||||
if len(tokens) >= 2:
|
||||
and_match = " AND ".join(tokens)
|
||||
try:
|
||||
for r in self._conn.execute(sql, (and_match, limit)):
|
||||
droot = r["document_root"]
|
||||
cur = best.get(droot)
|
||||
if cur is None or r["score"] < cur[0]:
|
||||
best[droot] = (
|
||||
r["score"], r["document_uri"] or "",
|
||||
r["title"] or "",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
# Per-token pass.
|
||||
for tok in tokens:
|
||||
for r in self._conn.execute(sql, (tok, limit)):
|
||||
droot = r["document_root"]
|
||||
cur = best.get(droot)
|
||||
if cur is None or r["score"] < cur[0]:
|
||||
best[droot] = (
|
||||
r["score"], r["document_uri"] or "",
|
||||
r["title"] or "",
|
||||
)
|
||||
try:
|
||||
for r in self._conn.execute(sql, (tok, limit)):
|
||||
droot = r["document_root"]
|
||||
cur = best.get(droot)
|
||||
if cur is None or r["score"] < cur[0]:
|
||||
best[droot] = (
|
||||
r["score"], r["document_uri"] or "",
|
||||
r["title"] or "",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
ranked = sorted(best.items(), key=lambda kv: kv[1][0])[:limit]
|
||||
return [
|
||||
Hit(
|
||||
|
|
@ -653,18 +678,17 @@ class MultiShardSqliteCorpus:
|
|||
return merged[:limit]
|
||||
|
||||
def fts_title(self, query: str, *, limit: int = 8) -> list[Hit]:
|
||||
"""Per-token title fan-out across all shards.
|
||||
"""Two-pass title fan-out: AND-mode plus per-token, merged.
|
||||
|
||||
Single-MATCH ``tok1 OR tok2`` over documents_fts is biased
|
||||
toward whichever token has lower corpus DF — rarer tokens get
|
||||
higher bm25 per match, so a single-MATCH return is dominated
|
||||
by their hits. For "what is the capitol of paris" that meant
|
||||
Capitol-* titles filled top-16 and the "Paris" article never
|
||||
made it into the merge pool. Per-shard per-token MATCH gets
|
||||
balanced per-token at the shard level, BUT MultiShard's
|
||||
generic _fanout then re-sorts globally by bm25 — undoing the
|
||||
balance. Fix: run per-token here too, giving each token its
|
||||
own slot in the final return.
|
||||
- **AND pass**: each shard runs a progressive-AND MATCH
|
||||
("tok1 AND tok2 AND …") so multi-token titles like
|
||||
"Spider-Man" rank where their multi-token content earns —
|
||||
single-token title floods don't crowd them out.
|
||||
|
||||
- **Per-token pass**: each token gets its own shard fan-out
|
||||
bucket; round-robin merge so every query token contributes
|
||||
even when one token's bm25 dominates (the "Paris" case —
|
||||
"capitol" alone would otherwise squeeze "Paris" out).
|
||||
"""
|
||||
from arborist.wallet.bucket import (
|
||||
_to_fts5, _FTS5_STOPWORDS, _FTS5_WORD_RE,
|
||||
|
|
@ -685,11 +709,6 @@ class MultiShardSqliteCorpus:
|
|||
if not tokens:
|
||||
return []
|
||||
|
||||
# Per-token: fan out across shards in parallel; each token
|
||||
# gets its OWN top-(limit*2) bucket. The final return is the
|
||||
# union of all tokens' buckets — guarantees every token has
|
||||
# surface area even when one token (e.g. "capitol", low-DF)
|
||||
# outranks another ("paris", higher-DF) on raw bm25.
|
||||
def _shard_token(item, token):
|
||||
sh_path, sc = item
|
||||
try:
|
||||
|
|
@ -698,6 +717,48 @@ class MultiShardSqliteCorpus:
|
|||
except Exception:
|
||||
return sh_path, []
|
||||
|
||||
# AND pass: progressive-AND single MATCH per shard. The shard's
|
||||
# own fts_title runs its own AND when len(tokens)>=2 + per-token,
|
||||
# so passing the original query in gets the full union — but we
|
||||
# explicitly want JUST the AND-mode hits in this bucket, so call
|
||||
# the shard with the joined AND expression directly (treating it
|
||||
# as a single MATCH token from the shard's perspective).
|
||||
and_bucket: list[Hit] = []
|
||||
if len(tokens) >= 2:
|
||||
and_match = " AND ".join(tokens)
|
||||
def _shard_and(item):
|
||||
sh_path, sc = item
|
||||
try:
|
||||
# Direct SQL query to bypass shard's own re-tokenizer.
|
||||
sql = (
|
||||
"SELECT d.document_root, d.document_uri, d.title, "
|
||||
" bm25(documents_fts) AS score "
|
||||
"FROM documents_fts "
|
||||
"JOIN documents d ON d.rowid = documents_fts.rowid "
|
||||
"WHERE documents_fts MATCH ? "
|
||||
"ORDER BY score LIMIT ?"
|
||||
)
|
||||
rows = list(sc._conn.execute(sql, (and_match, limit * 2)))
|
||||
return sh_path, rows
|
||||
except Exception:
|
||||
return sh_path, []
|
||||
with ThreadPoolExecutor(
|
||||
max_workers=max(1, len(self._shards))
|
||||
) as ex:
|
||||
futures = [ex.submit(_shard_and, item) for item in self._shards]
|
||||
for f in futures:
|
||||
sh_path, rows = f.result()
|
||||
for r in rows:
|
||||
and_bucket.append(Hit(
|
||||
document_root=r["document_root"],
|
||||
document_uri=r["document_uri"] or "",
|
||||
title=r["title"] or "",
|
||||
score=r["score"],
|
||||
shard_id=sh_path,
|
||||
))
|
||||
and_bucket.sort(key=lambda h: h.score)
|
||||
and_bucket = and_bucket[: limit * 2]
|
||||
|
||||
per_token_buckets: list[list[Hit]] = []
|
||||
with ThreadPoolExecutor(
|
||||
max_workers=max(1, len(self._shards))
|
||||
|
|
@ -721,22 +782,35 @@ class MultiShardSqliteCorpus:
|
|||
))
|
||||
bucket.sort(key=lambda h: h.score)
|
||||
per_token_buckets.append(bucket[: limit * 2])
|
||||
# Union: each token gets equal allocation. Dedupe by
|
||||
# document_root keeping best bm25 across tokens. The final
|
||||
# sort is per-bucket-aware: walk all buckets, take the
|
||||
# top-N from each round-robin, until we hit `limit`.
|
||||
# Merge: dedupe by document_root keeping best bm25 across
|
||||
# all buckets (AND + per-token). AND-mode usually wins
|
||||
# — its bm25 reflects matching MULTIPLE tokens which gives
|
||||
# a stronger score for the multi-token title.
|
||||
best: dict[str, Hit] = {}
|
||||
for h in and_bucket:
|
||||
cur = best.get(h.document_root)
|
||||
if cur is None or h.score < cur.score:
|
||||
best[h.document_root] = h
|
||||
for bucket in per_token_buckets:
|
||||
for h in bucket:
|
||||
cur = best.get(h.document_root)
|
||||
if cur is None or h.score < cur.score:
|
||||
best[h.document_root] = h
|
||||
# Round-robin across token buckets so each token contributes
|
||||
# representatives even when one token's bm25 dominates the
|
||||
# global sort. Within a bucket, walk by bm25-ascending order.
|
||||
bucket_iters = [iter(b) for b in per_token_buckets]
|
||||
seen_roots: set[str] = set()
|
||||
# Output: AND hits come first (best multi-token matches), then
|
||||
# round-robin across per-token buckets so no token's surface
|
||||
# area starves. The AND pass takes priority because its hits
|
||||
# are the strongest title-relevance signal we have.
|
||||
ranked: list[Hit] = []
|
||||
seen_roots: set[str] = set()
|
||||
for h in and_bucket:
|
||||
if len(ranked) >= limit:
|
||||
break
|
||||
if h.document_root in seen_roots:
|
||||
continue
|
||||
seen_roots.add(h.document_root)
|
||||
ranked.append(best[h.document_root])
|
||||
# Round-robin per-token fill.
|
||||
bucket_iters = [iter(b) for b in per_token_buckets]
|
||||
active = True
|
||||
while active and len(ranked) < limit:
|
||||
active = False
|
||||
|
|
@ -748,9 +822,6 @@ class MultiShardSqliteCorpus:
|
|||
if h.document_root in seen_roots:
|
||||
continue
|
||||
seen_roots.add(h.document_root)
|
||||
# Use the merged best-score (in case the same
|
||||
# doc surfaced from multiple tokens with better
|
||||
# bm25 elsewhere).
|
||||
ranked.append(best[h.document_root])
|
||||
break
|
||||
return ranked
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue