diff --git a/arborist/wallet/sidecar.py b/arborist/wallet/sidecar.py index f991218..21cc776 100644 --- a/arborist/wallet/sidecar.py +++ b/arborist/wallet/sidecar.py @@ -108,14 +108,68 @@ def fold_accents(text: str) -> str: ) +# Roman numeral <-> Arabic digit equivalence table for 1..20. +# A user typing 'final fantasy 7' should overlap a title 'Final +# Fantasy VII'. This is the same gap local query.py closes with its +# numeral-fold step. We only need 1..20 — beyond that, titles use +# Arabic digits in practice. +_NUMERAL_PAIRS: dict[str, str] = { + "1": "i", "i": "1", + "2": "ii", "ii": "2", + "3": "iii", "iii": "3", + "4": "iv", "iv": "4", + "5": "v", "v": "5", + "6": "vi", "vi": "6", + "7": "vii", "vii": "7", + "8": "viii","viii": "8", + "9": "ix", "ix": "9", + "10": "x", "x": "10", + "11": "xi", "xi": "11", + "12": "xii","xii": "12", + "13": "xiii","xiii": "13", + "14": "xiv","xiv": "14", + "15": "xv", "xv": "15", + "16": "xvi","xvi": "16", + "17": "xvii","xvii": "17", + "18": "xviii","xviii": "18", + "19": "xix","xix": "19", + "20": "xx", "xx": "20", +} + + +def numeral_expand(tokens) -> set: + """Return a set with each token's numeral-equivalent form added. + + 'final fantasy 7' tokens → {final, fantasy, 7, vii} + 'Final Fantasy VII' tokens → {final, fantasy, vii, 7} + Overlap of the two sets is 4 (or 3 if the duplicate folds), vs 2 + without the fold — enough to lift the right title above siblings + that share only 'final' + 'fantasy'. + """ + out = set(tokens) + for t in tokens: + alt = _NUMERAL_PAIRS.get(t) + if alt is not None: + out.add(alt) + return out + + def tokenize_text(text: str) -> list[str]: - """Lowercase word tokens, drop stopwords + 1-char tokens. - Accent-fold first so 'pokémon' indexes/queries as 'pokemon'.""" + """Lowercase word tokens, drop stopwords + 1-char tokens (except + numeric singletons like '7' which carry meaning — 'Final Fantasy 7' + needs '7' to survive sanitization so the numeral-fold can later + pair it with 'VII'). Accent-fold first so 'pokémon' indexes/queries + as 'pokemon'.""" folded = fold_accents(text) - return [ - t.lower() for t in _WORD_RE.findall(folded) - if t.lower() not in STOPWORDS and len(t) > 1 - ] + out: list[str] = [] + for raw in _WORD_RE.findall(folded): + t = raw.lower() + if t in STOPWORDS: + continue + if len(t) <= 1 and not t.isdigit(): + continue + out.append(t) + return out # --------------------------------------------------------------------------- @@ -458,6 +512,12 @@ class SidecarReader: terms = tokenize_text(query) if not terms: return [] + # Numeral expand so 'final fantasy 7' query also looks up 'vii' + # in the dict (built from chunk text where 'Final Fantasy VII' + # tokenizes to ['final', 'fantasy', 'vii']). The expansion is + # additive — original tokens stay, equivalents joined as OR + # so docs matching either form contribute to BM25. + terms = sorted(numeral_expand(terms)) N = self.docs_count if not hasattr(self, "_avg_doc_len"): # Compute once + cache; needs the full doc_lens. Lazy because @@ -519,12 +579,15 @@ class SidecarReader: if len(t) > 4 and t.endswith("s") and not t.endswith("ss"): return t[:-1] return t - query_stems = {_t_stem(t) for t in terms} + query_stems = numeral_expand({_t_stem(t) for t in terms}) for did in scores: - # Same accent-fold + tokenize on the title side so - # 'pokemon' (query) overlaps 'Pokémon' (title). + # Same accent-fold + tokenize + numeral-expand on the + # title side so 'pokemon'/'pokémon' AND '7'/'VII' both + # overlap correctly. title = fold_accents(self._doc_title(did).lower()) - title_tokens = {_t_stem(t) for t in _WORD_RE.findall(title)} + title_tokens = numeral_expand( + {_t_stem(t) for t in _WORD_RE.findall(title)} + ) overlap = len(query_stems & title_tokens) if overlap: scores[did] += title_boost * overlap