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