Train a small classifier and compare it with a chat model

import argparse
import json
from pathlib import Path

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline

from semantic_notes import api

LABELS = ("RETURN", "RENEW", "RESERVE")
REVIEW = "REVIEW"
OTHER_REQUESTS = {
    "valid": [
        ("Does the library sell headphones?", REVIEW),
        ("Can I return this book and book the projector?", REVIEW),
        ("Where are the study rooms?", REVIEW),
    ],
    "test": [
        ("What time does the library open?", REVIEW),
        ("I would like to buy a notebook.", REVIEW),
        ("Could I book the scanner and extend my book loan?", REVIEW),
    ],
}


def load_examples(path):
    examples = []
    for line in path.read_text(encoding="utf-8").splitlines():
        if line.strip():
            system, user, assistant = json.loads(line)["messages"]
            examples.append((system["content"], user["content"], assistant["content"]))
    if not examples:
        raise ValueError(f"No examples in {path}.")
    return examples


def ask_chat(instructions, text, allowed, constrained):
    payload = {
        "model": "qwen2.5:1.5b",
        "messages": [
            {"role": "system", "content": instructions},
            {"role": "user", "content": text},
        ],
        "options": {"temperature": 0, "num_ctx": 2048, "num_predict": 30},
        "stream": False,
    }
    if constrained:
        payload["format"] = {
            "type": "object",
            "properties": {"label": {"enum": allowed}},
            "required": ["label"],
            "additionalProperties": False,
        }
    content = api("/api/chat", payload)["message"]["content"].strip()
    if constrained:
        try:
            content = json.loads(content)["label"]
        except (KeyError, TypeError, ValueError):
            return "(invalid)"
    return content if isinstance(content, str) else "(invalid)"


def report(title, predictions, expected):
    correct = sum(guess == target for guess, target in zip(predictions, expected))
    reviewed = sum(guess == REVIEW for guess in predictions)
    print(f"{title}: {correct}/{len(expected)} correct, {reviewed} sent to review")
    for label in (*LABELS, REVIEW):
        count = sum(target == label for target in expected)
        if count:
            matches = sum(guess == target == label for guess, target in zip(predictions, expected))
            print(f"  {label}: {matches}/{count}")


def compare(args):
    train = load_examples(args.data / "train.jsonl")
    data = load_examples(args.data / f"{args.split}.jsonl")
    instructions = data[0][0]
    samples = [(user, label) for _, user, label in data]
    allowed = list(LABELS)
    if args.include_review:
        samples.extend(OTHER_REQUESTS[args.split])
        allowed.append(REVIEW)
        instructions = (
            instructions.replace("Reply with exactly one label.", "") +
            " Use REVIEW for other requests or for requests with multiple intents. "
            "Reply with exactly one label."
        )
    texts, expected = zip(*samples)
    pipeline = make_pipeline(
        TfidfVectorizer(ngram_range=(1, 2), sublinear_tf=True),
        LogisticRegression(C=4, max_iter=1000),
    )
    pipeline.fit([user for _, user, _ in train], [label for _, _, label in train])
    probabilities = pipeline.predict_proba(texts)
    predictions = []
    for row, target, text in zip(probabilities, expected, texts):
        index = row.argmax()
        score = row[index]
        label = str(pipeline.classes_[index]) if score >= args.threshold else REVIEW
        predictions.append(label)
        print(f"{target:7} | {label:7} {score:.2f} | {text}")
    report("Classifier", predictions, expected)
    if args.compare_chat:
        guesses = [ask_chat(instructions, text, allowed, args.json_chat) for text in texts]
        for target, guess, text in zip(expected, guesses, texts):
            print(f"{target:7} | {guess!r} | {text}")
        report("Chat model", guesses, expected)
        print("  Invalid labels:", sum(guess not in allowed for guess in guesses))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Compare a small classifier and an optional local chat baseline.")
    parser.add_argument("--data", type=Path, default=Path("library_data"))
    parser.add_argument("--split", choices=("valid", "test"), default="test")
    parser.add_argument("--threshold", type=float, default=0.0)
    parser.add_argument("--include-review", action="store_true")
    parser.add_argument("--compare-chat", action="store_true")
    parser.add_argument("--json-chat", action="store_true")
    args = parser.parse_args()
    if not 0 <= args.threshold <= 1:
        parser.error("Choose a threshold from 0 to 1.")
    if args.json_chat and not args.compare_chat:
        parser.error("--json-chat requires --compare-chat.")
    try:
        compare(args)
    except (OSError, KeyError, TypeError, ValueError) as error:
        raise SystemExit(str(error)) from None

import json
from pathlib import Path

SYSTEM = (
    "Classify this fictional library request. Use RETURN for giving back an item, "
    "RENEW for extending a loan, or RESERVE for booking equipment. "
    "Reply with exactly one label."
)

SPLITS = {
    "train": [
        ("Where can I drop off this book?", "RETURN"),
        ("I finished my book. Can I bring it back this evening?", "RETURN"),
        ("Which desk takes items back?", "RETURN"),
        ("Is there an after-hours drop box for my borrowed DVD?", "RETURN"),
        ("I am done with the atlas and want to give it back.", "RETURN"),
        ("How do I hand in the borrowed map?", "RETURN"),
        ("Can I keep this book for another week?", "RENEW"),
        ("Please extend my loan until next Thursday.", "RENEW"),
        ("The novel is due tomorrow and I need more time.", "RENEW"),
        ("Can you move the due date on my audiobook?", "RENEW"),
        ("How do I prolong my checkout for this item?", "RENEW"),
        ("I'd like the current borrowing period to last longer.", "RENEW"),
        ("Book the 3D printer for Saturday afternoon.", "RESERVE"),
        ("I need a time slot on the scanner.", "RESERVE"),
        ("Can I sign up for the sewing machine on Friday?", "RESERVE"),
        ("Please hold the projector for my presentation at noon.", "RESERVE"),
        ("I want to use the recording microphone at three.", "RESERVE"),
        ("Is there an open slot for the video camera tomorrow?", "RESERVE"),
    ],
    "valid": [
        ("Does the south desk accept borrowed board games back?", "RETURN"),
        ("Can my current checkout stay active through Monday?", "RENEW"),
        ("Could I schedule the button maker next Wednesday?", "RESERVE"),
    ],
    "test": [
        ("I found the library's cookbook in my bag. Where should I hand it over?", "RETURN"),
        ("Could my friend bring my finished magazine back for me?", "RETURN"),
        ("I am halfway through and my borrowing period ends today. May I have two more days?", "RENEW"),
        ("The checkout expires before my trip ends. Can I change the deadline?", "RENEW"),
        ("Please set aside the digitizing station for me after lunch.", "RESERVE"),
        ("I would like the last available time for the label printer.", "RESERVE"),
    ],
}

folder = Path("library_data")
folder.mkdir(exist_ok=False)
for name, examples in SPLITS.items():
    with (folder / f"{name}.jsonl").open("w", encoding="utf-8") as file:
        for request, label in examples:
            row = {"messages": [
                {"role": "system", "content": SYSTEM},
                {"role": "user", "content": request},
                {"role": "assistant", "content": label},
            ]}
            file.write(json.dumps(row) + "\n")
    print(f"{name}: {len(examples)} examples")

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()