Setting the file. One moment. Losses Cross Encoder
scripts/train_multi_vector_encoder_example.py
Python·240 lines·9 KB
16
- BatchSamplers.NO_DUPLICATES (critical for MNRL-family)
17- load_best_model_at_end with a retrieval metric
18- Auto model card + optional Hub push
19
20Runs identically in two modes:
21
22 # Local
23 pip install "sentence-transformers[train]>=6.0"
24 python train_multi_vector_encoder_example.py
25
26 # Or with uv (no explicit install needed)
27 uv run train_multi_vector_encoder_example.py
28
29 # Multi-GPU
30 accelerate launch train_multi_vector_encoder_example.py
31
32 # Hugging Face Jobs (paste the entire file contents as `script`)
33 hf_jobs("uv", {
34 "script": "<contents of this file>",
35 "flavor": "a10g-large",
36 "timeout": "3h",
37 "secrets": {"HF_TOKEN": "$HF_TOKEN"},
38 })
39
40Adjust MODEL_NAME, DATASET_NAME, OUTPUT_DIR, RUN_NAME at the top of the script.
41Default Hub push: at end of run, public, under your authenticated user as
42`{user}/{RUN_NAME}`, wrapped in try/except. To skip the push, comment out the
43push_to_hub call. For HF Jobs (ephemeral env), also enable in-trainer push:
44add `push_to_hub=True`, `hub_model_id=RUN_NAME`, `hub_strategy="every_save"`
45to TrainingArguments.
46"""
47
48from __future__ import annotations
49
50import argparse
51import logging
52import os
53from contextlib import nullcontext
54
55import torch
56from datasets import load_dataset
57
58from sentence_transformers import (
59 MultiVectorEncoder,
60 MultiVectorEncoderModelCardData,
61 MultiVectorEncoderTrainer,
62 MultiVectorEncoderTrainingArguments,
63)
64from sentence_transformers.base.sampler import BatchSamplers
65from sentence_transformers.multi_vector_encoder.evaluation import MultiVectorNanoBEIREvaluator
66from sentence_transformers.multi_vector_encoder.losses import MultiVectorMultipleNegativesRankingLoss
67
68
69def autocast_ctx():
70 """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
71 if not torch.cuda.is_available():
72 return nullcontext()
73 dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
74 return torch.autocast("cuda", dtype=dtype)
75
76
77def log_trackio_dashboard():
78 """Surface the Trackio dashboard URL so the user can watch training live."""
79 try:
80 from huggingface_hub import whoami
81
82 hf_user = whoami().get("name")
83 if hf_user:
84 logging.info(
85 f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
86 )
87 except Exception:
88 pass
89
90
91# Bare HF encoder. MVE grafts a fresh 128-dim token projection on top.
92MODEL_NAME = "answerdotai/ModernBERT-base"
93DATASET_NAME = "sentence-transformers/msmarco-bm25"
94DATASET_SUBSET = "triplet"
95TRAIN_SIZE = 50_000
96EVAL_SIZE = 1_000
97OUTPUT_DIR = "models/modernbert-base-msmarco-colbert"
98RUN_NAME = "modernbert-base-msmarco-colbert"
99SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"
100
101
102def setup_logging():
103 """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
104 os.makedirs("logs", exist_ok=True)
105 logging.basicConfig(
106 format="%(asctime)s - %(message)s",
107 datefmt="%Y-%m-%d %H:%M:%S",
108 level=logging.INFO,
109 handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
110 force=True,
111 )
112 for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
113 logging.getLogger(noisy).setLevel(logging.WARNING)
114 if torch.cuda.is_available():
115 torch.set_float32_matmul_precision("high") # TF32 on Ampere+, no quality loss
116
117
118def main() -> None:
119 parser = argparse.ArgumentParser()
120 parser.add_argument(
121 "--eval-only", type=str, default=None, help="Skip training, load this saved model and run only the evaluator."
122 )
123 cli, _ = parser.parse_known_args()
124
125 setup_logging()
126
127 if cli.eval_only:
128 logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
129 model = MultiVectorEncoder(cli.eval_only)
130 evaluator = MultiVectorNanoBEIREvaluator(batch_size=16)
131 with autocast_ctx():
132 evaluator(model)
133 return
134
135 logging.info(f"Loading base model: {MODEL_NAME}")
136 # Load in fp32, autocast bf16 during forward. torch_dtype=bfloat16 breaks the optimizer's precision.
137 model = MultiVectorEncoder(
138 MODEL_NAME,
139 model_card_data=MultiVectorEncoderModelCardData(
140 language="en",
141 license="apache-2.0",
142 model_name=f"ColBERT {MODEL_NAME.split('/')[-1]} trained on MS MARCO",
143 ),
144 model_kwargs={"torch_dtype": "float32"},
145 )
146
147 logging.info(f"Loading dataset: {DATASET_NAME} ({DATASET_SUBSET})")
148 train_size = 50 if SMOKE_TEST else TRAIN_SIZE
149 eval_size = 20 if SMOKE_TEST else EVAL_SIZE
150 full_dataset = load_dataset(DATASET_NAME, DATASET_SUBSET, split="train").select(range(train_size + eval_size))
151 dataset_dict = full_dataset.train_test_split(test_size=eval_size, seed=12)
152 train_dataset = dataset_dict["train"]
153 eval_dataset = dataset_dict["test"]
154 if SMOKE_TEST:
155 logging.info("SMOKE_TEST=1: trimmed dataset, will run max_steps=1 and skip Hub push")
156 logging.info(f" train: {len(train_dataset):,} examples")
157 logging.info(f" eval: {len(eval_dataset):,} examples")
158
159 # MNRL with in-batch negatives + explicit hard negatives. scale=1.0 (default) is correct for
160 # unnormalized MaxSim: do not copy scale=20.0 from bi-encoder MNRL. Length-normalized MeanMaxSim
161 # scoring instead wants a scale of roughly the average query length.
162 loss = MultiVectorMultipleNegativesRankingLoss(model=model)
163
164 # Cheap in-training evaluator on 3 datasets (drives `load_best_model_at_end`).
165 # The end-of-run evaluator below covers the full 13-dataset suite for the VERDICT delta.
166 evaluator = MultiVectorNanoBEIREvaluator(dataset_names=["msmarco", "nq", "fiqa2018"], batch_size=16)
167 test_evaluator = MultiVectorNanoBEIREvaluator(show_progress_bar=True, batch_size=16)
168 logging.info("Baseline evaluation (before training):")
169 with autocast_ctx():
170 # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
171 evaluator(model) # populates primary_metric on the training-time evaluator
172 # Baseline on the full suite so the VERDICT delta is apples-to-apples with the post-training score.
173 baseline_eval = test_evaluator(model)[test_evaluator.primary_metric]
174 metric_key = f"eval_{evaluator.primary_metric}"
175
176 args = MultiVectorEncoderTrainingArguments(
177 output_dir=OUTPUT_DIR,
178 num_train_epochs=1,
179 max_steps=1 if SMOKE_TEST else -1,
180 per_device_train_batch_size=32,
181 per_device_eval_batch_size=32,
182 learning_rate=3e-5,
183 weight_decay=0.01,
184 warmup_steps=0.05,
185 lr_scheduler_type="linear",
186 bf16=True,
187 batch_sampler=BatchSamplers.NO_DUPLICATES,
188 eval_strategy="steps",
189 eval_steps=0.1,
190 save_strategy="steps",
191 save_steps=0.1,
192 save_total_limit=2,
193 logging_steps=0.01,
194 logging_first_step=True,
195 load_best_model_at_end=True,
196 metric_for_best_model=metric_key,
197 greater_is_better=True,
198 report_to="none" if SMOKE_TEST else "trackio",
199 run_name=RUN_NAME,
200 seed=12,
201 )
202
203 trainer = MultiVectorEncoderTrainer(
204 model=model,
205 args=args,
206 train_dataset=train_dataset,
207 eval_dataset=eval_dataset,
208 loss=loss,
209 evaluator=evaluator,
210 )
211 if not SMOKE_TEST:
212 log_trackio_dashboard()
213 trainer.train()
214
215 logging.info("Post-training evaluation (full NanoBEIR suite):")
216 with autocast_ctx():
217 score = test_evaluator(model)[test_evaluator.primary_metric]
218 delta = score - baseline_eval
219 verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
220 logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")
221
222 final_dir = f"{OUTPUT_DIR}/final"
223 model.save_pretrained(final_dir)
224 logging.info(f"Saved final model to {final_dir}")
225
226 if SMOKE_TEST:
227 logging.info("SMOKE_TEST=1: skipping Hub push")
228 return
229
230 try:
231 commit_url = model.push_to_hub(RUN_NAME) # public by default, uses your authenticated user
232 logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
233 except Exception:
234 import traceback
235
236 logging.error(f"Hub push failed:\n{traceback.format_exc()}")
237
238
239if __name__ == "__main__":
240 main()