Train and fine-tune sentence-transformer embedding models

Acts as a router - not a manual - for training or fine-tuning sentence-transformers models across three types: SentenceTransformer (bi-encoder, maps input to a fixed-dim dense vector, for retrieval/similarity/clustering/classification/paraphrase mining/dedup/multimodal), CrossEncoder (reranker, scores query-passage pairs jointly for two-stage retrieval), and SparseEncoder (SPLADE, sparse vocabulary-space vectors for learned-sparse retrieval on inverted-index backends like Elasticsearch/OpenSearch/Lucene) - with tiebreaker phrases ("embedding"/"vector search" → SentenceTransformer, "rerank"/"two-stage" → CrossEncoder, "SPLADE"/"sparse" → SparseEncoder) and an explicit instruction to ask if still ambiguous. It explicitly forbids synthesizing a training script from the router file alone: the per-type production template (scripts/train_<type>_example.py) must be copied as the actual starting point, since it contains load-bearing scaffolding - an autocast helper, a model-card class, a logger-silencing list, force=True, seeding, TF32 config, version-compatible imports, and named-evaluator metric handling - that prior agent runs have repeatedly gotten wrong when rolling their own from a synthesized snippet. Required reading before writing any code is split into per-type files (losses mapped to data shape and incompatibilities like Cached* losses conflicting with gradient checkpointing; evaluators mapped to tasks and their metric_for_best_model key format; for SentenceTransformer specifically, model architecture and pooling rules) and cross-cutting files read regardless of type: training-argument knobs and precision rules (load fp32 and autocast bf16/fp16, never torch_dtype=bfloat16 directly), dataset column-matching and hard-negative-mining recipes, base-model selection including a flagged ModernBERT max_seq_length=8192 trap, and a symptom-indexed troubleshooting file to skim even on a healthy run since its "metrics don't improve" and "hub push fails" entries catch frequent bugs cheaply before they fire. Hardware, HF Jobs execution, and prompt-prefix (query:/passage: style) reference files load only when applicable. Defaults run locally unless hardware can't fit the job, execute a single run and only propose iteration if the result is weak or marginal, and push to the Hub at end-of-run wrapped in try/except (with additional in-trainer push enabled specifically on HF Jobs' ephemeral environment). The produced script must satisfy non-negotiable contracts: capture a baseline_eval score before trainer.train() starts; emit a single scrapeable end-of-run line in the exact format VERDICT: WIN|MARGINAL|REGRESSION | score=... | baseline=... | delta=...; silence noisy HTTP/HF-download loggers to WARNING; tee logs to a run-named file; wrap the final model.push_to_hub(...) in try/except; smoke-test with max_steps=1 on a tiny data slice before any long run; for CrossEncoder specifically, include EarlyStoppingCallback(patience>=3) since rerankers often peak mid-training and regress afterward; and for SparseEncoder specifically, log query/corpus active-dimension counts on the verdict line, since high nDCG with collapsed sparsity isn't actually a win. The seven-step workflow: identify the model type, load that type's required reading, copy the matching production template, replace the model/dataset/run names and loss/evaluator for the actual task (cross-checked against the loss and evaluator reference files), smoke-test, run, then log the experiment and propose iteration if the verdict is weak or marginal.

pip install "sentence-transformers[train]>=5.0"
pip install trackio
hf auth login

Use it when training or fine-tuning any sentence-transformers model - a bi-encoder embedding model, a cross-encoder reranker, or a sparse SPLADE encoder.

Input is a model type (or ambiguous request to be disambiguated), a base model, and a dataset. Output is a training script copied from the correct production template, customized with the right loss/evaluator/data shape, smoke-tested, run, and ending in a scrapeable VERDICT line plus a Hub-pushed model.

Uses the sentence-transformers[train] library, Hugging Face Hub authentication for model push, optional experiment trackers (trackio, wandb, tensorboard, mlflow), and Hugging Face Jobs for hardware that exceeds local capacity.

ML engineers training or fine-tuning embedding, reranking, or sparse retrieval models who want the correct production-grade training script scaffolding (baseline capture, smoke tests, verdict reporting, type-specific gotchas like CrossEncoder early stopping) rather than a synthesized script missing load-bearing details prior runs have gotten wrong.