Embedding Training

Embedding Training Pipeline

Fine-tune SentenceTransformer bi-encoders for retrieval with the embedding training method. This skill covers the end-to-end loop: obtain triplets → train (pick a base + adapter mode) → evaluate retrieval against a labeled corpus → read metrics → refine. Every step is a copy-pasteable CLI command or a YAML edit; there is no ad hoc Python.

The embedding method is a sibling to sft/kto/grpo/dpo — it is wired into the same method registry (shared/utilities/paths.py TRAINING_METHODS), the same recipe/local-run path (tuner.py local-run), and the same output layout (embedding_output/<timestamp>/...). What differs is the model family (encoder, not causal-LM), the loss (contrastive, not next-token), and the evaluation axis (corpus-level retrieval metrics, not per-completion correctness).

Quick Reference

The exact local-run / eval flags follow the same surface as the SFT path — verify against tuner/cli/parser.py or --help before scripting; do not guess flags from memory.

Method at a Glance

The Registry (SSOT)

Every base model the pipeline can train or load is a block in Trainers/embedding/configs/model_registry.yaml. Adding a model is one YAML block — no Python change (the config-driven rule). Each block describes how to LOAD, PROMPT, and ADAPT the model. The four seed models:

See reference/embedding-training.md for the full spec schema (pooling, lora_target_modules, matryoshka_dims, prompt_required) and how the dual loader picks the fast vs fallback path.

Adapter Modes

frozen_head carries a compare-to-ST-baseline gate (R6): before training the head, the frozen-base embeddings are asserted to match a plain SentenceTransformer baseline within tolerance (guards against Unsloth nonstandard-pooling drift). The smoke test enforces this.

There is no qlora mode in v1 — the enum is exactly {full, lora, frozen_head}. A qlora value in config raises ValueError ("deferred to a later phase").

Dataset Format (triplets / pairs)

Training data is JSONL. The loader (Trainers/embedding/src/data_loader.py) accepts field aliases and explodes multi-negative records into one row per negative:

{"query": "How do I reset my password?", "positive": "Open Settings → Security → Reset Password.", "negatives": ["Our refund policy allows returns within 30 days.", "Dark mode is under Appearance."]}
{"query": "What payment methods are accepted?", "positive": "We accept Visa, Mastercard, PayPal, and Apple Pay."}

• Anchor aliases: query / anchor / question. Positive aliases: positive / pos. Negative aliases: negative / negatives / neg (scalar or list).
• A negatives list explodes into one (anchor, positive, negative) row per negative (standard ST hard-negative shape).
• Pair records (no negatives) yield (anchor, positive) — but do NOT mix pair and triplet records in one file; the loader drops pair rows when any record has negatives.
• The spec's query_prompt is prepended to anchors and passage_prompt to positives/negatives. E5-style models require these prefixes.

Checked-in canonical fixtures live at Datasets/embedding/examples/: triplets_smoke.jsonl (training) and corpus.jsonl / queries.jsonl / qrels.jsonl (retrieval eval). See reference/triplet-data.md.

Retrieval Evaluation

Embedding models are scored corpus-level, not per-completion. The retrieval verifier embeds a document corpus and a query set, runs FAISS top-k nearest-neighbour retrieval, and aggregates recall@k / MRR / nDCG@k / MAP against qrels, applying a pass/warn/fail threshold ladder. Thresholds live in scenario YAML (config-driven), never in Python:

retrieval_config:
  corpus: Datasets/embedding/examples/corpus.jsonl
  queries: Datasets/embedding/examples/queries.jsonl
  qrels: Datasets/embedding/examples/qrels.jsonl
  metrics: [recall@10, mrr@10, ndcg@10]
  model: { registry_name: bge-base-en }   # or model.path to a trained adapter
  thresholds:
    min: { ndcg@10: 0.30, recall@10: 0.40 }   # below min → fail
    warn_margin: 0.05                          # within [min, min+margin) → warn
  primary_metric: ndcg@10

See reference/retrieval-eval.md for the metric-spec grammar (<metric>@<k>), the threshold ladder, and how to point the eval at a trained run.

The Smoke Recipe

Trainers/recipes/embedding_bge_base_smoke.yaml trains a small bge-base-en LoRA adapter for a few steps on the local-Docker path. It rides the modern Unsloth image (which ships transformers 4.57.1) and layers the embedding deps (sentence-transformers, faiss-cpu, datasets) via a setup.pip overlay — without re-pinning transformers and without touching the legacy SFT transformers 4.45.2 island.

python tuner.py local-run \
  --job-config Trainers/recipes/embedding_bge_base_smoke.yaml \
  --yes

Overlay pins are determined by the cloud smoke, not invented. The recipe's setup.pip entries are intentionally unpinned (sentence-transformers, faiss-cpu, datasets with TBD-PENDING-SMOKE-TEST comments). The repo has a documented history of image/numpy/transformers mismatches, so the exact pins must be captured empirically from the FIRST WORKING cloud run on the modern Unsloth image and only then written back as pkg==X.Y.Z. Do not hardcode speculative versions. See reference/embedding-training.md → "Overlay pins".

Progressive Reference

Cross-References

Tips

• Start with bge-base-en + lora — the cheapest path that reuses merge/upload. Switch to frozen_head for a CPU-only smoke.
• e5-base REQUIRES query/passage prompts; the loader applies them from the registry, so use the registry name, never a bare HF id, for E5.
• Heavy GPU training is NOT a CI gate; the CPU fallback path is what CI runs.
• Keep the smoke fast: a handful of max_steps is enough to prove the method wires end-to-end. Bump for a real run.
• Always --dry-run first to validate the dataset shape and config without training.