Build local hybrid search with SQLite FTS5, EmbeddingGemma 2 and Python

import sqlite3

print("SQLite", sqlite3.sqlite_version)            # contentless_delete needs 3.43 or newer
con = sqlite3.connect(":memory:")
con.execute("CREATE VIRTUAL TABLE t USING fts5(x)")   # raises 'no such module: fts5' if missing
print("FTS5 is available")

import re


def chunk_text(text: str, target: int = 1100, whole: int = 1600, overlap: int = 150) -> list[str]:
    """Pack sentences into chunks of about target characters, with a short overlap."""
    text = text.strip()
    if len(text) <= whole:                      # short texts stay in one piece
        return [text] if text else []
    pieces = []
    for block in re.split(r"\n\s*\n", text):    # blank lines first, then sentences
        for sent in re.findall(r"[^.!?:;\n]+[.!?:;\n]*", block):
            sent = sent.strip()
            while len(sent) > target:           # hard-cut very long sentences
                pieces.append(sent[:target])
                sent = sent[target - overlap:]
            if sent:
                pieces.append(sent)
    chunks, cur = [], ""
    for p in pieces:
        if cur and len(cur) + len(p) + 1 > target:
            chunks.append(cur)
            cur = cur[-overlap:] + " " + p      # carry a tail into the next chunk
        else:
            cur = f"{cur} {p}".strip()
    if cur:
        chunks.append(cur)
    return chunks

import re
import sqlite3

SCHEMA = """
CREATE TABLE IF NOT EXISTS files (id INTEGER PRIMARY KEY, path TEXT UNIQUE);
CREATE TABLE IF NOT EXISTS items (
    id INTEGER PRIMARY KEY, file_id INTEGER, heading TEXT, text TEXT);
CREATE VIRTUAL TABLE IF NOT EXISTS fts USING fts5(
    name, heading, text,
    content='', contentless_delete=1,           -- needs SQLite 3.43 or newer
    tokenize='unicode61 remove_diacritics 2', prefix='3');
"""

STOP = {"the", "and", "for", "with", "der", "die", "das", "und", "mit", "von", "ein"}


def open_db(path: str = ":memory:") -> sqlite3.Connection:
    con = sqlite3.connect(path)
    con.executescript(SCHEMA)
    return con


def path_words(path: str) -> str:
    """'2024/Mietvertrag_final.pdf' -> '2024 Mietvertrag final pdf'"""
    words = re.sub(r"[_\-./]+", " ", path)
    return re.sub(r"(?<=[a-z])(?=[A-Z])", " ", words).strip()


def add_file(con, path: str, chunks: list[tuple[str, str]]) -> list[int]:
    """chunks = [(heading, text), ...]. Returns the new item ids."""
    file_id = con.execute("INSERT INTO files(path) VALUES (?)", (path,)).lastrowid
    ids = []
    for heading, text in chunks:
        item_id = con.execute(
            "INSERT INTO items(file_id, heading, text) VALUES (?,?,?)",
            (file_id, heading, text)).lastrowid
        con.execute("INSERT INTO fts(rowid, name, heading, text) VALUES (?,?,?,?)",
                    (item_id, path_words(path), heading, text))
        ids.append(item_id)
    return ids


def fts_query(q: str) -> str:
    """OR of the query's words; words of 4+ letters also match as prefixes."""
    terms = [t for t in re.findall(r"\w+", q.lower()) if t not in STOP]
    terms = list(dict.fromkeys(terms))
    return " OR ".join(f'"{t}"*' if len(t) >= 4 and not t.isdigit() else f'"{t}"'
                       for t in terms)


def keyword_search(con, q: str, limit: int = 60) -> list[int]:
    match = fts_query(q)
    if not match:
        return []
    rows = con.execute(                          # bm25 is negative: smaller is better
        "SELECT rowid FROM fts WHERE fts MATCH ? ORDER BY bm25(fts, 4.0, 2.0, 1.0) LIMIT ?",
        (match, limit))
    return [r[0] for r in rows]

# Adapt to your runtime. pip install -U sentence-transformers transformers
import numpy as np
import torch
from sentence_transformers import SentenceTransformer

