Setting the file. One moment.
Train Cross Encoder Listwise Example · Train Sentence Transformers · huggingface/skills · Skills Docs
Back to the top of the page Losses Cross Encoder
scripts/train_cross_encoder_listwise_example.py
scripts/ train_cross_encoder_listwise_example.py
Python · 282 lines · 10 KB
16
relevance, and you want a stronger ranker than pointwise BCE.
17
18 Data shape: `(query, [doc_1, ..., doc_K], [score_1, ..., score_K])` per row.
19 This script builds it via `mine_hard_negatives(..., output_format="labeled-list")`
20 starting from `(question, answer)` pairs: each row gets the positive plus K
21 hard negatives, with binary scores (1 for positive, 0 for negatives).
22
23 CRITICAL: `activation_fn=nn.Identity()` is mandatory for LambdaLoss / ListNet /
24 ListMLE / PListMLE / RankNet / MarginMSE / MSE (anything that's not
25 `BinaryCrossEntropyLoss` or `CrossEntropyLoss`). The default `Sigmoid` (with
26 `num_labels=1`) saturates raw logits >5 to ~1.0 inside `predict()`, silently
27 collapsing eval ranking. See `../references/troubleshooting.md` ("CrossEncoder
28 eval nDCG crashes after distillation / listwise / pairwise training").
29
30 OOM recovery for LambdaLoss: drop `mini_batch_size` first (chunking inside the
31 loss preserves the K-list semantic), then `per_device_train_batch_size` paired
32 with `gradient_accumulation_steps`, then reduce K (the per-query candidate-list
33 length) only as a last resort. Lowering K changes the experiment.
34
35 Run locally:
36 pip install "sentence-transformers[train]>=5.0"
37 python train_cross_encoder_listwise_example.py
38
39 Multi-GPU:
40 accelerate launch train_cross_encoder_listwise_example.py
41
42 Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
43 """
44
45 from __future__ import annotations
46
47 import argparse
48 import logging
49 import os
50 from contextlib import nullcontext
51
52 import torch
53 import torch.nn as nn
54 from datasets import load_dataset, load_from_disk
55 from transformers import EarlyStoppingCallback
56
57 from sentence_transformers import (
58 CrossEncoder,
59 CrossEncoderModelCardData,
60 CrossEncoderTrainer,
61 CrossEncoderTrainingArguments,
62 SentenceTransformer,
63 )
64 from sentence_transformers.base.evaluation import SequentialEvaluator
65 from sentence_transformers.cross_encoder.evaluation import (
66 CrossEncoderNanoBEIREvaluator,
67 CrossEncoderRerankingEvaluator,
68 )
69 from sentence_transformers.cross_encoder.losses import LambdaLoss
70 from sentence_transformers.util import mine_hard_negatives
71
72
73 def autocast_ctx ():
74 """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
75 if not torch.cuda.is_available():
76 return nullcontext()
77 dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
78 return torch.autocast( "cuda" , dtype = dtype)
79
80
81 def log_trackio_dashboard ():
82 """Surface the Trackio dashboard URL so the user can watch training live."""
83 try :
84 from huggingface_hub import whoami
85
86 hf_user = whoami().get( "name" )
87 if hf_user:
88 logging.info(
89 f "Trackio dashboard (live training progress): https://huggingface.co/spaces/ { hf_user } /trackio"
90 )
91 except Exception :
92 pass
93
94
95 MODEL_NAME = "answerdotai/ModernBERT-base"
96 DATASET_NAME = "sentence-transformers/gooaq"
97 RETRIEVER_NAME = "sentence-transformers/static-retrieval-mrl-en-v1"
98 TRAIN_SIZE = 100_000
99 EVAL_SIZE = 1_000
100 NUM_NEGATIVES = 7 # K-1 negatives + 1 positive per row
101 EVAL_RERANK_DEPTH = 30 # candidates per query in the in-domain eval set
102 OUTPUT_DIR = "models/modernbert-gooaq-lambda"
103 RUN_NAME = "modernbert-gooaq-lambda"
104 HARD_NEG_CACHE = f "data/ { RUN_NAME } -hard-negatives"
105 HARD_EVAL_CACHE = f "data/ { RUN_NAME } -hard-eval"
106 SMOKE_TEST = os.environ.get( "SMOKE_TEST" ) == "1"
107
108
109 def setup_logging ():
110 """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
111 os.makedirs( "logs" , exist_ok = True )
112 logging.basicConfig(
113 format = " %(asctime)s - %(message)s " ,
114 datefmt = "%Y-%m- %d %H:%M:%S" ,
115 level = logging. INFO ,
116 handlers = [logging.StreamHandler(), logging.FileHandler( f "logs/ { RUN_NAME } .log" )],
117 force = True ,
118 )
119 for noisy in ( "httpx" , "httpcore" , "huggingface_hub" , "urllib3" , "filelock" , "fsspec" ):
120 logging.getLogger(noisy).setLevel(logging. WARNING )
121 if torch.cuda.is_available():
122 torch.set_float32_matmul_precision( "high" )
123
124
125 def main () -> None :
126 parser = argparse.ArgumentParser()
127 parser.add_argument(
128 "--eval-only" , type = str , default = None , help = "Skip training; load this saved model and run only the evaluator."
129 )
130 cli, _ = parser.parse_known_args()
131
132 setup_logging()
133
134 if cli.eval_only:
135 logging.info( f "Eval-only mode: loading model from { cli.eval_only } " )
136 model = CrossEncoder(cli.eval_only)
137 evaluator = CrossEncoderNanoBEIREvaluator( dataset_names = [ "msmarco" , "nfcorpus" , "nq" ])
138 with autocast_ctx():
139 evaluator(model)
140 return
141
142 logging.info( f "Loading base model: { MODEL_NAME } " )
143 model = CrossEncoder(
144 MODEL_NAME ,
145 num_labels = 1 ,
146 activation_fn = nn.Identity(), # Mandatory for LambdaLoss. Sigmoid would saturate eval logits.
147 model_card_data = CrossEncoderModelCardData(
148 language = "en" ,
149 license = "apache-2.0" ,
150 model_name = f " { MODEL_NAME .split( '/' )[ - 1 ] } reranker trained with LambdaLoss on GooAQ" ,
151 ),
152 )
153 # ModernBERT defaults to max_seq_length=8192, which allocates activation memory
154 # for 8192-token sequences regardless of input length. Pin to a (q, doc) cap.
155 model.max_seq_length = 512
156
157 full_dataset = load_dataset( DATASET_NAME , split = "train" ).select( range ( TRAIN_SIZE ))
158 split = full_dataset.train_test_split( test_size = EVAL_SIZE , seed = 12 )
159 train_pairs, eval_pairs = split[ "train" ], split[ "test" ]
160
161 if os.path.isdir( HARD_NEG_CACHE ) and os.path.isdir( HARD_EVAL_CACHE ):
162 logging.info( "Loading cached mined hard-negative datasets" )
163 hard_train = load_from_disk( HARD_NEG_CACHE )
164 hard_eval = load_from_disk( HARD_EVAL_CACHE )
165 else :
166 logging.info( f "Mining hard negatives with { RETRIEVER_NAME } " )
167 retriever = SentenceTransformer( RETRIEVER_NAME )
168 hard_train = mine_hard_negatives(
169 train_pairs,
170 retriever,
171 num_negatives = NUM_NEGATIVES ,
172 range_min = 10 ,
173 range_max = 100 ,
174 sampling_strategy = "top" ,
175 output_format = "labeled-list" , # Listwise: (query, [docs], [scores])
176 use_faiss = True ,
177 batch_size = 4096 ,
178 )
179 hard_eval = mine_hard_negatives(
180 eval_pairs,
181 retriever,
182 corpus = full_dataset[ "answer" ],
183 num_negatives = EVAL_RERANK_DEPTH ,
184 output_format = "n-tuple" ,
185 use_faiss = True ,
186 batch_size = 4096 ,
187 )
188 hard_train.save_to_disk( HARD_NEG_CACHE )
189 hard_eval.save_to_disk( HARD_EVAL_CACHE )
190 del retriever
191 torch.cuda.empty_cache()
192 if SMOKE_TEST :
193 logging.info( "SMOKE_TEST=1: trimmed mined datasets; will run max_steps=1 and skip Hub push" )
194 hard_train = hard_train.select( range ( min ( 50 , len (hard_train))))
195 hard_eval = hard_eval.select( range ( min ( 20 , len (hard_eval))))
196 logging.info( f " train: { len (hard_train) :,} rows | columns: { hard_train.column_names } " )
197
198 loss = LambdaLoss( model = model, mini_batch_size = 16 ) # mini_batch_size: drop first if OOM
199
200 nano_beir = CrossEncoderNanoBEIREvaluator( dataset_names = [ "msmarco" , "nfcorpus" , "nq" ])
201 # Pure reranker quality: positive is in `documents` and `always_rerank_positives=True` (default).
202 in_domain = CrossEncoderRerankingEvaluator(
203 samples = [
204 {
205 "query" : row[ "question" ],
206 "positive" : [row[ "answer" ]],
207 "documents" : [row[ "answer" ]] + [row[col] for col in hard_eval.column_names[ 2 :]],
208 }
209 for row in hard_eval
210 ],
211 batch_size = 64 ,
212 name = "gooaq-dev" ,
213 )
214 evaluator = SequentialEvaluator([in_domain, nano_beir])
215 logging.info( "Baseline evaluation:" )
216 with autocast_ctx():
217 baseline_eval = evaluator(model)[in_domain.primary_metric]
218
219 args = CrossEncoderTrainingArguments(
220 output_dir = OUTPUT_DIR ,
221 num_train_epochs = 1 ,
222 max_steps = 1 if SMOKE_TEST else - 1 ,
223 per_device_train_batch_size = 64 ,
224 per_device_eval_batch_size = 64 ,
225 learning_rate = 2e-5 ,
226 weight_decay = 0.01 ,
227 warmup_steps = 0.1 ,
228 lr_scheduler_type = "linear" ,
229 bf16 = True ,
230 eval_strategy = "steps" ,
231 eval_steps = 0.1 ,
232 save_strategy = "steps" ,
233 save_steps = 0.1 ,
234 save_total_limit = 2 ,
235 logging_steps = 0.01 ,
236 logging_first_step = True ,
237 load_best_model_at_end = True ,
238 metric_for_best_model = f "eval_ { in_domain.primary_metric } " , # in-domain reranker > NanoBEIR
239 greater_is_better = True ,
240 report_to = "none" if SMOKE_TEST else "trackio" ,
241 run_name = RUN_NAME ,
242 seed = 12 ,
243 )
244
245 trainer = CrossEncoderTrainer(
246 model = model,
247 args = args,
248 train_dataset = hard_train,
249 loss = loss,
250 evaluator = evaluator,
251 callbacks = [EarlyStoppingCallback( early_stopping_patience = 3 )],
252 )
253 if not SMOKE_TEST :
254 log_trackio_dashboard()
255 trainer.train()
256
257 logging.info( "Post-training evaluation:" )
258 with autocast_ctx():
259 score = evaluator(model)[in_domain.primary_metric]
260 delta = score - baseline_eval
261 verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
262 logging.info( f "VERDICT: { verdict } | score= { score :.4f} | baseline= { baseline_eval :.4f} | delta= { delta :+.4f} " )
263
264 final_dir = f " { OUTPUT_DIR } /final"
265 model.save_pretrained(final_dir)
266 logging.info( f "Saved final model to { final_dir } " )
267
268 if SMOKE_TEST :
269 logging.info( "SMOKE_TEST=1: skipping Hub push" )
270 return
271
272 try :
273 commit_url = model.push_to_hub( RUN_NAME )
274 logging.info( f "Pushed model to { commit_url.rsplit( '/commit/' , 1 )[ 0 ] } " )
275 except Exception :
276 import traceback
277
278 logging.error( f "Hub push failed: \n{ traceback.format_exc() } " )
279
280
281 if __name__ == "__main__" :
282 main()