scripts/train_embeddings.py

"""Domain-adapt the retrieval embedding model on the Zephyr corpus.

Why this is worth doing
----------------------
`all-MiniLM-L6-v2` is trained on general web text. Zephyr's vocabulary is not
general: "west", "overlay", "binding", "shim", "work queue", "SYS_INIT",
"CONFIG_" all carry meanings the base model has never seen in this sense. A
query about "a binding" retrieves prose about contracts and obligations before
it retrieves devicetree.

Fine-tuning on in-domain pairs pulls those senses apart.

What it trains on
-----------------
No hand-labelled data, and no synthetic questions from a generator that would
just teach the retriever a generator's phrasing. The pairs are mined from the
documents' own structure:

  (section heading in context)  <->  (that section's body)

A heading is what a reader would type to find the body under it. That is
exactly the query/passage relationship retrieval needs, and it is already
written by the Zephyr doc authors.

`MultipleNegativesRankingLoss` supplies the negatives: every other passage in
the batch. No mining pass, and the harder the batch, the better the signal.

Honest scope
------------
This trains the *retriever*, not the generator. It does not make Qwen Coder
know more about Zephyr — it makes the passages Qwen is handed more likely to be
the right ones. Those are different problems and both matter; see the README
for how they fit together.

    python scripts/train_embeddings.py --epochs 1
    python scripts/train_embeddings.py --eval-only        # baseline, no training
"""

from __future__ import annotations

import argparse
import json
import random
import sys
from pathlib import Path

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")

ROOT = Path(__file__).resolve().parent.parent
DEFAULT_INDEX = ROOT / "data" / "index"
DEFAULT_OUT = ROOT / "data" / "embedding-model"
BASE_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

# A heading has to carry some meaning to work as a query. "Overview",
# "Introduction" and "Example" appear hundreds of times against unrelated
# bodies, so training on them teaches the model that those words mean nothing —
# or worse, that every "Overview" body is interchangeable.
GENERIC_HEADINGS = {
    "overview", "introduction", "example", "examples", "usage", "notes", "note",
    "summary", "description", "requirements", "configuration", "background",
    "api reference", "references", "see also", "limitations", "implementation",
    "samples", "building", "running", "testing", "troubleshooting",
}


def load_chunks(index_dir: Path) -> list[dict]:
    path = index_dir / "chunks.jsonl"
    if not path.exists():
        raise SystemExit(f"no chunks at {path} - run scripts/build_index.py first")
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()]


def mine_pairs(chunks: list[dict]) -> list[tuple[str, str]]:
    """(query, passage) pairs from heading/body structure."""
    pairs: list[tuple[str, str]] = []
    for chunk in chunks:
        heading = (chunk.get("section") or "").strip()
        if not heading or heading.lower() in GENERIC_HEADINGS or len(heading) < 4:
            continue
        # A bare heading is ambiguous across subsystems: "Configuration" under
        # Bluetooth and under Kconfig are different queries. Qualify it with the
        # document title, which is what a reader searching would supply anyway.
        title = (chunk.get("title") or "").strip()
        query = f"{title}: {heading}" if title and title != heading else heading
        pairs.append((query, chunk["text"]))
    return pairs


def recall_at_k(model, queries: list[str], passages: list[str], k: int = 5) -> float:
    """Share of queries whose own passage is in the top k of the whole pool."""
    import numpy as np

    query_vectors = model.encode(
        queries, convert_to_numpy=True, normalize_embeddings=True, batch_size=64
    )
    passage_vectors = model.encode(
        passages, convert_to_numpy=True, normalize_embeddings=True, batch_size=64
    )
    similarity = query_vectors @ passage_vectors.T
    top = np.argsort(-similarity, axis=1)[:, :k]
    hits = sum(row_index in row for row_index, row in enumerate(top))
    return hits / len(queries)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--index", type=Path, default=DEFAULT_INDEX)
    parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--base", default=BASE_MODEL)
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--eval-size", type=int, default=400)
    parser.add_argument("--seed", type=int, default=13)
    parser.add_argument("--eval-only", action="store_true")
    args = parser.parse_args()

    chunks = load_chunks(args.index)
    pairs = mine_pairs(chunks)
    print(f"{len(chunks)} chunks -> {len(pairs)} training pairs")
    if len(pairs) < 200:
        raise SystemExit("too few usable pairs - check the corpus")

    random.Random(args.seed).shuffle(pairs)
    held_out = pairs[: args.eval_size]
    train_pairs = pairs[args.eval_size :]
    eval_queries = [q for q, _ in held_out]
    eval_passages = [p for _, p in held_out]

    from sentence_transformers import SentenceTransformer

    print(f"\nBaseline: {args.base}")
    baseline_model = SentenceTransformer(args.base)
    baseline = recall_at_k(baseline_model, eval_queries, eval_passages)
    print(f"  recall@5 over {len(held_out)} held-out pairs: {baseline:.3f}")

    if args.eval_only:
        return 0

    from sentence_transformers import InputExample, losses
    from torch.utils.data import DataLoader

    model = SentenceTransformer(args.base)
    examples = [InputExample(texts=[query, passage]) for query, passage in train_pairs]
    loader = DataLoader(examples, shuffle=True, batch_size=args.batch_size, drop_last=True)
    # In-batch negatives: every other passage in the batch is a negative for
    # this query. No mining pass, and it scales with batch size.
    loss = losses.MultipleNegativesRankingLoss(model)

    print(f"\nTraining on {len(examples)} pairs, {args.epochs} epoch(s)...")
    model.fit(
        train_objectives=[(loader, loss)],
        epochs=args.epochs,
        warmup_steps=int(len(loader) * 0.1),
        show_progress_bar=True,
    )

    tuned = recall_at_k(model, eval_queries, eval_passages)
    delta = tuned - baseline
    print(f"\n  baseline recall@5: {baseline:.3f}")
    print(f"  tuned    recall@5: {tuned:.3f}   ({delta:+.3f})")

    if delta <= 0:
        # Say so rather than shipping a model that is worse than the thing it
        # replaces. A negative result here is a real result.
        print("\n  Tuning did not improve retrieval. Not writing the model.")
        print("  Try more epochs, a larger batch, or better pairs before shipping this.")
        return 1

    args.out.mkdir(parents=True, exist_ok=True)
    model.save(str(args.out))
    (args.out / "eval.json").write_text(
        json.dumps(
            {
                "base_model": args.base,
                "train_pairs": len(examples),
                "eval_pairs": len(held_out),
                "epochs": args.epochs,
                "batch_size": args.batch_size,
                "recall_at_5_baseline": round(baseline, 4),
                "recall_at_5_tuned": round(tuned, 4),
            },
            indent=2,
        ),
        encoding="utf-8",
    )
    print(f"\nSaved to {args.out}")
    print("Rebuild the index with it:")
    print(f"  python scripts/build_index.py --model {args.out}")
    return 0


if __name__ == "__main__":
    sys.exit(main())