# bfloat16 where the hardware supports it, float32 everywhere else. Never float16.
dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float32

model = SentenceTransformer(
    "google/embeddinggemma-2",
    model_kwargs={"torch_dtype": dtype},
    config_kwargs={"vision_config": None, "audio_config": None},   # text only
)


def embed(texts: list[str]) -> np.ndarray:
    """texts already carry their prefix: QUERY + question, or doc_prompt(title, chunk)."""
    vecs = model.encode(texts, batch_size=16, normalize_embeddings=True)
    return np.asarray(vecs, dtype=np.float32)        # (n, 768), unit length

# Adapt to your runtime. Apple Silicon, from the mlx-community model card:
#   pip install "git+https://github.com/Blaizzy/mlx-vlm.git@3d87e88402f307efbf68e568971aa887ee7d9ed0"
#   pip install -U "mlx>=0.32.3" "transformers>=5.18.0"
import mlx.core as mx
import numpy as np
from mlx_vlm.embedding_loader import load_embedding_model
from mlx_vlm.utils import get_model_path
from transformers import AutoTokenizer

path = get_model_path("mlx-community/embeddinggemma-2-8bit")
model = load_embedding_model(path)
tokenizer = AutoTokenizer.from_pretrained(path)


def embed(texts: list[str]) -> np.ndarray:
    enc = tokenizer(texts, padding=True, truncation=True, max_length=2048, return_tensors="np")
    out = model(**{k: mx.array(v) for k, v in enc.items()}).text_embeds
    return np.array(out.astype(mx.float32))          # (n, 768), unit length

import hashlib

import numpy as np

QUERY = "task: search result | query: "


def doc_prompt(title: str, chunk: str) -> str:
    title = title.replace("|", " ")[:200] or "none"
    return f"title: {title} | text: {chunk}"


VECTOR_SCHEMA = """
CREATE TABLE IF NOT EXISTS vectors (item_id INTEGER PRIMARY KEY, v BLOB);  -- float16, 1536 bytes
CREATE TABLE IF NOT EXISTS vbits   (item_id INTEGER PRIMARY KEY, b BLOB);  -- sign bits, 96 bytes
CREATE TABLE IF NOT EXISTS vcache  (key BLOB PRIMARY KEY, v BLOB);         -- prompt hash -> vector
"""


def embed_cached(con, prompts: list[str], embed) -> np.ndarray:
    """Embed only prompts never seen before. Identical prompts (copies of a file) cost nothing."""
    keys = [hashlib.blake2b(p.encode(), digest_size=16).digest() for p in prompts]
    got = {}
    for k in set(keys):
        row = con.execute("SELECT v FROM vcache WHERE key=?", (k,)).fetchone()
        if row:
            got[k] = np.frombuffer(row[0], np.float16).astype(np.float32)
    todo = {k: p for k, p in zip(keys, prompts) if k not in got}
    if todo:
        vecs = embed(list(todo.values()))
        for k, v in zip(todo, vecs):
            v16 = (v / np.linalg.norm(v)).astype(np.float16)
            con.execute("INSERT OR REPLACE INTO vcache VALUES (?,?)", (k, v16.tobytes()))
            got[k] = v16.astype(np.float32)
    return np.stack([got[k] for k in keys])


def store_vector(con, item_id: int, v: np.ndarray) -> None:
    v = v / np.linalg.norm(v)
    con.execute("INSERT OR REPLACE INTO vectors VALUES (?,?)", (item_id, v.astype(np.float16).tobytes()))
    con.execute("INSERT OR REPLACE INTO vbits VALUES (?,?)", (item_id, np.packbits(v > 0).tobytes()))

import numpy as np

POPCOUNT = np.array([bin(i).count("1") for i in range(256)], dtype=np.uint8)


def load_bits(con):
    """Keep this in memory and reload it when the table changes."""
    rows = con.execute("SELECT item_id, b FROM vbits ORDER BY item_id").fetchall()
    ids = np.array([r[0] for r in rows], dtype=np.int64)
    bits = np.frombuffer(b"".join(r[1] for r in rows), np.uint8).reshape(len(rows), -1)
    return ids, bits


