Find documents across languages

import argparse

from semantic_notes import api, cosine

PASSAGES = {
    "renewal-zh": "如果没有其他读者预约，你可以把借来的图书续借十四天。",
    "hours-es": "La sala de lectura cierra a las seis de la tarde los sábados.",
    "scanner-zh": "设备编号 EQ-314 对应平板扫描仪。请在媒体服务台预约。",
    "tripod-en": "Borrow a camera tripod from the media desk. You can keep it for two days.",
    "returns-en": "Return books to the front desk or the after-hours book drop.",
    "workshop-es": "El taller de escritura empieza a las dos en la sala principal.",
}
CASES = [
    ("How long can I extend a book loan if nobody is waiting?", "renewal-zh"),
    ("周六阅览室几点关门？", "hours-es"),
    ("¿Qué equipo corresponde a EQ-314?", "scanner-zh"),
    ("¿Dónde puedo pedir prestado un trípode para la cámara?", "tripod-en"),
    ("晚上可以把书还到哪里？", "returns-en"),
    ("What time does the writing workshop begin?", "workshop-es"),
]


def embeddings(model, texts):
    values = api("/api/embed", {"model": model, "input": texts, "truncate": False})["embeddings"]
    if len(values) != len(texts):
        raise ValueError("Ollama returned the wrong number of embeddings.")
    return values


def evaluate(model):
    texts = list(PASSAGES.values())
    queries = [question for question, _ in CASES]
    if model == "embeddinggemma":
        texts = ["title: none | text: " + text for text in texts]
        queries = ["task: search result | query: " + query for query in queries]
    documents = dict(zip(PASSAGES, embeddings(model, texts)))
    questions = embeddings(model, queries)
    first = within_three = 0
    print(model)
    for (question, expected), vector in zip(CASES, questions):
        ranked = sorted(documents, key=lambda key: cosine(documents[key], vector), reverse=True)
        first += ranked[0] == expected
        within_three += expected in ranked[:3]
        print(f"  {question}\n    expected {expected}; top three: {', '.join(ranked[:3])}")
    print(f"  First result: {first}/{len(CASES)} | Top three: {within_three}/{len(CASES)}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Evaluate cross-language retrieval on six fictional notes.")
    parser.add_argument("model", choices=("all-minilm", "embeddinggemma", "both"), default="both", nargs="?")
    args = parser.parse_args()
    try:
        for name in ("all-minilm", "embeddinggemma") if args.model == "both" else (args.model,):
            evaluate(name)
    except (OSError, KeyError, ValueError) as error:
        raise SystemExit(str(error)) from None

import argparse
from contextlib import closing
from itertools import islice
import json
import math
from pathlib import Path
import sqlite3
from urllib.error import HTTPError
from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener

BASE = "http://127.0.0.1:11434"
MODEL = "all-minilm:latest"
MAX_PASSAGES = 200


class NoRedirects(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def api(endpoint, payload=None):
    request = Request(BASE + endpoint)
    if payload is not None:
        request.data = json.dumps(payload).encode("utf-8")
        request.add_header("Content-Type", "application/json")
    opener = build_opener(ProxyHandler({}), NoRedirects())
    with opener.open(request, timeout=120) as response:
        return json.load(response)


def model_digest():
    for model in api("/api/tags")["models"]:
        if model["name"] == MODEL:
            return model["digest"]
    raise ValueError(f"Pull {MODEL} with Ollama first.")


def embed(texts):
    result = api("/api/embed", {
        "model": MODEL,
        "input": texts,
        "truncate": False,
    })["embeddings"]
    if len(result) != len(texts):
        raise ValueError("Ollama returned the wrong number of embeddings.")
    return result


def cosine(left, right):
    if len(left) != len(right):
        raise ValueError("Vector dimensions changed. Rebuild the index.")
    length = math.hypot(*left) * math.hypot(*right)
    return sum(x * y for x, y in zip(left, right)) / length if length else 0


def build_index(folder, database):
    from index_notes import passages

    rows = list(islice(passages(folder), MAX_PASSAGES + 1))
    if not rows or len(rows) > MAX_PASSAGES:
        raise ValueError(f"Use 1–{MAX_PASSAGES} nonempty passages for this example.")
    digest = model_digest()
    stored = []
    for start in range(0, len(rows), 16):
        batch = rows[start:start + 16]
        vectors = embed([text for _, _, text in batch])
        stored.extend((*row, json.dumps(vector)) for row, vector in zip(batch, vectors))
    if model_digest() != digest:
        raise ValueError("The embedding model changed during indexing. Build the index again.")
    with closing(sqlite3.connect(database)) as connection, connection:
        connection.execute("CREATE TABLE IF NOT EXISTS model (name TEXT, digest TEXT)")
        connection.execute(
            "CREATE TABLE IF NOT EXISTS passages "
            "(source TEXT, number INTEGER, text TEXT, vector TEXT)"
        )
        connection.execute("DELETE FROM passages")
        connection.execute("DELETE FROM model")
        connection.execute("INSERT INTO model VALUES (?, ?)", (MODEL, digest))
        connection.executemany("INSERT INTO passages VALUES (?, ?, ?, ?)", stored)
    print(f"Saved {len(stored)} embeddings in {database}.")


def find_notes(database, query):
    if not query.strip() or not database.is_file():
        raise ValueError("Build the index and enter a nonempty question.")
    with closing(sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True)) as connection:
        connection.execute("BEGIN")
        identity = connection.execute("SELECT name, digest FROM model").fetchone()
        rows = connection.execute("SELECT source, number, text, vector FROM passages").fetchall()
    digest = model_digest()
    if identity != (MODEL, digest):
        raise ValueError("The embedding model differs from the one in the index. Rebuild it.")
    vector, = embed([query])
    if model_digest() != digest:
        raise ValueError("The embedding model changed during the query. Rebuild the index.")
    matches = [
        (cosine(json.loads(embedding), vector), source, number, text)
        for source, number, text, embedding in rows
    ]
    ranked = sorted(matches, key=lambda match: match[0], reverse=True)
    for score, source, number, text in ranked[:3]:
        print(f"{score:.3f} | {source}, passage {number}\n{text}\n")


def main():
    parser = argparse.ArgumentParser(description="Store local embeddings and search them later.")
    parser.add_argument("--db", type=Path, default=Path("notes-vectors.sqlite3"))
    commands = parser.add_subparsers(dest="command", required=True)
    build = commands.add_parser("build")
    build.add_argument("folder", type=Path)
    find = commands.add_parser("find")
    find.add_argument("query")
    args = parser.parse_args()

    try:
        if args.command == "build":
            build_index(args.folder, args.db)
        else:
            find_notes(args.db, args.query)
    except HTTPError as error:
        raise SystemExit(f"Ollama returned HTTP {error.code}: {error.read().decode()}") from None
    except (OSError, sqlite3.Error, ValueError) as error:
        raise SystemExit(str(error)) from None


if __name__ == "__main__":
    main()