Setting the file. One moment.
Train Sparse Encoder Distillation Example · Train Sentence Transformers · huggingface/skills · Skills Docs
Back to the top of the page Losses Cross Encoder
Next
Script Train Sparse Encoder Example
scripts/ train_sparse_encoder_distillation_example.py
Python · 262 lines · 9 KB
16
17 Data shape: `(query, positive, negative, score_diff)` where
18 `score_diff = teacher(q, pos) - teacher(q, neg)`. This script uses
19 `sentence-transformers/msmarco` (`bert-ensemble-margin-mse` subset) which has
20 precomputed teacher score diffs. To distill from your own cross-encoder
21 teacher, run a one-time teacher pass over your (q, pos, neg) triples and
22 store the per-row score diff.
23
24 Why distill SPLADE from a cross-encoder: SPLADE alone is hard to train from
25 contrastive labels because the FLOPS regularizer fights early-training signal.
26 Distilling from a strong cross-encoder gives the model a dense regression
27 target and reaches stronger nDCG faster than MNRL-only.
28
29 Run locally:
30 pip install "sentence-transformers[train]>=5.0"
31 python train_sparse_encoder_distillation_example.py
32
33 Multi-GPU:
34 accelerate launch train_sparse_encoder_distillation_example.py
35
36 Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
37 """
38
39 from __future__ import annotations
40
41 import argparse
42 import logging
43 import os
44 from contextlib import nullcontext
45
46 import torch
47 from datasets import load_dataset, load_from_disk
48
49 from sentence_transformers import (
50 SparseEncoder,
51 SparseEncoderModelCardData,
52 SparseEncoderTrainer,
53 SparseEncoderTrainingArguments,
54 )
55 from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
56 from sentence_transformers.sparse_encoder.losses import SparseMarginMSELoss, SpladeLoss
57
58
59 def autocast_ctx ():
60 """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
61 if not torch.cuda.is_available():
62 return nullcontext()
63 dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
64 return torch.autocast( "cuda" , dtype = dtype)
65
66
67 def log_trackio_dashboard ():
68 """Surface the Trackio dashboard URL so the user can watch training live."""
69 try :
70 from huggingface_hub import whoami
71
72 hf_user = whoami().get( "name" )
73 if hf_user:
74 logging.info(
75 f "Trackio dashboard (live training progress): https://huggingface.co/spaces/ { hf_user } /trackio"
76 )
77 except Exception :
78 pass
79
80
81 MODEL_NAME = "Luyu/co-condenser-marco" # MS MARCO-tuned MLM base. Very strong starting point
82 DATASET_NAME = "sentence-transformers/msmarco"
83 DATASET_SUBSET = "bert-ensemble-margin-mse"
84 TRAIN_SIZE = 100_000
85 EVAL_SIZE = 5_000
86 OUTPUT_DIR = "models/splade-msmarco-distilled"
87 RUN_NAME = "splade-msmarco-distilled"
88 DATA_CACHE = f "data/ { RUN_NAME } -resolved"
89
90 QUERY_REGULARIZER_WEIGHT = 0.1 # higher than contrastive recipe. Distillation tolerates more sparsity pressure
91 DOCUMENT_REGULARIZER_WEIGHT = 0.08
92 SMOKE_TEST = os.environ.get( "SMOKE_TEST" ) == "1"
93
94
95 def setup_logging ():
96 """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
97 os.makedirs( "logs" , exist_ok = True )
98 logging.basicConfig(
99 format = " %(asctime)s - %(message)s " ,
100 datefmt = "%Y-%m- %d %H:%M:%S" ,
101 level = logging. INFO ,
102 handlers = [logging.StreamHandler(), logging.FileHandler( f "logs/ { RUN_NAME } .log" )],
103 force = True ,
104 )
105 for noisy in ( "httpx" , "httpcore" , "huggingface_hub" , "urllib3" , "filelock" , "fsspec" ):
106 logging.getLogger(noisy).setLevel(logging. WARNING )
107 if torch.cuda.is_available():
108 torch.set_float32_matmul_precision( "high" )
109
110
111 def load_resolved_dataset ():
112 """Load (query, positive, negative, score) rows. The MSMARCO subset is keyed by
113 passage_id / query_id. Resolve to text once and cache to disk so reruns skip the work."""
114 if os.path.isdir( DATA_CACHE ):
115 logging.info( f "Loading cached resolved dataset from { DATA_CACHE } " )
116 return load_from_disk( DATA_CACHE )
117
118 logging.info( f "Resolving { DATASET_NAME } / { DATASET_SUBSET } ids -> text (one-time, cached)" )
119 corpus_ds = load_dataset( DATASET_NAME , "corpus" , split = "train" )
120 corpus = dict ( zip (corpus_ds[ "passage_id" ], corpus_ds[ "passage" ]))
121 queries_ds = load_dataset( DATASET_NAME , "queries" , split = "train" )
122 queries = dict ( zip (queries_ds[ "query_id" ], queries_ds[ "query" ]))
123 raw = load_dataset( DATASET_NAME , DATASET_SUBSET , split = "train" ).select( range ( TRAIN_SIZE + EVAL_SIZE ))
124
125 def id_to_text (batch):
126 return {
127 "query" : [queries[qid] for qid in batch[ "query_id" ]],
128 "positive" : [corpus[pid] for pid in batch[ "positive_id" ]],
129 "negative" : [corpus[pid] for pid in batch[ "negative_id" ]],
130 "score" : batch[ "score" ],
131 }
132
133 resolved = raw.map(id_to_text, batched = True , remove_columns = [ "query_id" , "positive_id" , "negative_id" ])
134 resolved.save_to_disk( DATA_CACHE )
135 return resolved
136
137
138 def main () -> None :
139 parser = argparse.ArgumentParser()
140 parser.add_argument(
141 "--eval-only" , type = str , default = None , help = "Skip training; load this saved model and run only the evaluator."
142 )
143 cli, _ = parser.parse_known_args()
144
145 setup_logging()
146
147 if cli.eval_only:
148 logging.info( f "Eval-only mode: loading model from { cli.eval_only } " )
149 model = SparseEncoder(cli.eval_only)
150 evaluator = SparseNanoBEIREvaluator( dataset_names = [ "msmarco" , "nfcorpus" , "nq" ])
151 with autocast_ctx():
152 evaluator(model)
153 return
154
155 logging.info( f "Loading base model: { MODEL_NAME } " )
156 model = SparseEncoder(
157 MODEL_NAME ,
158 model_card_data = SparseEncoderModelCardData(
159 language = "en" ,
160 license = "apache-2.0" ,
161 model_name = f "SPLADE from { MODEL_NAME .split( '/' )[ - 1 ] } distilled from MS MARCO ensemble" ,
162 ),
163 )
164 model.max_seq_length = 256
165
166 resolved = load_resolved_dataset()
167 if SMOKE_TEST :
168 logging.info( "SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push" )
169 resolved = resolved.select( range ( min ( 70 , len (resolved))))
170 eval_size = 20 if SMOKE_TEST else EVAL_SIZE
171 split = resolved.train_test_split( test_size = eval_size, seed = 12 )
172 train_dataset = split[ "train" ]
173 eval_dataset = split[ "test" ]
174 logging.info( f " train: { len (train_dataset) :,} rows | eval: { len (eval_dataset) :,} rows" )
175 logging.info( f " columns: { train_dataset.column_names } " )
176
177 loss = SpladeLoss(
178 model = model,
179 loss = SparseMarginMSELoss( model = model),
180 query_regularizer_weight = QUERY_REGULARIZER_WEIGHT ,
181 document_regularizer_weight = DOCUMENT_REGULARIZER_WEIGHT ,
182 )
183
184 evaluator = SparseNanoBEIREvaluator( dataset_names = [ "msmarco" , "nfcorpus" , "nq" ])
185 logging.info( "Baseline evaluation (fill-mask base scores near zero, confirms pipeline):" )
186 with autocast_ctx():
187 # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
188 baseline_result = evaluator(model)
189 baseline_eval = baseline_result[evaluator.primary_metric]
190 metric_key = f "eval_ { evaluator.primary_metric } "
191
192 args = SparseEncoderTrainingArguments(
193 output_dir = OUTPUT_DIR ,
194 num_train_epochs = 1 ,
195 max_steps = 1 if SMOKE_TEST else - 1 ,
196 per_device_train_batch_size = 16 ,
197 per_device_eval_batch_size = 16 ,
198 learning_rate = 2e-5 ,
199 weight_decay = 0.01 ,
200 warmup_steps = 0.1 ,
201 lr_scheduler_type = "linear" ,
202 bf16 = True ,
203 eval_strategy = "steps" ,
204 eval_steps = 0.1 ,
205 save_strategy = "steps" ,
206 save_steps = 0.1 ,
207 save_total_limit = 2 ,
208 logging_steps = 0.01 ,
209 logging_first_step = True ,
210 load_best_model_at_end = True ,
211 metric_for_best_model = metric_key,
212 greater_is_better = True ,
213 report_to = "none" if SMOKE_TEST else "trackio" ,
214 run_name = RUN_NAME ,
215 seed = 12 ,
216 )
217
218 trainer = SparseEncoderTrainer(
219 model = model,
220 args = args,
221 train_dataset = train_dataset,
222 eval_dataset = eval_dataset,
223 loss = loss,
224 evaluator = evaluator,
225 )
226 if not SMOKE_TEST :
227 log_trackio_dashboard()
228 trainer.train()
229
230 logging.info( "Post-training evaluation:" )
231 with autocast_ctx():
232 result = evaluator(model)
233 score = result[evaluator.primary_metric]
234 delta = score - baseline_eval
235 verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
236 # Active-dim keys come back name-prefixed (e.g. "NanoBEIR_..._query_active_dims"). Suffix-match for compat.
237 qad = next ((v for k, v in result.items() if k.endswith( "query_active_dims" )), "n/a" )
238 cad = next ((v for k, v in result.items() if k.endswith( "corpus_active_dims" )), "n/a" )
239 logging.info(
240 f "VERDICT: { verdict } | score= { score :.4f} | baseline= { baseline_eval :.4f} | delta= { delta :+.4f} "
241 f "| query_active= { qad } corpus_active= { cad } "
242 )
243
244 final_dir = f " { OUTPUT_DIR } /final"
245 model.save_pretrained(final_dir)
246 logging.info( f "Saved final model to { final_dir } " )
247
248 if SMOKE_TEST :
249 logging.info( "SMOKE_TEST=1: skipping Hub push" )
250 return
251
252 try :
253 commit_url = model.push_to_hub( RUN_NAME )
254 logging.info( f "Pushed model to { commit_url.rsplit( '/commit/' , 1 )[ 0 ] } " )
255 except Exception :
256 import traceback
257
258 logging.error( f "Hub push failed: \n{ traceback.format_exc() } " )
259
260
261 if __name__ == "__main__" :
262 main()