def rescore(con, ids, qv: np.ndarray, k: int):
    """Exact cosine on the float16 vectors. ids=None scans every vector."""
    if ids is None:
        rows = con.execute("SELECT item_id, v FROM vectors").fetchall()
    else:
        marks = ",".join("?" * len(ids))
        rows = con.execute(f"SELECT item_id, v FROM vectors WHERE item_id IN ({marks})",
                           [int(i) for i in ids]).fetchall()
    M = np.stack([np.frombuffer(v, np.float16).astype(np.float32) for _, v in rows])
    sims = (M @ qv) / np.linalg.norm(M, axis=1)
    return [(rows[i][0], float(sims[i])) for i in np.argsort(-sims)[:k]]


def vector_search(con, index, qv: np.ndarray, k: int = 60, oversample: int = 16):
    """1-bit Hamming shortlist of k * oversample items, then exact re-scoring."""
    ids, bits = index
    if len(ids) == 0:
        return []
    dist = POPCOUNT[bits ^ np.packbits(qv > 0)].sum(axis=1, dtype=np.int32)
    n = min(len(ids), k * oversample)
    cand = ids[np.argpartition(dist, n - 1)[:n]] if n < len(ids) else ids
    return rescore(con, cand, qv, k)


def overlap_at(con, index, queries, k: int = 10, oversample: int = 16) -> float:
    """Share of the exact top-k that the 1-bit path also returns. Run it on your own vectors."""
    hits = 0
    for qv in queries:
        exact = {i for i, _ in rescore(con, None, qv, k)}
        fast = {i for i, _ in vector_search(con, index, qv, k, oversample)}
        hits += len(exact & fast)
    return hits / (k * len(queries))

# Step 5, check. Needs VECTOR_SCHEMA and store_vector (step 4), load_bits and overlap_at (step 5).
import sqlite3

import numpy as np


def unit(a: np.ndarray) -> np.ndarray:
    return a / np.linalg.norm(a, axis=1, keepdims=True)


def overlap_demo(X: np.ndarray, Q: np.ndarray, oversamples=(1, 4, 16)) -> dict:
    con = sqlite3.connect(":memory:")
    con.executescript(VECTOR_SCHEMA)
    for i, v in enumerate(X, start=1):
        store_vector(con, i, v)
    index = load_bits(con)
    return {o: round(overlap_at(con, index, list(Q), 10, o), 2) for o in oversamples}


rng = np.random.default_rng(7)
n, clusters, dims = 30_000, 300, 768
centers = rng.standard_normal((clusters, dims)).astype(np.float32)
X = unit(centers[rng.integers(0, clusters, n)] + 0.9 * rng.standard_normal((n, dims)).astype(np.float32))
Q = unit(X[rng.integers(0, n, 60)] + 1.5 * rng.standard_normal((60, dims)).astype(np.float32) / np.sqrt(dims))
print("clustered:", overlap_demo(X, Q))

R = unit(rng.standard_normal((n, dims)).astype(np.float32))
QR = unit(R[rng.integers(0, n, 40)] + 0.7 * rng.standard_normal((40, dims)).astype(np.float32) / np.sqrt(dims))
print("random:   ", overlap_demo(R, QR, (1, 16)))

# Prints, with seed 7 and NumPy 2.5.3 (the values move with the seed):
# clustered: {1: 0.4, 4: 0.81, 16: 1.0}
# random:    {1: 0.17, 16: 0.42}

# Optional: let sqlite-vec do the 1-bit shortlist. pip install sqlite-vec
# Needs a Python whose sqlite3 module can load extensions.
import sqlite_vec


def use_sqlite_vec(con) -> None:
    con.enable_load_extension(True)
    sqlite_vec.load(con)
    con.enable_load_extension(False)
    con.execute("CREATE VIRTUAL TABLE IF NOT EXISTS vbin USING vec0(embedding bit[768])")


def store_bits(con, item_id: int, v: np.ndarray) -> None:
    con.execute("INSERT INTO vbin(rowid, embedding) VALUES (?, vec_bit(?))",
                (item_id, np.packbits(v > 0).tobytes()))


