Upload bm25_index.py with huggingface_hub

bm25_index.py

@@ -0,0 +1,111 @@
+"""Lexical (keyword) index over the same 158,945 passages, for the dense-vs-hybrid comparison.
+
+Burmese has no spaces between words, so passages are split into syllables with the standard syllable-break
+rule (Ye Kyaw Thu's sylbreak); each passage is indexed on its syllables and syllable bigrams (bigrams
+approximate words), plus lowercased Latin words and digits. Scores are BM25 (k1 = 1.5, b = 0.75), precomputed
+into a sparse matrix so a query is a sum of columns.
+
+    python bm25_index.py                    # build data/bm25/ from data/index/meta_*.jsonl (≈ 5–10 min on the laptop)
+    python bm25_index.py --query "ကျောင်းဆရာ လစာ"   # try a query
+
+    from bm25_index import BM25
+    bm = BM25()                             # loads data/bm25/
+    scores = bm.scores("question")          # numpy array, one score per passage (same order as the dense index)
+"""
+import argparse, json, os, re, time
+import numpy as np
+import scipy.sparse as sp
+from common import DATA, INDEX_DIR, embed_text_of, citation_of
+
+BM25_DIR = os.path.join(DATA, "bm25")
+
+# --- Burmese syllable breaking (sylbreak rule) ---------------------------------------------------------------
+_CONS = "က-အ"; _SS = "္"; _ATHAT = "်"
+_OTHER = "ဣဤဥဦဧဩဪဿ၌၍၏"
+_BREAK = re.compile(r"((?<!" + _SS + r")[" + _CONS + r"](?!့?[" + _ATHAT + _SS + r"])|[" + _OTHER + r"])")   # ့ may sit between consonant and ်
+_LATIN = re.compile(r"[A-Za-z]+|[0-9]+|[၀-၉]+")
+_MY_RUN = re.compile(r"[က-႟]+")
+
+def syllables(run):
+    return [s for s in _BREAK.sub(r" \1", run).split() if s]
+
+def tokenize(text):
+    """Syllables + syllable bigrams for Burmese runs; lowercased words/numbers for Latin and digits."""
+    toks = []
+    for m in re.finditer(r"[က-႟]+|[A-Za-z]+|[0-9]+", text):
+        s = m.group()
+        if s[0] >= "က":
+            sy = syllables(s)
+            toks.extend(sy); toks.extend(a + b for a, b in zip(sy, sy[1:]))
+        else:
+            toks.append(s.lower())
+    return toks
+
+# --- build -----------------------------------------------------------------------------------------------------
+def build(k1=1.5, b=0.75, min_df=2):
+    os.makedirs(BM25_DIR, exist_ok=True)
+    vocab = {}; R, C, V = [], [], []; rows = []; cols = []; vals = []; dl = []
+    t0 = time.time(); n = 0
+    import glob
+    def flush():
+        R.append(np.array(rows, dtype=np.int32)); C.append(np.array(cols, dtype=np.int32)); V.append(np.array(vals, dtype=np.float32))
+        rows.clear(); cols.clear(); vals.clear()
+    for mp in sorted(glob.glob(os.path.join(INDEX_DIR, "meta_*.jsonl"))):
+        for line in open(mp, encoding="utf-8"):
+            r = json.loads(line)
+            counts = {}
+            for t in tokenize(embed_text_of(r)):
+                tid = vocab.get(t)
+                if tid is None: tid = vocab[t] = len(vocab)
+                counts[tid] = counts.get(tid, 0) + 1
+            dl.append(sum(counts.values()))
+            rows.extend([n] * len(counts)); cols.extend(counts.keys()); vals.extend(counts.values())
+            n += 1
+            if n % 2000 == 0: flush()
+            if n % 20000 == 0: print(f"  {n:,} passages · {len(vocab):,} terms · {time.time()-t0:.0f}s", flush=True)
+    flush()
+    tf = sp.csr_matrix((np.concatenate(V), (np.concatenate(R), np.concatenate(C))), shape=(n, len(vocab)))
+    del R, C, V
+    dl = np.array(dl, dtype=np.float32); avgdl = dl.mean()
+    df = np.diff(tf.tocsc().indptr).astype(np.float32)
+    keep = df >= min_df                                   # drop terms seen in a single passage (mostly noise bigrams)
+    inv = {i: t for t, i in vocab.items()}
+    new_ids = -np.ones(len(vocab), dtype=np.int64); new_ids[keep] = np.arange(int(keep.sum()))
+    vocab = {inv[i]: int(new_ids[i]) for i in np.flatnonzero(keep)}
+    tf = tf[:, keep]; df = df[keep]
+    idf = np.log(1 + (n - df + 0.5) / (df + 0.5))
+    tf = tf.tocoo()
+    denom = tf.data + k1 * (1 - b + b * dl[tf.row] / avgdl)
+    w = (idf[tf.col] * tf.data * (k1 + 1) / denom).astype(np.float32)
+    W = sp.csc_matrix((w, (tf.row, tf.col)), shape=tf.shape)
+    sp.save_npz(os.path.join(BM25_DIR, "weights.npz"), W)
+    json.dump({"vocab": vocab, "n": n, "k1": k1, "b": b, "avgdl": float(avgdl), "min_df": min_df},
+              open(os.path.join(BM25_DIR, "vocab.json"), "w"), ensure_ascii=False)
+    print(f"built: {n:,} passages · {len(vocab):,} terms (df ≥ {min_df}) · {W.nnz:,} weights · {time.time()-t0:.0f}s → {BM25_DIR}")
+
+class BM25:
+    def __init__(self, path=BM25_DIR):
+        meta = json.load(open(os.path.join(path, "vocab.json"), encoding="utf-8"))
+        self.vocab = meta["vocab"]; self.n = meta["n"]
+        self.W = sp.load_npz(os.path.join(path, "weights.npz")).tocsc()
+
+    def scores(self, query):
+        ids = [self.vocab[t] for t in tokenize(query) if t in self.vocab]
+        if not ids: return np.zeros(self.n, dtype=np.float32)
+        return np.asarray(self.W[:, ids].sum(axis=1)).ravel().astype(np.float32)
+
+    def search(self, query, k=5):
+        s = self.scores(query); top = np.argpartition(-s, k)[:k]
+        return sorted(((float(s[j]), int(j)) for j in top), reverse=True)
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser(); ap.add_argument("--query"); ap.add_argument("--k", type=int, default=5)
+    a = ap.parse_args()
+    if a.query:
+        from search import load_meta
+        bm = BM25(); meta = load_meta()
+        t = time.time(); hits = bm.search(a.query, a.k)
+        print(f"({time.time()-t:.2f}s)  tokens: {tokenize(a.query)[:20]}")
+        for r, (s, j) in enumerate(hits, 1): print(f"{r}. [{s:.2f}] {citation_of(meta[j])}\n   {meta[j]['text'][:200]}")
+    else:
+        build()