Spaces:
Runtime error
Runtime error
| """Ranked search over a notebook's chat memory: BM25, written in place. | |
| grep finds exact strings but cannot rank and misses synonyms. Each topic file carries keywords with synonyms, and | |
| every searchable unit (a saved quote or the summary) is scored together with its topic's label, gist and keywords, so | |
| a search for "fee" can find a chat that only said "pricing". A notebook's chats are small, so this runs in | |
| milliseconds from the files already on disk, with no index service.""" | |
| import math | |
| import re | |
| from collections import Counter | |
| from dataclasses import dataclass | |
| from app.chat_memory import files | |
| K1 = 1.5 | |
| B = 0.75 | |
| MAX_RESULTS = 8 | |
| CUTOFF = 0.5 # drop results scoring under this share of the best one | |
| SNIPPET_CHARS = 240 | |
| _WORD = re.compile(r"[a-z0-9]+") | |
| _STOP = frozenset("a an and are as at be by for from has have i in is it its me my of on or that the this to was " | |
| "we were what when which who with you your about did do does said say".split()) | |
| def tokenize(text: str) -> list[str]: | |
| out = [] | |
| for w in _WORD.findall(text.lower()): | |
| if w in _STOP: | |
| continue | |
| out.append(w[:-1] if len(w) > 3 and w.endswith("s") and not w.endswith("ss") else w) | |
| return out | |
| class Unit: | |
| path: str # chat area relative: {chat}/{slug}.md | |
| line: int # 1 based line in that file | |
| text: str # what is shown | |
| topic: str | |
| last: str # the topic's last active date | |
| tokens: list[str] | |
| def units_for(path: str, text: str) -> list[Unit]: | |
| topic = files.parse_topic(text) | |
| if topic is None: | |
| return [] | |
| lines = text.splitlines() | |
| context = tokenize(f"{topic.label} {topic.gist} {' '.join(topic.keywords)}") # document expansion | |
| units: list[Unit] = [] | |
| def add(line: int, shown: str) -> None: | |
| units.append(Unit(path, line, shown, topic.label, topic.last, tokenize(shown) + context)) | |
| summary_line = next((i for i, ln in enumerate(lines, start=1) if ln.strip() == "## Summary"), 0) | |
| if topic.summary: | |
| add(summary_line + 2 if summary_line else 1, topic.summary) | |
| head = {q.message_id: 0 for q in topic.verbatim} | |
| for i, ln in enumerate(lines, start=1): | |
| for mid in head: | |
| if ln.startswith(f"[{mid} "): | |
| head[mid] = i | |
| for q in topic.verbatim: | |
| add(head[q.message_id] + 1 if head[q.message_id] else 1, q.text) | |
| return units | |
| class Bm25: | |
| def __init__(self, units: list[Unit]): | |
| self.units = units | |
| self.tf = [Counter(u.tokens) for u in units] | |
| n = max(len(units), 1) | |
| self.avg = sum(len(u.tokens) for u in units) / n if units else 0.0 | |
| df: Counter[str] = Counter() | |
| for c in self.tf: | |
| df.update(c.keys()) | |
| self.idf = {t: math.log(1 + (n - d + 0.5) / (d + 0.5)) for t, d in df.items()} | |
| def search(self, query: str, since: str = "") -> list[tuple[float, Unit]]: | |
| terms = list(dict.fromkeys(tokenize(query))) | |
| scored = [] | |
| for unit, tf in zip(self.units, self.tf, strict=True): | |
| if since and unit.last < since: | |
| continue | |
| length = len(unit.tokens) or 1 | |
| score = 0.0 | |
| for t in terms: | |
| f = tf.get(t, 0) | |
| if f: | |
| score += self.idf[t] * f * (K1 + 1) / (f + K1 * (1 - B + B * length / (self.avg or 1))) | |
| if score > 0: | |
| scored.append((score, unit)) | |
| scored.sort(key=lambda s: s[0], reverse=True) | |
| return scored | |
| def cut(scored: list[tuple[float, Unit]]) -> list[tuple[float, Unit]]: | |
| """Keep the clear winners: results within CUTOFF of the best score, at most MAX_RESULTS, at least one.""" | |
| if not scored: | |
| return [] | |
| top = scored[0][0] | |
| return [s for s in scored if s[0] >= top * CUTOFF][:MAX_RESULTS] | |
| _cache: dict[tuple, Bm25] = {} | |
| def index_for(key: tuple, load_units) -> Bm25: | |
| """The BM25 index for one notebook, rebuilt only when its files change (the key holds the manifest hashes).""" | |
| if key not in _cache: | |
| if len(_cache) >= 8: | |
| _cache.pop(next(iter(_cache))) | |
| _cache[key] = Bm25(load_units()) | |
| return _cache[key] | |