def vector_search_vec(con, qv: np.ndarray, k: int = 60, oversample: int = 16):
    cand = [r[0] for r in con.execute(
        "SELECT rowid FROM vbin WHERE embedding MATCH vec_bit(?) AND k = ?",
        (np.packbits(qv > 0).tobytes(), k * oversample))]
    return rescore(con, cand, qv, k) if cand else []

def rrf(ranked: dict[str, list[int]], weights: dict[str, float], k: int = 60) -> dict[int, float]:
    """Weighted reciprocal rank fusion. ranked maps a list name to item ids, best first."""
    score: dict[int, float] = {}
    for name, ids in ranked.items():
        w = weights.get(name, 1.0)
        for rank, item in enumerate(ids, start=1):
            score[item] = score.get(item, 0.0) + w / (k + rank)
    return score


WEIGHTS = {"keyword": 1.0, "meaning": 1.0}   # start equal, tune on your own questions later


def by_file(scores: dict[int, float], item_file: dict[int, int]):
    """Roll chunk scores up to files: best + 0.3 * second + 0.1 * third."""
    per_file: dict[int, list[tuple[float, int]]] = {}
    for item, s in scores.items():
        per_file.setdefault(item_file[item], []).append((s, item))
    out = []
    for file_id, hits in per_file.items():
        hits.sort(reverse=True)
        total = sum(w * s for w, (s, _) in zip((1.0, 0.3, 0.1), hits))
        out.append((file_id, total, hits[0][1]))      # file, score, best chunk
    return sorted(out, key=lambda t: -t[1])

def search(con, index, query: str, embed, limit: int = 10, weights=WEIGHTS) -> list[dict]:
    qv = embed([QUERY + query])[0]
    qv = qv / np.linalg.norm(qv)
    ranked = {
        "keyword": keyword_search(con, query),
        "meaning": [i for i, _ in vector_search(con, index, qv)],
    }
    scores = rrf(ranked, weights)
    if not scores:
        return []
    marks = ",".join("?" * len(scores))
    item_file = dict(con.execute(f"SELECT id, file_id FROM items WHERE id IN ({marks})", list(scores)))
    paths = dict(con.execute("SELECT id, path FROM files"))
    results = []
    for file_id, score, best in by_file(scores, item_file)[:limit]:
        text = con.execute("SELECT text FROM items WHERE id=?", (best,)).fetchone()[0]
        results.append({"path": paths[file_id], "score": score, "text": text})
    return results

def gated_rerank(query, results, score_fn, top=10, gate=0.8, k=10, w_retrieval=0.75):
    """Re-rank only near-ties. score_fn(query, texts) returns one number per text."""
    n = min(top, len(results))
    if n < 2 or results[1]["score"] < gate * results[0]["score"]:
        return results                           # retrieval is decisive: skip the model
    rr = score_fn(query, [f'{r["path"]}\n{r["text"]}' for r in results[:n]])
    rr_rank = {i: r for r, i in enumerate(sorted(range(n), key=lambda i: -rr[i]))}
    fused = {i: w_retrieval / (k + i) + 1.0 / (k + rr_rank[i]) for i in range(n)}
    order = sorted(range(n), key=lambda i: -fused[i])
    return [results[i] for i in order] + results[n:]

# Adapt to your runtime. pip install -U sentence-transformers
from sentence_transformers import CrossEncoder

reranker = CrossEncoder(
    "Qwen/Qwen3-Reranker-0.6B",
    prompts={"files": "Given a search query over a person's documents, "
                      "judge whether the document is what they are looking for"},
    default_prompt_name="files",
)


def score_fn(query: str, texts: list[str]) -> list[float]:
    # Raw yes/no logit differences. Only their order is used.
    # Cutting by characters is a rough stand-in for cutting at about 128 tokens.
    return reranker.predict([(query, t[:600]) for t in texts]).tolist()

GOLD = [   # your own questions; "expect" holds path fragments that count as a right answer
    {"q": "tomato soup recipe", "expect": ["recipes"]},
    {"q": "flight booking Lisbon", "expect": ["lisbon", "trip"]},
]


