# SkillPatch skill: train-sentence-transformers

A router skill for training or fine-tuning sentence-transformers models, covering three model types: SentenceTransformer (bi-encoder), CrossEncoder (reranker), and SparseEncoder (SPLADE). It guides agents to the correct reference documents and production script templates based on the task, and covers loss selection, hard-negative mining, evaluators, distillation, LoRA, Matryoshka, and Hugging Face Hub publishing. Acts as a structured decision-making entry point rather than a standalone manual.

You (the agent) were given this URL and asked to use this skill. This is a **public** skill — no SkillPatch account, API key, or prior setup is required. Two ways to use it:

**1. Use it right now** — the skill's complete file tree (SKILL.md + all reference files) is inlined below; read `SKILL.md` first, then follow it, consulting the other files as it directs.

**2. Install the exact package onto disk** (recommended if you can run a shell — this reproduces the skill byte-for-byte, including any binary assets that can't be inlined):

```bash
mkdir -p .claude/skills/train-sentence-transformers
curl -sSL https://skillpatch.dev/install_skill/train-sentence-transformers | tar -xz -C .claude/skills/
```

(`.claude/skills/` is Claude Code's convention; use whatever directory your agent loads skills from.)


---

## Skill files (28)

- `SKILL.md`
- `references/base_model_selection.md`
- `references/dataset_formats.md`
- `references/evaluators_cross_encoder.md`
- `references/evaluators_sentence_transformer.md`
- `references/evaluators_sparse_encoder.md`
- `references/hardware_guide.md`
- `references/hf_jobs_execution.md`
- `references/losses_cross_encoder.md`
- `references/losses_sentence_transformer.md`
- `references/losses_sparse_encoder.md`
- `references/model_architectures.md`
- `references/prompts_and_instructions.md`
- `references/training_args.md`
- `references/troubleshooting.md`
- `scripts/mine_hard_negatives.py`
- `scripts/train_cross_encoder_distillation_example.py`
- `scripts/train_cross_encoder_example.py`
- `scripts/train_cross_encoder_listwise_example.py`
- `scripts/train_sentence_transformer_distillation_example.py`
- `scripts/train_sentence_transformer_example.py`
- `scripts/train_sentence_transformer_make_multilingual_example.py`
- `scripts/train_sentence_transformer_matryoshka_example.py`
- `scripts/train_sentence_transformer_multi_dataset_example.py`
- `scripts/train_sentence_transformer_static_embedding_example.py`
- `scripts/train_sentence_transformer_with_lora_example.py`
- `scripts/train_sparse_encoder_distillation_example.py`
- `scripts/train_sparse_encoder_example.py`


### `SKILL.md`

````markdown
---
name: train-sentence-transformers
description: Train or fine-tune sentence-transformers models across `SentenceTransformer` (bi-encoder; dense or static embedding model; for retrieval, similarity, clustering, classification, paraphrase mining, dedup, multimodal), `CrossEncoder` (reranker; pair scoring for two-stage retrieval / pair classification), and `SparseEncoder` (SPLADE, sparse embedding model; for learned-sparse retrieval). Covers loss selection, hard-negative mining, evaluators, distillation, LoRA, Matryoshka, and Hugging Face Hub publishing. Use for any sentence-transformers training task.
---

# Train a sentence-transformers Model

**This SKILL.md is a router, not a manual.** It tells you which references and example scripts to load for your task. The actual content — recommended losses, evaluators, training-script structure, model selection, training-arg knobs, troubleshooting — lives in `references/` and `scripts/`.

**Do not synthesize a training script from this file alone.** Open the per-type production template (`scripts/train_<type>_example.py`) and copy it as your starting point. The templates contain load-bearing scaffolding (autocast helper, model-card class, logger silencing list, `force=True`, `seed`, TF32, version-compatible imports, named-evaluator metric handling) that prior agent runs have repeatedly missed when rolling their own from a synthesized snippet.

## 1. Identify the model type

| Tag | Class | What it does | When to pick |
|---|---|---|---|
| **[SentenceTransformer]** | `SentenceTransformer` (bi-encoder) | Maps each input to a fixed-dim dense vector | Retrieval, similarity, clustering, classification, paraphrase mining, dedup |
| **[CrossEncoder]** | `CrossEncoder` (reranker) | Scores `(query, passage)` pairs jointly | Two-stage retrieval (rerank top-100 from bi-encoder), pair classification |
| **[SparseEncoder]** | `SparseEncoder` (SPLADE) | Sparse vectors over the vocabulary | Learned-sparse retrieval, inverted-index backends (Elasticsearch / OpenSearch / Lucene) |

Tiebreakers when the request is ambiguous: "embedding model" / "vector search" / "similarity" → **[SentenceTransformer]**. "rerank" / "ranker" / "two-stage" → **[CrossEncoder]**. "SPLADE" / "sparse" / "inverted index" → **[SparseEncoder]**. If still unclear, ask.

## 2. Required reading

**Read these in full before writing any code. Do not triage by perceived relevance.**

### Per-type — always required

**[SentenceTransformer]**
- `references/losses_sentence_transformer.md` — loss-to-data-shape mapping; `BatchSamplers.NO_DUPLICATES` requirement for MNRL-family; `Cached*` ↔ `gradient_checkpointing` incompatibility.
- `references/evaluators_sentence_transformer.md` — evaluator-to-task mapping; `metric_for_best_model` key construction (named vs unnamed); per-evaluator `primary_metric` values.
- `references/model_architectures.md` — encoder vs decoder vs static vs Router pipelines; pooling rules (mean / cls / lasttoken); auto-mean-pooling behavior for fresh-start MLM bases.
- `scripts/train_sentence_transformer_example.py` — production template; copy this as your starting point.

**[CrossEncoder]**
- `references/losses_cross_encoder.md` — pointwise / pairwise / listwise / distillation; `pos_weight` derivation; `activation_fn=Identity()` mandatory for non-BCE losses (silent eval-rank collapse otherwise).
- `references/evaluators_cross_encoder.md` — `CrossEncoderRerankingEvaluator` recipe; named-evaluator key format `eval_{name}_{primary_metric}`.
- `scripts/train_cross_encoder_example.py` — production template; copy this as your starting point.

**[SparseEncoder]**
- `references/losses_sparse_encoder.md` — `SpladeLoss` wrapper requirement; FLOPS regularizer weights; smoke-test active-dim ramp behavior.
- `references/evaluators_sparse_encoder.md` — `SparseNanoBEIREvaluator` (English-only) and the in-domain alternative; `eval_{name}_{primary_metric}` key format.
- `scripts/train_sparse_encoder_example.py` — production template; copy this as your starting point.

### Cross-cutting — always required (regardless of task)

- `references/training_args.md` — `TrainingArguments` knobs, precision rules (load fp32 + autocast bf16/fp16; never `torch_dtype=bfloat16`), `warmup_steps` (float) vs deprecated `warmup_ratio`, `save_steps` must be a multiple of `eval_steps` for `load_best_model_at_end`, schedulers, HPO, tracker, resume, hub-push variants.
- `references/dataset_formats.md` — column-matching rules (label name auto-detection; column-order-not-name); reshaping recipes; hard-negative mining options.
- `references/base_model_selection.md` — discovery commands; per-type model namespaces; ModernBERT-family `max_seq_length=8192` trap; `datasets >= 4` script-loader rejection; non-English starting-point shortcuts.
- `references/troubleshooting.md` — symptom-indexed failure recipes. Skim the section headings on every run, even a healthy one; the "Metrics don't improve" and "Hub push fails" entries cover bugs that bite frequently and are cheaper to recognize before they fire than to debug after.

### Cross-cutting — load when applicable

- `references/hardware_guide.md` — VRAM sizing, multi-GPU, FSDP / DeepSpeed, HF Jobs flavors. Required for >24GB models, multi-GPU, or HF Jobs runs.
- `references/hf_jobs_execution.md` — required when running on HF Jobs.
- `references/prompts_and_instructions.md` — required when using prompt-tuned bases (E5, BGE, GTE, Qwen3-Embedding, Instructor, Nomic, etc.) or adding `query: ` / `passage: ` style prefixes.

### Variant scripts (open when the task matches)
- **[SentenceTransformer]** `scripts/train_sentence_transformer_<matryoshka|multi_dataset|with_lora|distillation|make_multilingual|static_embedding>_example.py`.
- **[CrossEncoder]** `scripts/train_cross_encoder_<distillation|listwise>_example.py`.
- **[SparseEncoder]** `scripts/train_sparse_encoder_distillation_example.py`.
- Hard-negative mining CLI — `scripts/mine_hard_negatives.py`.

## 3. Defaults

Override only if the user specifies otherwise:
- **Local execution.** Pitch HF Jobs only if local hardware can't fit the job.
- **Single run.** After it completes, propose experimentation if the user would benefit (weak/marginal verdict, "see how high you can push it" framing, etc.). Iteration rules in `references/training_args.md` (Experimentation section).
- **Public Hub push at end-of-run, wrapped in try-except.** On HF Jobs (ephemeral env) ALSO enable in-trainer push (`push_to_hub=True` + `hub_strategy="every_save"`); details in `references/hf_jobs_execution.md`.

## 4. Constraints the produced script must satisfy

These are non-negotiable contracts. Implementation lives in the production templates and references — do not reinvent.

- Capture the pre-training evaluator score as `baseline_eval` **before** `trainer.train()`.
- Emit a single end-of-run line: `VERDICT: WIN|MARGINAL|REGRESSION | score=... | baseline=... | delta=...`. A monitor scrapes for this.
- Silence `httpx`, `httpcore`, `huggingface_hub`, `urllib3`, `filelock`, `fsspec` to WARNING (otherwise HF download URLs flood the agent's context).
- Tee logs to `logs/{RUN_NAME}.log`.
- End with `model.push_to_hub(...)` wrapped in `try/except`.
- Smoke-test before any long run (`max_steps=1` + tiny dataset slice). The production templates show one common pattern (`SMOKE_TEST` env var).
- **[CrossEncoder]** Include `EarlyStoppingCallback(patience>=3)` — CE rerankers often peak mid-training and regress.
- **[SparseEncoder]** Log `query_active_dims` / `corpus_active_dims` on the verdict line; high nDCG with collapsed sparsity is not a win. The keys come back name-prefixed (e.g. `..._query_active_dims`); use suffix matching to pluck them — see the SPARSE production template for the exact pattern.

## 5. Workflow

1. Identify the model type (§1). Ask if ambiguous.
2. Load the §2 required-reading files for that type.
3. Open `scripts/train_<type>_example.py` and copy it as your starting point.
4. Replace `MODEL_NAME`, `DATASET_NAME`, `RUN_NAME`, the loss, and the evaluator with the user's task. Cross-check loss/data-shape match against `references/losses_<type>.md`; cross-check the `metric_for_best_model` key against `references/evaluators_<type>.md` (named evaluators format the key as `eval_{name}_{primary_metric}`).
5. Smoke-test (`max_steps=1`).
6. Run.
7. After the run, append to `logs/experiments.md` and propose iteration if the verdict is weak/marginal.

## Prerequisites

```bash
pip install "sentence-transformers[train]>=5.0"        # add [train,image] / [audio] / [video] for [SentenceTransformer] multimodal
pip install trackio                                    # optional tracker; or wandb / tensorboard / mlflow
hf auth login                                          # or set HF_TOKEN with write scope (for Hub push)
```

GPU strongly recommended. CPU works only for demos and `[SentenceTransformer]` `StaticEmbedding`.

````


### `references/base_model_selection.md`

````markdown
# Base Model Selection

Leaderboards rotate every few months; don't trust any hardcoded "best" pick. Discover current options live — run **both** sort orders since most-downloaded surfaces proven options and trending surfaces recent SOTA that may not have download volume yet.

## Discovery commands

**[BI]**:
```bash
hf models list --filter sentence-transformers --sort downloads --limit 20
hf models list --filter sentence-transformers --sort trending  --limit 20
```

**[CE]**:
```bash
hf models list --filter sentence-transformers --filter text-ranking --sort downloads --limit 20
hf models list --filter sentence-transformers --filter text-ranking --sort trending  --limit 20
```

**[SPARSE]**:
```bash
hf models list --filter sentence-transformers --filter sparse-encoder --sort downloads --limit 20
hf models list --filter sentence-transformers --filter sparse-encoder --sort trending  --limit 20
```

Optional language narrowing (any type): add `--filter <language-code>`. Not all multilingual models tag each language, so missing matches doesn't mean the model can't handle that language — re-run without the filter to compare.

```bash
hf models card <id> --text                        # confirm dimensions, max_seq_length, license, languages
```

Cross-check the [MTEB leaderboard](https://huggingface.co/spaces/mteb/leaderboard) (pick the relevant tab) before committing to a multi-hour run.

## [BI] Bi-Encoder

Continue from an existing retriever beats fresh-start + 100k–500k pairs. Common namespaces as of 2026-Q2 (verify against discovery commands — the field rotates):

- **English encoder retrievers**: `sentence-transformers/all-*` (MiniLM-L6-v2, mpnet-base-v2 still the most-downloaded models on the Hub), `BAAI/bge-*-en-v1.5`, `nomic-ai/nomic-embed-text-v1.5`, `mixedbread-ai/mxbai-embed-large-v1`, `Alibaba-NLP/gte-*`, `Snowflake/snowflake-arctic-embed-*`, `jinaai/jina-embeddings-v5-text-small` / `-nano`, `microsoft/harrier-oss-v1-270m` / `-0.6b`.
- **Multilingual encoder retrievers**: `sentence-transformers/paraphrase-multilingual-*`, `intfloat/multilingual-e5-*`, `ibm-granite/granite-embedding-*-multilingual-r2`, `google/embeddinggemma-300m`, `voyageai/voyage-4-nano`.
- **Long documents (8k+)**: `nomic-ai/modernbert-embed-*`, `answerdotai/ModernBERT-large`.
- **Decoder LLM retrievers** (multilingual; **need last-token pooling**): `Qwen/Qwen3-Embedding-*` (0.6B / 4B / 8B), `Qwen/Qwen3-VL-Embedding-*` (multimodal), `codefuse-ai/F2LLM-v2-*`.
- **Fresh-start English** (≥500k pairs + domain-fit reason): `microsoft/mpnet-base`, `answerdotai/ModernBERT-base`, `google-bert/bert-base-uncased`, `jhu-clsp/ettin-encoder-*` (17m / 32m / 68m / 150m / 400m / 1b — paired ModernBERT encoder family).
- **Fresh-start multilingual**: `FacebookAI/xlm-roberta-base` (MLM-only, needs contrastive training), `microsoft/mdeberta-v3-base`, `jhu-clsp/mmBERT-base` / `-small`.
- **CPU / small footprint** (`StaticEmbedding`): `StaticEmbedding(tokenizer, embedding_dim=...)`. **Model size = `vocab_size × dim × 4 bytes`** — pick a small-vocab tokenizer or you get a giant model: 30k-vocab `bert-base-uncased` × 128 dim ≈ 15 MB; **250k-vocab `paraphrase-multilingual-MiniLM-L12-v2` × 256 dim ≈ 256 MB**. Random init needs 1M+ pairs; warm-start (`StaticEmbedding.from_distillation(...)`) helps under ~100k pairs.

Architecture variants (encoder / decoder / static / Router), pooling rules, and decoder-vs-encoder setup paths: `model_architectures.md`.

**ModernBERT-family bases default to `max_seq_length=8192`.** That allocates activation memory for 8192-token sequences regardless of your data length and silently drives Windows VRAM into "shared memory" spillover. After loading any ModernBERT / mmBERT / Ettin / gte-modernbert / nomic-modernbert base, **explicitly set `model.max_seq_length = 256` (or 512 for documents)** unless you actually need long context.

## [CE] Cross-Encoder

Continue from an existing reranker beats fresh-start + 100k–500k pairs in most domains; default to this unless you have a strong reason otherwise. Common namespaces as of 2026-Q2:

- **English encoder rerankers**: `cross-encoder/ms-marco-*`, `BAAI/bge-reranker-*`, `mixedbread-ai/mxbai-rerank-*-v1` / `-v2`, `Alibaba-NLP/gte-reranker-modernbert-*`, `ibm-granite/granite-embedding-reranker-english-*`.
- **Multilingual encoder rerankers**: `cross-encoder/mmarco-*`, `BAAI/bge-reranker-v2-m3`, `Alibaba-NLP/gte-multilingual-reranker-*`, `ibm-granite/granite-embedding-reranker-multilingual-*`.
- **Decoder LLM rerankers** (multilingual; `num_labels=1` last-token-style scoring): `Qwen/Qwen3-Reranker-*` (0.6B / 4B / 8B), `Qwen/Qwen3-VL-Reranker-*` (multimodal).
- **Fresh-start**: `microsoft/MiniLM-L12-H384-uncased`, `answerdotai/ModernBERT-base` / `-large`, `jhu-clsp/ettin-encoder-*`, `FacebookAI/xlm-roberta-base` (multilingual), `microsoft/mdeberta-v3-base` (multilingual), `jhu-clsp/mmBERT-base` / `-small` (multilingual). Pass `num_labels >= 2` for classification cross-encoders.

Encoder-only bases are still the latency-efficient default (bidirectional attention is well-suited to the reranking use case at small parameter counts), but decoder LLM rerankers are now competitive at the top of MTEB Reranking when latency / memory budget allows.

**Minimum dataset:** 500k+ labeled `(query, passage, label)` tuples for production; 10k–100k labeled pairs for continue-training on domain data. Low-resource languages may have less than 10k labeled pairs; in that case lean on a multilingual base's pretraining and accept a noisier signal.

**"Small" multilingual is ~100M+ params**, not 17M-50M like the English small models. mMiniLMv2-L12-H384 (~117M) is roughly the small-end for usable multilingual rerankers.

## [SPARSE] Sparse Encoder (SPLADE)

SPLADE requires a fill-mask / `AutoModelForMaskedLM`-compatible checkpoint. Encoder-only MLM models work out of the box; **decoder LLMs do not**.

- **Continue from existing SPLADE — English**: `naver/splade-*` (the canonical family), `opensearch-project/opensearch-neural-sparse-encoding-*` (incl. `-doc-v2-distill`, `-doc-v3-distill` / `-doc-v3-gte`), `prithivida/Splade_PP_en_v*`, `ibm-granite/granite-embedding-30m-sparse`.
- **Continue from existing SPLADE — multilingual**: `opensearch-project/opensearch-neural-sparse-encoding-multilingual-v1`.
- **Fresh-start English** (≥500k pairs): any encoder with an MLM head — `distilbert/distilbert-base-uncased`, `google-bert/bert-base-uncased`. Pure `AutoModel` checkpoints without MLM won't work. Discover MLM bases: `hf models list --filter fill-mask --sort downloads --limit 20`.
- **Fresh-start multilingual**: `FacebookAI/xlm-roberta-base` (has MLM head). For other multilingual MLM bases: add `--filter <language-code>`.

**Minimum dataset:** 500k+ triplets (with mined hard negatives) for a competitive SPLADE; 50k+ triplets for domain adaptation on existing SPLADE.

## Cross-cutting tips

- **Non-English retrieval starting points** (when language tag returns 0 results): check `intfloat/multilingual_e5_train_data` for parallel pair data; MIRACL via the `sentence-transformers/miracl` mirror for multilingual retrieval; mMARCO via `unicamp-dl/mmarco` (14 languages, parquet-backed).
- **Avoid script-based dataset loaders.** `datasets >= 4` rejects them with `RuntimeError: Dataset scripts are no longer supported`. Look for parquet-backed mirrors (e.g. `sentence-transformers/miracl` instead of `miracl/miracl`).
- **`hf datasets sql` requires DuckDB** (`pip install duckdb`). Without it, fall back to `python -c "from datasets import load_dataset; ds = load_dataset('<id>', ...); print(ds.column_names, ds[0])"`.

````


### `references/dataset_formats.md`

````markdown
# Dataset Formats

This reference covers: how datasets map to losses, how to reshape data when it doesn't fit, and how to mine hard negatives.

## The two rules

From the sentence-transformers training overview:

1. If the loss requires a label, the dataset must have a column named **`label`, `labels`, `score`, or `scores`**. Any column with one of these names is the label.
2. All other columns are **inputs**. The loss defines how many input columns it expects. Column **names don't matter; order does**.

Example: `CoSENTLoss` expects 2 inputs + a float label. A dataset with columns `["premise", "hypothesis", "score"]` works. A dataset with `["score", "premise", "hypothesis"]` does not — reorder first.

## Per-loss data shapes

The per-type loss references (`losses_sentence_transformer.md`, `losses_cross_encoder.md`, `losses_sparse_encoder.md`) are the canonical mappings from data shape to loss. Cross-cutting recipe bits those tables don't show:

- **`CosineSimilarityLoss`** wants `score` normalized to `[0, 1]`; `CoSENTLoss` / `AnglELoss` are pairwise-ranking and ignore absolute scale, so on `stsb` (raw 0-5) divide by 5 only when using cosine-similarity.
- **`BatchAllTripletLoss` / `BatchHardTripletLoss` / `BatchSemiHardTripletLoss`** need `batch_sampler=BatchSamplers.GROUP_BY_LABEL` so multiple samples per label appear in the same batch.
- **`MSELoss` (distillation)** label is the teacher's full embedding vector (a list of floats), not a scalar score.
- **`MarginMSELoss` (distillation)** label is `teacher_score(q, pos) - teacher_score(q, neg)`, precomputed per row.
- **N-tuple shape** for MNRL `(anchor, positive, negative_1, negative_2, ..., negative_N)` (1-indexed) is produced by `mine_hard_negatives(..., output_format="n-tuple")`; "labeled-list" output_format produces the CrossEncoder listwise shape.

## Reshaping operations

If your data doesn't fit the loss's expected shape:

### Reorder columns

```python
# Columns are ["hypothesis", "premise", "score"] but CoSENTLoss expects premise first.
dataset = dataset.select_columns(["premise", "hypothesis", "score"])
```

### Rename label column

```python
# Your label is called "relevance" but ST wants "label".
dataset = dataset.rename_column("relevance", "label")
```

### Drop extra columns

```python
# ST will treat every non-label column as an input. Drop metadata.
dataset = dataset.remove_columns(["source_id", "created_at", "language"])
```

### Convert dtypes

```python
# Label is str, need float for CoSENTLoss.
dataset = dataset.map(lambda x: {"label": float(x["label"])})
```

## Hard-negative mining

`mine_hard_negatives` (in `sentence_transformers.util`) produces a training dataset with mined negatives using a retriever. Hard negatives are the single highest-leverage lever for retrieval-model quality.

### Basic usage

```python
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import mine_hard_negatives

retriever = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

mined = mine_hard_negatives(
    dataset=train_pairs,                # has (anchor, positive) or (q, a) columns
    model=retriever,
    num_negatives=5,
    range_min=0, range_max=100,         # rank window to sample hard negatives from
    sampling_strategy="top",            # "top" = rank-1 hardest; "random" = random in window
    output_format="n-tuple",            # "triplet" | "n-tuple" | "labeled-pair" | "labeled-list"
    use_faiss=True,
)
```

### Output formats

- `"triplet"` — `(anchor, positive, negative)` triplets. One row per `(query, negative)` pair.
- `"n-tuple"` — `(anchor, positive, negative_1, negative_2, ..., negative_N)` (1-indexed) — one row per query.
- `"labeled-pair"` — `(anchor, text, label)` with `label=1` for positives and `label=0` for negatives. Good for `BinaryCrossEntropyLoss`.
- `"labeled-list"` — `(anchor, texts, labels)` — one row per query with a list of candidates. Good for listwise losses.

### Filtering false negatives

If the retriever returns "negatives" that are actually relevant, they become false negatives and hurt training. Filter them:

```python
mined = mine_hard_negatives(
    dataset=train_pairs,
    model=retriever,
    cross_encoder=CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2"),   # score candidates
    num_negatives=5,
    max_score=0.9,                      # drop candidates scoring above 0.9
    relative_margin=0.05,               # require neg_score < pos_score * (1 - 0.05)
    absolute_margin=0.2,                # require neg_score < pos_score - 0.2
    output_format="n-tuple",
    use_faiss=True,
)
```

Use **either** `relative_margin` or `absolute_margin`, usually not both. `max_score` is independently useful as a hard ceiling.

### CLI

`scripts/mine_hard_negatives.py` is a CLI wrapper — see it for a ready-to-run command.

## Choosing the right `range_min` / `range_max`

`range_max=None` is the default; pass an integer to cap how far down the ranked list to sample from.

- `range_min=0`, `range_max=100` — sample from the top-100 retrieved. Good default.
- `range_min=10`, `range_max=100` — skip the top-10 (often contains true positives). Safer if you lack a cross-encoder.
- `range_min=0`, `range_max=1000` — wider net, more diverse negatives, slower.
- `sampling_strategy="top"` — always pick the rank-1 hardest. Maximum training signal per row.
- `sampling_strategy="random"` — pick randomly within the range. More robust if your retriever is itself noisy.

## Quick Hub-side dataset checks

`hf datasets sql "SELECT * FROM 'hf://datasets/<id>/<split>' LIMIT 5"` streams rows via DuckDB without `load_dataset(...)` — the fastest way to confirm column names match your loss before a full validation run. `hf datasets info <id>` shows config / splits / size; `hf datasets card <id> --text` renders the README.

## Gotchas

- **`remove_unused_columns=True` (default)**: the trainer drops columns that aren't passed to the model's forward. Usually fine, but if you rely on a custom collator that uses metadata columns, set `remove_unused_columns=False`.
- **Floats stored as strings after CSV load**: `load_dataset("csv", ...)` keeps columns as strings by default. Cast with `.map(lambda x: {"label": float(x["label"])})`.
- **Mined hard negatives with `include_positives=True`** include the positive as a negative in the output list — only useful when you're building an evaluator or want to measure rank of the positive. For training, leave it `False`.

````


### `references/evaluators_cross_encoder.md`

````markdown
# Evaluators (Cross-Encoder)

All cross-encoder evaluators live in `sentence_transformers.cross_encoder.evaluation`.

## Choosing the right evaluator

| Task | Evaluator |
|---|---|
| Rerank retrieval results (nDCG@k on BM25 top-N) — fast default | `CrossEncoderNanoBEIREvaluator` |
| Rerank with custom candidates per query | `CrossEncoderRerankingEvaluator` |
| Binary / multi-class pair classification | `CrossEncoderClassificationEvaluator` |
| Continuous pair scoring (STS-style) | `CrossEncoderCorrelationEvaluator` |

Wrap multiple in `SequentialEvaluator` (from `sentence_transformers.base.evaluation`) to track them together:

```python
from sentence_transformers.base.evaluation import SequentialEvaluator
evaluator = SequentialEvaluator([nano_beir_eval, custom_rerank_eval])
```

## The default: `CrossEncoderNanoBEIREvaluator`

Analog of `NanoBEIREvaluator` for rerankers. Takes BM25 top-100 for each NanoBEIR query and measures how well the cross-encoder re-ranks them.

```python
from sentence_transformers.cross_encoder.evaluation import CrossEncoderNanoBEIREvaluator

evaluator = CrossEncoderNanoBEIREvaluator(
    dataset_names=["msmarco", "nfcorpus", "nq"],   # default: 11 of 13 NanoBEIR datasets (excludes "arguana", "touche2020")
    batch_size=64,
    rerank_k=100,                                   # rerank the BM25 top-K
)
```

Output key for `metric_for_best_model`: **`eval_NanoBEIR_R100_mean_ndcg@10`**. The `R100` signals "rerank top-100"; if you change `rerank_k`, the prefix changes (e.g. `R50`).

Each individual dataset contributes `eval_Nano{DatasetName}_R100_ndcg@10` (e.g. `eval_NanoMSMARCO_R100_ndcg@10`) too.

## Custom reranking with your own candidates

Use when you have query + positive + distractor candidates that aren't part of NanoBEIR:

```python
from sentence_transformers.cross_encoder.evaluation import CrossEncoderRerankingEvaluator

samples = [
    {"query": "...", "positive": ["the gold answer"], "documents": ["...", "...", ...]}
    for ...
]

evaluator = CrossEncoderRerankingEvaluator(
    samples=samples,
    batch_size=64,
    name="my-rerank",
    always_rerank_positives=False,   # default is True; override to False for realistic eval
)
```

- `always_rerank_positives=True` (the library default) forces the positive into the candidate pool even when the retriever missed it. The reranker is graded only on candidates it can actually score, so the metric reflects pure reranker quality.
- `always_rerank_positives=False`: the positive is only reranked if it's already in `documents`. If the retriever missed it, the rank counts as N+1. This reflects end-to-end retriever+reranker quality. A positive the retriever missed is lost, regardless of reranker skill.

Output key: `eval_{name}_ndcg@10`, `eval_{name}_map`, `eval_{name}_mrr@10`.

## Classification-style cross-encoders

### `CrossEncoderClassificationEvaluator`

Works for both binary (`num_labels=1`) and multi-class (`num_labels>=2`) cross-encoders. Branches internally:
- `num_labels=1`: binary mode. Sweeps thresholds to report accuracy, F1, precision, recall, and **average_precision** (primary).
- `num_labels>=2`: multi-class mode (e.g. NLI: entailment / neutral / contradiction). Reports **f1_macro** (primary), f1_micro, f1_weighted, and per-class precision / recall.

```python
from sentence_transformers.cross_encoder.evaluation import CrossEncoderClassificationEvaluator

evaluator = CrossEncoderClassificationEvaluator(
    sentence_pairs=[(premise, hypothesis), ...],
    labels=[0, 1, 2, ...],
    batch_size=64,
    name="nli-dev",
)
```

Output keys (binary, `num_labels=1`): `eval_{name}_accuracy`, `eval_{name}_f1`, `eval_{name}_average_precision` (primary).
Output keys (multi-class, `num_labels>=2`): `eval_{name}_f1_macro` (primary), `eval_{name}_f1_micro`, `eval_{name}_f1_weighted`.

### `CrossEncoderCorrelationEvaluator`

For continuous-score cross-encoders (like an STS cross-encoder outputting a similarity score). Reports Pearson/Spearman vs. gold scores.

```python
from sentence_transformers.cross_encoder.evaluation import CrossEncoderCorrelationEvaluator

evaluator = CrossEncoderCorrelationEvaluator(
    sentence_pairs=[(a, b), ...],
    scores=[0.4, 0.8, ...],
    name="stsb-dev",
)
```

Output keys: `eval_{name}_spearman`, `eval_{name}_pearson`.

## Writing `metric_for_best_model`

Pattern: `f"eval_{evaluator.primary_metric}"`. Inspect after construction: `print(evaluator.primary_metric)`. Common values:
- `eval_NanoBEIR_R100_mean_ndcg@10` — `CrossEncoderNanoBEIREvaluator` default
- `eval_{name}_ndcg@10` — `CrossEncoderRerankingEvaluator`
- `eval_{name}_average_precision` — `CrossEncoderClassificationEvaluator` (binary, `num_labels=1`)
- `eval_{name}_f1_macro` — `CrossEncoderClassificationEvaluator` (multi-class, `num_labels>=2`)
- `eval_{name}_spearman` — `CrossEncoderCorrelationEvaluator`

## Gotchas

- **Always run `evaluator(model)` once before training** — pre-training baseline. Tiny post-training delta means the loss/data/base is wrong.
- `CrossEncoderClassificationEvaluator` accepts both `num_labels=1` (binary, primary `average_precision`) and `num_labels>=2` (multi-class, primary `f1_macro`); `CrossEncoderCorrelationEvaluator` requires `num_labels=1`.
- The default `dataset_names=None` excludes `arguana` and `touche2020` (Argument-Retrieval task differs from the rest); pass `dataset_names=list(DATASET_NAME_TO_HUMAN_READABLE)` from `sentence_transformers.cross_encoder.evaluation.nano_beir` to actually run all 13.
- Subset NanoBEIR datasets during training (3–4) to keep eval cheap; run the broader set post-training.

````


### `references/evaluators_sentence_transformer.md`

````markdown
# Evaluators (Bi-Encoder)

All bi-encoder evaluators live in `sentence_transformers.sentence_transformer.evaluation`.

## Choosing the right evaluator

| Task | Evaluator |
|---|---|
| Retrieval (nDCG, MRR, Recall) — fast default | `NanoBEIREvaluator` |
| Retrieval on your own corpus / qrels | `InformationRetrievalEvaluator` |
| STS / continuous similarity | `EmbeddingSimilarityEvaluator` |
| Binary classification | `BinaryClassificationEvaluator` |
| Triplet accuracy | `TripletEvaluator` |
| Reranking (from retrieval candidates) | `RerankingEvaluator` |
| MSE vs. teacher (distillation) | `MSEEvaluator`, `MSEEvaluatorFromDataFrame` |
| Paraphrase mining | `ParaphraseMiningEvaluator` |
| Translation (cross-lingual alignment) | `TranslationEvaluator` |
| Label accuracy (classification during training) | `LabelAccuracyEvaluator` |

Wrap multiple evaluators in `SequentialEvaluator` to track all of them together:

```python
from sentence_transformers.sentence_transformer.evaluation import SequentialEvaluator
evaluator = SequentialEvaluator([evaluator1, evaluator2, evaluator3])
```

## The big three

### `NanoBEIREvaluator` (retrieval)

Small, fast subset of BEIR. Typically runs in <1 minute on a mid-range GPU. Default choice for retrieval training.

```python
from sentence_transformers.sentence_transformer.evaluation import NanoBEIREvaluator

evaluator = NanoBEIREvaluator(
    dataset_names=["msmarco", "nfcorpus", "nq"],   # default: all 13 NanoBEIR datasets
    batch_size=128,
    show_progress_bar=False,
)
```

- Default dataset list covers 13 tasks; pick a subset for speed during training.
- Output key for `metric_for_best_model`: **`eval_NanoBEIR_mean_cosine_ndcg@10`** (bi-encoder default = cosine similarity).

### `EmbeddingSimilarityEvaluator` (STS-style)

Computes Pearson/Spearman correlation between model cosine similarities and gold labels.

```python
from sentence_transformers.sentence_transformer.evaluation import EmbeddingSimilarityEvaluator
from sentence_transformers.util.similarity import SimilarityFunction

evaluator = EmbeddingSimilarityEvaluator(
    sentences1=stsb["sentence1"],
    sentences2=stsb["sentence2"],
    scores=stsb["score"],
    main_similarity=SimilarityFunction.COSINE,
    name="sts-dev",
)
```

- `main_similarity` can be `COSINE`, `DOT_PRODUCT`, `EUCLIDEAN`, `MANHATTAN`.
- `name` is used in the output key: `eval_sts-dev_spearman_cosine`, `eval_sts-dev_pearson_cosine`, etc.

### `InformationRetrievalEvaluator` (full retrieval)

Use when you have your **own** corpus + queries + qrels (not one of the NanoBEIR tasks).

```python
from sentence_transformers.sentence_transformer.evaluation import InformationRetrievalEvaluator

evaluator = InformationRetrievalEvaluator(
    queries={qid: query_text for qid, query_text in ...},
    corpus={doc_id: doc_text for doc_id, doc_text in ...},
    relevant_docs={qid: {doc_id, ...} for qid in ...},   # qid -> set of relevant doc_ids
    name="my-retrieval",
    mrr_at_k=[10],
    ndcg_at_k=[10],
    accuracy_at_k=[1, 5, 10],
    precision_recall_at_k=[1, 5, 10],
    map_at_k=[100],
    show_progress_bar=False,
    batch_size=64,
)
```

Output keys: `eval_{name}_cosine_ndcg@10`, `eval_{name}_cosine_mrr@10`, etc.

Heavy for large corpora — each eval encodes the full corpus. Don't run it every 100 steps. Use `NanoBEIREvaluator` for frequent evaluation during training and reserve full IR for milestones / post-training.

## Other bi-encoder evaluators

### `BinaryClassificationEvaluator`

For labeled pair classification (e.g. duplicate detection, entailment as binary). Reports accuracy, F1, precision/recall, AP. Supports all distance metrics — finds the best threshold per metric.

### `TripletEvaluator`

For `(anchor, positive, negative)` triplets. Reports the fraction of triplets where the positive is closer to the anchor than the negative.

### `RerankingEvaluator`

For custom re-ranking datasets: you provide candidates per query, the evaluator computes MAP and MRR. Good for measuring retrieval-quality on a held-out set.

### `MSEEvaluator` / `MSEEvaluatorFromDataFrame`

For distillation setups. Compares student embeddings against teacher embeddings (or teacher scores), reports MSE.

### `ParaphraseMiningEvaluator`

For paraphrase-mining tasks. Given a corpus of labeled paraphrase pairs, computes mining quality (F1 at various thresholds).

### `TranslationEvaluator`

For cross-lingual / `make_multilingual`-style alignment checking. Measures whether the student aligns sentences across languages.

### `LabelAccuracyEvaluator`

For a `SoftmaxLoss`-trained classifier head. Reports accuracy on held-out data.

## Writing `metric_for_best_model`

Pattern: `f"eval_{evaluator.primary_metric}"`. Inspect after construction: `print(evaluator.primary_metric)`. Common values:
- `eval_NanoBEIR_mean_cosine_ndcg@10` — `NanoBEIREvaluator`
- `eval_sts-dev_spearman_cosine` — `EmbeddingSimilarityEvaluator(name="sts-dev")`
- `eval_{name}_cosine_ndcg@10` — `InformationRetrievalEvaluator(name=...)`

## Multi-dimensional evaluation (Matryoshka)

For Matryoshka-trained models, evaluate at each target dimension:

```python
per_dim_evaluators = [
    EmbeddingSimilarityEvaluator(
        sentences1=..., sentences2=..., scores=...,
        main_similarity=SimilarityFunction.COSINE,
        name=f"sts-dev-{dim}",
        truncate_dim=dim,
    ) for dim in [768, 512, 256, 128, 64]
]
evaluator = SequentialEvaluator(per_dim_evaluators, main_score_function=lambda scores: scores[0])
```

The first evaluator's score drives `load_best_model_at_end`.

## Gotchas

- **Always run `evaluator(model)` once before training** — pre-training baseline. If the post-training delta is tiny, the loss/data/base is wrong.
- Don't run `InformationRetrievalEvaluator` with a large corpus (>100k docs) at frequent `eval_steps` — use `NanoBEIREvaluator` during training, reserve full IR for end-of-training.
- `greater_is_better=True` is the default; right for nDCG / MRR / accuracy.

````


### `references/evaluators_sparse_encoder.md`

````markdown
# Evaluators (Sparse Encoder)

All sparse-encoder evaluators live in `sentence_transformers.sparse_encoder.evaluation`. They mirror the bi-encoder versions with a `Sparse` prefix and default to **dot product** similarity (cosine on sparse vectors is less meaningful).

## Choosing the right evaluator

| Task | Evaluator |
|---|---|
| Retrieval (nDCG, MRR, Recall) — fast default | `SparseNanoBEIREvaluator` |
| Retrieval on your own corpus / qrels | `SparseInformationRetrievalEvaluator` |
| STS / continuous similarity | `SparseEmbeddingSimilarityEvaluator` |
| Binary classification | `SparseBinaryClassificationEvaluator` |
| Triplet accuracy | `SparseTripletEvaluator` |
| Reranking (from retrieval candidates) | `SparseRerankingEvaluator` |
| MSE vs. teacher (distillation) | `SparseMSEEvaluator` |
| Translation (cross-lingual alignment) | `SparseTranslationEvaluator` |
| Hybrid BM25 + sparse retrieval | `ReciprocalRankFusionEvaluator` |

Wrap multiple in `SequentialEvaluator` (from `sentence_transformers.base.evaluation`):

```python
from sentence_transformers.base.evaluation import SequentialEvaluator
evaluator = SequentialEvaluator([sparse_nano_beir, my_custom_ir])
```

## The default: `SparseNanoBEIREvaluator`

Small, fast subset of BEIR adapted for sparse retrieval. Typical runtime <1 minute on a mid-range GPU.

```python
from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator

evaluator = SparseNanoBEIREvaluator(
    dataset_names=["msmarco", "nfcorpus", "nq"],   # default: all 13 NanoBEIR datasets
    batch_size=32,
    show_progress_bar=False,
)
```

Output key for `metric_for_best_model`: **`eval_NanoBEIR_mean_dot_ndcg@10`** (sparse defaults to dot product).

### Sparsity tracking

Unlike the dense variant, the sparse evaluator also reports **active dimension counts** so you can monitor sparsity during training:

- `query_active_dims` — non-zero entries per query vector
- `document_active_dims` — non-zero entries per document vector

A healthy SPLADE checkpoint typically shows ~30–50 active dims for queries and ~150–250 for documents. If these drift toward the vocab size (~30k), the FLOPS regularization isn't doing its job — raise `query_regularizer_weight` / `document_regularizer_weight` in `SpladeLoss`.

## Retrieval on your own corpus

### `SparseInformationRetrievalEvaluator`

Same shape as the dense version but operates on sparse vectors internally:

```python
from sentence_transformers.sparse_encoder.evaluation import SparseInformationRetrievalEvaluator

evaluator = SparseInformationRetrievalEvaluator(
    queries={qid: text for qid, text in ...},
    corpus={doc_id: text for doc_id, text in ...},
    relevant_docs={qid: {doc_id, ...} for qid in ...},
    name="my-sparse-ir",
    ndcg_at_k=[10],
    mrr_at_k=[10],
    accuracy_at_k=[1, 5, 10],
    map_at_k=[100],
    batch_size=32,
)
```

Output keys: `eval_{name}_dot_ndcg@10`, `eval_{name}_dot_mrr@10`, etc. Also reports active-dims.

Heavy for large corpora. Use `SparseNanoBEIREvaluator` during training; reserve full IR for post-training.

## Hybrid retrieval

### `ReciprocalRankFusionEvaluator`

Measures the performance of combining your sparse encoder with BM25 (or any other retriever) via reciprocal-rank fusion. Useful when shipping a hybrid system is the actual deployment target.

## Other sparse evaluators

### `SparseEmbeddingSimilarityEvaluator`

STS-style. Computes Pearson/Spearman between sparse vector similarities and gold labels. Uses dot product by default.

### `SparseBinaryClassificationEvaluator`

For labeled pair classification with sparse embeddings.

### `SparseTripletEvaluator`

For `(anchor, positive, negative)` triplets — reports fraction where the positive is closer than the negative (by dot product).

### `SparseRerankingEvaluator`

For custom re-ranking with sparse embeddings. Same semantics as the dense `RerankingEvaluator`.

### `SparseMSEEvaluator`

For distillation setups. Compares sparse student embeddings against teacher outputs.

### `SparseTranslationEvaluator`

For cross-lingual / `make_multilingual`-style alignment checking with sparse embeddings.

## Writing `metric_for_best_model`

Pattern: `f"eval_{evaluator.primary_metric}"`. Inspect after construction: `print(evaluator.primary_metric)`. Common values:
- `eval_NanoBEIR_mean_dot_ndcg@10` — `SparseNanoBEIREvaluator` default
- `eval_{name}_dot_ndcg@10` — `SparseInformationRetrievalEvaluator`
- `eval_{name}_spearman_dot` — `SparseEmbeddingSimilarityEvaluator`

## Gotchas

- **Always run `evaluator(model)` once before training** — confirms the pipeline works (a fill-mask base scores ~0 on retrieval until trained).
- Sparse evaluators default to dot product; cosine on sparse vectors isn't meaningful.
- Don't compare dense and sparse metrics directly — different scales (cosine ∈ [-1, 1] vs. dot ∈ [0, ∞)).
- Always check `query_active_dims` / `document_active_dims` — thousands of active dims per doc means the FLOPS regularizer is mistuned, even if nDCG looks OK.

````


### `references/hardware_guide.md`

````markdown
# Hardware Guide

Training embedding models is memory-bound more often than compute-bound.

## If you hit OOM

Try in this order:

1. **Reduce `per_device_train_batch_size`**. Raise `gradient_accumulation_steps` to keep the effective batch size for regression losses. (For MNRL, effective batch via grad-accum is **not** equivalent — see point 3.)
2. **Enable `gradient_checkpointing=True`**. ~30% slower, ~40% less activation memory. Incompatible with `Cached*` losses.
3. **Switch to a `Cached*` loss**:
   - `CachedMultipleNegativesRankingLoss(model, mini_batch_size=32)` — forwards in mini-batches, accumulates the contrastive loss over the full batch. Can simulate batch sizes of 1024+ on a 24GB GPU.
   - `CachedSpladeLoss(model, loss=..., mini_batch_size=16)` — same trick for sparse.
   - `CachedGISTEmbedLoss(model, guide_model, mini_batch_size=32)` — GIST variant.
4. **Enable PEFT / LoRA** for decoder models >1B. `LoraConfig(r=64, lora_alpha=128, task_type="FEATURE_EXTRACTION")`. See `../scripts/train_sentence_transformer_with_lora_example.py` (docstring covers when to use, hyperparams, QLoRA, sharing).
5. **Move to multi-GPU**. See below.
6. **Shorten sequences**. If truncating to 128 is already sufficient for your task, set `max_seq_length` on the transformer module.

## Multi-GPU

`sentence-transformers` uses `accelerate` under the hood. Distributed training works without code changes.

### Data parallel (DDP)

Launch:

```bash
accelerate launch train.py

# or explicitly:
accelerate launch --multi_gpu --num_processes=4 train.py
```

`per_device_train_batch_size` stays per-GPU. Effective batch size scales linearly. MNRL's in-batch negatives remain **per-device**, not global, unless you pass `gather_across_devices=True` to losses that support it (`MultipleNegativesRankingLoss`, `CachedMultipleNegativesRankingLoss`, the symmetric variants, `GISTEmbedLoss`, `CachedGISTEmbedLoss`, `SparseMultipleNegativesRankingLoss`).

### FSDP / DeepSpeed

For models >3B, use `accelerate config` to enable FSDP or DeepSpeed ZeRO. Both are supported — `sentence-transformers` doesn't require any code changes, only the launch config.

```bash
accelerate config                   # interactive; choose FSDP or DeepSpeed
accelerate launch train.py
```

With FSDP full-shard: a 7B model trains on 4×24GB GPUs that would OOM on any single one of them.

**FSDP caveats** (from the [distributed training docs](https://sbert.net/docs/sentence_transformer/training/distributed.html)):
- **Evaluators don't run under FSDP** as of writing — the eval hooks call `model.encode()` which FSDP-wrapped modules can't service mid-training. Plan to evaluate post-training on a single-GPU load of the final checkpoint instead, or train with DDP if you need mid-training evaluation.
- **Layer wrapping must be specified**, e.g. `fsdp_config={"transformer_layer_cls_to_wrap": "BertLayer"}` (substitute the right layer class for your model: `BertLayer`, `LlamaDecoderLayer`, `Qwen2DecoderLayer`, etc.). Without this FSDP sharding can silently misbehave.
- **Slower than DDP** for models that fit on a single GPU — only reach for FSDP when you actually need the memory savings.

DeepSpeed ZeRO-2/3 is an alternative with its own config; works identically at the `accelerate config` level.

## Effective batch size for contrastive losses

For `MultipleNegativesRankingLoss` and its variants, **batch size is a quality knob**, not just a speed knob. Bigger batches = more in-batch negatives = richer gradients.

The effective pool of in-batch negatives per anchor:

| Setup | In-batch negatives per anchor |
|---|---|
| Single GPU, batch 64 | 63 |
| 4× DDP, per-device batch 64 | 63 local only by default; 255 with `MultipleNegativesRankingLoss(model, gather_across_devices=True)` |
| Single GPU, CachedMNRL, mini_batch 32, batch 256 | 255 |
| 4× DDP, CachedMNRL, per-device 256 | 255 local; 1023 with `gather_across_devices=True` |

For large corpora (retrieval), push toward 512+ effective negatives. For small, clean datasets (STS), 64 is plenty.

## Precision choice by GPU

| GPU generation | Recommended |
|---|---|
| T4, V100, GTX 1xxx, RTX 2xxx | `fp16=True` |
| RTX 3xxx, A10G, A100, L4 | `bf16=True` |
| RTX 4xxx, H100, B200 | `bf16=True` (or fp8 on H100 via specific kernels — not default) |
| Apple M-series / ROCm | MPS/ROCm support is variable; `fp16` or `fp32` most reliable |

bf16 is more numerically stable and almost always preferred when available.

## Hugging Face Jobs flavor guide

Hugging Face Jobs requires a Pro/Team/Enterprise plan. Pricing is approximate and subject to change — see the [Jobs pricing page](https://huggingface.co/docs/huggingface_hub/guides/jobs).

| Flavor | Memory | Typical use | Est. $/hr |
|---|---|---|---|
| `cpu-basic` | ~2 GB | Dataset prep, validation, hard-neg mining (small) | <$0.10 |
| `cpu-upgrade` | ~4 GB | Same, slightly bigger | $0.10 |
| `t4-small` | 16 GB | Demos, MiniLM/DistilBERT with small batches | ~$0.75 |
| `t4-medium` | 16 GB | MiniLM / DistilBERT with larger batch | ~$1.50 |
| `l4x1` | 24 GB | BERT-base, MPNet, ModernBERT-base | ~$2.50 |
| `a10g-small` | 24 GB | BERT-base to BERT-large | ~$3.50 |
| `a10g-large` | 48 GB | ModernBERT-large, Qwen3-0.6B | ~$5.00 |
| `a10g-largex2` | 96 GB (2× 48GB) | Mid-size multi-GPU | ~$10 |
| `a100-large` | 80 GB | Large models or big contrastive batches | ~$10–12 |
| `h100` | 80 GB | Biggest single-GPU | ~$12 |
| `h100x8` | 640 GB | LLM-scale distributed | ~$96 |

Defaults by base model:
- MiniLM / DistilBERT -> `t4-small`
- BERT-base / MPNet / ModernBERT-base -> `a10g-small` or `l4x1`
- BERT-large / ModernBERT-large -> `a10g-large`
- Qwen3-0.6B decoder base -> `a10g-large`
- 1B+ decoder bases with LoRA -> `a10g-large` or `a100-large`

Always start one flavor **smaller** than you think you need: OOM on Jobs is cheap ($0.50–$5 for a failed run). Underprovisioned is better than overprovisioned for the first attempt. When budgeting `timeout`, add **20–30% buffer** for model loading, checkpoint saving, and Hub push.

````


### `references/hf_jobs_execution.md`

````markdown
# Hugging Face Jobs Execution

Run training on Hugging Face's managed GPUs without provisioning any local infrastructure. The same training script runs locally and on Jobs — this reference covers only the Jobs-specific concerns.

## Prerequisites

- Hugging Face account with a **Pro, Team, or Enterprise** plan. Jobs are paid.
- `HF_TOKEN` with **write** permission. Log in once locally with `hf auth login` (the modern command from the `hf` CLI; the older `huggingface-cli login` still works but is deprecated).
- Access to the `hf_jobs()` MCP tool, or the `hf` CLI (`curl -LsSf https://hf.co/cli/install.sh | bash -s`).

## The three submission paths

### 1. Inline script via MCP (recommended in Claude Code)

Pass the full training script as `script`. Dependencies come from the PEP 723 header.

```python
hf_jobs("uv", {
    "script": """
# /// script
# requires-python = ">=3.10"
# dependencies = ["sentence-transformers[train]>=5.0", "trackio"]
# ///

# <full training script content>
""",
    "flavor": "a10g-large",
    "timeout": "3h",
    "secrets": {"HF_TOKEN": "$HF_TOKEN"},
})
```

### 2. Script-from-URL via MCP

Upload the script to the Hub (as a model or dataset repo file) or a Gist, then reference by URL:

```python
hf_jobs("uv", {
    "script": "https://huggingface.co/USERNAME/scripts/resolve/main/train_bi_encoder.py",
    "flavor": "a10g-large",
    "timeout": "3h",
    "secrets": {"HF_TOKEN": "$HF_TOKEN"},
})
```

Local file paths (`./train.py`, `/path/to/train.py`) **do not work** — Jobs run in isolated containers without access to your filesystem.

### 3. CLI

```bash
hf jobs uv run \
    --flavor a10g-large \
    --timeout 3h \
    --secrets HF_TOKEN \
    "https://huggingface.co/USERNAME/scripts/resolve/main/train.py"
```

Syntax gotchas:
- Command order is `hf jobs uv run`, **not** `hf jobs run uv`.
- Flags (`--flavor`, `--timeout`, `--secrets`) go **before** the script URL.
- `--secrets` (plural), not `--secret`.

## Required script modifications for Jobs

Add these to your `TrainingArguments`:

```python
args = SentenceTransformerTrainingArguments(
    ...,
    push_to_hub=True,
    hub_model_id="your-username/my-model",
    hub_strategy="every_save",        # push each checkpoint; timeout-safe
    save_strategy="steps",
    save_steps=0.1,                   # 10 saves/pushes per epoch; scales with dataset size
)
```

Why each matters:

| Argument | Why |
|---|---|
| `push_to_hub=True` | The Jobs container is destroyed after the job finishes. Without Hub push, all weights are lost. |
| `hub_model_id` | Required to identify the destination repo. |
| `hub_strategy="every_save"` | Default, but worth being deliberate about on Jobs: each checkpoint is pushed as it's written, so a timeout leaves all completed checkpoints on the Hub. `"end"` only pushes once `trainer.train()` returns, so a timeout loses everything. |
| `save_strategy="steps"` + `save_steps=0.1` | Checkpoints must actually be saved for `hub_strategy="every_save"` to push them. Fractional `0.1` = save every 10% of training, auto-scales with dataset size. |

## Secrets

Secrets are environment variables injected into the Jobs container. They never appear in logs and are not part of the script.

| Secret | Required when |
|---|---|
| `HF_TOKEN` | Always, for Hub push. Also covers Trackio auth. |
| `WANDB_API_KEY` | Using `report_to="wandb"`. |
| `MLFLOW_TRACKING_URI`, `MLFLOW_TRACKING_TOKEN` | Using MLflow with a remote server. |

The `$HF_TOKEN` syntax in the job config references the value from your local environment at submission time — the literal string `$HF_TOKEN` is replaced with your token's value. Never hardcode tokens in the script itself.

Trackio (the default tracker in this skill) uses `HF_TOKEN` for auth, so no extra secrets are needed. Only switch to the W&B / MLflow rows above if you're using those trackers.

## Timeout

Default is **30 minutes**, which is too short for almost any real training. Set explicitly:

```python
"timeout": "2h"       # 2 hours
"timeout": "90m"      # 90 minutes
"timeout": "1.5h"     # 90 minutes
"timeout": 7200       # seconds, as integer
```

Rule: **estimated training time × 1.3**. The extra buffer covers model loading, dataset caching, checkpoint saving, and Hub push.

On timeout, the container is killed immediately. Only data on the Hub (`hub_strategy="every_save"` saves you here) or in persistent volumes survives.

## Dataset caching

Hugging Face datasets are cached at `~/.cache/huggingface/datasets` by default — **inside the container**, which is destroyed after the job. Each Jobs run re-downloads the dataset.

For large datasets (>5 GB), this matters. Options:

- **Persistent `/data` volume** (Jobs feature, check current documentation): set `HF_DATASETS_CACHE=/data/datasets` so caches persist across jobs.
- **Pre-cache locally, push to Hub**: if the dataset is on Hub already, nothing to do. If it's local-only, `dataset.push_to_hub(...)` once so subsequent jobs load from Hub.

## Monitoring a running job

```bash
hf jobs ps [--all]                        # running (or all) jobs
hf jobs inspect <job-id>                  # full config + status
hf jobs logs <job-id> [--follow|--tail N] # tail or stream
hf jobs cancel <job-id>
hf jobs hardware                          # list flavors + hourly rates
```

`hf jobs logs <id> --follow` under `Bash run_in_background` pairs nicely with a `Monitor` watching for the `VERDICT:` line emitted by your training script's verdict block.

MCP equivalents (signatures may vary by server version — check the actual
tool listing): `hf_jobs("ps")`, `hf_jobs("logs", {"job_id": ...})`,
`hf_jobs("cancel", {"job_id": ...})`.

For recurring runs, `hf jobs scheduled uv run "<cron>" <script> ...`
schedules; `hf jobs scheduled ps/suspend/delete` manages.

## Common failures

### "Model not found on Hub" after a successful-looking run

The run succeeded but `push_to_hub` was not enabled. The container is gone; the weights are gone.

Fix: always set `push_to_hub=True` + `hub_model_id=...` + `secrets={"HF_TOKEN": "$HF_TOKEN"}`.

### Tracker not connecting

- **Trackio:** `HF_TOKEN` missing or lacks write permission. Add `"secrets": {"HF_TOKEN": "$HF_TOKEN"}` and make sure the token has write access.
- **W&B:** `WANDB_API_KEY` missing. Add `"secrets": {"HF_TOKEN": "$HF_TOKEN", "WANDB_API_KEY": "$WANDB_API_KEY"}`.

### OOM on first step

Flavor too small. Move up one tier (see `hardware_guide.md`).

### Training starts but eval hangs forever

`eval_strategy="steps"` with no `eval_dataset`. Always provide an eval dataset, or set `eval_strategy="no"`.

### Dataset download times out

Large dataset or slow cold-cache. Increase `timeout` or pre-cache to a persistent volume.

### `CachedMultipleNegativesRankingLoss` + `gradient_checkpointing=True` crash

The cached losses are incompatible with gradient checkpointing. Disable `gradient_checkpointing`.

After submission, the MCP returns a job ID. Monitor with `hf_jobs("logs", {"job_id": ...})` when you want an update — don't poll in a tight loop. End-to-end submission templates live in `scripts/train_sentence_transformer_example.py` / `scripts/train_cross_encoder_example.py` / `scripts/train_sparse_encoder_example.py`; wrap the script contents in the inline pattern from §1 above.

````


### `references/losses_cross_encoder.md`

````markdown
# Cross-Encoder Losses (reranker)

All losses live in `sentence_transformers.cross_encoder.losses`.

Cross-encoder losses fall into three families: **pointwise** (score each pair independently), **pairwise** (score pairs of pairs), and **listwise** (score a whole ranked list at once). Plus contrastive/distillation variants for specific setups.

## Top-line decision table

| You have | Use | Family |
|---|---|---|
| `(query, passage, label)` with `label ∈ {0, 1}` or `[0, 1]` | `BinaryCrossEntropyLoss` | Pointwise |
| `(query, passage, class_id)` multi-class | `CrossEntropyLoss` | Pointwise |
| `(query, passage)` with implicit positives, want contrastive | `CachedMultipleNegativesRankingLoss` | Contrastive |
| `(query, passages, labels)` (one row per query, with parallel lists of candidate passages and their relevance scores) | `LambdaLoss` | Listwise |
| Same listwise shape, want well-tested simpler loss | `ListNetLoss` or `ListMLELoss` | Listwise |
| `(query, positive, negative)` pairwise | `RankNetLoss` | Pairwise |
| `(query, passage, teacher_score)` distillation from stronger reranker | `MSELoss` or `MarginMSELoss` | Distillation |

## Pointwise losses

### `BinaryCrossEntropyLoss`

**The default cross-encoder loss.** For `(query, passage, label)` where label is 0/1 or a continuous [0, 1] relevance score.

```python
loss = BinaryCrossEntropyLoss(model, pos_weight=torch.tensor(5.0))
```

- `pos_weight`: upweight positives when your dataset is imbalanced (e.g. 1 positive + N hard negatives). A good default is `pos_weight = num_hard_negatives`.
- Handles both binary labels and continuous relevance scores — if you have graded relevance (0.0, 0.5, 1.0) BCE still works.

### `CrossEntropyLoss`

Multi-class classification over passages. Use with `CrossEncoder("...", num_labels=N)` where N is the number of classes (e.g. 3 for NLI: entailment / neutral / contradiction).

## Contrastive losses (for reranker training)

### `MultipleNegativesRankingLoss` (cross-encoder)

Contrastive training mirroring bi-encoder MNRL. Applies in-batch negatives: for each `(query, positive)`, every other `positive` in the batch is used as a negative.

- **Data**: `(query, positive)` or `(query, positive, negative_0, negative_1, ...)`.
- For CrossEncoder training, **prefer `BinaryCrossEntropyLoss` with mined hard negatives** as the default; MNRL is the fallback when you only have `(query, positive)` pairs and want to avoid a separate mining pass.

### `CachedMultipleNegativesRankingLoss`

Cached variant (GradCache) — decouples per-device batch size from effective in-batch negatives, same as the bi-encoder cached version.

```python
loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=16)
```

- Incompatible with `gradient_checkpointing=True`.
- Key choice for training rerankers with effective batch 256+ on a single GPU.

## Distillation losses

> **`activation_fn=nn.Identity()` is mandatory** for all distillation, listwise, and pairwise losses below — only `BinaryCrossEntropyLoss` and `CrossEntropyLoss` tolerate the default `Sigmoid`. The loss sees raw logits during training, but the model's `activation_fn` is applied at evaluation via `predict()`; the default `Sigmoid` (with `num_labels=1`) saturates raw logits >5 to ~1.0 inside `predict()`, silently collapsing eval ranking (training loss looks healthy while nDCG crashes from ~0.59 to ~0.14). Construct as `CrossEncoder("...", num_labels=1, activation_fn=torch.nn.Identity())`. See `troubleshooting.md` ("CrossEncoder eval nDCG crashes after distillation / listwise / pairwise training") for the failure-mode walkthrough.

### `MSELoss` (cross-encoder)

Regress the student's output score to the teacher's score. Construct the model with `activation_fn=nn.Identity()` (see callout above).

- **Data**: `(query, passage, teacher_score)`.
- The teacher is typically a larger/stronger cross-encoder. Precompute its scores once, store as the label.

### `MarginMSELoss` (cross-encoder)

Regress the **difference** between positive and negative scores against the teacher's margin. Typically gives better distillation than plain MSE.

- **Data**: `(query, positive, negative, score_diff)` where `score_diff = teacher_score(query, positive) - teacher_score(query, negative)`.
- Popular recipe for MS MARCO-style distillation.

## Listwise losses

All listwise losses expect the dataset to have per-query **lists** of candidate passages and their scores — typically via a collator that groups rows by query.

> Listwise and pairwise losses below also require `activation_fn=nn.Identity()` (see the Distillation-section callout above).

### `LambdaLoss`

The state-of-the-art listwise ranking loss. Optimizes a surrogate of nDCG via weighted pairwise comparisons.

**Data shape**: `(query, [doc1...docN], [score1...scoreN])` per row; One query, a list of candidate documents, and a parallel list of relevance scores. Use `mine_hard_negatives(..., output_format="labeled-list", ...)` to build this from `(query, positive)` pairs.

```python
import torch.nn as nn
from sentence_transformers.cross_encoder.losses import LambdaLoss, NDCGLoss2PPScheme

model = CrossEncoder("...", num_labels=1, activation_fn=nn.Identity())
loss = LambdaLoss(model, weighting_scheme=NDCGLoss2PPScheme())
```

- Strong default when you have multiple candidates per query and graded relevance.
- `weighting_scheme`: `NDCGLoss2PPScheme` (default), `NDCGLoss2Scheme`, `LambdaRankScheme`. The `NDCGLoss2PPScheme` default was shown to reach the strongest performance in the original LambdaLoss paper.

#### LambdaLoss-specific operational notes

- **OOM recovery order**: drop `mini_batch_size` first (chunking inside the loss preserves the K-list semantic), then `per_device_train_batch_size` paired with `gradient_accumulation_steps`, then reduce K (the per-query candidate-list length) only as a last resort. Lowering K changes the experiment; the others don't.
- **Tiny loss at large K is expected.** With `NDCGLoss2PPScheme`, the loss normalizes by the discount-weighted pair count — at K=128 the loss can scale to ~`1e-4` numerically. That's not "training broken"; read `eval_NanoBEIR_R100_mean_ndcg@10` (or your equivalent) instead of the training loss to judge progress.
- For very large K (>=128), the O(K²) weight buffer per query is materialized outside of the forward chunking, so even small `mini_batch_size` may not be enough. Consider scoring strategy: top-K hard negatives often beat random K.

### `ListNetLoss`

Treats the ranking as a probability distribution (via softmax) and minimizes cross-entropy against the teacher distribution.

### `ListMLELoss`

Maximum-likelihood estimation over permutations. Simpler than LambdaLoss; decent default.

### `PListMLELoss`

Position-aware ListMLE — weights higher-ranked items more heavily. Often outperforms plain ListMLE on top-k metrics.

### `RankNetLoss`

Pairwise classification: for each pair of candidates, predict which is ranked higher via cross-entropy.

- Simpler and faster than LambdaLoss.
- Scales quadratically with list length; not great for long candidate lists (>20).

### `ADRMSELoss`

Alternative listwise formulation (Approx Discounted Rank MSE) from the Rank-DistiLLM paper. Same data shape as `LambdaLoss`. In practice LambdaLoss is the stronger default; `RankNetLoss` was reported as marginally more effective (~0.002 nDCG@10) than ADRMSELoss on the paper's LLM-distillation setup, and LambdaLoss generally beats both.

Hard-negative mining is essential for any contrastive reranker (random negatives teach nothing). See `dataset_formats.md` (Hard-negative mining section) and `../scripts/mine_hard_negatives.py`.

## Gotchas

- **`BinaryCrossEntropyLoss` without `pos_weight`** when you have 5+ hard negatives per positive: the loss under-weights the positive signal. Set `pos_weight=num_hard_negatives`.
- **`CachedMultipleNegativesRankingLoss` + `gradient_checkpointing=True`**: crash. Pick one.
- **Listwise losses with wildly different list lengths per query**: some losses don't handle ragged lists well. Pad or truncate to a fixed length.
- **`MarginMSELoss` without precomputed teacher score diffs**: this loss needs the `score_diff` label column populated from a teacher pass; it does not compute the teacher internally.
- **Training a cross-encoder with `num_labels=1` then using `CrossEntropyLoss`**: mismatch — BCE is for num_labels=1, CE is for num_labels>=2.
- **Default `Sigmoid` activation under any non-BCE loss**: silently destroys eval ranking. Pass `activation_fn=torch.nn.Identity()` to `CrossEncoder(...)` for distillation, listwise, and pairwise losses (everything except BCE/CE). See the callouts in the Distillation and Listwise sections.
- **Custom CE head writing to the wrong feature key**: a custom scoring head must populate `features["scores"]` (not `features["sentence_embedding"]`). Otherwise `CrossEncoder.predict()` raises `KeyError: 'scores'` at inference time, even though training succeeds.
- **Loading a CE with a custom-class head fails from a different script**: if you defined a `class ClassifierHead(nn.Module)` inline in `train.py` and saved the model, `modules.json` records `__main__.ClassifierHead`. Loading via `CrossEncoder("path")` from any other entry point raises `ImportError: Module '__main__' does not define a 'ClassifierHead' attribute`. Either move the class into an importable file (`my_pkg/heads.py`), build the same shape from stock ST modules (`Dense + LayerNorm + Dense`), or document that the model is only loadable from the same script.

````


### `references/losses_sentence_transformer.md`

````markdown
# Bi-Encoder Losses (SentenceTransformer)

All losses live in `sentence_transformers.sentence_transformer.losses`.

Losses are grouped by data shape. The #1 rule: **pick a loss that matches your data**, not the other way around.

## Top-line decision table

| You have | Use |
|---|---|
| `(anchor, positive)` pairs | `MultipleNegativesRankingLoss` (or Cached variant for large batches) |
| `(anchor, positive, negative)` triplets | `MultipleNegativesRankingLoss` — it handles triplets natively |
| `(text1, text2, score)` with `score ∈ [-1, 1]` or `[0, 1]` | `CoSENTLoss` (strongly recommended) |
| `(text1, text2, label)` with `label ∈ {0, 1}` | `OnlineContrastiveLoss` |
| `(text, class_id)` single-column with integer class | `BatchAllTripletLoss` |
| `(query, positive, negative, score_diff)` | `MarginMSELoss` (distillation) |
| `(text, teacher_embedding)` | `MSELoss` (embedding distillation) |
| Want multiple output dims from one training | Wrap any of the above in `MatryoshkaLoss` |
| No labels at all, just sentences | `DenoisingAutoEncoderLoss` or `ContrastiveTensionLossInBatchNegatives` |

## Contrastive losses (pairs + triplets, no labels)

### `MultipleNegativesRankingLoss` (MNRL)

The default bi-encoder loss. Uses **in-batch negatives**: every other `positive` in the batch acts as a negative for the current `anchor`.

```python
loss = MultipleNegativesRankingLoss(model, scale=20.0)  # similarity_fct defaults to cos_sim
```

- **Data**: `(anchor, positive)` or `(anchor, positive, negative)`. More columns = more explicit hard negatives per row.
- **Scale**: temperature. Default `scale=20.0` multiplies similarities by 20 (equivalent to softmax temperature 0.05). Tune only if cosine similarities end up saturated.
- **Critical**: set `batch_sampler=BatchSamplers.NO_DUPLICATES` on training args. Otherwise duplicate anchors create false negatives.
- **Tip**: batch size matters a lot — more in-batch negatives = better gradients.

### `CachedMultipleNegativesRankingLoss`

Same loss, but with gradient caching (GradCache): forwards in mini-batches but computes the contrastive loss over the full batch. Use this when you want effective batch size of 256+ but your GPU can only fit 32 forwards.

```python
loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=32)
```

- **Incompatible with `gradient_checkpointing=True`**.
- Set `mini_batch_size` to whatever `per_device_train_batch_size` would be if you couldn't use this. Then crank the actual `per_device_train_batch_size` to what you want the effective batch to be (256+, 1024+).

### `MultipleNegativesSymmetricRankingLoss`

MNRL computed bidirectionally — scores positives from both (anchor -> positive) and (positive -> anchor) directions. Slightly better on retrieval tasks where the "anchor" and "positive" distinctions are soft (paraphrase, deduplication).

### `CachedMultipleNegativesSymmetricRankingLoss`

Cached variant of the above.

### `GISTEmbedLoss`

Like MNRL, but uses a **guide model** (a separate pretrained Sentence Transformer) to **filter out false negatives** before computing the contrastive loss. The guide model scores each potential negative; if it looks too similar to the positive, it's excluded.

```python
guide = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
loss = GISTEmbedLoss(model, guide=guide)
```

- Expensive: guide model does a forward pass per batch.
- Strong when your in-batch negatives are noisy (e.g. small corpus with many near-duplicates).

### `CachedGISTEmbedLoss`

Cached + GIST.

### `MegaBatchMarginLoss`

In-batch margin-based triplet: for each anchor, find the hardest negative in the batch and apply a margin loss. Older pattern, usually outperformed by MNRL.

### `TripletLoss`

Classic triplet margin loss on explicit `(anchor, positive, negative)`. Uses a fixed margin and the hardest in-batch is not considered — only the provided triplet.

```python
loss = TripletLoss(model, distance_metric=TripletDistanceMetric.EUCLIDEAN, triplet_margin=5)
```

- Simpler than MNRL; less powerful when the batch has useful negatives.
- Good for cases where you have *trusted* pre-mined triplets and want to avoid in-batch noise.

## Batch-triplet losses (single column + integer label)

These mine triplets *within* the batch from samples sharing a label.

### `BatchAllTripletLoss`

For each anchor, form triplets with all positive/negative combinations in the batch. Max signal per batch.

```python
loss = BatchAllTripletLoss(model, margin=5)
```

- **Data**: single text column + integer `label`. Needs multiple samples per label in each batch (set `batch_sampler=BatchSamplers.GROUP_BY_LABEL`).

### `BatchHardTripletLoss`

Same, but only the single hardest positive + hardest negative per anchor.

### `BatchSemiHardTripletLoss`

Semi-hard mining: negatives harder than the positive but easier than the margin. Often more stable than fully-hard.

### `BatchHardSoftMarginTripletLoss`

Variant with a soft margin (log-sum-exp) instead of a fixed margin hinge.

**When to use batch-triplet losses**: classification-style datasets (labels are class IDs, not pair annotations). E.g. "train an embedder where samples from the same class are close."

## Scored regression losses (labeled pairs, float score)

### `CoSENTLoss`

**The recommended regression loss** for `(text1, text2, score)`. Trains on pairwise ranking: for any two pairs `(a, b)` and `(c, d)` with `score(a,b) > score(c,d)`, the model should score `(a, b)` higher. Much better than squared error.

```python
loss = CoSENTLoss(model, scale=20.0)
```

- **Data**: `(text1, text2, float_score)`. Labels can be `[0, 1]` or `[-1, 1]`.
- Works well with small datasets (STS-B, 5k pairs).

### `AnglELoss`

Similar to CoSENT but uses angle-based optimization in complex space. Sometimes outperforms CoSENT on tasks with fine-grained similarity gradations. Strong alternative.

### `CosineSimilarityLoss`

Squared-error loss on cosine similarity: `mse(cos(text1, text2), label)`. Simpler than CoSENT, usually worse. Keep for legacy / reproducibility.

## Contrastive labeled losses (labeled pairs, binary label)

### `ContrastiveLoss`

For `(text1, text2, label)` where `label ∈ {0, 1}`. Minimizes distance for positives; pushes negatives past a margin.

```python
loss = ContrastiveLoss(model, margin=0.5, distance_metric=SiameseDistanceMetric.COSINE_DISTANCE)
```

### `OnlineContrastiveLoss`

Same setup but **ignores "easy" pairs** (positives already close, negatives already far) and only optimizes hard ones. Much more robust to label noise.

```python
loss = OnlineContrastiveLoss(model, margin=0.5)
```

**Preferred over `ContrastiveLoss`** for most practical labeled pair datasets.

### `SoftmaxLoss`

A classifier head on concatenated `(u, v, |u-v|)` embeddings, trained with cross-entropy. Useful when you have NLI-style multi-class labels (entailment / neutral / contradiction) and want a categorical loss. Historically important (it trained the first popular sentence embedding models) but **generally outperformed by MNRL**.

## Distillation losses

### `MSELoss`

Regress the student's embedding to match a teacher's embedding.

- **Data**: `(text, teacher_embedding)`. The teacher embedding is a fixed vector per row.
- Use when distilling a smaller bi-encoder from a larger one.

### `MarginMSELoss`

For `(query, positive, negative, score_diff)`: minimize `mse(student_score_diff, teacher_score_diff)`. The teacher is typically a cross-encoder that produced the score differences.

- **Data**: 3 text columns + 1 float column.
- Workhorse of dense retriever training from cross-encoder teachers (ms-marco distillation).

### `DistillKLDivLoss`

KL-divergence distillation: student's softmax distribution over candidates should match teacher's.

- **Data**: `(query, passages[], teacher_scores[])`.
- Good for list-wise distillation when you have multiple candidates per query.

See `../scripts/train_sentence_transformer_distillation_example.py` for the end-to-end pattern (its docstring covers Embedding MSE / Margin MSE / Listwise KL with full recipes).

## Regularizer / wrapper losses

These don't have their own data shape — they wrap another loss and add a regularization objective.

### `MatryoshkaLoss`

Train **once**, deploy at any of several dimensions. Wraps any loss and computes it at multiple truncated dimensions, adding them weighted.

```python
base_loss = MultipleNegativesRankingLoss(model)
loss = MatryoshkaLoss(
    model,
    base_loss,
    matryoshka_dims=[768, 512, 256, 128, 64],
    matryoshka_weights=[1, 1, 1, 1, 1],   # relative weighting per dim
)
```

- At inference, `SentenceTransformer(..., truncate_dim=128)` gives 128-dim output with ~95% of full quality.
- Default matryoshka_dims: pick the dims you want to deploy at. Smaller dims improve compression at the cost of slight quality drop.

### `Matryoshka2dLoss`

2D-Matryoshka: reduce dimension **and** number of transformer layers in a single wrapper. Internally composes `MatryoshkaLoss` + `AdaptiveLayerLoss`, so you only need this one (don't wrap it in AdaptiveLayerLoss yourself). Deploy at any (dim, layer) pair at inference.

### `AdaptiveLayerLoss`

Wrap any loss; adds a term that trains each of the transformer's layers to be a valid exit point. Deploy with fewer layers at inference for faster encoding.

```python
loss = AdaptiveLayerLoss(
    model,
    base_loss,
    n_layers_per_step=1,
    last_layer_weight=1.0,
    prior_layers_weight=1.0,
)
```

### `GlobalOrthogonalRegularizationLoss` (GOR)

Stand-alone regularizer (not a wrapper, despite living in this section). Penalizes embedding pairs whose dot product deviates from orthogonality, encouraging the model to spread embeddings across the full vector space. Use it alongside a primary contrastive loss by summing the two outputs in your own training step; can help with downstream retrieval diversity.

## Unsupervised losses

### `DenoisingAutoEncoderLoss` (TSDAE)

Sentence-level denoising autoencoder: corrupt a sentence (drop tokens), force the model to reconstruct it. Pretraining-style — useful for domain adaptation when you have unlabeled in-domain sentences.

### `ContrastiveTensionLoss`

Unsupervised contrastive: two copies of the model encode the same sentence; they should agree. Pure self-supervised.

### `ContrastiveTensionLossInBatchNegatives`

CT with in-batch negatives. Stronger than vanilla CT.

## Gotchas

- **`MultipleNegativesRankingLoss` without `BatchSamplers.NO_DUPLICATES`** will include duplicate anchors in the same batch, destroying training signal. Always set the sampler.
- **Any `Cached*` loss + `gradient_checkpointing=True` = crash.** Pick one.
- **`TripletLoss` with bad negatives** (too easy) = loss hits zero fast and model stops learning. Mine hard negatives first.
- **Matryoshka wrapping `CachedMultipleNegativesRankingLoss`**: supported, but the cached-loss's mini-batch semantics apply to the base loss only. Think twice before combining.

````


### `references/losses_sparse_encoder.md`

````markdown
# Sparse-Encoder Losses (SPLADE)

All losses live in `sentence_transformers.sparse_encoder.losses`.

This reference targets the **SPLADE** architecture (Transformer + SpladePooling). The sparse-encoder package also exports `CSRLoss` and `CSRReconstructionLoss` for the CSR architecture (Transformer + Pooling + SparseAutoEncoder); those are out of scope here — see the sbert.net docs if you're training a CSR model.

Choosing a loss means (a) pick a base loss (contrastive, regression, distillation) and (b) wrap it in `SpladeLoss` to add FLOPS regularization.

## Top-line decision table

| You have | Use |
|---|---|
| `(anchor, positive)` or triplet, SPLADE architecture | `SpladeLoss(loss=SparseMultipleNegativesRankingLoss(model), ...)` |
| Same, want effective batch size of 256+ | `CachedSpladeLoss(...)` |
| `(text1, text2, score)` labeled pairs | `SparseCoSENTLoss` or `SparseCosineSimilarityLoss` |
| Distillation from cross-encoder teacher | `SparseMarginMSELoss` |
| Listwise distillation | `SparseDistillKLDivLoss` |
| Explicit triplet | `SparseTripletLoss` |

## The core wrapper: `SpladeLoss`

`SpladeLoss` adds **FLOPS regularization** on top of another sparse loss. FLOPS regularization penalizes non-zero activations, keeping embeddings genuinely sparse.

```python
loss = SpladeLoss(
    model=model,
    loss=SparseMultipleNegativesRankingLoss(model=model),
    query_regularizer_weight=5e-5,
    document_regularizer_weight=3e-5,
)
```

- `query_regularizer_weight`: how much to penalize non-zero terms in query embeddings.
- `document_regularizer_weight`: same for documents.
- Typical range: 1e-5 to 1e-4. Higher = sparser embeddings, lower recall; lower = denser, possibly better recall.
- `SparseEncoderTrainer` automatically registers a `SpladeRegularizerWeightSchedulerCallback` whenever the loss is a `SpladeLoss`. The callback ramps the weights from 0 up to the target over the first ~33% of training; the default shape is `SchedulerType.QUADRATIC` (not linear). The ramp length and shape are configured on the callback (`SpladeRegularizerWeightSchedulerCallback(loss=..., warmup_ratio=..., scheduler_type=...)`), not on `SpladeLoss`; to override, instantiate the callback yourself and pass it via `callbacks=[...]`. This ramp is important; starting with full regularization from step 0 kills learning.

Use `CachedSpladeLoss` for the GradCache variant.

## Contrastive losses (no labels)

### `SparseMultipleNegativesRankingLoss`

Sparse analog of bi-encoder MNRL. In-batch contrastive.

```python
inner = SparseMultipleNegativesRankingLoss(model=model)
loss = SpladeLoss(model=model, loss=inner, query_regularizer_weight=5e-5, document_regularizer_weight=3e-5)
```

- **Always wrap in `SpladeLoss`** for SPLADE architectures.
- Set `batch_sampler=BatchSamplers.NO_DUPLICATES` on training args.

### `SparseTripletLoss`

Classic triplet margin loss on explicit `(anchor, positive, negative)`.

## Labeled regression losses

### `SparseCoSENTLoss`

Pairwise ranking loss for `(text1, text2, score)`. Mirrors bi-encoder `CoSENTLoss`.

### `SparseCosineSimilarityLoss`

MSE on cosine similarity. Simpler, usually worse than CoSENT.

### `SparseAnglELoss`

Angle-based loss in complex space. Alternative to CoSENT.

## Distillation losses

### `SparseMSELoss`

Embedding MSE. Student sparse embedding should match teacher embedding.

- **Data**: `(text, teacher_embedding)`.
- Teacher can be a dense bi-encoder or another sparse model.

### `SparseMarginMSELoss`

Margin MSE from a cross-encoder teacher.

- **Data**: `(query, positive, negative, score_diff)` where `score_diff = teacher_score(query, positive) - teacher_score(query, negative)`.
- Typical recipe for training SPLADE from cross-encoder labels (ms-marco distillation).
- Wrap in `SpladeLoss(model, loss=SparseMarginMSELoss(model), ...)` for SPLADE.

### `SparseDistillKLDivLoss`

Listwise KL-div distillation — student's softmax distribution over candidates should match teacher's.

## Independent regularizer

### `FlopsLoss`

Standalone FLOPS regularizer. Usually you use this via `SpladeLoss`, not directly.

For regularizer-weight tuning and dense-output recovery, see `troubleshooting.md` ("SPLADE embeddings are dense"). MLM-head requirement: `base_model_selection.md` (SPARSE section). Active-dim sparsity targets and how to monitor them: `evaluators_sparse_encoder.md` (Sparsity tracking).

## Gotchas

- **`SparseMultipleNegativesRankingLoss` without `SpladeLoss` wrapping on a SPLADE model**: no FLOPS regularization -> dense outputs defeating the purpose of SPLADE. Always wrap.
- **`CachedSpladeLoss` + `gradient_checkpointing=True`**: crash. Pick one.
- **Starting training with full FLOPS regularization at step 0**: the model outputs zero everywhere and gets stuck. The built-in scheduler avoids this — don't override it unless you know why.
- **`query_regularizer_weight` == `document_regularizer_weight`**: usually wrong. Queries should be sparser than documents (fewer terms per query). Since higher regularization drives more zeros, give the query weight the larger value. `query_regularizer_weight=5e-5`, `document_regularizer_weight=3e-5` is a good starting ratio.

````


### `references/model_architectures.md`

````markdown
# Model Architectures (SentenceTransformer)

The `SentenceTransformer` class is a `torch.nn.Sequential` of modules. The common shape is `Transformer` + `Pooling` (+ optional `Normalize` / `Dense`), but four distinct architecture families are supported and the right choice depends on the task.

## The four architecture families

| Family | Backbone | Pooling | Use case |
|---|---|---|---|
| **Encoder (bidirectional)** | BERT, RoBERTa, DeBERTa, MPNet, ModernBERT, XLM-R | `mean` (default) or `cls` | Short/medium text, general default |
| **Decoder (causal LLM)** | Qwen, Llama, Mistral, Gemma | `lasttoken` | Long context, instruction-tunable, larger quality ceiling |
| **Static embeddings** | `StaticEmbedding` module | N/A | CPU-only, <10MB, extremely fast |
| **Multimodal / Router** | VLM backbones or composed encoders | depends | Text + image / audio / video |

Below, each family with concrete setup.

## Encoder models (the default)

The historical default and still usually the right choice for text embeddings.

```python
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("microsoft/mpnet-base")
# Auto-constructs: Transformer(feature-extraction) -> Pooling(mean).
```

When `SentenceTransformer("<checkpoint>")` is called with a raw HF encoder, it auto-wraps the transformer and adds `Pooling(..., pooling_mode="mean")`.

To customize pooling or add modules:

```python
from sentence_transformers import SentenceTransformer
from sentence_transformers.sentence_transformer.modules import Normalize, Pooling, Transformer

transformer = Transformer("answerdotai/ModernBERT-base")
pooling = Pooling(transformer.get_embedding_dimension(), pooling_mode="cls")    # or "mean", "lasttoken", ...
model = SentenceTransformer(modules=[transformer, pooling, Normalize()])
```

**Pooling modes**:
- `mean` (default) — average of token embeddings, masked to the attention mask. Strongest default.
- `cls` — embedding of the `[CLS]` token. Works if the base was CLS-pretrained.
- `max` — element-wise max across tokens. Rare.
- `mean_sqrt_len_tokens` — mean scaled by √seq_len. Empirically helps on some tasks.
- `weightedmean` — token-position-weighted mean. Useful for decoder bases as a non-last-token alternative.
- `lasttoken` — embedding of the last token. Required for causal-LM bases (see decoder section below).

Don't switch pooling mid-training. Pick once.

## Decoder / causal LLM models

Strong at long context, instruction following, multilingual. Memory-hungry — typically LoRA-trained rather than full fine-tuned.

**Two setup paths depending on whether the model was already adapted for embeddings:**

```python
# Path A: already-adapted embedding checkpoint (ships with the right modules):
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B")   # just works

# Path B: raw decoder LLM, build the pipeline manually:
from sentence_transformers import SentenceTransformer
from sentence_transformers.sentence_transformer.modules import Normalize, Pooling, Transformer

transformer = Transformer(
    "Qwen/Qwen2.5-0.5B",
    transformer_task="text-generation",         # critical: causal attention, no bidirectional
    processor_kwargs={"padding_side": "left"},  # last-token pooling wants left-padding
)
pooling = Pooling(transformer.get_embedding_dimension(), pooling_mode="lasttoken")
model = SentenceTransformer(modules=[transformer, pooling, Normalize()])
```

Skipping `transformer_task="text-generation"` or `pooling_mode="lasttoken"` on a raw decoder gives embeddings that look plausible until you benchmark.

**Why last-token pooling:** causal attention means only the last token has seen the full sequence. Mean-pooling a causal model averages embeddings that only saw prefixes — the result doesn't represent the whole input.

For training decoder bases:
- Learning rate: typically `1e-4` or higher (not `2e-5` like encoders).
- LoRA is almost always the right choice for >1B-param bases; see `../scripts/train_sentence_transformer_with_lora_example.py` (its docstring covers when to use, hyperparams, QLoRA for 7B+, and adapter sharing).

## Static embeddings

`StaticEmbedding` skips the transformer entirely — each token maps to a pre-computed vector via a lookup table. No attention, no contextualization.

**When to use:**
- CPU inference, no GPU, browser / edge / on-device deployment.
- Need <10MB model size.
- Latency budget <1ms per embedding.
- Have >1M training pairs (contextualization is replaced by per-token optimization; this takes data).

**When NOT to use:**
- Task needs contextual understanding (polysemy, syntax, long-range dependencies).
- You have <100k training pairs — the model won't learn enough.

**Setup:**

```python
from sentence_transformers import SentenceTransformer
from sentence_transformers.sentence_transformer.modules import StaticEmbedding
from tokenizers import Tokenizer

tokenizer = Tokenizer.from_pretrained("google-bert/bert-base-uncased")
static_embedding = StaticEmbedding(tokenizer, embedding_dim=512)
model = SentenceTransformer(modules=[static_embedding])
```

Train with `MultipleNegativesRankingLoss` on a large contrastive dataset (1M+ pairs).

**Warm starts vs. random init** — with **>1M training samples**, random-init beats `StaticEmbedding.from_model2vec(...)` or `.from_distillation(...)` warm starts. With smaller datasets, warm starts help.

```python
# For smaller datasets (<100k), warm-start:
static_embedding = StaticEmbedding.from_model2vec("minishlab/potion-base-8M")
# or:
static_embedding = StaticEmbedding.from_distillation("sentence-transformers/all-MiniLM-L6-v2", vocabulary=list(tokenizer.get_vocab().keys()))
```

See `../scripts/train_sentence_transformer_static_embedding_example.py` for a runnable end-to-end recipe (random init + MNRL + Matryoshka + bf16 + lr=2e-1) and the [Static Embeddings blog post](https://huggingface.co/blog/static-embeddings) for benchmarks.

## Multimodal via VLM backbone

Modern vision-language models can be loaded directly and produce joint text+image embeddings:

```python
from sentence_transformers import SentenceTransformer

model = SentenceTransformer(
    "Qwen/Qwen3-VL-Embedding-2B",
    model_kwargs={"attn_implementation": "flash_attention_2"},  # do NOT set torch_dtype here; see training_args.md
    processor_kwargs={"min_pixels": 28 * 28, "max_pixels": 600 * 600},
)

# Check which modalities this model supports:
print(model.modalities)
# ['text', 'image', 'video', 'message']
```

Training data can mix text, PIL images, image paths/URLs, audio, and mixed-modality dicts like `{"image": <PIL>, "text": "describe this"}`. The data collator handles preprocessing via the model's `preprocess` method.

Install multimodal extras: `pip install "sentence-transformers[image]"` (or `[audio]`, `[video]`).

**Precision**: load in fp32 and pass `bf16=True` (or `fp16=True`) to TrainingArguments — autocast handles the inference path. Don't set `torch_dtype="bfloat16"` in `model_kwargs`: it puts Adam state in bf16 and silently degrades quality (see `training_args.md`).

## Multimodal via Router

Instead of one VLM backbone, compose separate encoders per modality:

```python
from sentence_transformers import SentenceTransformer
from sentence_transformers.sentence_transformer.modules import Dense, Pooling, Router, Transformer

# Text encoder
text_encoder = Transformer("sentence-transformers/all-MiniLM-L6-v2")
text_pooling = Pooling(text_encoder.get_embedding_dimension(), pooling_mode="mean")
# Project text to match image encoder's dim
text_projection = Dense(text_encoder.get_embedding_dimension(), 768)

# Image encoder (SigLIP outputs pooled embeddings directly)
image_encoder = Transformer("google/siglip2-base-patch16-224")

router = Router(
    sub_modules={
        "text": [text_encoder, text_pooling, text_projection],
        "image": [image_encoder],
    },
)
model = SentenceTransformer(modules=[router])
```

**Warning**: Router-based models have unaligned embedding spaces at init — you must train to align them. Use a `Dense` projection layer when dimensions differ. Task-based routing (different encoders for queries vs. documents) is also supported via `route_mappings`; see the `Router` docstring.

## Gotchas

- **Decoder base with mean pooling**: silently produces garbage embeddings. Always use `lasttoken`.
- **Router multimodal without training**: the separate encoders' embedding spaces are unaligned at init. Don't expect useful cross-modal similarity until you've trained with a loss that aligns the spaces.
- **StaticEmbedding with fewer than 100k pairs**: the model won't learn enough. Either warm-start via `from_model2vec` / `from_distillation`, or use a regular encoder.
- **Large VLM backbones on consumer GPUs**: combine LoRA + `attn_implementation="flash_attention_2"`. With LoRA only, you can additionally pass `torch_dtype="bfloat16"` — the bf16 base weights are frozen, so the Adam-state precision concern from the precision rule above doesn't apply (the LoRA adapter stays fp32, so its optimizer state stays fp32). Without LoRA, follow the precision rule: keep weights fp32 and rely on `bf16=True` autocast.

````


### `references/prompts_and_instructions.md`

````markdown
# Prompts and Instructions

Modern embedding models (E5, BGE, GTE, Qwen3-Embedding, Nomic, Instructor, etc.) use **prompts** / **instructions** at encode time — short prefixes like `"query: "`, `"passage: "`, or `"Represent this sentence for retrieval:"` — to signal task intent.

In sentence-transformers, prompts work during both **training** and **inference**, and the library keeps them aligned automatically.

The prompt is literally prepended to the input text before tokenization.

## Model-level prompts

Set prompts on the model so `encode()`, `encode_query()`, `encode_document()` use them automatically:

```python
model = SentenceTransformer("your-base-model")
model.prompts = {
    "query":    "query: ",
    "document": "passage: ",
}
model.default_prompt_name = "document"        # used when none is specified
```

At inference:

```python
q_emb = model.encode_query("What is the capital of France?")
# Internally: encodes "query: What is the capital of France?"

d_emb = model.encode_document(["Paris is the capital of France.", "Berlin is the capital of Germany."])
# Internally: encodes "passage: Paris..." etc.

# Explicit prompt_name:
emb = model.encode(["some text"], prompt_name="query")
```

## Training with prompts

Set prompts on the training args so they're applied automatically to the right columns during training. Four shapes, increasingly specific:

### 1. Single prompt (applied everywhere)

```python
args = SentenceTransformerTrainingArguments(
    ...,
    prompts="Represent this sentence for similarity: ",
)
```

Every input column gets the prefix. Simplest and often sufficient for single-task training.

### 2. Per-column prompts

```python
prompts = {
    "anchor":   "query: ",
    "positive": "passage: ",
    "negative": "passage: ",
}
```

Keys are **column names**. Different prefixes for different roles.

### 3. Per-dataset prompts (multi-dataset)

```python
prompts = {
    "all-nli": "Classify the entailment relationship: ",
    "stsb":    "Score semantic similarity: ",
}
```

Keys are **dataset names** (matching the `train_dataset` dict keys).

### 4. Per-dataset + per-column

```python
prompts = {
    "all-nli": "",                                          # all-nli: no prompt
    "msmarco": {                                            # msmarco: per-column
        "query":    "query: ",
        "positive": "passage: ",
        "negative": "passage: ",
    },
}
```

## Cross-encoder specifics

Cross-encoders restrict prompts to single-value or per-dataset only (no per-column). This is because cross-encoders take text pairs, not separate columns.

```python
args = CrossEncoderTrainingArguments(
    ...,
    prompts="Rank this passage for the query: ",   # single prompt
)
# or
args = CrossEncoderTrainingArguments(
    ...,
    prompts={"msmarco": "...", "gooaq": "..."},    # per-dataset
)
```

## Sparse encoders (SPLADE)

Sparse encoders support prompts with the same four shapes as bi-encoders (added in v5.x). Same `model.prompts = {...}` API.

## Pooling and prompt tokens (bi-encoder)

When using mean/last-token pooling with prompt prefixes, the prompt tokens themselves can either:
- **Be included** in the pooled embedding (default behavior). Simpler, what most models do.
- **Be excluded** — only the actual content tokens are pooled.

To exclude, flip `include_prompt` on the Pooling module — found via `isinstance` rather than a fixed index, since multimodal / Router pipelines don't always put Pooling at position 1:

```python
from sentence_transformers.sentence_transformer.modules import Pooling

for module in model:
    if isinstance(module, Pooling):
        module.include_prompt = False
```

Or use the helper if your model exposes one (e.g. `SparseEncoder.set_pooling_include_prompt(False)`).

**When to exclude**: when the prompt is purely a task signal and you don't want it to dilute the semantic representation. Most E5 / BGE / GTE models leave prompts included.

During training with `prompts=`, the `include_prompt` setting is respected; the trainer also tracks `include_prompt_lengths` internally so pooling skips the prompt tokens correctly when `include_prompt=False`.

## Inference alignment

If you trained with `prompts={"query": "query: ", ...}`, you must:

1. **Save the prompts on the model**: `model.prompts = args.prompts` before `save_pretrained`, OR let the library do this automatically via the model card data.
2. **Use the same prompts at inference**: call `encode_query()` / `encode_document()` and the saved prompts are applied.

If the saved model's `config_sentence_transformers.json` has `prompts` + `default_prompt_name`, anyone who loads it gets the correct prompts for free.

## Instruction tuning (different from prompts)

Some models (Instructor, Qwen3-Embedding) use **instructions** — longer descriptions like `"Represent the biomedical query for retrieving relevant passages:"`. In sentence-transformers these are modeled the same way: just set them as prompts.

The model-card convention is to list available instructions/prompts in a dict:

```python
model.prompts = {
    "msmarco-query":    "Represent the query for MS MARCO retrieval: ",
    "msmarco-doc":      "Represent the passage for MS MARCO retrieval: ",
    "sts":              "Represent the sentence for semantic similarity: ",
    "classification":   "Represent the sentence for topic classification: ",
}
```

Users call `model.encode(["..."], prompt_name="msmarco-query")`.

## Gotchas

- **Training with prompts but forgetting to set them on the model before `save_pretrained`**: the saved model won't know about the prompts. Users will encode *without* the prefix and get bad results. Fix: `model.prompts = args.prompts` before saving, or use `SentenceTransformerModelCardData(...)` with the prompt info.
- **Using `encode_query()` at inference but didn't train with a "query" prompt**: the method will just call regular `encode` (no prefix applied), which is fine. But the documentation implies you use it, and users may be confused.
- **Trailing spaces matter**: `"query: "` vs `"query:"` — the former has a trailing space before the real text. Inspect your training data to know which one your model was trained with.
- **Mixing prompted and non-prompted data in multi-dataset training** is fine as long as your `prompts` dict covers it per-dataset (e.g. `{"dataset_a": "query: ", "dataset_b": ""}`).
- **`include_prompt=False` with small datasets**: the model may under-fit because you've effectively shortened every input. Usually leave as True.

````


### `references/training_args.md`

````markdown
# Training Arguments

`SentenceTransformerTrainingArguments`, `CrossEncoderTrainingArguments`, and `SparseEncoderTrainingArguments` all inherit from Hugging Face's `TrainingArguments`, so 95% of the arguments are the same.

This reference covers the arguments that actually matter for embedding-model training.

## The recommended default set

Start with this and adjust only what you have a reason to change:

```python
from sentence_transformers import SentenceTransformerTrainingArguments
from sentence_transformers.base.sampler import BatchSamplers

args = SentenceTransformerTrainingArguments(
    output_dir="models/my-model",

    # Duration
    num_train_epochs=1,
    # max_steps=10_000,                    # alternative to epochs

    # Batch size
    per_device_train_batch_size=64,
    per_device_eval_batch_size=64,
    gradient_accumulation_steps=1,

    # Optimizer
    learning_rate=2e-5,
    warmup_steps=0.1,                      # transformers v5.2 deprecated `warmup_ratio`; pass the ratio as a float directly to `warmup_steps`
    lr_scheduler_type="linear",
    weight_decay=0.0,

    # Precision
    bf16=True,                             # fp16=True on older GPUs (T4, V100)

    # Sampler (bi-encoder + sparse-encoder)
    batch_sampler=BatchSamplers.NO_DUPLICATES,

    # Eval + checkpointing
    eval_strategy="steps",
    eval_steps=0.1,                        # fraction: 10 evals/epoch, scales with dataset size
    save_strategy="steps",
    save_steps=0.1,                        # keep aligned with eval_steps for load_best_model_at_end
    save_total_limit=2,
    load_best_model_at_end=True,
    metric_for_best_model="eval_NanoBEIR_mean_cosine_ndcg@10",
    greater_is_better=True,

    # Logging
    logging_steps=0.01,                    # fraction: ~100 log lines/epoch
    logging_first_step=True,
    run_name="my-model",
    report_to="trackio",                   # or "wandb", "tensorboard", "mlflow", "none"
)
```

## Duration

- `num_train_epochs` — most common. 1 for large datasets (>500k), 3–10 for small.
- `max_steps` — use instead of epochs when you want a fixed compute budget. Overrides `num_train_epochs`.
- For huge datasets where 1 epoch is overkill, pick `max_steps` matching your compute plan.

## Batch size

Effective batch size = `per_device_train_batch_size × num_gpus × gradient_accumulation_steps`.

Rules of thumb:
- **Contrastive losses (MNRL, GIST, SMNRL)**: push `per_device_train_batch_size` as high as VRAM allows. Larger in-batch negatives = better gradients. 64–256 typical.
- **Regression losses (CoSENTLoss, CosineSimilarityLoss, etc.)**: batch size matters less. 16–64 works.
- **Cross-encoders**: batch size is less critical for quality. 32–128 typical.
- If you can't fit the desired per-device batch, use `gradient_accumulation_steps` to simulate it — but for MNRL-family losses this **does not** provide the same benefit as a real batch (in-batch negatives are still only per-device). Use `CachedMultipleNegativesRankingLoss` instead.

## Learning rate and schedule

- `2e-5` is the safe default for full fine-tuning of BERT-family encoders.
- `1e-4` to `5e-4` for LoRA / PEFT adapters.
- `2e-1` for training a `StaticEmbedding` model from scratch (much higher than transformers because each token is a free-floating vector with no upstream gradients).
- `lr_scheduler_type="linear"` with `warmup_steps=0.1` is standard (a float `< 1` is interpreted as a fraction of total steps). `"cosine"` works equally well; `"constant_with_warmup"` is fine for very short runs. The legacy `warmup_ratio` was deprecated in transformers v5.2 in favor of `warmup_steps` accepting floats; passing `warmup_ratio=...` still works but emits a DeprecationWarning.
- If loss goes NaN, **drop LR first** before anything else.

## Precision

**The non-negotiable rule:** load the model in **fp32** (the default — don't pass `torch_dtype=torch.bfloat16` to the model constructor or `model_kwargs`). Use the `bf16=True` / `fp16=True` flags below to enable **autocast**, not a weight cast. The trainer keeps the model and optimizer state in fp32 and autocasts activations to bf16/fp16 at forward/backward time. This preserves Adam's full-precision moments while giving you most of the bf16 throughput.

Casting weights to bf16 *before* the optimizer is created puts the Adam state (`exp_avg`, `exp_avg_sq`) in bf16 too — bf16's 7-bit mantissa is too coarse for small gradient moments and you get silent quality regressions across runs.

| Flag | When to use |
|---|---|
| `bf16=True` | Ampere (A10G, A100, 3090) and newer (Hopper, Ada). Preferred when supported — more numerically stable than fp16. Activations only; weights stay fp32. |
| `fp16=True` | Older GPUs (T4, V100, 2080, Titan V). Be prepared to drop LR or enable loss scaling if you see NaN. Activations only; weights stay fp32. |
| Neither | fp32 throughout. Slow; only useful for debugging numerical issues. |

Do not set both `bf16=True` and `fp16=True`.

**Evaluator calls outside the trainer** (typically the pre-training baseline + a final post-training pass) don't get the trainer's autocast. Wrap them manually for the speedup — and note that the wrap is **only strictly required when the model uses `attn_implementation="flash_attention_2"`**, since FA2 kernels need bf16/fp16 inputs to function. Without FA2 the wrap is a throughput optimization, not a correctness requirement:

```python
import torch
from contextlib import nullcontext

def autocast_ctx():
    if not torch.cuda.is_available():
        return nullcontext()
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype)

with autocast_ctx():
    evaluator(model)            # baseline
trainer.train()
with autocast_ctx():
    evaluator(model)            # post-training
```

**FlashAttention 2** wants bf16/fp16 inputs but doesn't require bf16 weights. Pass `model_kwargs={"attn_implementation": "flash_attention_2"}` *without* `torch_dtype`, and let `bf16=True` autocast feed bf16 activations to FA2. Weights stay fp32, optimizer state stays fp32.

## Batch sampler (bi-encoder + sparse-encoder)

`batch_sampler=BatchSamplers.NO_DUPLICATES` is critical for contrastive losses. Without it, the same (anchor, positive) can appear multiple times in one batch, turning legitimate positives into false negatives.

Use `BatchSamplers.NO_DUPLICATES` for MNRL / SMNRL / CachedMNRL / GIST (the default recommendation); `BatchSamplers.GROUP_BY_LABEL` for batch-triplet losses (`BatchAllTripletLoss`, `BatchHardTripletLoss`); `BatchSamplers.NO_DUPLICATES_HASHED` only for very large datasets where per-batch string comparison gets slow.

For multi-dataset training, the analogous `MultiDatasetBatchSamplers` class controls how to draw from each dataset (`ROUND_ROBIN`, `PROPORTIONAL`). Under DDP, each dataset is sharded automatically per process — no extra config needed; set `multi_dataset_batch_sampler=...` once and it works for 1-GPU and N-GPU runs identically.

## Evaluation and checkpointing

```python
eval_strategy="steps",
eval_steps=0.1,                        # evaluate every 10% of training
save_strategy="steps",
save_steps=0.1,                        # save at the same cadence (required for load_best_model_at_end)
save_total_limit=2,
load_best_model_at_end=True,
metric_for_best_model="eval_<EvaluatorName>_<metric>",
greater_is_better=True,
```

**Prefer fractional values over absolute step counts.** `eval_steps=0.1` / `save_steps=0.1` / `logging_steps=0.01` are interpreted as fractions of total training steps (10 evals per epoch, 100 log lines per epoch) and auto-scale when the dataset size or epoch count changes. The HF Trainer converts a `float < 1` to `int(total_steps * fraction)` at init time, so the same config works whether you're training on 10k or 10M rows — no need to recompute absolute step counts each time.

Use an absolute integer (e.g. `eval_steps=500`) only when you have a specific reason: comparing runs at known step counts, or when `max_steps` is set to an unusual value that makes fractions awkward.

Non-negotiable rules:
1. `save_steps` must be a multiple of `eval_steps` (or equal) when `load_best_model_at_end=True`, so the best-eval checkpoint is actually on disk. Matching them is the simplest path (e.g. both at `0.1`).
2. If `eval_strategy="steps"` and you don't pass `eval_dataset`, training hangs. Either provide an eval dataset or set `eval_strategy="no"`.
3. `metric_for_best_model` must match the exact key the evaluator writes. The pattern is usually `f"eval_{evaluator.primary_metric}"`. Common values:
   - `NanoBEIREvaluator` (bi-encoder, cosine): `eval_NanoBEIR_mean_cosine_ndcg@10`
   - `SparseNanoBEIREvaluator` (sparse, dot): `eval_NanoBEIR_mean_dot_ndcg@10`
   - `CrossEncoderNanoBEIREvaluator` (rerank from BM25 top-100): `eval_NanoBEIR_R100_mean_ndcg@10`
   - `EmbeddingSimilarityEvaluator(name="sts-dev")`: `eval_sts-dev_spearman_cosine`

## Early stopping

Add `EarlyStoppingCallback` via `callbacks=[...]`:

```python
from transformers import EarlyStoppingCallback

trainer = SentenceTransformerTrainer(
    ..., callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
)
```

This requires `load_best_model_at_end=True` and `metric_for_best_model=...` to be set.

`early_stopping_patience=3` means "stop if the best metric hasn't improved in 3 consecutive eval rounds." Use `early_stopping_threshold=0.001` to require a minimum improvement.

**When it actually matters:**
- **Cross-encoders**: strongly recommended. CE rerankers typically peak mid-training and then degrade — the best checkpoint is rarely the last. Early stopping saves compute *and* guards against quality regression.
- **Bi-encoders and sparse encoders**: usually plateau rather than regress, so early stopping fires much less often. `load_best_model_at_end=True` alone gives you the right final model; adding the callback is a belt-and-suspenders safety net.

## Resuming training

`trainer.train(resume_from_checkpoint=True)` resumes from the latest checkpoint in `output_dir`. Pass a specific path to resume from a specific step: `resume_from_checkpoint="models/my-model/checkpoint-500"`.

State that persists across resumption: optimizer, scheduler, random seeds, trainer step counter. State that does **not**: dataset iteration order for `IterableDataset` — if you use streaming datasets, you must handle resumption yourself.

## Hub push

`push_to_hub=True` + `hub_model_id="your-username/my-model"` + `hub_strategy="every_save"` is the standard pattern. On HF Jobs, also pass `secrets={"HF_TOKEN": "$HF_TOKEN"}` on the job submission. The four `hub_strategy` values: `"every_save"` (each checkpoint, mandatory for HF Jobs), `"end"` (final only), `"checkpoint"` (latest, overwrite), `"all_checkpoints"` (each as a separate commit).

## Logging

```python
logging_steps=0.01,             # fraction: ~100 log lines per epoch (use an int for a fixed cadence)
logging_first_step=True,        # log before any training; useful sanity check
logging_dir=None,               # defaults to output_dir/runs
report_to="trackio",            # or ["trackio", "tensorboard"] for multiple; "none" disables all
run_name="meaningful-name",     # shown in the tracker UI
```

**Tracker recommendation:**
- **Trackio** (default) for solo / small-team work: zero friction beyond `HF_TOKEN`. Auto-creates a Space at `https://huggingface.co/spaces/<your-username>/trackio` on the first run; subsequent runs append and group by `run_name`.
- **W&B** for larger teams or sweep / report features. `pip install wandb && wandb login` (or set `WANDB_API_KEY`).
- **TensorBoard** for air-gapped environments. No remote dashboard.
- **MLflow** when it's already the org standard.

For trackio sweeps / ablations, use `trackio.init(project="...", name="...", group="v1", config={...})` before training to group related runs side-by-side. Without `trackio.init()`, defaults are derived from `run_name` and HF username.

**Tracker gotchas:**
- `report_to="all"` enables every installed integration (usually more than you want); `"none"` disables everything (the current `transformers` default). Always set explicitly.
- Trackio on HF Jobs without `secrets={"HF_TOKEN": "$HF_TOKEN"}` fails silently. W&B on HF Jobs needs `WANDB_API_KEY` in `secrets`.
- The HF Trainer logs only on rank 0 under DDP; custom logging in your script may need explicit rank checks to avoid duplicate writes.

## Memory-saving arguments

```python
gradient_checkpointing=True,    # trades compute for memory. ~30% slowdown, ~40% less memory.
gradient_checkpointing_kwargs={"use_reentrant": False},
torch_empty_cache_steps=1000,   # periodically clear PyTorch allocator cache
dataloader_num_workers=2,       # parallel data loading; 2-4 is usually enough
dataloader_pin_memory=True,
```

Do **not** combine `gradient_checkpointing=True` with any `Cached*` loss — they conflict.

## Hyperparameter search

`trainer.hyperparameter_search(...)` is supported for all three trainers via Hugging Face's `Trainer` API (uses Optuna, Ray Tune, Sigopt, or W&B as backends).

Minimal example:

```python
def model_init(trial):
    return SentenceTransformer("microsoft/mpnet-base")

def hp_space(trial):
    return {
        "learning_rate": trial.suggest_float("learning_rate", 1e-6, 1e-4, log=True),
        "num_train_epochs": trial.suggest_int("num_train_epochs", 1, 3),
        "per_device_train_batch_size": trial.suggest_categorical(
            "per_device_train_batch_size", [32, 64, 128]
        ),
    }

trainer = SentenceTransformerTrainer(
    model=None, model_init=model_init,
    args=args, train_dataset=train_dataset, eval_dataset=eval_dataset,
    loss=lambda model: MultipleNegativesRankingLoss(model),   # function that takes model -> loss
    evaluator=evaluator,
)

best_run = trainer.hyperparameter_search(
    hp_space=hp_space,
    direction="maximize",
    n_trials=10,
    backend="optuna",
)
print(best_run)
```

Install a backend: `pip install optuna` (or `ray[tune]`).

HPO is expensive. Don't reach for it until a single manually-tuned run is working end-to-end. For most production models, picking a reasonable LR from the range above and tuning batch size is enough.

## Multi-task training args (brief)

When training on a dict of datasets with a dict of losses, add:

```python
multi_dataset_batch_sampler=MultiDatasetBatchSamplers.PROPORTIONAL,  # or ROUND_ROBIN
```

See `../scripts/train_sentence_transformer_multi_dataset_example.py` (docstring covers per-dataset losses, single-loss + DatasetDict variant, samplers, gotchas).

## Don't

- Don't set `eval_strategy="epoch"` without setting `save_strategy="epoch"` — checkpoint/eval alignment matters for `load_best_model_at_end`.
- Don't set `remove_unused_columns=False` unless you have a custom collator that consumes metadata columns the loss doesn't see. The default (True) is safer — it drops unused columns automatically.
- Don't set `seed` to verify reproducibility and then expect bit-for-bit identical runs on different GPUs or across PyTorch versions — exact reproducibility across hardware is not guaranteed.
- Don't tune `adam_beta1` / `adam_beta2` / `adam_epsilon` unless you have a specific reason. The defaults are fine for 99% of cases.

````


### `references/troubleshooting.md`

````markdown
# Troubleshooting

Common failures in sentence-transformers training, with root causes and fixes. Organized by symptom.

## Out of memory (OOM)

**Symptom:** `torch.cuda.OutOfMemoryError` or `CUDA out of memory`.

**Fixes in order:**

1. Reduce `per_device_train_batch_size`. For MNRL, compensate via `CachedMultipleNegativesRankingLoss(model, mini_batch_size=...)` with a large outer batch.
2. Set `gradient_accumulation_steps` to accumulate gradients over multiple smaller batches (for non-contrastive losses).
3. Enable `gradient_checkpointing=True`. ~30% slower, ~40% less activation memory. **Incompatible with `Cached*` losses.**
4. Use `bf16=True` instead of fp32 (if your GPU supports Ampere+).
5. Reduce `max_seq_length` on the transformer module: `model[0].max_seq_length = 128`.
6. Use LoRA (`peft`) for large decoder bases. See `../scripts/train_sentence_transformer_with_lora_example.py` (docstring covers when to use, hyperparams, QLoRA, gotchas).
7. Move to a bigger GPU or multi-GPU (`accelerate launch`).

**See also**: `hardware_guide.md` for VRAM estimates by batch size.

## Loss is NaN or Inf

**Symptom:** training loss printed as `nan` or `inf`, or suddenly jumps to a huge value.

**Fixes:**

1. **Drop learning rate** first. Try `5e-6` or `1e-6`.
2. Enable `warmup_steps=0.1` (a float `< 1` is interpreted as a fraction of total steps), or set an absolute `warmup_steps=500`.
3. Switch fp16 -> bf16 (more numerically stable). If stuck on fp16 (older GPU), try disabling fp16 for a sanity run.
4. Check the dataset for broken rows: empty strings, NaN label values, mismatched column counts. Inspect a few rows with `print(dataset[:5])`.
5. For custom losses or unusual base models (very long context), consider adding `max_grad_norm=1.0` to the training args.
6. Check tokenization: some tokenizers emit no tokens for certain unicode / whitespace-only inputs, which can cause downstream NaN in mean pooling.

## Metrics don't improve / are at baseline

**Symptom:** eval metrics after training are the same as before.

**Fixes:**

1. **Re-run dataset inspector** with `--loss <your-loss>`. Most likely cause: column order wrong, label column not detected, or loss expects a different shape.
2. For retrieval: **your negatives are too easy**. Mine hard negatives (`scripts/mine_hard_negatives.py`).
3. Check `metric_for_best_model` — if the key doesn't match what the evaluator writes, the trainer silently uses the final checkpoint (not the best).
4. Confirm `BatchSamplers.NO_DUPLICATES` is set for MNRL-style losses. Without it, the in-batch negative signal is corrupted.
5. Base model is wrong for the task (e.g. decoder-only LLM for short-text STS). Try a different base.
6. Learning rate too low. Default `2e-5` works for encoders; LoRA wants `1e-4+`; static embeddings want about `2e-1` (much higher than transformers).
7. Dataset too small for the chosen loss. Contrastive losses need >10k pairs to be meaningful.

## Training hangs at the first eval

**Symptom:** training starts, then hangs indefinitely at the first evaluation step.

**Fix:** you set `eval_strategy="steps"` or `"epoch"` but either didn't pass `eval_dataset` or passed an empty one. Either provide an `eval_dataset` or set `eval_strategy="no"`.

## Training hangs at startup (multi-GPU)

**Symptom:** `accelerate launch` runs, prints "Found X GPUs" + model-loading messages, then hangs.

**Fixes:**

1. If using a custom dataset class, ensure `__len__` is implemented and returns a consistent length across processes.
2. If using `batch_sampler=BatchSamplers.NO_DUPLICATES` with a very small dataset relative to the world size, batches may not form. Use a larger dataset or smaller per-device batch.
3. Check for mismatched PyTorch / CUDA versions across nodes. `nvidia-smi` + `python -c "import torch; print(torch.version.cuda)"` on each node.
4. NCCL timeout. Set `NCCL_TIMEOUT=300` (seconds) env var.

## `CachedMultipleNegativesRankingLoss` crashes

**Symptom:** error like "element 0 of tensors does not require grad", or a cryptic autograd error.

**Fix:** you have `gradient_checkpointing=True`. The cached losses do their own forward/backward orchestration; gradient checkpointing conflicts with it. Disable `gradient_checkpointing`.

Same applies to `CachedSpladeLoss`, `CachedGISTEmbedLoss`, and any other `Cached*` loss.

## Hub push fails

**Symptom:** `HTTPError: 401` or `403` during `push_to_hub`.

**Fixes:**

1. Run `hf auth whoami`. If it fails, run `hf auth login`.
2. Token needs **write** permission. Regenerate from https://huggingface.co/settings/tokens.
3. Repo either must exist and you have write access, or `hub_private_repo=True`/`False` is set so the library can create it.
4. On HF Jobs: pass `secrets={"HF_TOKEN": "$HF_TOKEN"}` on the job submission.

## Tracker (Trackio / W&B / TensorBoard) not logging

**Symptom:** training runs but no metrics appear in the tracker UI.

**Fixes:**

- Trackio: confirm `pip install trackio` succeeded. No login step — trackio uses your `HF_TOKEN` (set by `hf auth login` or the `HF_TOKEN` env var). On HF Jobs, `HF_TOKEN` must be in `secrets`.
- W&B: confirm `wandb login` succeeded, or `WANDB_API_KEY` env var is set. On HF Jobs, `WANDB_API_KEY` must be in `secrets`.
- `report_to` not set to the right tracker: `report_to="trackio"` (or `"wandb"`, or a list like `["trackio", "tensorboard"]`).
- TensorBoard: check `logging_dir` (defaults to `output_dir/runs/<timestamp>`); point TB at the parent.

## Base model loads but `encode` produces garbage

**Symptom:** `model.encode(["test"])` returns constant vectors, or all-zeros, or NaN.

**Fixes:**

1. You loaded a classification fine-tune as a base (e.g. a BERT fine-tuned on SQuAD). The CLS head is a QA head, not a usable pooling layer. Use the underlying pretrained encoder (`bert-base-uncased`), not the task-specific checkpoint.
2. For decoder models: you're using mean pooling on causal attention. Switch to last-token pooling.
3. For SPLADE: the model's MLM head isn't initialized properly. Make sure the base has `AutoModelForMaskedLM` compatibility.

## Dataset loads but eval is trivially correct

**Symptom:** eval metric is 1.0 (perfect) from the first evaluation step.

**Fix:** your eval set overlaps the train set. Check `dataset.train_test_split(test_size=...)` was called correctly, or that the Hub dataset's `train` vs. `dev` splits are actually disjoint.

## Model-card generation fails

**Symptom:** warning about model-card generation, or `README.md` is missing after `save_pretrained`.

**Fixes:**

- `codecarbon` may be attempting to write emissions and failing. Set `CODECARBON_LOG_LEVEL=error` or uninstall codecarbon.
- Some training state couldn't be serialized (custom objects with unusual types). Pass `model_card_data=SentenceTransformerModelCardData(...)` with explicit fields to bypass inference.

## `num_labels` mismatch on CrossEncoder

**Symptom:** `BinaryCrossEntropyLoss` raises about mismatched dims, or `CrossEntropyLoss` complains.

**Fix:** `num_labels=1` goes with `BinaryCrossEntropyLoss`; `num_labels>=2` goes with `CrossEntropyLoss`. Set `CrossEncoder("...", num_labels=1)` for BCE.

## CrossEncoder eval nDCG crashes after distillation / listwise / pairwise training

**Symptom:** training loss looks healthy, baseline eval looked fine, but post-training eval nDCG drops a lot (e.g. 0.59 → 0.14). Every checkpoint after the first eval is below baseline.

**Root cause:** the default `Sigmoid` activation function on `CrossEncoder(num_labels=1, ...)` saturates raw logits >5 to ~1.0. Distillation/listwise/pairwise losses (`MSELoss`, `MarginMSELoss`, `LambdaLoss`, `RankNetLoss`, `ListNetLoss`, `ListMLELoss`, `PListMLELoss`, `ADRMSELoss`) operate on raw logits — once the model learns to push positives past the saturation point, ranking information is lost.

**Fix:** construct the model with an `Identity` activation:

```python
import torch.nn as nn
model = CrossEncoder("...", num_labels=1, activation_fn=nn.Identity())
```

Keep the default `Sigmoid` only for `BinaryCrossEntropyLoss` (which uses BCE-with-logits internally and wants sigmoidable inputs).

## LambdaLoss training loss is tiny (e.g. 1e-4): is training broken?

**Symptom:** with `LambdaLoss(model, weighting_scheme=NDCGLoss2PPScheme())` and large `K` (>=64), training loss prints in the 1e-3 to 1e-5 range.

**Root cause:** `NDCGLoss2PPScheme` normalizes by the discount-weighted pair count, which scales roughly with K. The numeric magnitude of the loss isn't the signal you should track.

**Fix:** ignore training loss for LambdaLoss; watch the eval metric (`eval_NanoBEIR_R100_mean_ndcg@10` or your `metric_for_best_model`) instead. If eval is moving in the right direction and loss is "too small", that's expected.

## LambdaLoss OOM: what to drop first

**Symptom:** `CUDA out of memory` during a `LambdaLoss` training step, especially at K>=64.

**Recovery order:**

1. **Lower `mini_batch_size`** on the loss first. The internal forward chunking preserves the K-list semantic — this is the cheapest knob and doesn't change the experiment.
2. Lower `per_device_train_batch_size` and compensate with `gradient_accumulation_steps` to keep total batch fixed.
3. **Reduce K (the candidate-list length per query) only as a last resort.** Lowering K changes what the loss is computing; this is an experimental change, not a memory tweak.

For very large K (>=128), `NDCGLoss2PPScheme` materializes O(K²) weight buffers per query that aren't covered by the forward chunking, so even small `mini_batch_size` may not be enough — at that point K is the right knob.

## Loading a model with a custom inline `nn.Module` raises `ImportError`

**Symptom:** training and saving via `train.py` works fine; loading the saved model from any other script (`predict.py`, a notebook, or another package) fails with:

```
ImportError: Module "__main__" does not define a "ClassifierHead" attribute
```

**Root cause:** when a custom `nn.Module` is defined inline in a script, Python records its qualname as `__main__.ClassifierHead`. ST writes that qualname into `modules.json` at save time. Loading from a different entry point can't resolve `__main__.ClassifierHead` — the class doesn't exist in the loader's `__main__`.

**Fixes (any one):**

1. Move the class to an importable module: `from my_pkg.heads import ClassifierHead`. Save again so `modules.json` records `my_pkg.heads.ClassifierHead`.
2. Build the same shape using stock ST modules (`Dense + LayerNorm + Dense`) instead of a custom class — those are always loadable.
3. Document that the model is only loadable from the same script (acceptable for one-off experiments, not for shipping models).

## SPLADE embeddings are dense (not sparse)

**Symptom:** after training, `(embedding != 0).sum(dim=-1)` is in the thousands, not ~30-250.

**Fixes:**

1. You're missing the `SpladeLoss` wrapper. `SparseMultipleNegativesRankingLoss` alone doesn't add FLOPS regularization. Wrap: `SpladeLoss(model, loss=inner_loss, query_regularizer_weight=5e-5, document_regularizer_weight=3e-5)`.
2. Regularizer weights too low. Increase to 1e-4 or higher.
3. Scheduler wiped out early. `SpladeRegularizerWeightSchedulerCallback` ramps weights from 0 to target over the first ~33% of training (`warmup_ratio=1/3` by default; configurable). On very short runs this never reaches target. Either train longer or set `query_regularizer_weight` / `document_regularizer_weight` higher to compensate.

## `ValueError: The dataset has ... columns but the loss expects N`

**Fix:** column count mismatch. Drop extra columns (`dataset.remove_columns([...])`) or reorder with `select_columns`. Names don't matter; count and order do.

## Encoding on CPU is painfully slow

**Symptom:** a few dozen sentences per second instead of 1000s.

**Fixes:**

- Ensure the model is on GPU: `model.to("cuda")` (or `device_map="auto"` when loading).
- For genuine CPU inference (no GPU available), consider switching to a `StaticEmbedding`-based model — ~1000x faster on CPU than a transformer.

## Hub model loads, but `model.encode(prompt_name="query")` acts like no prompt was applied

**Fix:** the saved model doesn't have `prompts` in `config_sentence_transformers.json`. When training, set `model.prompts = args.prompts` before `save_pretrained`, or use `SentenceTransformerModelCardData(prompts=...)`.

## `accelerate launch` runs only on one GPU

**Fix:** run `accelerate config` first, set number of GPUs and precision. Or pass explicitly: `accelerate launch --multi_gpu --num_processes=4 train.py`.

## Cached loss + PEFT adapter fails to backprop

**Symptom:** "None of the inputs have requires_grad=True" when using Cached* loss with a LoRA adapter.

**Fix:** after `add_adapter`, call:

```python
model.transformers_model.enable_input_require_grads()
```

This ensures gradient flow through the frozen base + trainable adapter.

## Related reference docs

- `training_args.md` (shared) — the args that affect all of the above.
- `hardware_guide.md` (shared) — VRAM sizing and multi-GPU.
- `dataset_formats.md` (shared) — column/loss validation.
- `losses_sentence_transformer.md` / `losses_cross_encoder.md` / `losses_sparse_encoder.md` (per-model-type catalogs) — loss-specific quirks.

````


### `scripts/mine_hard_negatives.py`

```
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers[train]>=5.0",
#     "datasets>=2.19.0",
# ]
# ///
"""Mine hard negatives as a pre-training step for contrastive losses.

Hard negatives are the single highest-leverage lever for retrieval quality.
This script is a thin, CLI-friendly wrapper around
`sentence_transformers.util.mine_hard_negatives`.

Typical workflow:
1. Start from a dataset of (anchor, positive) pairs.
2. Pick a retriever model (can be your current base model or a stronger one).
3. Run this script to produce a new dataset with N mined negatives per anchor.
4. Train with MultipleNegativesRankingLoss or CachedMultipleNegativesRankingLoss.

Usage:
    python mine_hard_negatives.py \\
        --dataset sentence-transformers/gooaq \\
        --model sentence-transformers/all-MiniLM-L6-v2 \\
        --num-negatives 5 \\
        --output-path data/gooaq-hard-negatives

    # Mine from a separate document corpus (recommended for production):
    python mine_hard_negatives.py \\
        --dataset sentence-transformers/gooaq \\
        --model sentence-transformers/all-MiniLM-L6-v2 \\
        --corpus-dataset sentence-transformers/wikipedia-en-passages \\
        --corpus-column text \\
        --num-negatives 5 \\
        --output-path data/gooaq-hn-wiki

    # With a cross-encoder as an "oracle" to filter negatives by score:
    python mine_hard_negatives.py \\
        --dataset sentence-transformers/gooaq \\
        --model sentence-transformers/all-MiniLM-L6-v2 \\
        --cross-encoder cross-encoder/ms-marco-MiniLM-L-6-v2 \\
        --num-negatives 5 \\
        --max-score 0.9 \\
        --relative-margin 0.05 \\
        --output-path data/gooaq-hn-filtered

    # Push the mined dataset to the Hub:
    python mine_hard_negatives.py \\
        --dataset sentence-transformers/gooaq --model ... --num-negatives 5 \\
        --push-to-hub your-username/gooaq-hard-negatives

Key options:
    --num-negatives     How many hard negatives to mine per anchor (default 3).
    --range-min/max     Which retrieval-rank window to sample from (default 0..100).
    --sampling-strategy "top" (rank-1 hardest) or "random" (within the window).
    --relative-margin   Require that negative_score < positive_score * (1 - margin).
    --max-score         Filter candidates above this score (likely false negatives).
    --cross-encoder     Use a cross-encoder to re-score candidates before filtering.
    --corpus-dataset    Mine from a separate document pool instead of the input
                        dataset's positives. Recommended for production: typical
                        retrieval corpora (Wikipedia, MSMARCO passages) are far
                        larger than your training-pair pool, giving harder negatives.

See the `mine_hard_negatives` API reference for full semantics and all flags.
"""

from __future__ import annotations

import argparse
import logging
import sys

from datasets import load_dataset

from sentence_transformers import CrossEncoder, SentenceTransformer
from sentence_transformers.util import mine_hard_negatives

logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)
for _noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
    logging.getLogger(_noisy).setLevel(logging.WARNING)


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--dataset", required=True)
    p.add_argument("--subset", default=None)
    p.add_argument("--split", default="train")
    p.add_argument("--model", required=True, help="Retriever / bi-encoder used to score candidates")
    p.add_argument("--cross-encoder", default=None, help="Optional CrossEncoder to re-score and filter")
    p.add_argument("--anchor-column", default=None)
    p.add_argument("--positive-column", default=None)
    p.add_argument(
        "--num-negatives",
        type=int,
        default=3,
        help="Number of hard negatives to mine per anchor. Default 3 matches the library.",
    )
    p.add_argument("--range-min", type=int, default=0)
    p.add_argument("--range-max", type=int, default=100)
    p.add_argument("--sampling-strategy", choices=["top", "random"], default="top")
    p.add_argument(
        "--max-score", type=float, default=None, help="Drop candidates scoring above this (likely false negatives)"
    )
    p.add_argument("--min-score", type=float, default=None)
    p.add_argument("--absolute-margin", type=float, default=None)
    p.add_argument("--relative-margin", type=float, default=None)
    p.add_argument(
        "--output-format",
        choices=["triplet", "n-tuple", "labeled-pair", "labeled-list"],
        default="triplet",
    )
    p.add_argument("--include-positives", action="store_true")
    p.add_argument("--output-scores", action="store_true")
    p.add_argument("--batch-size", type=int, default=32)
    p.add_argument("--use-faiss", action="store_true")
    p.add_argument(
        "--corpus-dataset",
        default=None,
        help="Optional Hub dataset id or local path for a separate document pool to mine from. "
        "If unset, mines negatives from the input dataset's positives.",
    )
    p.add_argument("--corpus-subset", default=None, help="Subset of --corpus-dataset (optional)")
    p.add_argument("--corpus-split", default="train", help="Split of --corpus-dataset (default 'train')")
    p.add_argument(
        "--corpus-column", default="text", help="Text column to extract from --corpus-dataset (default 'text')"
    )
    p.add_argument("--output-path", default=None, help="Local directory to save the mined dataset to")
    p.add_argument("--push-to-hub", default=None, help="Hub repo id to push the mined dataset to (optional)")
    p.add_argument("--private", action="store_true", help="Push as a private repo")
    return p


def main() -> int:
    args = build_parser().parse_args()

    dataset = (
        load_dataset(args.dataset, args.subset, split=args.split)
        if args.subset
        else load_dataset(args.dataset, split=args.split)
    )
    print(f"Loaded {len(dataset):,} rows from {args.dataset} (split={args.split})")

    model = SentenceTransformer(args.model)
    cross_encoder = CrossEncoder(args.cross_encoder) if args.cross_encoder else None

    corpus = None
    if args.corpus_dataset:
        corpus_ds = (
            load_dataset(args.corpus_dataset, args.corpus_subset, split=args.corpus_split)
            if args.corpus_subset
            else load_dataset(args.corpus_dataset, split=args.corpus_split)
        )
        if args.corpus_column not in corpus_ds.column_names:
            raise SystemExit(
                f"--corpus-column '{args.corpus_column}' not in {args.corpus_dataset} columns: "
                f"{corpus_ds.column_names}"
            )
        corpus = list(corpus_ds[args.corpus_column])
        print(f"Loaded corpus: {len(corpus):,} documents from {args.corpus_dataset}.{args.corpus_column}")

    mined = mine_hard_negatives(
        dataset=dataset,
        model=model,
        corpus=corpus,
        cross_encoder=cross_encoder,
        anchor_column_name=args.anchor_column,
        positive_column_name=args.positive_column,
        num_negatives=args.num_negatives,
        range_min=args.range_min,
        range_max=args.range_max,
        sampling_strategy=args.sampling_strategy,
        max_score=args.max_score,
        min_score=args.min_score,
        absolute_margin=args.absolute_margin,
        relative_margin=args.relative_margin,
        output_format=args.output_format,
        include_positives=args.include_positives,
        output_scores=args.output_scores,
        batch_size=args.batch_size,
        use_faiss=args.use_faiss,
    )
    print(f"Mined dataset: {len(mined):,} rows | columns: {mined.column_names}")

    if args.output_path:
        mined.save_to_disk(args.output_path)
        print(f"Saved to {args.output_path}")

    if args.push_to_hub:
        mined.push_to_hub(args.push_to_hub, private=args.private)
        print(f"Pushed to https://huggingface.co/datasets/{args.push_to_hub}")

    if not args.output_path and not args.push_to_hub:
        print("No --output-path or --push-to-hub provided; nothing persisted.")

    return 0


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

```


### `scripts/train_cross_encoder_distillation_example.py`

```
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers[train]>=5.0",
#     "datasets>=2.19.0",
#     "accelerate>=0.26.0",
#     "trackio",
# ]
# ///
"""CrossEncoder distillation: train a small reranker from a stronger one's scores.

Uses MarginMSELoss on `(query, positive, negative, score_diff)` where
`score_diff = teacher(q, pos) - teacher(q, neg)`. Workhorse of MS MARCO-style
distillation: the student learns the teacher's score gaps without ever calling
the teacher at training time (precomputed score diffs).

For Embedding-MSE / Listwise-KL variants and the broader pattern map, see the
sibling `train_sentence_transformer_distillation_example.py` docstring.

CRITICAL: `activation_fn=nn.Identity()` is mandatory. The default `Sigmoid` (with
`num_labels=1`) saturates raw logits >5 to ~1.0 inside `predict()` at eval time,
silently collapsing eval ranking (training loss stays healthy while nDCG drops
from e.g. ~0.59 to ~0.14). See `../references/troubleshooting.md` ("CrossEncoder
eval nDCG crashes after distillation / listwise / pairwise training").

Data: this script uses `sentence-transformers/msmarco` (`bert-ensemble-margin-mse`
subset), which has precomputed teacher score diffs per (q, pos, neg) row. To
distill from your own teacher, replace the dataset loading with a one-time
teacher pass over your (q, pos, neg) triples and store `score_diff = teacher_pos
- teacher_neg` as the label.

Run locally:
    pip install "sentence-transformers[train]>=5.0"
    python train_cross_encoder_distillation_example.py

Multi-GPU:
    accelerate launch train_cross_encoder_distillation_example.py

Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
"""

from __future__ import annotations

import argparse
import logging
import os
from contextlib import nullcontext

import torch
import torch.nn as nn
from datasets import load_dataset, load_from_disk
from transformers import EarlyStoppingCallback

from sentence_transformers import (
    CrossEncoder,
    CrossEncoderModelCardData,
    CrossEncoderTrainer,
    CrossEncoderTrainingArguments,
)
from sentence_transformers.cross_encoder.evaluation import CrossEncoderNanoBEIREvaluator
from sentence_transformers.cross_encoder.losses import MarginMSELoss


def autocast_ctx():
    """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
    if not torch.cuda.is_available():
        return nullcontext()
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype)


def log_trackio_dashboard():
    """Surface the Trackio dashboard URL so the user can watch training live."""
    try:
        from huggingface_hub import whoami

        hf_user = whoami().get("name")
        if hf_user:
            logging.info(
                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
            )
    except Exception:
        pass


MODEL_NAME = "microsoft/MiniLM-L12-H384-uncased"
DATASET_NAME = "sentence-transformers/msmarco"
DATASET_SUBSET = "bert-ensemble-margin-mse"
TRAIN_SIZE = 100_000
EVAL_SIZE = 5_000
OUTPUT_DIR = "models/minilm-msmarco-distilled"
RUN_NAME = "minilm-msmarco-distilled"
DATA_CACHE = f"data/{RUN_NAME}-resolved"
SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"


def setup_logging():
    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
    os.makedirs("logs", exist_ok=True)
    logging.basicConfig(
        format="%(asctime)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.INFO,
        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
        force=True,
    )
    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision("high")


def load_resolved_dataset():
    """Load (query, positive, negative, score) rows. The MSMARCO subset is keyed by
    passage_id / query_id; resolve to text once and cache to disk so reruns skip the work."""
    if os.path.isdir(DATA_CACHE):
        logging.info(f"Loading cached resolved dataset from {DATA_CACHE}")
        return load_from_disk(DATA_CACHE)

    logging.info(f"Resolving {DATASET_NAME}/{DATASET_SUBSET} ids -> text (one-time, cached)")
    corpus_ds = load_dataset(DATASET_NAME, "corpus", split="train")
    corpus = dict(zip(corpus_ds["passage_id"], corpus_ds["passage"]))
    queries_ds = load_dataset(DATASET_NAME, "queries", split="train")
    queries = dict(zip(queries_ds["query_id"], queries_ds["query"]))
    raw = load_dataset(DATASET_NAME, DATASET_SUBSET, split="train").select(range(TRAIN_SIZE + EVAL_SIZE))

    def id_to_text(batch):
        return {
            "query": [queries[qid] for qid in batch["query_id"]],
            "positive": [corpus[pid] for pid in batch["positive_id"]],
            "negative": [corpus[pid] for pid in batch["negative_id"]],
            "score": batch["score"],
        }

    resolved = raw.map(id_to_text, batched=True, remove_columns=["query_id", "positive_id", "negative_id"])
    resolved.save_to_disk(DATA_CACHE)
    return resolved


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
    )
    cli, _ = parser.parse_known_args()

    setup_logging()

    if cli.eval_only:
        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
        model = CrossEncoder(cli.eval_only)
        evaluator = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
        with autocast_ctx():
            evaluator(model)
        return

    logging.info(f"Loading base model: {MODEL_NAME}")
    model = CrossEncoder(
        MODEL_NAME,
        num_labels=1,
        activation_fn=nn.Identity(),  # Mandatory for distillation losses.
        model_card_data=CrossEncoderModelCardData(
            language="en",
            license="apache-2.0",
            model_name=f"{MODEL_NAME.split('/')[-1]} reranker distilled from MS MARCO ensemble",
        ),
    )

    resolved = load_resolved_dataset()
    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
        resolved = resolved.select(range(min(70, len(resolved))))
    eval_size = 20 if SMOKE_TEST else EVAL_SIZE
    split = resolved.train_test_split(test_size=eval_size, seed=12)
    train_dataset = split["train"]
    eval_dataset = split["test"]
    logging.info(f"  train: {len(train_dataset):,} rows | eval: {len(eval_dataset):,} rows")
    logging.info(f"  columns: {train_dataset.column_names}")

    loss = MarginMSELoss(model)

    evaluator = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
    logging.info("Baseline evaluation:")
    with autocast_ctx():
        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
        baseline_eval = evaluator(model)[evaluator.primary_metric]
    metric_key = f"eval_{evaluator.primary_metric}"

    args = CrossEncoderTrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=1,
        max_steps=1 if SMOKE_TEST else -1,
        per_device_train_batch_size=32,
        per_device_eval_batch_size=32,
        learning_rate=8e-6,  # Lower than typical 2e-5; distillation regression converges faster
        weight_decay=0.01,
        warmup_steps=0.1,
        lr_scheduler_type="linear",
        bf16=True,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.1,
        save_total_limit=2,
        logging_steps=0.01,
        logging_first_step=True,
        load_best_model_at_end=True,
        metric_for_best_model=metric_key,
        greater_is_better=True,
        report_to="none" if SMOKE_TEST else "trackio",
        run_name=RUN_NAME,
        seed=12,
    )

    trainer = CrossEncoderTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        loss=loss,
        evaluator=evaluator,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],  # CE rerankers peak mid-training
    )
    if not SMOKE_TEST:
        log_trackio_dashboard()
    trainer.train()

    logging.info("Post-training evaluation:")
    with autocast_ctx():
        score = evaluator(model)[evaluator.primary_metric]
    delta = score - baseline_eval
    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
    logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")

    final_dir = f"{OUTPUT_DIR}/final"
    model.save_pretrained(final_dir)
    logging.info(f"Saved final model to {final_dir}")

    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: skipping Hub push")
        return

    try:
        commit_url = model.push_to_hub(RUN_NAME)
        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
    except Exception:
        import traceback

        logging.error(f"Hub push failed:\n{traceback.format_exc()}")


if __name__ == "__main__":
    main()

```


### `scripts/train_cross_encoder_example.py`

```
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers[train]>=5.0",
#     "datasets>=2.19.0",
#     "accelerate>=0.26.0",
#     "trackio",
# ]
# ///
"""Production-ready cross-encoder (reranker) training template.

Demonstrates:
- BinaryCrossEntropyLoss on (query, passage, label) pointwise data
- Hard-negative mining (`mine_hard_negatives` with `output_format="labeled-pair"`)
  to produce the labeled training data BCE needs, starting from (question, answer)
  pairs
- `pos_weight=num_negatives` to offset the positive/negative imbalance
- CrossEncoderNanoBEIREvaluator for retrieval reranking metrics
- load_best_model_at_end on the retrieval metric
- Auto model card + optional Hub push

Run locally:
    pip install "sentence-transformers[train]>=5.0"
    python train_cross_encoder_example.py

Multi-GPU:
    accelerate launch train_cross_encoder_example.py

Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
"""

from __future__ import annotations

import argparse
import logging
import os
from contextlib import nullcontext

import torch
from datasets import load_dataset, load_from_disk
from transformers import EarlyStoppingCallback

from sentence_transformers import (
    CrossEncoder,
    CrossEncoderModelCardData,
    CrossEncoderTrainer,
    CrossEncoderTrainingArguments,
    SentenceTransformer,
)
from sentence_transformers.cross_encoder.evaluation import CrossEncoderNanoBEIREvaluator
from sentence_transformers.cross_encoder.losses import BinaryCrossEntropyLoss
from sentence_transformers.util import mine_hard_negatives


def autocast_ctx():
    """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
    if not torch.cuda.is_available():
        return nullcontext()
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype)


def log_trackio_dashboard():
    """Surface the Trackio dashboard URL so the user can watch training live."""
    try:
        from huggingface_hub import whoami

        hf_user = whoami().get("name")
        if hf_user:
            logging.info(
                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
            )
    except Exception:
        pass


MODEL_NAME = "microsoft/MiniLM-L12-H384-uncased"
DATASET_NAME = "sentence-transformers/gooaq"
RETRIEVER_NAME = "sentence-transformers/static-retrieval-mrl-en-v1"
TRAIN_SIZE = 100_000
EVAL_SIZE = 1_000
NUM_NEGATIVES = 5
OUTPUT_DIR = "models/minilm-gooaq-ce"
RUN_NAME = "minilm-gooaq-ce"
HARD_NEG_CACHE = f"data/{RUN_NAME}-hard-negatives"  # delete this dir to remine
SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"


def setup_logging():
    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
    os.makedirs("logs", exist_ok=True)
    logging.basicConfig(
        format="%(asctime)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.INFO,
        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
        force=True,
    )
    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision("high")  # TF32 on Ampere+, no quality loss


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
    )
    cli, _ = parser.parse_known_args()

    setup_logging()

    if cli.eval_only:
        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
        model = CrossEncoder(cli.eval_only)
        evaluator = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
        with autocast_ctx():
            evaluator(model)
        return

    logging.info(f"Loading base model: {MODEL_NAME}")
    model = CrossEncoder(
        MODEL_NAME,
        num_labels=1,
        model_card_data=CrossEncoderModelCardData(
            language="en",
            license="apache-2.0",
            model_name=f"{MODEL_NAME.split('/')[-1]} reranker finetuned on GooAQ",
        ),
    )

    logging.info(f"Loading dataset: {DATASET_NAME}")
    train_size = 50 if SMOKE_TEST else TRAIN_SIZE
    eval_size = 20 if SMOKE_TEST else EVAL_SIZE
    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
    pairs = load_dataset(DATASET_NAME, split="train").select(range(train_size + eval_size))

    if os.path.isdir(HARD_NEG_CACHE):
        logging.info(f"Loading cached mined hard negatives from {HARD_NEG_CACHE}")
        labeled = load_from_disk(HARD_NEG_CACHE)
    else:
        logging.info(f"Mining hard negatives with {RETRIEVER_NAME}")
        retriever = SentenceTransformer(RETRIEVER_NAME)
        labeled = mine_hard_negatives(
            dataset=pairs,
            model=retriever,
            num_negatives=NUM_NEGATIVES,
            range_min=10,  # skip the top-10 (likely to contain true positives)
            range_max=100,
            sampling_strategy="top",
            output_format="labeled-pair",
            use_faiss=True,
        )
        labeled.save_to_disk(HARD_NEG_CACHE)
        logging.info(f"Saved mined dataset to {HARD_NEG_CACHE} (delete to remine)")
        del retriever
        torch.cuda.empty_cache()
    # EVAL_SIZE here counts labeled-pair rows, not distinct queries: each
    # query contributes 1 positive + NUM_NEGATIVES negatives, so e.g. 1000
    # rows is approximately 1000 / (1 + NUM_NEGATIVES) distinct queries.
    split = labeled.train_test_split(test_size=eval_size, seed=12)
    train_dataset = split["train"]
    eval_dataset = split["test"]
    logging.info(f"  train: {len(train_dataset):,} rows | eval: {len(eval_dataset):,} rows")
    logging.info(f"  columns: {train_dataset.column_names}")

    # pos_weight = negatives / positives, derived from the actual label distribution
    # so it stays correct if rows get filtered or the mining ratio drifts.
    n_pos = sum(1 for label in train_dataset["label"] if label > 0.5)
    n_neg = len(train_dataset) - n_pos
    pos_weight_value = n_neg / max(n_pos, 1)
    logging.info(f"  positives: {n_pos:,} | negatives: {n_neg:,} | pos_weight: {pos_weight_value:.2f}")
    loss = BinaryCrossEntropyLoss(model, pos_weight=torch.tensor(pos_weight_value))

    evaluator = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
    logging.info("Baseline evaluation:")
    with autocast_ctx():
        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
        baseline_eval = evaluator(model)[evaluator.primary_metric]
    metric_key = f"eval_{evaluator.primary_metric}"

    args = CrossEncoderTrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=1,
        max_steps=1 if SMOKE_TEST else -1,
        per_device_train_batch_size=64,
        per_device_eval_batch_size=64,
        learning_rate=2e-5,
        weight_decay=0.01,
        warmup_steps=0.1,
        lr_scheduler_type="linear",
        bf16=True,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.1,
        save_total_limit=2,
        logging_steps=0.01,
        logging_first_step=True,
        load_best_model_at_end=True,
        metric_for_best_model=metric_key,
        greater_is_better=True,
        report_to="none" if SMOKE_TEST else "trackio",
        run_name=RUN_NAME,
        seed=12,
    )

    # EarlyStoppingCallback earns its keep for cross-encoders: CE rerankers
    # typically peak mid-training and then degrade, so stopping at the best
    # eval checkpoint is load-bearing (unlike bi-encoders, which tend to
    # plateau rather than regress).
    trainer = CrossEncoderTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        loss=loss,
        evaluator=evaluator,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
    )
    if not SMOKE_TEST:
        log_trackio_dashboard()
    trainer.train()

    logging.info("Post-training evaluation:")
    with autocast_ctx():
        score = evaluator(model)[evaluator.primary_metric]
    delta = score - baseline_eval
    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
    logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")

    final_dir = f"{OUTPUT_DIR}/final"
    model.save_pretrained(final_dir)
    logging.info(f"Saved final model to {final_dir}")

    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: skipping Hub push")
        return

    try:
        commit_url = model.push_to_hub(RUN_NAME)
        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
    except Exception:
        import traceback

        logging.error(f"Hub push failed:\n{traceback.format_exc()}")


if __name__ == "__main__":
    main()

```


### `scripts/train_cross_encoder_listwise_example.py`

```
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers[train]>=5.0",
#     "datasets>=2.19.0",
#     "accelerate>=0.26.0",
#     "trackio",
# ]
# ///
"""CrossEncoder listwise training with LambdaLoss.

LambdaLoss is the state-of-the-art listwise ranking loss — it optimizes a
surrogate of nDCG via weighted pairwise comparisons over a per-query candidate
list. Use this when you have multiple candidates per query with graded
relevance, and you want a stronger ranker than pointwise BCE.

Data shape: `(query, [doc_1, ..., doc_K], [score_1, ..., score_K])` per row.
This script builds it via `mine_hard_negatives(..., output_format="labeled-list")`
starting from `(question, answer)` pairs: each row gets the positive plus K
hard negatives, with binary scores (1 for positive, 0 for negatives).

CRITICAL: `activation_fn=nn.Identity()` is mandatory for LambdaLoss / ListNet /
ListMLE / PListMLE / RankNet / MarginMSE / MSE — anything that's not
`BinaryCrossEntropyLoss` or `CrossEntropyLoss`. The default `Sigmoid` (with
`num_labels=1`) saturates raw logits >5 to ~1.0 inside `predict()`, silently
collapsing eval ranking. See `../references/troubleshooting.md` ("CrossEncoder
eval nDCG crashes after distillation / listwise / pairwise training").

OOM recovery for LambdaLoss: drop `mini_batch_size` first (chunking inside the
loss preserves the K-list semantic), then `per_device_train_batch_size` paired
with `gradient_accumulation_steps`, then reduce K (the per-query candidate-list
length) only as a last resort. Lowering K changes the experiment.

Run locally:
    pip install "sentence-transformers[train]>=5.0"
    python train_cross_encoder_listwise_example.py

Multi-GPU:
    accelerate launch train_cross_encoder_listwise_example.py

Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
"""

from __future__ import annotations

import argparse
import logging
import os
from contextlib import nullcontext

import torch
import torch.nn as nn
from datasets import load_dataset, load_from_disk
from transformers import EarlyStoppingCallback

from sentence_transformers import (
    CrossEncoder,
    CrossEncoderModelCardData,
    CrossEncoderTrainer,
    CrossEncoderTrainingArguments,
    SentenceTransformer,
)
from sentence_transformers.base.evaluation import SequentialEvaluator
from sentence_transformers.cross_encoder.evaluation import (
    CrossEncoderNanoBEIREvaluator,
    CrossEncoderRerankingEvaluator,
)
from sentence_transformers.cross_encoder.losses import LambdaLoss
from sentence_transformers.util import mine_hard_negatives


def autocast_ctx():
    """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
    if not torch.cuda.is_available():
        return nullcontext()
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype)


def log_trackio_dashboard():
    """Surface the Trackio dashboard URL so the user can watch training live."""
    try:
        from huggingface_hub import whoami

        hf_user = whoami().get("name")
        if hf_user:
            logging.info(
                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
            )
    except Exception:
        pass


MODEL_NAME = "answerdotai/ModernBERT-base"
DATASET_NAME = "sentence-transformers/gooaq"
RETRIEVER_NAME = "sentence-transformers/static-retrieval-mrl-en-v1"
TRAIN_SIZE = 100_000
EVAL_SIZE = 1_000
NUM_NEGATIVES = 7  # K-1 negatives + 1 positive per row
EVAL_RERANK_DEPTH = 30  # candidates per query in the in-domain eval set
OUTPUT_DIR = "models/modernbert-gooaq-lambda"
RUN_NAME = "modernbert-gooaq-lambda"
HARD_NEG_CACHE = f"data/{RUN_NAME}-hard-negatives"
HARD_EVAL_CACHE = f"data/{RUN_NAME}-hard-eval"
SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"


def setup_logging():
    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
    os.makedirs("logs", exist_ok=True)
    logging.basicConfig(
        format="%(asctime)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.INFO,
        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
        force=True,
    )
    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision("high")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
    )
    cli, _ = parser.parse_known_args()

    setup_logging()

    if cli.eval_only:
        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
        model = CrossEncoder(cli.eval_only)
        evaluator = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
        with autocast_ctx():
            evaluator(model)
        return

    logging.info(f"Loading base model: {MODEL_NAME}")
    model = CrossEncoder(
        MODEL_NAME,
        num_labels=1,
        activation_fn=nn.Identity(),  # Mandatory for LambdaLoss; Sigmoid would saturate eval logits.
        model_card_data=CrossEncoderModelCardData(
            language="en",
            license="apache-2.0",
            model_name=f"{MODEL_NAME.split('/')[-1]} reranker trained with LambdaLoss on GooAQ",
        ),
    )
    # ModernBERT defaults to max_seq_length=8192, which allocates activation memory
    # for 8192-token sequences regardless of input length. Pin to a (q, doc) cap.
    model.max_seq_length = 512

    full_dataset = load_dataset(DATASET_NAME, split="train").select(range(TRAIN_SIZE))
    split = full_dataset.train_test_split(test_size=EVAL_SIZE, seed=12)
    train_pairs, eval_pairs = split["train"], split["test"]

    if os.path.isdir(HARD_NEG_CACHE) and os.path.isdir(HARD_EVAL_CACHE):
        logging.info("Loading cached mined hard-negative datasets")
        hard_train = load_from_disk(HARD_NEG_CACHE)
        hard_eval = load_from_disk(HARD_EVAL_CACHE)
    else:
        logging.info(f"Mining hard negatives with {RETRIEVER_NAME}")
        retriever = SentenceTransformer(RETRIEVER_NAME)
        hard_train = mine_hard_negatives(
            train_pairs,
            retriever,
            num_negatives=NUM_NEGATIVES,
            range_min=10,
            range_max=100,
            sampling_strategy="top",
            output_format="labeled-list",  # Listwise: (query, [docs], [scores])
            use_faiss=True,
            batch_size=4096,
        )
        hard_eval = mine_hard_negatives(
            eval_pairs,
            retriever,
            corpus=full_dataset["answer"],
            num_negatives=EVAL_RERANK_DEPTH,
            output_format="n-tuple",
            use_faiss=True,
            batch_size=4096,
        )
        hard_train.save_to_disk(HARD_NEG_CACHE)
        hard_eval.save_to_disk(HARD_EVAL_CACHE)
        del retriever
        torch.cuda.empty_cache()
    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: trimmed mined datasets; will run max_steps=1 and skip Hub push")
        hard_train = hard_train.select(range(min(50, len(hard_train))))
        hard_eval = hard_eval.select(range(min(20, len(hard_eval))))
    logging.info(f"  train: {len(hard_train):,} rows | columns: {hard_train.column_names}")

    loss = LambdaLoss(model=model, mini_batch_size=16)  # mini_batch_size: drop first if OOM

    nano_beir = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
    # Pure reranker quality: positive is in `documents` and `always_rerank_positives=True` (default).
    in_domain = CrossEncoderRerankingEvaluator(
        samples=[
            {
                "query": row["question"],
                "positive": [row["answer"]],
                "documents": [row["answer"]] + [row[col] for col in hard_eval.column_names[2:]],
            }
            for row in hard_eval
        ],
        batch_size=64,
        name="gooaq-dev",
    )
    evaluator = SequentialEvaluator([in_domain, nano_beir])
    logging.info("Baseline evaluation:")
    with autocast_ctx():
        baseline_eval = evaluator(model)[in_domain.primary_metric]

    args = CrossEncoderTrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=1,
        max_steps=1 if SMOKE_TEST else -1,
        per_device_train_batch_size=64,
        per_device_eval_batch_size=64,
        learning_rate=2e-5,
        weight_decay=0.01,
        warmup_steps=0.1,
        lr_scheduler_type="linear",
        bf16=True,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.1,
        save_total_limit=2,
        logging_steps=0.01,
        logging_first_step=True,
        load_best_model_at_end=True,
        metric_for_best_model=f"eval_{in_domain.primary_metric}",  # in-domain reranker > NanoBEIR
        greater_is_better=True,
        report_to="none" if SMOKE_TEST else "trackio",
        run_name=RUN_NAME,
        seed=12,
    )

    trainer = CrossEncoderTrainer(
        model=model,
        args=args,
        train_dataset=hard_train,
        loss=loss,
        evaluator=evaluator,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
    )
    if not SMOKE_TEST:
        log_trackio_dashboard()
    trainer.train()

    logging.info("Post-training evaluation:")
    with autocast_ctx():
        score = evaluator(model)[in_domain.primary_metric]
    delta = score - baseline_eval
    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
    logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")

    final_dir = f"{OUTPUT_DIR}/final"
    model.save_pretrained(final_dir)
    logging.info(f"Saved final model to {final_dir}")

    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: skipping Hub push")
        return

    try:
        commit_url = model.push_to_hub(RUN_NAME)
        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
    except Exception:
        import traceback

        logging.error(f"Hub push failed:\n{traceback.format_exc()}")


if __name__ == "__main__":
    main()

```


### `scripts/train_sentence_transformer_distillation_example.py`

```
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers[train]>=5.0",
#     "datasets>=2.19.0",
#     "accelerate>=0.26.0",
#     "trackio",
# ]
# ///
"""Distillation template: train a small student to match a stronger teacher's embeddings.

This script implements **embedding MSE** (Pattern 1 below): pre-compute teacher
embeddings, train the student to minimize MSE against them. Yields a smaller /
faster student typically reaching 95-99% of teacher quality.

Three distillation patterns:
- Pattern 1 (this script): `(text, teacher_embedding)` + `MSELoss`. Cheapest,
  most data-efficient. Use when student + teacher are both bi-encoders with the
  same output dim and you have a pile of unlabeled text.
- Pattern 2: `(query, positive, negative, score_diff)` + `MarginMSELoss` from a
  CrossEncoder teacher's score differences. Workhorse of ms-marco distillation.
  See `../references/losses_sentence_transformer.md` (MarginMSELoss section).
- Pattern 3: `(query, positive, neg_1, ..., neg_n, labels)` + `DistillKLDivLoss`
  to preserve the full teacher distribution. More data-hungry; natural fit when
  distilling from an ensemble of rerankers.

Mismatched dims: if the student's dim is smaller than the teacher's, MSELoss
fails. Add a PCA-init `Dense` projection so the teacher matches the student:

    from sklearn.decomposition import PCA
    from sentence_transformers.sentence_transformer.modules import Dense
    pca = PCA(n_components=student.get_embedding_dimension())
    pca.fit(teacher.encode(sentences[:20_000], convert_to_numpy=True))
    dense = Dense(
        in_features=teacher.get_embedding_dimension(),
        out_features=student.get_embedding_dimension(),
        bias=False, activation_function=torch.nn.Identity(),
    )
    dense.linear.weight = torch.nn.Parameter(torch.from_numpy(pca.components_).float())
    teacher.add_module("dense", dense)

Distilling to a CrossEncoder student: construct with `activation_fn=nn.Identity()`
or eval ranking collapses silently. Every non-BCE CE loss expects raw logits
during training, but the model's `activation_fn` runs at eval time inside
`predict()`. Default `Sigmoid` (when `num_labels=1`) saturates raw logits >5 to
~1.0, dropping nDCG from e.g. ~0.59 to ~0.14 with healthy-looking training loss.
Applies to all CE distillation / listwise / pairwise losses; see SKILL.md
Directive 7 ([CE]).

Layer pruning shortcut for Pattern 1: copy the teacher, delete layers (often
keeps 99%+ of quality at a fraction of the layers), then distill with MSELoss:

    from copy import deepcopy
    student = deepcopy(teacher)
    layers = student.transformers_model.encoder.layer  # BERT/MPNet/DistilBERT
    student.transformers_model.encoder.layer = torch.nn.ModuleList(
        [layers[0], layers[3], layers[6], layers[9]]
    )
    student.transformers_model.config.num_hidden_layers = 4

Tips: pre-compute teacher outputs once and cache (`dataset.save_to_disk`); LR
1e-4 (higher than the usual 2e-5; the target is dense regression); 1 epoch is
usually enough; the student inherits the teacher's weaknesses, so pick a
teacher strong on YOUR task; if the teacher expects an instruction prefix,
include it during teacher encoding so the student's target matches inference.

For multilingual student distillation (extend an English teacher to other
languages without in-language supervised data), see `train_sentence_transformer_make_multilingual_example.py`.
"""

from __future__ import annotations

import argparse
import logging
import os
from contextlib import nullcontext

import torch
from datasets import Dataset, load_dataset

from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerModelCardData,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)
from sentence_transformers.sentence_transformer.evaluation import EmbeddingSimilarityEvaluator
from sentence_transformers.sentence_transformer.losses import MSELoss
from sentence_transformers.sentence_transformer.modules import Normalize
from sentence_transformers.util.similarity import SimilarityFunction


def autocast_ctx():
    """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
    if not torch.cuda.is_available():
        return nullcontext()
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype)


def log_trackio_dashboard():
    """Surface the Trackio dashboard URL so the user can watch training live."""
    try:
        from huggingface_hub import whoami

        hf_user = whoami().get("name")
        if hf_user:
            logging.info(
                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
            )
    except Exception:
        pass


TEACHER_MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"
STUDENT_MODEL_NAME = "distilbert/distilbert-base-uncased"

CORPUS_DATASET = "sentence-transformers/all-nli"
CORPUS_SUBSET = "pair"

TRAIN_SIZE = 50_000
EVAL_SIZE = 1_000
OUTPUT_DIR = "models/distilbert-distilled-from-mpnet"
RUN_NAME = "distilbert-distill-from-mpnet"

TEACHER_ENCODE_BATCH_SIZE = 256
TRAIN_BATCH_SIZE = 128
SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"


def setup_logging():
    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
    os.makedirs("logs", exist_ok=True)
    logging.basicConfig(
        format="%(asctime)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.INFO,
        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
        force=True,
    )
    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision("high")  # TF32 on Ampere+, no quality loss


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
    )
    cli, _ = parser.parse_known_args()

    setup_logging()

    if cli.eval_only:
        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
        model = SentenceTransformer(cli.eval_only)
        stsb = load_dataset("sentence-transformers/stsb", split="validation")
        evaluator = EmbeddingSimilarityEvaluator(
            sentences1=stsb["sentence1"],
            sentences2=stsb["sentence2"],
            scores=stsb["score"],
            main_similarity=SimilarityFunction.COSINE,
            name="sts-dev",
        )
        with autocast_ctx():
            evaluator(model)
        return

    logging.info(f"Loading teacher: {TEACHER_MODEL_NAME}")
    teacher = SentenceTransformer(TEACHER_MODEL_NAME)

    logging.info(f"Loading student: {STUDENT_MODEL_NAME}")
    student = SentenceTransformer(
        STUDENT_MODEL_NAME,
        model_card_data=SentenceTransformerModelCardData(
            language="en",
            license="apache-2.0",
            model_name=f"{STUDENT_MODEL_NAME.split('/')[-1]} distilled from {TEACHER_MODEL_NAME.split('/')[-1]}",
        ),
    )
    # Match the teacher's final Normalize. MSELoss against unit-norm targets fights student
    # outputs at norm ~5-10 and can silently regress
    if any(isinstance(m, Normalize) for m in teacher) and not any(isinstance(m, Normalize) for m in student):
        student.append(Normalize())

    if student.get_embedding_dimension() != teacher.get_embedding_dimension():
        raise SystemExit(
            f"Student dim ({student.get_embedding_dimension()}) != teacher dim "
            f"({teacher.get_embedding_dimension()}). Plain MSELoss requires matching dims. "
            "Either pick a student with matching dim, or add a Dense projection layer "
            "(PCA-initialized from teacher embeddings). See the 'MISMATCHED EMBEDDING DIMS' "
            "section in this script's docstring."
        )

    logging.info(f"Loading corpus: {CORPUS_DATASET} ({CORPUS_SUBSET})")
    train_size = 50 if SMOKE_TEST else TRAIN_SIZE
    eval_size = 20 if SMOKE_TEST else EVAL_SIZE
    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
    raw = load_dataset(CORPUS_DATASET, CORPUS_SUBSET, split="train")
    sentences = list(dict.fromkeys(s for row in raw for s in (row["anchor"], row["positive"]) if isinstance(s, str)))
    sentences = sentences[: train_size + eval_size]
    train_sentences = sentences[:train_size]
    eval_sentences = sentences[train_size : train_size + eval_size]

    logging.info(f"Encoding {len(train_sentences):,} training sentences with the teacher (may take a while)")
    teacher_train = teacher.encode(
        train_sentences, batch_size=TEACHER_ENCODE_BATCH_SIZE, convert_to_numpy=True, show_progress_bar=True
    )

    logging.info(f"Encoding {len(eval_sentences):,} eval sentences with the teacher")
    teacher_eval = teacher.encode(
        eval_sentences, batch_size=TEACHER_ENCODE_BATCH_SIZE, convert_to_numpy=True, show_progress_bar=True
    )

    train_dataset = Dataset.from_dict({"sentence": train_sentences, "label": teacher_train.tolist()})
    eval_dataset = Dataset.from_dict({"sentence": eval_sentences, "label": teacher_eval.tolist()})

    logging.info(f"Building training dataset ({len(train_dataset):,}) and eval dataset ({len(eval_dataset):,})")

    loss = MSELoss(model=student)

    logging.info("Setting up STS-B evaluator for quality tracking")
    stsb = load_dataset("sentence-transformers/stsb", split="validation")
    evaluator = EmbeddingSimilarityEvaluator(
        sentences1=stsb["sentence1"],
        sentences2=stsb["sentence2"],
        scores=stsb["score"],
        main_similarity=SimilarityFunction.COSINE,
        name="sts-dev",
    )
    logging.info("Teacher performance:")
    evaluator(teacher)
    logging.info("Student performance before distillation:")
    # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
    baseline_eval = evaluator(student)[evaluator.primary_metric]
    metric_key = f"eval_{evaluator.primary_metric}"

    args = SentenceTransformerTrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=1,
        max_steps=1 if SMOKE_TEST else -1,
        per_device_train_batch_size=TRAIN_BATCH_SIZE,
        per_device_eval_batch_size=TRAIN_BATCH_SIZE,
        learning_rate=1e-4,
        weight_decay=0.01,
        warmup_steps=0.1,
        bf16=True,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.1,
        save_total_limit=2,
        logging_steps=0.01,
        logging_first_step=True,
        load_best_model_at_end=True,
        metric_for_best_model=metric_key,
        greater_is_better=True,
        report_to="none" if SMOKE_TEST else "trackio",
        run_name=RUN_NAME,
        seed=12,
    )

    trainer = SentenceTransformerTrainer(
        model=student,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        loss=loss,
        evaluator=evaluator,
    )
    if not SMOKE_TEST:
        log_trackio_dashboard()
    trainer.train()

    logging.info("Student performance after distillation:")
    score = evaluator(student)[evaluator.primary_metric]
    delta = score - baseline_eval
    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
    logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")

    final_dir = f"{OUTPUT_DIR}/final"
    student.save_pretrained(final_dir)
    logging.info(f"Saved to {final_dir}")

    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: skipping Hub push")
        return

    try:
        commit_url = student.push_to_hub(RUN_NAME)
        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
    except Exception:
        import traceback

        logging.error(f"Hub push failed:\n{traceback.format_exc()}")


if __name__ == "__main__":
    main()

```


### `scripts/train_sentence_transformer_example.py`

```
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers[train]>=5.0",
#     "datasets>=2.19.0",
#     "accelerate>=0.26.0",
#     "trackio",
# ]
# ///
"""Production-ready bi-encoder (SentenceTransformer) training template.

This script demonstrates a recommended setup:
- MultipleNegativesRankingLoss on (anchor, positive, negative) triplets
- NanoBEIREvaluator for retrieval metrics during training
- BatchSamplers.NO_DUPLICATES (critical for MNRL)
- load_best_model_at_end with a retrieval metric
- Auto model card + optional Hub push

Runs identically in two modes:

    # Local
    pip install "sentence-transformers[train]>=5.0"
    python train_sentence_transformer_example.py

    # Or with uv (no explicit install needed)
    uv run train_sentence_transformer_example.py

    # Multi-GPU
    accelerate launch train_sentence_transformer_example.py

    # Hugging Face Jobs (paste the entire file contents as `script`)
    hf_jobs("uv", {
        "script": "<contents of this file>",
        "flavor": "a10g-large",
        "timeout": "3h",
        "secrets": {"HF_TOKEN": "$HF_TOKEN"},
    })

Adjust MODEL_NAME, DATASET_NAME, OUTPUT_DIR, RUN_NAME at the top of the script.
Default Hub push: at end of run, public, under your authenticated user as
`{user}/{RUN_NAME}`, wrapped in try/except. To skip the push, comment out the
push_to_hub call. For HF Jobs (ephemeral env), also enable in-trainer push:
add `push_to_hub=True`, `hub_model_id=RUN_NAME`, `hub_strategy="every_save"`
to TrainingArguments.
"""

from __future__ import annotations

import argparse
import logging
import os
from contextlib import nullcontext

import torch
from datasets import load_dataset

from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerModelCardData,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)
from sentence_transformers.base.sampler import BatchSamplers
from sentence_transformers.sentence_transformer.evaluation import NanoBEIREvaluator
from sentence_transformers.sentence_transformer.losses import MultipleNegativesRankingLoss


def autocast_ctx():
    """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
    if not torch.cuda.is_available():
        return nullcontext()
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype)


def log_trackio_dashboard():
    """Surface the Trackio dashboard URL so the user can watch training live."""
    try:
        from huggingface_hub import whoami

        hf_user = whoami().get("name")
        if hf_user:
            logging.info(
                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
            )
    except Exception:
        pass


MODEL_NAME = "microsoft/mpnet-base"
DATASET_NAME = "sentence-transformers/all-nli"
DATASET_SUBSET = "triplet"
TRAIN_SIZE = 50_000
EVAL_SIZE = 1_000
OUTPUT_DIR = "models/mpnet-base-all-nli"
RUN_NAME = "mpnet-base-all-nli"
SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"


def setup_logging():
    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
    os.makedirs("logs", exist_ok=True)
    logging.basicConfig(
        format="%(asctime)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.INFO,
        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
        force=True,
    )
    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision("high")  # TF32 on Ampere+, no quality loss


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
    )
    cli, _ = parser.parse_known_args()

    setup_logging()

    if cli.eval_only:
        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
        model = SentenceTransformer(cli.eval_only)
        evaluator = NanoBEIREvaluator()
        with autocast_ctx():
            evaluator(model)
        return

    logging.info(f"Loading base model: {MODEL_NAME}")
    model = SentenceTransformer(
        MODEL_NAME,
        model_card_data=SentenceTransformerModelCardData(
            language="en",
            license="apache-2.0",
            model_name=f"{MODEL_NAME.split('/')[-1]} finetuned on AllNLI",
        ),
    )

    logging.info(f"Loading dataset: {DATASET_NAME} ({DATASET_SUBSET})")
    train_size = 50 if SMOKE_TEST else TRAIN_SIZE
    eval_size = 20 if SMOKE_TEST else EVAL_SIZE
    train_dataset = load_dataset(DATASET_NAME, DATASET_SUBSET, split="train").select(range(train_size))
    eval_dataset = load_dataset(DATASET_NAME, DATASET_SUBSET, split="dev").select(range(eval_size))
    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
    logging.info(f"  train: {len(train_dataset):,} examples")
    logging.info(f"  eval:  {len(eval_dataset):,} examples")

    loss = MultipleNegativesRankingLoss(model)

    evaluator = NanoBEIREvaluator()
    logging.info("Baseline evaluation (before training):")
    with autocast_ctx():
        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
        baseline_eval = evaluator(model)[evaluator.primary_metric]
    metric_key = f"eval_{evaluator.primary_metric}"

    args = SentenceTransformerTrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=1,
        max_steps=1 if SMOKE_TEST else -1,
        per_device_train_batch_size=64,
        per_device_eval_batch_size=64,
        learning_rate=2e-5,
        weight_decay=0.01,
        warmup_steps=0.1,
        lr_scheduler_type="linear",
        bf16=True,
        batch_sampler=BatchSamplers.NO_DUPLICATES,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.1,
        save_total_limit=2,
        logging_steps=0.01,
        logging_first_step=True,
        load_best_model_at_end=True,
        metric_for_best_model=metric_key,
        greater_is_better=True,
        report_to="none" if SMOKE_TEST else "trackio",
        run_name=RUN_NAME,
        seed=12,
    )

    trainer = SentenceTransformerTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        loss=loss,
        evaluator=evaluator,
    )
    if not SMOKE_TEST:
        log_trackio_dashboard()
    trainer.train()

    logging.info("Post-training evaluation:")
    with autocast_ctx():
        score = evaluator(model)[evaluator.primary_metric]
    delta = score - baseline_eval
    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
    logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")

    final_dir = f"{OUTPUT_DIR}/final"
    model.save_pretrained(final_dir)
    logging.info(f"Saved final model to {final_dir}")

    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: skipping Hub push")
        return

    try:
        commit_url = model.push_to_hub(RUN_NAME)  # public by default; uses your authenticated user
        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
    except Exception:
        import traceback

        logging.error(f"Hub push failed:\n{traceback.format_exc()}")


if __name__ == "__main__":
    main()

```


### `scripts/train_sentence_transformer_make_multilingual_example.py`

```
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers[train]>=5.0",
#     "datasets>=2.19.0",
#     "accelerate>=0.26.0",
#     "trackio",
# ]
# ///
"""Multilingual teacher-student distillation: extend an English bi-encoder to other languages.

The trick (Reimers & Gurevych 2020, https://huggingface.co/papers/2004.09813):
the teacher embeds English; the multilingual student is trained so that BOTH
the English sentence AND its translation map to the SAME teacher embedding.
Cross-lingual retrieval works out of the box because translations sit near
each other in the joint space.

Use when you have a strong English bi-encoder and want a multilingual version
but lack in-language supervised data. If you DO have in-language labeled data,
train directly with MNRL / CoSENTLoss on it; that usually wins on in-language
tasks.

Data: parallel `(english, non_english)` pairs. `sentence-transformers/parallel-sentences-*`
covers many corpora (talks, europarl, tatoeba, wikimatrix, opensubtitles, jw300,
news-commentary, ...) with `{src}-{tgt}` subsets. ~500k pairs per language is
plenty.

Picks:
- Teacher: any English bi-encoder you want a multilingual copy of
  (all-mpnet-base-v2, all-MiniLM-L6-v2, BAAI/bge-base-en-v1.5, intfloat/e5-base-v2).
- Student: must be multilingual (xlm-roberta-base, paraphrase-multilingual-MiniLM-L12-v2,
  microsoft/mdeberta-v3-base).
- Student dim must match teacher dim, otherwise add a PCA-init Dense projection
  (see train_sentence_transformer_distillation_example.py).
"""

from __future__ import annotations

import argparse
import logging
import os
from contextlib import nullcontext

import numpy as np
import torch
from datasets import DatasetDict, load_dataset

from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerModelCardData,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)
from sentence_transformers.sentence_transformer.evaluation import (
    MSEEvaluator,
    SequentialEvaluator,
    TranslationEvaluator,
)
from sentence_transformers.sentence_transformer.losses import MSELoss
from sentence_transformers.sentence_transformer.modules import Normalize


def autocast_ctx():
    """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
    if not torch.cuda.is_available():
        return nullcontext()
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype)


def log_trackio_dashboard():
    """Surface the Trackio dashboard URL so the user can watch training live."""
    try:
        from huggingface_hub import whoami

        hf_user = whoami().get("name")
        if hf_user:
            logging.info(
                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
            )
    except Exception:
        pass


TEACHER_MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"
STUDENT_MODEL_NAME = "FacebookAI/xlm-roberta-base"

PARALLEL_DATASET = "sentence-transformers/parallel-sentences-talks"
SOURCE_LANGUAGE = "en"
TARGET_LANGUAGES = ("de", "es", "fr", "it")

MAX_SENTENCES_PER_LANGUAGE = 200_000
EVAL_SENTENCES_PER_LANGUAGE = 1_000
STUDENT_MAX_SEQ_LENGTH = 128

OUTPUT_DIR = "models/xlm-roberta-multilingual-from-mpnet"
RUN_NAME = "xlm-roberta-multilingual-from-mpnet"

TEACHER_ENCODE_BATCH_SIZE = 256
TRAIN_BATCH_SIZE = 64
SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"


def setup_logging():
    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
    os.makedirs("logs", exist_ok=True)
    logging.basicConfig(
        format="%(asctime)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.INFO,
        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
        force=True,
    )
    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision("high")  # TF32 on Ampere+, no quality loss


def load_parallel_data() -> tuple[DatasetDict, DatasetDict]:
    """Load (en, target_lang) parallel pairs for each target language as a DatasetDict."""
    train_dict = DatasetDict()
    eval_dict = DatasetDict()
    for tgt in TARGET_LANGUAGES:
        subset = f"{SOURCE_LANGUAGE}-{tgt}"
        try:
            train_ds = load_dataset(PARALLEL_DATASET, subset, split="train")
        except Exception as exc:
            logging.error(f"Could not load {PARALLEL_DATASET}/{subset}: {exc}")
            continue
        if len(train_ds) > MAX_SENTENCES_PER_LANGUAGE:
            train_ds = train_ds.select(range(MAX_SENTENCES_PER_LANGUAGE))

        try:
            eval_ds = load_dataset(PARALLEL_DATASET, subset, split="dev").select(range(EVAL_SENTENCES_PER_LANGUAGE))
        except Exception:
            split = train_ds.train_test_split(test_size=EVAL_SENTENCES_PER_LANGUAGE, shuffle=True, seed=12)
            train_ds, eval_ds = split["train"], split["test"]

        train_dict[subset] = train_ds
        eval_dict[subset] = eval_ds
    if not train_dict:
        raise SystemExit(f"No language subsets loaded from {PARALLEL_DATASET}. Check TARGET_LANGUAGES.")
    return train_dict, eval_dict


def build_evaluator(eval_dict: DatasetDict, teacher: SentenceTransformer) -> SequentialEvaluator:
    """Per-language MSE + TranslationEvaluator. `main_score_function` averages
    translation accuracies only; MSE (`negative_mse * 100`) is on a different
    scale and would break the verdict threshold if mixed in."""
    sub_evaluators = []
    for subset, ds in eval_dict.items():
        sub_evaluators.append(
            MSEEvaluator(
                source_sentences=ds["english"],
                target_sentences=ds["non_english"],
                name=subset,
                teacher_model=teacher,
                batch_size=TEACHER_ENCODE_BATCH_SIZE,
            )
        )
        sub_evaluators.append(
            TranslationEvaluator(
                source_sentences=ds["english"],
                target_sentences=ds["non_english"],
                name=subset,
                batch_size=TEACHER_ENCODE_BATCH_SIZE,
            )
        )
    # Sub-evaluators alternate MSE / Translation per language; scores[1::2] are the translation accuracies.
    return SequentialEvaluator(sub_evaluators, main_score_function=lambda scores: float(np.mean(scores[1::2])))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
    )
    cli, _ = parser.parse_known_args()

    setup_logging()

    if cli.eval_only:
        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
        student = SentenceTransformer(cli.eval_only)
        teacher = SentenceTransformer(TEACHER_MODEL_NAME)
        _, eval_dict = load_parallel_data()
        evaluator = build_evaluator(eval_dict, teacher)
        with autocast_ctx():
            evaluator(student)
        return

    logging.info(f"Loading teacher: {TEACHER_MODEL_NAME}")
    teacher = SentenceTransformer(TEACHER_MODEL_NAME)

    logging.info(f"Loading student: {STUDENT_MODEL_NAME}")
    student = SentenceTransformer(
        STUDENT_MODEL_NAME,
        model_card_data=SentenceTransformerModelCardData(
            language=[SOURCE_LANGUAGE, *TARGET_LANGUAGES],
            license="apache-2.0",
            model_name=f"{STUDENT_MODEL_NAME.split('/')[-1]} multilingual from {TEACHER_MODEL_NAME.split('/')[-1]}",
        ),
    )
    student.max_seq_length = STUDENT_MAX_SEQ_LENGTH
    # Match the teacher's final Normalize. MSELoss against unit-norm targets fights student
    # outputs at norm ~5-10 and can silently regress
    if any(isinstance(m, Normalize) for m in teacher) and not any(isinstance(m, Normalize) for m in student):
        student.append(Normalize())

    if student.get_embedding_dimension() != teacher.get_embedding_dimension():
        raise SystemExit(
            f"Student dim ({student.get_embedding_dimension()}) != teacher dim "
            f"({teacher.get_embedding_dimension()}). MSELoss requires matching dims. "
            "Either pick a student with matching dim, or add a Dense projection layer "
            "(see train_sentence_transformer_distillation_example.py 'MISMATCHED EMBEDDING DIMS')."
        )

    logging.info("Loading parallel data")
    train_dict, eval_dict = load_parallel_data()
    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: trimming each language subset; will run max_steps=1 and skip Hub push")
        train_dict = DatasetDict({k: v.select(range(min(50, len(v)))) for k, v in train_dict.items()})
        eval_dict = DatasetDict({k: v.select(range(min(20, len(v)))) for k, v in eval_dict.items()})

    def attach_teacher_label(batch):
        return {
            "english": batch["english"],
            "non_english": batch["non_english"],
            "label": teacher.encode(batch["english"], batch_size=TEACHER_ENCODE_BATCH_SIZE, show_progress_bar=False),
        }

    column_names = list(train_dict.values())[0].column_names
    logging.info("Encoding training English sentences with teacher (cached on disk if you save_to_disk)")
    train_dict = train_dict.map(attach_teacher_label, batched=True, batch_size=10_000, remove_columns=column_names)
    eval_dict = eval_dict.map(attach_teacher_label, batched=True, batch_size=10_000, remove_columns=column_names)

    loss = MSELoss(model=student)

    evaluator = build_evaluator(eval_dict, teacher)
    logging.info("Student baseline (before training):")
    with autocast_ctx():
        baseline_eval = evaluator(student)["sequential_score"]

    args = SentenceTransformerTrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=3,
        max_steps=1 if SMOKE_TEST else -1,
        per_device_train_batch_size=TRAIN_BATCH_SIZE,
        per_device_eval_batch_size=TRAIN_BATCH_SIZE,
        learning_rate=2e-5,
        weight_decay=0.01,
        warmup_steps=0.1,
        bf16=True,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.1,
        save_total_limit=2,
        logging_steps=0.01,
        logging_first_step=True,
        load_best_model_at_end=True,
        metric_for_best_model="eval_sequential_score",
        greater_is_better=True,
        report_to="none" if SMOKE_TEST else "trackio",
        run_name=RUN_NAME,
        seed=12,
    )

    trainer = SentenceTransformerTrainer(
        model=student,
        args=args,
        train_dataset=train_dict,
        eval_dataset=eval_dict,
        loss=loss,
        evaluator=evaluator,
    )
    if not SMOKE_TEST:
        log_trackio_dashboard()
    trainer.train()

    logging.info("Final student evaluation:")
    with autocast_ctx():
        score = evaluator(student)["sequential_score"]
    delta = score - baseline_eval
    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
    logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")

    final_dir = f"{OUTPUT_DIR}/final"
    student.save_pretrained(final_dir)
    logging.info(f"Saved to {final_dir}")

    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: skipping Hub push")
        return

    try:
        commit_url = student.push_to_hub(RUN_NAME)
        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
    except Exception:
        import traceback

        logging.error(f"Hub push failed:\n{traceback.format_exc()}")


if __name__ == "__main__":
    main()

```


### `scripts/train_sentence_transformer_matryoshka_example.py`

```
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers[train]>=5.0",
#     "datasets>=2.19.0",
#     "accelerate>=0.26.0",
#     "trackio",
# ]
# ///
"""Matryoshka (MRL) training: train once, deploy at multiple embedding dimensions.

MatryoshkaLoss wraps a base loss and optimizes it at several truncated dimensions
simultaneously. At inference, load with `truncate_dim=<target>` to get that size
with ~95% of full-dim quality.

Typical use: train at [768, 512, 256, 128, 64], deploy at 128 for 6x smaller
index + 6x faster ANN with minimal quality loss.

Run locally:
    pip install "sentence-transformers[train]>=5.0"
    python train_sentence_transformer_matryoshka_example.py
"""

from __future__ import annotations

import argparse
import logging
import os
from contextlib import nullcontext

import torch
from datasets import load_dataset

from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)
from sentence_transformers.base.sampler import BatchSamplers
from sentence_transformers.sentence_transformer.evaluation import (
    EmbeddingSimilarityEvaluator,
    NanoBEIREvaluator,
    SequentialEvaluator,
)
from sentence_transformers.sentence_transformer.losses import MatryoshkaLoss, MultipleNegativesRankingLoss
from sentence_transformers.util.similarity import SimilarityFunction


def autocast_ctx():
    """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
    if not torch.cuda.is_available():
        return nullcontext()
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype)


def log_trackio_dashboard():
    """Surface the Trackio dashboard URL so the user can watch training live."""
    try:
        from huggingface_hub import whoami

        hf_user = whoami().get("name")
        if hf_user:
            logging.info(
                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
            )
    except Exception:
        pass


MODEL_NAME = "microsoft/mpnet-base"
MATRYOSHKA_DIMS = [768, 512, 256, 128, 64]
OUTPUT_DIR = "models/mpnet-matryoshka"
RUN_NAME = "mpnet-matryoshka"
SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"


def setup_logging():
    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
    os.makedirs("logs", exist_ok=True)
    logging.basicConfig(
        format="%(asctime)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.INFO,
        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
        force=True,
    )
    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision("high")  # TF32 on Ampere+, no quality loss


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
    )
    cli, _ = parser.parse_known_args()

    setup_logging()

    if cli.eval_only:
        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
        model = SentenceTransformer(cli.eval_only)
        evaluator = NanoBEIREvaluator()
        with autocast_ctx():
            evaluator(model)
        return

    model = SentenceTransformer(MODEL_NAME)

    train_size = 50 if SMOKE_TEST else 50_000
    eval_size = 20 if SMOKE_TEST else 1_000
    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
    train_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split="train").select(range(train_size))
    eval_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split="dev").select(range(eval_size))

    inner_loss = MultipleNegativesRankingLoss(model)
    loss = MatryoshkaLoss(model, inner_loss, matryoshka_dims=MATRYOSHKA_DIMS)

    stsb = load_dataset("sentence-transformers/stsb", split="validation")
    per_dim_evaluators = [
        EmbeddingSimilarityEvaluator(
            sentences1=stsb["sentence1"],
            sentences2=stsb["sentence2"],
            scores=stsb["score"],
            main_similarity=SimilarityFunction.COSINE,
            name=f"sts-dev-{dim}",
            truncate_dim=dim,
        )
        for dim in MATRYOSHKA_DIMS
    ]
    evaluator = SequentialEvaluator(
        [*per_dim_evaluators, NanoBEIREvaluator()],
        main_score_function=lambda scores: scores[0],
    )
    logging.info("Baseline evaluation (before training):")
    with autocast_ctx():
        # Must run before deriving metric_key: each sub-evaluator mutates its primary_metric to add the name_ prefix.
        baseline_result = evaluator(model)
    # Drive on the first per-dim evaluator's metric (matches main_score_function above).
    metric_key = f"eval_{per_dim_evaluators[0].primary_metric}"
    baseline_eval = baseline_result[per_dim_evaluators[0].primary_metric]

    args = SentenceTransformerTrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=1,
        max_steps=1 if SMOKE_TEST else -1,
        per_device_train_batch_size=128,
        per_device_eval_batch_size=128,
        learning_rate=2e-5,
        weight_decay=0.01,
        warmup_steps=0.1,
        bf16=True,
        batch_sampler=BatchSamplers.NO_DUPLICATES,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.1,
        save_total_limit=2,
        logging_steps=0.01,
        logging_first_step=True,
        load_best_model_at_end=True,
        metric_for_best_model=metric_key,
        greater_is_better=True,
        report_to="none" if SMOKE_TEST else "trackio",
        run_name=RUN_NAME,
        seed=12,
    )

    trainer = SentenceTransformerTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        loss=loss,
        evaluator=evaluator,
    )
    if not SMOKE_TEST:
        log_trackio_dashboard()
    trainer.train()

    logging.info("Post-training evaluation:")
    with autocast_ctx():
        score = evaluator(model)[per_dim_evaluators[0].primary_metric]
    delta = score - baseline_eval
    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
    logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")

    final_dir = f"{OUTPUT_DIR}/final"
    model.save_pretrained(final_dir)
    logging.info(f"Saved to {final_dir}")
    logging.info(f"To use at a specific dimension, load with: SentenceTransformer({final_dir!r}, truncate_dim=128)")

    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: skipping Hub push")
        return

    try:
        commit_url = model.push_to_hub(RUN_NAME)
        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
    except Exception:
        import traceback

        logging.error(f"Hub push failed:\n{traceback.format_exc()}")


if __name__ == "__main__":
    main()

```


### `scripts/train_sentence_transformer_multi_dataset_example.py`

```
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers[train]>=5.0",
#     "datasets>=2.19.0",
#     "accelerate>=0.26.0",
#     "trackio",
# ]
# ///
"""Multi-dataset / multi-task training: train one model on several datasets at once.

Pass `train_dataset` (and optionally `eval_dataset`) as dicts. Pass `loss` as
either a dict keyed by the same names (one loss per dataset, this script's
default — Variant A) or a single loss instance applied to every dataset
(Variant B). Dict keys are arbitrary but must match exactly across all three
dicts; they show up in log output as `loss_all-nli=...`, `loss_stsb=...`.

Three reasons to use multi-dataset training:
- Multi-task: combine datasets with different signals (retrieval + STS +
  classification) into a single general-purpose embedder.
- Data augmentation: add a supplementary dataset (STS labels alongside your
  main retrieval pairs) as a regularizer.
- Domain coverage: train on several domains at once rather than sequentially.

Variant A (this script): different shapes per dataset, so each needs its own
matching loss. Pass `loss` as a dict; the trainer dispatches per-dataset and
mixing loss arities (MNRL with 3 inputs + CoSENTLoss with 2+label) is fine.

Variant B: same shape, same loss, but you want each mini-batch drawn from a
single domain (so MNRL in-batch negatives stay in-domain and remain genuinely
hard). Pass ONE loss and a dict of datasets:

    train_datasets = {"medical": medical_pairs, "legal": legal_pairs, "code": code_pairs}
    loss = MultipleNegativesRankingLoss(model)
    trainer = SentenceTransformerTrainer(model=model, args=args,
        train_dataset=train_datasets, loss=loss, ...)

The multi-dataset batch sampler draws each batch from a single dataset, so a
3-domain MNRL run gets in-domain negatives by construction. Counter-intuitive
benefit: DatasetDict can outperform `concatenate_datasets` even with losses
that don't share across the batch (e.g. LambdaLoss in cross-encoder training).

Multi-dataset samplers:
- `PROPORTIONAL` (default): sample from each dataset in proportion to its size.
  Every row is seen ~once per epoch. Bias toward the largest dataset.
- `ROUND_ROBIN`: alternate evenly; training stops when the SMALLEST is
  exhausted. Equal screen-time per task.
Common pattern: `PROPORTIONAL` for 1 epoch, then `ROUND_ROBIN` for a second
if a smaller task's loss is still decreasing.

Per-dataset prompts (bi-encoder, sparse-encoder): pass `prompts={"all-nli": "",
"stsb": "Represent ...: ", "msmarco": {"query": "query: ", "positive":
"passage: ", ...}}` to TrainingArguments. The nested per-column form works for
bi-encoder and sparse-encoder; cross-encoders support single-value or
per-dataset only. See `../references/prompts_and_instructions.md`.

Eval metric aggregation: with a dict `eval_dataset`, each dataset's loss is
logged separately (`eval_loss_all-nli`, `eval_loss_stsb`). The evaluator runs
on the full model, so its metrics aren't per-dataset unless you wrap a
`SequentialEvaluator` with per-dataset sub-evaluators. Set
`metric_for_best_model` to a single evaluator metric, NOT a per-dataset loss.

Gotchas: keys must match EXACTLY across all three dicts (train/eval/loss) or
training fails at step 0; `NO_DUPLICATES` + `PROPORTIONAL` works (deduplicates
within each batch regardless of source dataset); `ROUND_ROBIN` with uneven
dataset sizes means `num_train_epochs=N` is N passes over the SMALLEST — use
`PROPORTIONAL` or `max_steps` if you want N passes over the largest.
"""

from __future__ import annotations

import argparse
import logging
import os
from contextlib import nullcontext

import torch
from datasets import load_dataset

from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)
from sentence_transformers.base.sampler import BatchSamplers
from sentence_transformers.sentence_transformer.evaluation import (
    EmbeddingSimilarityEvaluator,
    NanoBEIREvaluator,
)
from sentence_transformers.sentence_transformer.losses import (
    CoSENTLoss,
    MultipleNegativesRankingLoss,
)
from sentence_transformers.util.similarity import SimilarityFunction


def autocast_ctx():
    """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
    if not torch.cuda.is_available():
        return nullcontext()
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype)


def log_trackio_dashboard():
    """Surface the Trackio dashboard URL so the user can watch training live."""
    try:
        from huggingface_hub import whoami

        hf_user = whoami().get("name")
        if hf_user:
            logging.info(
                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
            )
    except Exception:
        pass


RUN_NAME = "mpnet-nli-stsb"
SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"


def setup_logging():
    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
    os.makedirs("logs", exist_ok=True)
    logging.basicConfig(
        format="%(asctime)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.INFO,
        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
        force=True,
    )
    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision("high")  # TF32 on Ampere+, no quality loss


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
    )
    cli, _ = parser.parse_known_args()

    setup_logging()

    if cli.eval_only:
        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
        model = SentenceTransformer(cli.eval_only)
        evaluator = NanoBEIREvaluator()
        with autocast_ctx():
            evaluator(model)
        return

    model = SentenceTransformer("microsoft/mpnet-base")

    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: trimmed datasets; will run max_steps=1 and skip Hub push")
    nli_train_size = 50 if SMOKE_TEST else 50_000
    nli_eval_size = 20 if SMOKE_TEST else 500
    nli_train = load_dataset("sentence-transformers/all-nli", "triplet", split="train").select(range(nli_train_size))
    stsb_train = load_dataset("sentence-transformers/stsb", split="train")
    if SMOKE_TEST:
        stsb_train = stsb_train.select(range(min(50, len(stsb_train))))
    nli_eval = load_dataset("sentence-transformers/all-nli", "triplet", split="dev").select(range(nli_eval_size))
    stsb_eval = load_dataset("sentence-transformers/stsb", split="validation")
    if SMOKE_TEST:
        stsb_eval = stsb_eval.select(range(min(20, len(stsb_eval))))

    train_datasets = {"all-nli": nli_train, "stsb": stsb_train}
    eval_datasets = {"all-nli": nli_eval, "stsb": stsb_eval}

    losses = {
        "all-nli": MultipleNegativesRankingLoss(model),
        "stsb": CoSENTLoss(model),
    }

    evaluator = EmbeddingSimilarityEvaluator(
        sentences1=stsb_eval["sentence1"],
        sentences2=stsb_eval["sentence2"],
        scores=stsb_eval["score"],
        main_similarity=SimilarityFunction.COSINE,
        name="sts-dev",
    )
    logging.info("Baseline evaluation (before training):")
    with autocast_ctx():
        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
        baseline_eval = evaluator(model)[evaluator.primary_metric]
    metric_key = f"eval_{evaluator.primary_metric}"

    # multi_dataset_batch_sampler defaults to PROPORTIONAL (samples each dataset
    # in proportion to its size). To force equal alternation between datasets:
    #   from sentence_transformers.base.sampler import MultiDatasetBatchSamplers
    #   ... multi_dataset_batch_sampler=MultiDatasetBatchSamplers.ROUND_ROBIN ...
    args = SentenceTransformerTrainingArguments(
        output_dir="models/mpnet-nli-stsb",
        num_train_epochs=1,
        max_steps=1 if SMOKE_TEST else -1,
        per_device_train_batch_size=32,
        per_device_eval_batch_size=32,
        learning_rate=2e-5,
        weight_decay=0.01,
        warmup_steps=0.1,
        bf16=True,
        batch_sampler=BatchSamplers.NO_DUPLICATES,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.1,
        save_total_limit=2,
        logging_steps=0.01,
        logging_first_step=True,
        load_best_model_at_end=True,
        metric_for_best_model=metric_key,
        greater_is_better=True,
        report_to="none" if SMOKE_TEST else "trackio",
        run_name="mpnet-nli-stsb",
        seed=12,
    )

    trainer = SentenceTransformerTrainer(
        model=model,
        args=args,
        train_dataset=train_datasets,
        eval_dataset=eval_datasets,
        loss=losses,
        evaluator=evaluator,
    )
    if not SMOKE_TEST:
        log_trackio_dashboard()
    trainer.train()

    logging.info("Post-training evaluation:")
    with autocast_ctx():
        score = evaluator(model)[evaluator.primary_metric]
    delta = score - baseline_eval
    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
    logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")

    model.save_pretrained("models/mpnet-nli-stsb/final")

    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: skipping Hub push")
        return

    try:
        commit_url = model.push_to_hub(RUN_NAME)
        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
    except Exception:
        import traceback

        logging.error(f"Hub push failed:\n{traceback.format_exc()}")


if __name__ == "__main__":
    main()

```


### `scripts/train_sentence_transformer_static_embedding_example.py`

```
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers[train]>=5.0",
#     "datasets>=2.19.0",
#     "accelerate>=0.26.0",
#     "trackio",
#     "tokenizers>=0.20",
#     "model2vec",  # only needed when WARMSTART=True (StaticEmbedding.from_model2vec)
# ]
# ///
"""Train a StaticEmbedding model on a contrastive dataset.

StaticEmbedding is a token-bag model: a per-token embedding table averaged over
the tokens of an input. No transformer, no attention. Inference is ~20x faster
on GPU and ~80x faster on CPU than a small encoder, with surprisingly competitive
quality on retrieval benchmarks when trained on >=1M contrastive pairs.

Two init paths via `WARMSTART` constant:
- `WARMSTART=False` (default): random init. Use with >=1M contrastive pairs;
  reaches a higher ceiling than warm-start when given enough data. The default
  dataset below (GooAQ, ~3M pairs) is comfortably in this regime.
- `WARMSTART=True`: `StaticEmbedding.from_model2vec(...)` — distil from a
  model2vec checkpoint. Flip to True if you swap in a smaller dataset (<1M
  pairs); converges faster and reaches better quality at lower data scales.

Demonstrates:
- MultipleNegativesRankingLoss wrapped in MatryoshkaLoss for nested embedding dims
- Large batch size (1024+) with a high LR (~2e-1 for random init, ~5e-2 for warm-
  start) since the loss surface for a token-bag is much flatter than for a
  pretrained encoder
- BatchSamplers.NO_DUPLICATES (load-bearing for in-batch negatives with duplicated
  anchors)
- NanoBEIREvaluator at full embedding dim
- Auto model card + optional Hub push

Run locally (CPU works for inference, but training needs a GPU for batch=1024+):
    pip install "sentence-transformers[train]>=5.0"
    python train_sentence_transformer_static_embedding_example.py

Multi-GPU:
    accelerate launch train_sentence_transformer_static_embedding_example.py

Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).

References:
- HF blog post: https://huggingface.co/blog/static-embeddings
- Module docs: sentence_transformers.sentence_transformer.modules.StaticEmbedding
"""

from __future__ import annotations

import argparse
import logging
import os
from contextlib import nullcontext

import datasets
import torch
from datasets import load_dataset
from tokenizers import Tokenizer

from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerModelCardData,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)
from sentence_transformers.base.sampler import BatchSamplers
from sentence_transformers.sentence_transformer.evaluation import NanoBEIREvaluator
from sentence_transformers.sentence_transformer.losses import (
    MatryoshkaLoss,
    MultipleNegativesRankingLoss,
)
from sentence_transformers.sentence_transformer.modules import StaticEmbedding


def autocast_ctx():
    """bf16/fp16 autocast for evaluator calls outside the trainer."""
    if not torch.cuda.is_available():
        return nullcontext()
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype)


def log_trackio_dashboard():
    """Surface the Trackio dashboard URL so the user can watch training live."""
    try:
        from huggingface_hub import whoami

        hf_user = whoami().get("name")
        if hf_user:
            logging.info(
                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
            )
    except Exception:
        pass


TOKENIZER_NAME = "google-bert/bert-base-uncased"
EMBEDDING_DIM = 1024
MATRYOSHKA_DIMS = [1024, 512, 256, 128, 64, 32]  # ordered largest-first per MatryoshkaLoss

# False: random init (recommended for >=1M pairs; reaches a higher ceiling).
# True: warm-start from a model2vec checkpoint (recommended for <1M pairs).
# Default False because the example dataset (GooAQ, ~3M pairs) is well above the threshold.
WARMSTART = False
WARMSTART_MODEL2VEC = "minishlab/potion-base-8M"

OUTPUT_DIR = "models/static-embedding-bert-uncased"
RUN_NAME = "static-embedding-bert-uncased"
SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"


def setup_logging():
    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
    os.makedirs("logs", exist_ok=True)
    logging.basicConfig(
        format="%(asctime)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.INFO,
        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
        force=True,
    )
    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision("high")


def load_pair_dataset() -> datasets.Dataset:
    """Load a contrastive-pair dataset for training.

    StaticEmbedding starts from random initialization, so it needs *a lot* of
    contrastive signal to converge. GooAQ alone provides ~3M (question, answer)
    pairs, comfortably over the >=1M threshold below which a warm-start would
    beat random init. For stronger production models, concatenate more sources
    (NaturalQuestions, MSMARCO, MIRACL, etc.) and shuffle, in the same family of
    sources used in `sentence-transformers/static-retrieval-mrl-en-v1`.
    """
    return (
        load_dataset("sentence-transformers/gooaq", split="train")
        .rename_columns({"question": "anchor", "answer": "positive"})
        .select_columns(["anchor", "positive"])
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
    )
    cli, _ = parser.parse_known_args()

    setup_logging()

    if cli.eval_only:
        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
        model = SentenceTransformer(cli.eval_only)
        evaluator = NanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
        with autocast_ctx():
            evaluator(model)
        return

    if WARMSTART:
        logging.info(f"Warm-starting StaticEmbedding from model2vec: {WARMSTART_MODEL2VEC}")
        # `StaticEmbedding.from_distillation("<bi-encoder>", vocabulary=...)` is the
        # alternative warm-start path (distil from a stronger teacher's vectors); pick
        # one. model2vec is faster to load and converges quickly on smaller datasets.
        static_embedding = StaticEmbedding.from_model2vec(WARMSTART_MODEL2VEC)
    else:
        logging.info(f"Random-init StaticEmbedding from {TOKENIZER_NAME} tokenizer (dim={EMBEDDING_DIM})")
        tokenizer = Tokenizer.from_pretrained(TOKENIZER_NAME)
        static_embedding = StaticEmbedding(tokenizer, embedding_dim=EMBEDDING_DIM)
    model = SentenceTransformer(
        modules=[static_embedding],
        model_card_data=SentenceTransformerModelCardData(
            language="en",
            license="apache-2.0",
            model_name=f"Static embedding ({EMBEDDING_DIM}d) trained on contrastive pairs",
        ),
    )

    logging.info("Loading + concatenating training datasets")
    full = load_pair_dataset()
    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
        full = full.select(range(min(200, len(full))))
    eval_size = 20 if SMOKE_TEST else 10_000
    split = full.train_test_split(test_size=eval_size, seed=12)
    train_dataset = split["train"]
    eval_dataset = split["test"]
    logging.info(f"  train: {len(train_dataset):,} rows | eval: {len(eval_dataset):,} rows")
    logging.info(f"  columns: {train_dataset.column_names}")

    inner = MultipleNegativesRankingLoss(model)
    loss = MatryoshkaLoss(model, inner, matryoshka_dims=MATRYOSHKA_DIMS)

    evaluator = NanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
    logging.info("Baseline evaluation (random init scores near zero; warm-start scores 0.3+):")
    with autocast_ctx():
        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
        baseline_eval = evaluator(model)[evaluator.primary_metric]
    metric_key = f"eval_{evaluator.primary_metric}"

    args = SentenceTransformerTrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=1,
        max_steps=1 if SMOKE_TEST else -1,
        per_device_train_batch_size=2048,
        per_device_eval_batch_size=2048,
        learning_rate=5e-2
        if WARMSTART
        else 2e-1,  # warm-start needs less LR; both far higher than encoder fine-tuning
        weight_decay=0.0,  # weight decay on a token-bag is usually harmful
        warmup_steps=0.1,
        lr_scheduler_type="linear",
        bf16=True,
        batch_sampler=BatchSamplers.NO_DUPLICATES,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.1,
        save_total_limit=2,
        logging_steps=0.01,
        logging_first_step=True,
        load_best_model_at_end=True,
        metric_for_best_model=metric_key,
        greater_is_better=True,
        report_to="none" if SMOKE_TEST else "trackio",
        run_name=RUN_NAME,
        seed=12,
    )

    trainer = SentenceTransformerTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        loss=loss,
        evaluator=evaluator,
    )
    if not SMOKE_TEST:
        log_trackio_dashboard()
    trainer.train()

    logging.info("Post-training evaluation:")
    with autocast_ctx():
        score = evaluator(model)[evaluator.primary_metric]
    delta = score - baseline_eval
    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
    logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")

    final_dir = f"{OUTPUT_DIR}/final"
    model.save_pretrained(final_dir)
    logging.info(f"Saved final model to {final_dir}")

    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: skipping Hub push")
        return

    try:
        commit_url = model.push_to_hub(RUN_NAME)
        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
    except Exception:
        import traceback

        logging.error(f"Hub push failed:\n{traceback.format_exc()}")


if __name__ == "__main__":
    main()

```


### `scripts/train_sentence_transformer_with_lora_example.py`

```
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers[train]>=5.0",
#     "peft>=0.7.0",
#     "datasets>=2.19.0",
#     "accelerate>=0.26.0",
#     "trackio",
# ]
# ///
"""LoRA / PEFT adapter training: memory-efficient fine-tuning.

Instead of training all parameters, LoRA injects small low-rank adapter
matrices and trains only those (~50-100x fewer trainable params). Works out
of the box with sentence-transformers via `peft`.

Use LoRA when: the base is a large decoder (Qwen3, Llama, Mistral, Gemma) and
a full fine-tune is VRAM-prohibitive; you want multiple task-specific adapters
on one base; you're nudging an existing strong retriever (E5-mistral,
Qwen3-Embedding) on domain data. Skip LoRA for small encoders (BERT-base,
MiniLM — full fine-tune is tractable and usually better) or tiny datasets
(<1k pairs — adapter rank becomes the bottleneck).

This template defaults to BERT-base for portability. Swap `MODEL_NAME` to a
decoder backbone for the use case LoRA actually shines on.

Architecture variants:
- Bi-encoder (this script): `TaskType.FEATURE_EXTRACTION`.
- Sparse-encoder: same pattern; `SparseEncoder` supports `add_adapter`. See
  `examples/sparse_encoder/training/peft/train_splade_gooaq_peft.py`.
- Cross-encoder: use `TaskType.SEQ_CLS` for `num_labels >= 1`. Community
  examples are sparse; smoke-test with `max_steps=1` first.

Key hyperparameters:
- `r` (rank): 8-128. Bigger = more capacity + memory. 64 is a strong default.
- `lora_alpha`: typically 2 x r (some teams use 1 x r for stability).
- `lora_dropout`: 0.05-0.1; raise to 0.1 for small datasets.
- `target_modules=None` auto-picks attention modules; pass
  `["q_proj", "k_proj", "v_proj", "o_proj"]` (attention) or
  `["gate_proj", "up_proj", "down_proj"]` (MLP) for explicit control.
- `modules_to_save=["pooler"]` for CLS-pooled bases — the pooler Dense should
  be trained too, not adapted.
- LR is HIGHER than full fine-tune: 1e-4 to 5e-4 for LoRA vs. 2e-5 full.

Rough memory savings on a 0.6B base (bf16, batch 64, seq 128): full fine-tune
~24 GB, LoRA r=64 ~10 GB (~12M trainable, 2%), LoRA r=16 ~8 GB (~3M, 0.5%).
Bigger savings on 7B+ models.

Saving / sharing: `model.save_pretrained("dir")` writes ONLY the adapter (few
MB) plus a reference to the base model. Loaders call the same one-liner;
`peft` is invoked and the base downloaded on demand. For a merged model that
loads without `peft` (needed for vLLM-style servers), call
`model.transformers_model.merge_and_unload()` then `save_pretrained` /
`push_to_hub`.

Swapping adapters at inference (the main multi-task deployment win):
    model = SentenceTransformer("base-model")
    model.load_adapter("adapter-a", adapter_name="a")
    model.load_adapter("adapter-b", adapter_name="b")
    model.set_adapter("a"); emb_a = model.encode([...])

QLoRA (4-bit base + LoRA) for 7B+ on consumer GPUs:
    from transformers import BitsAndBytesConfig
    bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                             bnb_4bit_compute_dtype=torch.bfloat16)
    model = SentenceTransformer("Qwen/Qwen3-Embedding-7B",
                                 model_kwargs={"quantization_config": bnb})
    model.add_adapter(LoraConfig(r=64, lora_alpha=128, ...))
`pip install bitsandbytes` first. Linux-only for bitsandbytes (Windows: the
fork or WSL).

Known issues:
- LoRA PEFT on Qwen2.5-VL / paligemma / gemma3 / internvl / aya_vision under
  transformers v5: `AutoModel.from_pretrained(peft_path)` crashes with
  `KeyError: 'qwen2_vl'`. Pin transformers to 4.x or wait for the upstream fix.
- `gradient_checkpointing=True` + LoRA: usually works; if you hit "None of
  the inputs have requires_grad=True", call
  `model.transformers_model.enable_input_require_grads()` after `add_adapter`.
- `add_adapter` before pooling: when building from scratch (not loading a
  pre-assembled checkpoint), call `add_adapter` AFTER
  `SentenceTransformer(modules=[...])` is complete.

Common gotchas: LR still at 2e-5 (LoRA needs higher); forgetting to merge for
vLLM-style servers (they don't load `peft`); `r=8` too small for retrievers
trained on millions of pairs (try 32 or 64); `modules_to_save` missing the
pooler on CLS-pooled bases.
"""

from __future__ import annotations

import argparse
import logging
import os
from contextlib import nullcontext

import torch
from datasets import load_dataset
from peft import LoraConfig, TaskType

from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerModelCardData,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)
from sentence_transformers.base.sampler import BatchSamplers
from sentence_transformers.sentence_transformer.evaluation import NanoBEIREvaluator
from sentence_transformers.sentence_transformer.losses import CachedMultipleNegativesRankingLoss


def autocast_ctx():
    """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
    if not torch.cuda.is_available():
        return nullcontext()
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype)


def log_trackio_dashboard():
    """Surface the Trackio dashboard URL so the user can watch training live."""
    try:
        from huggingface_hub import whoami

        hf_user = whoami().get("name")
        if hf_user:
            logging.info(
                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
            )
    except Exception:
        pass


MODEL_NAME = "google-bert/bert-base-uncased"
OUTPUT_DIR = "models/bert-base-gooaq-lora"
RUN_NAME = "bert-base-gooaq-lora"
SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"


def setup_logging():
    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
    os.makedirs("logs", exist_ok=True)
    logging.basicConfig(
        format="%(asctime)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.INFO,
        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
        force=True,
    )
    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision("high")  # TF32 on Ampere+, no quality loss


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
    )
    cli, _ = parser.parse_known_args()

    setup_logging()

    if cli.eval_only:
        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
        model = SentenceTransformer(cli.eval_only)
        evaluator = NanoBEIREvaluator()
        with autocast_ctx():
            evaluator(model)
        return

    model = SentenceTransformer(
        MODEL_NAME,
        model_card_data=SentenceTransformerModelCardData(
            language="en",
            license="apache-2.0",
            model_name=f"{MODEL_NAME.split('/')[-1]} LoRA adapter on GooAQ",
        ),
    )

    peft_config = LoraConfig(
        task_type=TaskType.FEATURE_EXTRACTION,
        inference_mode=False,
        r=64,
        lora_alpha=128,
        lora_dropout=0.1,
    )
    model.add_adapter(peft_config)
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    logging.info(f"trainable params: {trainable:,} / {total:,} ({100 * trainable / total:.2f}%)")

    full = load_dataset("sentence-transformers/gooaq", split="train")
    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
        full = full.select(range(min(200, len(full))))
    eval_size = 20 if SMOKE_TEST else 10_000
    train_cap = 50 if SMOKE_TEST else 500_000
    split = full.train_test_split(test_size=eval_size, seed=12)
    train_dataset = split["train"].select(range(min(train_cap, len(split["train"]))))
    eval_dataset = split["test"]
    logging.info(f"train={len(train_dataset):,} eval={len(eval_dataset):,}")

    loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=32)

    evaluator = NanoBEIREvaluator()
    logging.info("Baseline:")
    with autocast_ctx():
        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
        baseline_eval = evaluator(model)[evaluator.primary_metric]
    metric_key = f"eval_{evaluator.primary_metric}"

    args = SentenceTransformerTrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=1,
        max_steps=1 if SMOKE_TEST else -1,
        per_device_train_batch_size=512,
        per_device_eval_batch_size=512,
        learning_rate=1e-4,
        weight_decay=0.01,
        warmup_steps=0.1,
        bf16=True,
        batch_sampler=BatchSamplers.NO_DUPLICATES,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.1,
        save_total_limit=2,
        logging_steps=0.01,
        logging_first_step=True,
        load_best_model_at_end=True,
        metric_for_best_model=metric_key,
        greater_is_better=True,
        report_to="none" if SMOKE_TEST else "trackio",
        run_name=RUN_NAME,
        seed=12,
    )

    trainer = SentenceTransformerTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        loss=loss,
        evaluator=evaluator,
    )
    if not SMOKE_TEST:
        log_trackio_dashboard()
    trainer.train()

    logging.info("Post-training evaluation:")
    with autocast_ctx():
        score = evaluator(model)[evaluator.primary_metric]
    delta = score - baseline_eval
    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
    logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")

    model.save_pretrained(f"{OUTPUT_DIR}/final")

    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: skipping Hub push")
        return

    try:
        commit_url = model.push_to_hub(RUN_NAME)
        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
    except Exception:
        import traceback

        logging.error(f"Hub push failed:\n{traceback.format_exc()}")


if __name__ == "__main__":
    main()

```


### `scripts/train_sparse_encoder_distillation_example.py`

```
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers[train]>=5.0",
#     "datasets>=2.19.0",
#     "accelerate>=0.26.0",
#     "trackio",
# ]
# ///
"""SPLADE distillation from a cross-encoder teacher.

Trains a SPLADE sparse retriever to match a stronger cross-encoder's score
gaps via `SparseMarginMSELoss` wrapped in `SpladeLoss` (the FLOPS regularizer
is non-negotiable; without it embeddings collapse to dense).

Data shape: `(query, positive, negative, score_diff)` where
`score_diff = teacher(q, pos) - teacher(q, neg)`. This script uses
`sentence-transformers/msmarco` (`bert-ensemble-margin-mse` subset) which has
precomputed teacher score diffs. To distill from your own cross-encoder
teacher, run a one-time teacher pass over your (q, pos, neg) triples and
store the per-row score diff.

Why distill SPLADE from a cross-encoder: SPLADE alone is hard to train from
contrastive labels because the FLOPS regularizer fights early-training signal;
distilling from a strong cross-encoder gives the model a dense regression
target and reaches stronger nDCG faster than MNRL-only.

Run locally:
    pip install "sentence-transformers[train]>=5.0"
    python train_sparse_encoder_distillation_example.py

Multi-GPU:
    accelerate launch train_sparse_encoder_distillation_example.py

Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
"""

from __future__ import annotations

import argparse
import logging
import os
from contextlib import nullcontext

import torch
from datasets import load_dataset, load_from_disk

from sentence_transformers import (
    SparseEncoder,
    SparseEncoderModelCardData,
    SparseEncoderTrainer,
    SparseEncoderTrainingArguments,
)
from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
from sentence_transformers.sparse_encoder.losses import SparseMarginMSELoss, SpladeLoss


def autocast_ctx():
    """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
    if not torch.cuda.is_available():
        return nullcontext()
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype)


def log_trackio_dashboard():
    """Surface the Trackio dashboard URL so the user can watch training live."""
    try:
        from huggingface_hub import whoami

        hf_user = whoami().get("name")
        if hf_user:
            logging.info(
                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
            )
    except Exception:
        pass


MODEL_NAME = "Luyu/co-condenser-marco"  # MS MARCO-tuned MLM base; very strong starting point
DATASET_NAME = "sentence-transformers/msmarco"
DATASET_SUBSET = "bert-ensemble-margin-mse"
TRAIN_SIZE = 100_000
EVAL_SIZE = 5_000
OUTPUT_DIR = "models/splade-msmarco-distilled"
RUN_NAME = "splade-msmarco-distilled"
DATA_CACHE = f"data/{RUN_NAME}-resolved"

QUERY_REGULARIZER_WEIGHT = 0.1  # higher than contrastive recipe; distillation tolerates more sparsity pressure
DOCUMENT_REGULARIZER_WEIGHT = 0.08
SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"


def setup_logging():
    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
    os.makedirs("logs", exist_ok=True)
    logging.basicConfig(
        format="%(asctime)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.INFO,
        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
        force=True,
    )
    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision("high")


def load_resolved_dataset():
    """Load (query, positive, negative, score) rows. The MSMARCO subset is keyed by
    passage_id / query_id; resolve to text once and cache to disk so reruns skip the work."""
    if os.path.isdir(DATA_CACHE):
        logging.info(f"Loading cached resolved dataset from {DATA_CACHE}")
        return load_from_disk(DATA_CACHE)

    logging.info(f"Resolving {DATASET_NAME}/{DATASET_SUBSET} ids -> text (one-time, cached)")
    corpus_ds = load_dataset(DATASET_NAME, "corpus", split="train")
    corpus = dict(zip(corpus_ds["passage_id"], corpus_ds["passage"]))
    queries_ds = load_dataset(DATASET_NAME, "queries", split="train")
    queries = dict(zip(queries_ds["query_id"], queries_ds["query"]))
    raw = load_dataset(DATASET_NAME, DATASET_SUBSET, split="train").select(range(TRAIN_SIZE + EVAL_SIZE))

    def id_to_text(batch):
        return {
            "query": [queries[qid] for qid in batch["query_id"]],
            "positive": [corpus[pid] for pid in batch["positive_id"]],
            "negative": [corpus[pid] for pid in batch["negative_id"]],
            "score": batch["score"],
        }

    resolved = raw.map(id_to_text, batched=True, remove_columns=["query_id", "positive_id", "negative_id"])
    resolved.save_to_disk(DATA_CACHE)
    return resolved


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
    )
    cli, _ = parser.parse_known_args()

    setup_logging()

    if cli.eval_only:
        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
        model = SparseEncoder(cli.eval_only)
        evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
        with autocast_ctx():
            evaluator(model)
        return

    logging.info(f"Loading base model: {MODEL_NAME}")
    model = SparseEncoder(
        MODEL_NAME,
        model_card_data=SparseEncoderModelCardData(
            language="en",
            license="apache-2.0",
            model_name=f"SPLADE from {MODEL_NAME.split('/')[-1]} distilled from MS MARCO ensemble",
        ),
    )
    model.max_seq_length = 256

    resolved = load_resolved_dataset()
    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
        resolved = resolved.select(range(min(70, len(resolved))))
    eval_size = 20 if SMOKE_TEST else EVAL_SIZE
    split = resolved.train_test_split(test_size=eval_size, seed=12)
    train_dataset = split["train"]
    eval_dataset = split["test"]
    logging.info(f"  train: {len(train_dataset):,} rows | eval: {len(eval_dataset):,} rows")
    logging.info(f"  columns: {train_dataset.column_names}")

    loss = SpladeLoss(
        model=model,
        loss=SparseMarginMSELoss(model=model),
        query_regularizer_weight=QUERY_REGULARIZER_WEIGHT,
        document_regularizer_weight=DOCUMENT_REGULARIZER_WEIGHT,
    )

    evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
    logging.info("Baseline evaluation (fill-mask base scores near zero, confirms pipeline):")
    with autocast_ctx():
        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
        baseline_result = evaluator(model)
        baseline_eval = baseline_result[evaluator.primary_metric]
    metric_key = f"eval_{evaluator.primary_metric}"

    args = SparseEncoderTrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=1,
        max_steps=1 if SMOKE_TEST else -1,
        per_device_train_batch_size=16,
        per_device_eval_batch_size=16,
        learning_rate=2e-5,
        weight_decay=0.01,
        warmup_steps=0.1,
        lr_scheduler_type="linear",
        bf16=True,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.1,
        save_total_limit=2,
        logging_steps=0.01,
        logging_first_step=True,
        load_best_model_at_end=True,
        metric_for_best_model=metric_key,
        greater_is_better=True,
        report_to="none" if SMOKE_TEST else "trackio",
        run_name=RUN_NAME,
        seed=12,
    )

    trainer = SparseEncoderTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        loss=loss,
        evaluator=evaluator,
    )
    if not SMOKE_TEST:
        log_trackio_dashboard()
    trainer.train()

    logging.info("Post-training evaluation:")
    with autocast_ctx():
        result = evaluator(model)
        score = result[evaluator.primary_metric]
    delta = score - baseline_eval
    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
    # Active-dim keys come back name-prefixed (e.g. "NanoBEIR_..._query_active_dims"); suffix-match for compat.
    qad = next((v for k, v in result.items() if k.endswith("query_active_dims")), "n/a")
    cad = next((v for k, v in result.items() if k.endswith("corpus_active_dims")), "n/a")
    logging.info(
        f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f} "
        f"| query_active={qad} corpus_active={cad}"
    )

    final_dir = f"{OUTPUT_DIR}/final"
    model.save_pretrained(final_dir)
    logging.info(f"Saved final model to {final_dir}")

    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: skipping Hub push")
        return

    try:
        commit_url = model.push_to_hub(RUN_NAME)
        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
    except Exception:
        import traceback

        logging.error(f"Hub push failed:\n{traceback.format_exc()}")


if __name__ == "__main__":
    main()

```


### `scripts/train_sparse_encoder_example.py`

```
#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers[train]>=5.0",
#     "datasets>=2.19.0",
#     "accelerate>=0.26.0",
#     "trackio",
# ]
# ///
"""Production-ready sparse-encoder (SPLADE) training template.

Demonstrates:
- SpladeLoss wrapping SparseMultipleNegativesRankingLoss
- FLOPS regularization (`query_regularizer_weight` / `document_regularizer_weight`)
- SparseNanoBEIREvaluator for sparse retrieval metrics
- load_best_model_at_end on the retrieval metric

Base model must expose a masked-LM head; any `AutoModelForMaskedLM`-compatible
checkpoint works (DistilBERT, BERT, MiniLM MLM variants, existing SPLADE models).

Run locally:
    pip install "sentence-transformers[train]>=5.0"
    python train_sparse_encoder_example.py

Multi-GPU:
    accelerate launch train_sparse_encoder_example.py

Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
"""

from __future__ import annotations

import argparse
import logging
import os
from contextlib import nullcontext

import torch
from datasets import load_dataset

from sentence_transformers import (
    SparseEncoder,
    SparseEncoderModelCardData,
    SparseEncoderTrainer,
    SparseEncoderTrainingArguments,
)
from sentence_transformers.base.sampler import BatchSamplers
from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
from sentence_transformers.sparse_encoder.losses import (
    SparseMultipleNegativesRankingLoss,
    SpladeLoss,
)


def autocast_ctx():
    """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
    if not torch.cuda.is_available():
        return nullcontext()
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    return torch.autocast("cuda", dtype=dtype)


def log_trackio_dashboard():
    """Surface the Trackio dashboard URL so the user can watch training live."""
    try:
        from huggingface_hub import whoami

        hf_user = whoami().get("name")
        if hf_user:
            logging.info(
                f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
            )
    except Exception:
        pass


MODEL_NAME = "distilbert/distilbert-base-uncased"
DATASET_NAME = "sentence-transformers/gooaq"
TRAIN_SIZE = 100_000
EVAL_SIZE = 1_000
OUTPUT_DIR = "models/distilbert-splade-gooaq"
RUN_NAME = "distilbert-splade-gooaq"

QUERY_REGULARIZER_WEIGHT = 5e-5
DOCUMENT_REGULARIZER_WEIGHT = 3e-5
SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"


def setup_logging():
    """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
    os.makedirs("logs", exist_ok=True)
    logging.basicConfig(
        format="%(asctime)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.INFO,
        handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
        force=True,
    )
    for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
        logging.getLogger(noisy).setLevel(logging.WARNING)
    if torch.cuda.is_available():
        torch.set_float32_matmul_precision("high")  # TF32 on Ampere+, no quality loss


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
    )
    cli, _ = parser.parse_known_args()

    setup_logging()

    if cli.eval_only:
        logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
        model = SparseEncoder(cli.eval_only)
        evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
        with autocast_ctx():
            evaluator(model)
        return

    logging.info(f"Loading base model: {MODEL_NAME}")
    # Prompts are optional for SPLADE: most BERT-style MLM bases don't need them.
    # If you're starting from a CSR (`Transformer + Pooling + SparseAutoEncoder`)
    # base like `tomaarsen/csr-mxbai-embed-large-v1-nq` that *was* trained with
    # prompts, mirror them here to preserve quality:
    #   prompts={"query": "Represent this sentence for similarity: ", "document": ""},
    #   default_prompt_name="document",
    model = SparseEncoder(
        MODEL_NAME,
        model_card_data=SparseEncoderModelCardData(
            language="en",
            license="apache-2.0",
            model_name=f"SPLADE from {MODEL_NAME.split('/')[-1]} trained on GooAQ",
        ),
    )

    logging.info(f"Loading dataset: {DATASET_NAME}")
    train_size = 50 if SMOKE_TEST else TRAIN_SIZE
    eval_size = 20 if SMOKE_TEST else EVAL_SIZE
    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
    full = load_dataset(DATASET_NAME, split="train")
    split = full.train_test_split(test_size=eval_size, seed=12)
    train_dataset = split["train"].select(range(min(train_size, len(split["train"]))))
    eval_dataset = split["test"]
    logging.info(f"  train: {len(train_dataset):,} rows | eval: {len(eval_dataset):,} rows")
    logging.info(f"  columns: {train_dataset.column_names}")

    loss = SpladeLoss(
        model=model,
        loss=SparseMultipleNegativesRankingLoss(model=model),
        query_regularizer_weight=QUERY_REGULARIZER_WEIGHT,
        document_regularizer_weight=DOCUMENT_REGULARIZER_WEIGHT,
    )

    evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
    logging.info("Baseline evaluation:")
    with autocast_ctx():
        # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
        baseline_result = evaluator(model)
        baseline_eval = baseline_result[evaluator.primary_metric]
    metric_key = f"eval_{evaluator.primary_metric}"

    args = SparseEncoderTrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=1,
        max_steps=1 if SMOKE_TEST else -1,
        per_device_train_batch_size=32,
        per_device_eval_batch_size=32,
        learning_rate=2e-5,
        weight_decay=0.01,
        warmup_steps=0.1,
        lr_scheduler_type="linear",
        bf16=True,
        batch_sampler=BatchSamplers.NO_DUPLICATES,
        eval_strategy="steps",
        eval_steps=0.1,
        save_strategy="steps",
        save_steps=0.1,
        save_total_limit=2,
        logging_steps=0.01,
        logging_first_step=True,
        load_best_model_at_end=True,
        metric_for_best_model=metric_key,
        greater_is_better=True,
        report_to="none" if SMOKE_TEST else "trackio",
        run_name=RUN_NAME,
        seed=12,
    )

    trainer = SparseEncoderTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        loss=loss,
        evaluator=evaluator,
    )
    if not SMOKE_TEST:
        log_trackio_dashboard()
    trainer.train()

    logging.info("Post-training evaluation:")
    with autocast_ctx():
        result = evaluator(model)
        score = result[evaluator.primary_metric]
    delta = score - baseline_eval
    verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
    # Active-dim keys come back name-prefixed (e.g. "NanoBEIR_..._query_active_dims"); suffix-match for compat.
    qad = next((v for k, v in result.items() if k.endswith("query_active_dims")), "n/a")
    cad = next((v for k, v in result.items() if k.endswith("corpus_active_dims")), "n/a")
    logging.info(
        f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f} "
        f"| query_active={qad} corpus_active={cad}"
    )

    final_dir = f"{OUTPUT_DIR}/final"
    model.save_pretrained(final_dir)
    logging.info(f"Saved final model to {final_dir}")

    if SMOKE_TEST:
        logging.info("SMOKE_TEST=1: skipping Hub push")
        return

    try:
        commit_url = model.push_to_hub(RUN_NAME)
        logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
    except Exception:
        import traceback

        logging.error(f"Hub push failed:\n{traceback.format_exc()}")


if __name__ == "__main__":
    main()

```