Setting the file. One moment.
Train Sentence Transformer Multi Dataset Example · Train Sentence Transformers · huggingface/skills · Skills Docs
Back to the top of the page Losses Cross Encoder
scripts/ train_sentence_transformer_multi_dataset_example.py
Python · 258 lines · 10 KB
16
(Variant B). Dict keys are arbitrary but must match exactly across all three
17 dicts. They show up in log output as `loss_all-nli=...`, `loss_stsb=...`.
18
19 Three reasons to use multi-dataset training:
20 - Multi-task: combine datasets with different signals (retrieval + STS +
21 classification) into a single general-purpose embedder.
22 - Data augmentation: add a supplementary dataset (STS labels alongside your
23 main retrieval pairs) as a regularizer.
24 - Domain coverage: train on several domains at once rather than sequentially.
25
26 Variant A (this script): different shapes per dataset, so each needs its own
27 matching loss. Pass `loss` as a dict. The trainer dispatches per-dataset and
28 mixing loss arities (MNRL with 3 inputs + CoSENTLoss with 2+label) is fine.
29
30 Variant B: same shape, same loss, but you want each mini-batch drawn from a
31 single domain (so MNRL in-batch negatives stay in-domain and remain genuinely
32 hard). Pass ONE loss and a dict of datasets:
33
34 train_datasets = {"medical": medical_pairs, "legal": legal_pairs, "code": code_pairs}
35 loss = MultipleNegativesRankingLoss(model)
36 trainer = SentenceTransformerTrainer(model=model, args=args,
37 train_dataset=train_datasets, loss=loss, ...)
38
39 The multi-dataset batch sampler draws each batch from a single dataset, so a
40 3-domain MNRL run gets in-domain negatives by construction. Counter-intuitive
41 benefit: DatasetDict can outperform `concatenate_datasets` even with losses
42 that don't share across the batch (e.g. LambdaLoss in cross-encoder training).
43
44 Multi-dataset samplers:
45 - `PROPORTIONAL` (default): sample from each dataset in proportion to its size.
46 Every row is seen ~once per epoch. Bias toward the largest dataset.
47 - `ROUND_ROBIN`: alternate evenly. Training stops when the SMALLEST is
48 exhausted. Equal screen-time per task.
49 Common pattern: `PROPORTIONAL` for 1 epoch, then `ROUND_ROBIN` for a second
50 if a smaller task's loss is still decreasing.
51
52 Per-dataset prompts (bi-encoder, sparse-encoder): pass `prompts={"all-nli": "",
53 "stsb": "Represent ...: ", "msmarco": {"query": "query: ", "positive":
54 "passage: ", ...}}` to TrainingArguments. The nested per-column form works for
55 bi-encoder and sparse-encoder. Cross-encoders support single-value or
56 per-dataset only. See `../references/prompts_and_instructions.md`.
57
58 Eval metric aggregation: with a dict `eval_dataset`, each dataset's loss is
59 logged separately (`eval_loss_all-nli`, `eval_loss_stsb`). The evaluator runs
60 on the full model, so its metrics aren't per-dataset unless you wrap a
61 `SequentialEvaluator` with per-dataset sub-evaluators. Set
62 `metric_for_best_model` to a single evaluator metric, NOT a per-dataset loss.
63
64 Gotchas: keys must match EXACTLY across all three dicts (train/eval/loss) or
65 training fails at step 0. `NO_DUPLICATES` + `PROPORTIONAL` works (deduplicates
66 within each batch regardless of source dataset). `ROUND_ROBIN` with uneven
67 dataset sizes means `num_train_epochs=N` is N passes over the SMALLEST. Use
68 `PROPORTIONAL` or `max_steps` if you want N passes over the largest.
69 """
70
71 from __future__ import annotations
72
73 import argparse
74 import logging
75 import os
76 from contextlib import nullcontext
77
78 import torch
79 from datasets import load_dataset
80
81 from sentence_transformers import (
82 SentenceTransformer,
83 SentenceTransformerTrainer,
84 SentenceTransformerTrainingArguments,
85 )
86 from sentence_transformers.base.sampler import BatchSamplers
87 from sentence_transformers.sentence_transformer.evaluation import (
88 EmbeddingSimilarityEvaluator,
89 NanoBEIREvaluator,
90 )
91 from sentence_transformers.sentence_transformer.losses import (
92 CoSENTLoss,
93 MultipleNegativesRankingLoss,
94 )
95 from sentence_transformers.util.similarity import SimilarityFunction
96
97
98 def autocast_ctx ():
99 """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
100 if not torch.cuda.is_available():
101 return nullcontext()
102 dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
103 return torch.autocast( "cuda" , dtype = dtype)
104
105
106 def log_trackio_dashboard ():
107 """Surface the Trackio dashboard URL so the user can watch training live."""
108 try :
109 from huggingface_hub import whoami
110
111 hf_user = whoami().get( "name" )
112 if hf_user:
113 logging.info(
114 f "Trackio dashboard (live training progress): https://huggingface.co/spaces/ { hf_user } /trackio"
115 )
116 except Exception :
117 pass
118
119
120 RUN_NAME = "mpnet-nli-stsb"
121 SMOKE_TEST = os.environ.get( "SMOKE_TEST" ) == "1"
122
123
124 def setup_logging ():
125 """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
126 os.makedirs( "logs" , exist_ok = True )
127 logging.basicConfig(
128 format = " %(asctime)s - %(message)s " ,
129 datefmt = "%Y-%m- %d %H:%M:%S" ,
130 level = logging. INFO ,
131 handlers = [logging.StreamHandler(), logging.FileHandler( f "logs/ { RUN_NAME } .log" )],
132 force = True ,
133 )
134 for noisy in ( "httpx" , "httpcore" , "huggingface_hub" , "urllib3" , "filelock" , "fsspec" ):
135 logging.getLogger(noisy).setLevel(logging. WARNING )
136 if torch.cuda.is_available():
137 torch.set_float32_matmul_precision( "high" ) # TF32 on Ampere+, no quality loss
138
139
140 def main () -> None :
141 parser = argparse.ArgumentParser()
142 parser.add_argument(
143 "--eval-only" , type = str , default = None , help = "Skip training; load this saved model and run only the evaluator."
144 )
145 cli, _ = parser.parse_known_args()
146
147 setup_logging()
148
149 if cli.eval_only:
150 logging.info( f "Eval-only mode: loading model from { cli.eval_only } " )
151 model = SentenceTransformer(cli.eval_only)
152 evaluator = NanoBEIREvaluator()
153 with autocast_ctx():
154 evaluator(model)
155 return
156
157 model = SentenceTransformer( "microsoft/mpnet-base" )
158
159 if SMOKE_TEST :
160 logging.info( "SMOKE_TEST=1: trimmed datasets; will run max_steps=1 and skip Hub push" )
161 nli_train_size = 50 if SMOKE_TEST else 50_000
162 nli_eval_size = 20 if SMOKE_TEST else 500
163 nli_train = load_dataset( "sentence-transformers/all-nli" , "triplet" , split = "train" ).select( range (nli_train_size))
164 stsb_train = load_dataset( "sentence-transformers/stsb" , split = "train" )
165 if SMOKE_TEST :
166 stsb_train = stsb_train.select( range ( min ( 50 , len (stsb_train))))
167 nli_eval = load_dataset( "sentence-transformers/all-nli" , "triplet" , split = "dev" ).select( range (nli_eval_size))
168 stsb_eval = load_dataset( "sentence-transformers/stsb" , split = "validation" )
169 if SMOKE_TEST :
170 stsb_eval = stsb_eval.select( range ( min ( 20 , len (stsb_eval))))
171
172 train_datasets = { "all-nli" : nli_train, "stsb" : stsb_train}
173 eval_datasets = { "all-nli" : nli_eval, "stsb" : stsb_eval}
174
175 losses = {
176 "all-nli" : MultipleNegativesRankingLoss(model),
177 "stsb" : CoSENTLoss(model),
178 }
179
180 evaluator = EmbeddingSimilarityEvaluator(
181 sentences1 = stsb_eval[ "sentence1" ],
182 sentences2 = stsb_eval[ "sentence2" ],
183 scores = stsb_eval[ "score" ],
184 main_similarity = SimilarityFunction. COSINE ,
185 name = "sts-dev" ,
186 )
187 logging.info( "Baseline evaluation (before training):" )
188 with autocast_ctx():
189 # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
190 baseline_eval = evaluator(model)[evaluator.primary_metric]
191 metric_key = f "eval_ { evaluator.primary_metric } "
192
193 # multi_dataset_batch_sampler defaults to PROPORTIONAL (samples each dataset
194 # in proportion to its size). To force equal alternation between datasets:
195 # from sentence_transformers.base.sampler import MultiDatasetBatchSamplers
196 # ... multi_dataset_batch_sampler=MultiDatasetBatchSamplers.ROUND_ROBIN ...
197 args = SentenceTransformerTrainingArguments(
198 output_dir = "models/mpnet-nli-stsb" ,
199 num_train_epochs = 1 ,
200 max_steps = 1 if SMOKE_TEST else - 1 ,
201 per_device_train_batch_size = 32 ,
202 per_device_eval_batch_size = 32 ,
203 learning_rate = 2e-5 ,
204 weight_decay = 0.01 ,
205 warmup_steps = 0.1 ,
206 bf16 = True ,
207 batch_sampler = BatchSamplers. NO_DUPLICATES ,
208 eval_strategy = "steps" ,
209 eval_steps = 0.1 ,
210 save_strategy = "steps" ,
211 save_steps = 0.1 ,
212 save_total_limit = 2 ,
213 logging_steps = 0.01 ,
214 logging_first_step = True ,
215 load_best_model_at_end = True ,
216 metric_for_best_model = metric_key,
217 greater_is_better = True ,
218 report_to = "none" if SMOKE_TEST else "trackio" ,
219 run_name = "mpnet-nli-stsb" ,
220 seed = 12 ,
221 )
222
223 trainer = SentenceTransformerTrainer(
224 model = model,
225 args = args,
226 train_dataset = train_datasets,
227 eval_dataset = eval_datasets,
228 loss = losses,
229 evaluator = evaluator,
230 )
231 if not SMOKE_TEST :
232 log_trackio_dashboard()
233 trainer.train()
234
235 logging.info( "Post-training evaluation:" )
236 with autocast_ctx():
237 score = evaluator(model)[evaluator.primary_metric]
238 delta = score - baseline_eval
239 verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
240 logging.info( f "VERDICT: { verdict } | score= { score :.4f} | baseline= { baseline_eval :.4f} | delta= { delta :+.4f} " )
241
242 model.save_pretrained( "models/mpnet-nli-stsb/final" )
243
244 if SMOKE_TEST :
245 logging.info( "SMOKE_TEST=1: skipping Hub push" )
246 return
247
248 try :
249 commit_url = model.push_to_hub( RUN_NAME )
250 logging.info( f "Pushed model to { commit_url.rsplit( '/commit/' , 1 )[ 0 ] } " )
251 except Exception :
252 import traceback
253
254 logging.error( f "Hub push failed: \n{ traceback.format_exc() } " )
255
256
257 if __name__ == "__main__" :
258 main()