def first_hit(results: list[dict], expect: list[str]):
    for rank, r in enumerate(results, start=1):
        if any(frag.lower() in r["path"].lower() for frag in expect):
            return rank
    return None                                  # not in the results at all


def evaluate(search_fn, gold, depth: int = 10) -> list:
    """One rank (or None) per question. search_fn(question) -> results, best first."""
    return [first_hit(search_fn(g["q"])[:depth], g["expect"]) for g in gold]


def report(ranks: list) -> dict:
    n = len(ranks)

    def hit(k):
        return sum(1 for r in ranks if r is not None and r <= k) / n

    return {"hit@1": hit(1), "hit@3": hit(3), "hit@10": hit(10),
            "mrr": sum(1 / r for r in ranks if r) / n}

import json
import time


def accept(before: dict, after: dict, weights=None, min_gain=0.01, max_loss=0.01):
    """before/after: gold-set name -> ranks from evaluate(). Returns (ok, reason)."""
    weights = weights or {name: 1.0 for name in before}
    gain = sum(weights[n] * (report(after[n])["mrr"] - report(before[n])["mrr"])
               for n in before) / sum(weights.values())
    for n in before:
        loss = report(before[n])["mrr"] - report(after[n])["mrr"]
        if loss > max_loss:
            return False, f"{n} lost {loss:.3f} MRR"
        if any(b == 1 and a != 1 for b, a in zip(before[n], after[n])):
            return False, f"{n} lost a rank-1 hit"
    return gain >= min_gain, f"weighted MRR {gain:+.3f}"


def tune(run_eval, current: dict, step=0.2, log="quality.jsonl") -> dict:
    """run_eval(weights) -> {set name: ranks}. Try one weight at a time, keep the best accepted change."""
    before, best, best_gain, best_why = run_eval(current), None, 0.0, ""
    for name in current:
        for delta in (-step, step):
            cand = {**current, name: round(max(0.05, current[name] + delta), 2)}
            after = run_eval(cand)
            ok, why = accept(before, after)
            gain = sum(report(after[n])["mrr"] - report(before[n])["mrr"] for n in before) / len(before)
            if ok and gain > best_gain:
                best, best_gain, best_why = cand, gain, why
    entry = {"when": time.strftime("%Y-%m-%d %H:%M:%S"), "previous": current,
             "action": "applied" if best else "kept", "candidate": best, "reason": best_why or "no change passed"}
    with open(log, "a") as f:
        f.write(json.dumps(entry) + "\n")
    return best or current


def rollback(log="quality.jsonl"):
    """Weights from before the most recent applied change, or None."""
    entries = [json.loads(line) for line in open(log)]
    applied = [e for e in entries if e["action"] == "applied"]
    return applied[-1]["previous"] if applied else None

# Step 2, check. German spellings against the step 2 schema. Needs open_db, add_file and keyword_search.
con = open_db()
add_file(con, "notes/a.txt", [("", "Herr Müller wohnt in der Straße")])
add_file(con, "notes/b.txt", [("", "Herr Muller wohnt in der Strasse")])
add_file(con, "notes/c.txt", [("", "Der Mietvertrag läuft")])
add_file(con, "notes/d.txt", [("", "Der Untermietvertrag läuft")])


def paths_for(query: str) -> list[str]:
    ids = keyword_search(con, query)
    if not ids:
        return []
    marks = ",".join("?" * len(ids))
    rows = con.execute(
        f"SELECT DISTINCT f.path FROM items i JOIN files f ON f.id = i.file_id "
        f"WHERE i.id IN ({marks}) ORDER BY f.path", ids)
    return [r[0] for r in rows]


for q in ["Müller", "Muller", "Mueller", "Straße", "Strasse", "miet", "Vertrag"]:
    print(f"{q:8} -> {paths_for(q)}")

# Prints, on SQLite 3.53.4:
# Müller   -> ['notes/a.txt', 'notes/b.txt']
# Muller   -> ['notes/a.txt', 'notes/b.txt']
# Mueller  -> []
# Straße   -> ['notes/a.txt']
# Strasse  -> ['notes/b.txt']
# miet     -> ['notes/c.txt']
# Vertrag  -> []