Setting the file. One moment. Losses Cross Encoder
scripts/train_cross_encoder_example.py
scripts/train_cross_encoder_example.py
Python·253 lines·9 KB
to produce the labeled training data BCE needs, starting from (question, answer)
17 pairs
18- `pos_weight=num_negatives` to offset the positive/negative imbalance
19- CrossEncoderNanoBEIREvaluator for retrieval reranking metrics
20- load_best_model_at_end on the retrieval metric
21- Auto model card + optional Hub push
22
23Run locally:
24 pip install "sentence-transformers[train]>=5.0"
25 python train_cross_encoder_example.py
26
27Multi-GPU:
28 accelerate launch train_cross_encoder_example.py
29
30Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
31"""
32
33from __future__ import annotations
34
35import argparse
36import logging
37import os
38from contextlib import nullcontext
39
40import torch
41from datasets import load_dataset, load_from_disk
42from transformers import EarlyStoppingCallback
43
44from sentence_transformers import (
45 CrossEncoder,
46 CrossEncoderModelCardData,
47 CrossEncoderTrainer,
48 CrossEncoderTrainingArguments,
49 SentenceTransformer,
50)
51from sentence_transformers.cross_encoder.evaluation import CrossEncoderNanoBEIREvaluator
52from sentence_transformers.cross_encoder.losses import BinaryCrossEntropyLoss
53from sentence_transformers.util import mine_hard_negatives
54
55
56def autocast_ctx():
57 """bf16/fp16 autocast for evaluator calls outside the trainer (which has its own autocast)."""
58 if not torch.cuda.is_available():
59 return nullcontext()
60 dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
61 return torch.autocast("cuda", dtype=dtype)
62
63
64def log_trackio_dashboard():
65 """Surface the Trackio dashboard URL so the user can watch training live."""
66 try:
67 from huggingface_hub import whoami
68
69 hf_user = whoami().get("name")
70 if hf_user:
71 logging.info(
72 f"Trackio dashboard (live training progress): https://huggingface.co/spaces/{hf_user}/trackio"
73 )
74 except Exception:
75 pass
76
77
78MODEL_NAME = "microsoft/MiniLM-L12-H384-uncased"
79DATASET_NAME = "sentence-transformers/gooaq"
80RETRIEVER_NAME = "sentence-transformers/static-retrieval-mrl-en-v1"
81TRAIN_SIZE = 100_000
82EVAL_SIZE = 1_000
83NUM_NEGATIVES = 5
84OUTPUT_DIR = "models/minilm-gooaq-ce"
85RUN_NAME = "minilm-gooaq-ce"
86HARD_NEG_CACHE = f"data/{RUN_NAME}-hard-negatives" # delete this dir to remine
87SMOKE_TEST = os.environ.get("SMOKE_TEST") == "1"
88
89
90def setup_logging():
91 """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
92 os.makedirs("logs", exist_ok=True)
93 logging.basicConfig(
94 format="%(asctime)s - %(message)s",
95 datefmt="%Y-%m-%d %H:%M:%S",
96 level=logging.INFO,
97 handlers=[logging.StreamHandler(), logging.FileHandler(f"logs/{RUN_NAME}.log")],
98 force=True,
99 )
100 for noisy in ("httpx", "httpcore", "huggingface_hub", "urllib3", "filelock", "fsspec"):
101 logging.getLogger(noisy).setLevel(logging.WARNING)
102 if torch.cuda.is_available():
103 torch.set_float32_matmul_precision("high") # TF32 on Ampere+, no quality loss
104
105
106def main() -> None:
107 parser = argparse.ArgumentParser()
108 parser.add_argument(
109 "--eval-only", type=str, default=None, help="Skip training; load this saved model and run only the evaluator."
110 )
111 cli, _ = parser.parse_known_args()
112
113 setup_logging()
114
115 if cli.eval_only:
116 logging.info(f"Eval-only mode: loading model from {cli.eval_only}")
117 model = CrossEncoder(cli.eval_only)
118 evaluator = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
119 with autocast_ctx():
120 evaluator(model)
121 return
122
123 logging.info(f"Loading base model: {MODEL_NAME}")
124 model = CrossEncoder(
125 MODEL_NAME,
126 num_labels=1,
127 model_card_data=CrossEncoderModelCardData(
128 language="en",
129 license="apache-2.0",
130 model_name=f"{MODEL_NAME.split('/')[-1]} reranker finetuned on GooAQ",
131 ),
132 )
133
134 logging.info(f"Loading dataset: {DATASET_NAME}")
135 train_size = 50 if SMOKE_TEST else TRAIN_SIZE
136 eval_size = 20 if SMOKE_TEST else EVAL_SIZE
137 if SMOKE_TEST:
138 logging.info("SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push")
139 pairs = load_dataset(DATASET_NAME, split="train").select(range(train_size + eval_size))
140
141 if os.path.isdir(HARD_NEG_CACHE):
142 logging.info(f"Loading cached mined hard negatives from {HARD_NEG_CACHE}")
143 labeled = load_from_disk(HARD_NEG_CACHE)
144 else:
145 logging.info(f"Mining hard negatives with {RETRIEVER_NAME}")
146 retriever = SentenceTransformer(RETRIEVER_NAME)
147 labeled = mine_hard_negatives(
148 dataset=pairs,
149 model=retriever,
150 num_negatives=NUM_NEGATIVES,
151 range_min=10, # skip the top-10 (likely to contain true positives)
152 range_max=100,
153 sampling_strategy="top",
154 output_format="labeled-pair",
155 use_faiss=True,
156 )
157 labeled.save_to_disk(HARD_NEG_CACHE)
158 logging.info(f"Saved mined dataset to {HARD_NEG_CACHE} (delete to remine)")
159 del retriever
160 torch.cuda.empty_cache()
161 # EVAL_SIZE here counts labeled-pair rows, not distinct queries: each
162 # query contributes 1 positive + NUM_NEGATIVES negatives, so e.g. 1000
163 # rows is approximately 1000 / (1 + NUM_NEGATIVES) distinct queries.
164 split = labeled.train_test_split(test_size=eval_size, seed=12)
165 train_dataset = split["train"]
166 eval_dataset = split["test"]
167 logging.info(f" train: {len(train_dataset):,} rows | eval: {len(eval_dataset):,} rows")
168 logging.info(f" columns: {train_dataset.column_names}")
169
170 # pos_weight = negatives / positives, derived from the actual label distribution
171 # so it stays correct if rows get filtered or the mining ratio drifts.
172 n_pos = sum(1 for label in train_dataset["label"] if label > 0.5)
173 n_neg = len(train_dataset) - n_pos
174 pos_weight_value = n_neg / max(n_pos, 1)
175 logging.info(f" positives: {n_pos:,} | negatives: {n_neg:,} | pos_weight: {pos_weight_value:.2f}")
176 loss = BinaryCrossEntropyLoss(model, pos_weight=torch.tensor(pos_weight_value))
177
178 evaluator = CrossEncoderNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"])
179 logging.info("Baseline evaluation:")
180 with autocast_ctx():
181 # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
182 baseline_eval = evaluator(model)[evaluator.primary_metric]
183 metric_key = f"eval_{evaluator.primary_metric}"
184
185 args = CrossEncoderTrainingArguments(
186 output_dir=OUTPUT_DIR,
187 num_train_epochs=1,
188 max_steps=1 if SMOKE_TEST else -1,
189 per_device_train_batch_size=64,
190 per_device_eval_batch_size=64,
191 learning_rate=2e-5,
192 weight_decay=0.01,
193 warmup_steps=0.1,
194 lr_scheduler_type="linear",
195 bf16=True,
196 eval_strategy="steps",
197 eval_steps=0.1,
198 save_strategy="steps",
199 save_steps=0.1,
200 save_total_limit=2,
201 logging_steps=0.01,
202 logging_first_step=True,
203 load_best_model_at_end=True,
204 metric_for_best_model=metric_key,
205 greater_is_better=True,
206 report_to="none" if SMOKE_TEST else "trackio",
207 run_name=RUN_NAME,
208 seed=12,
209 )
210
211 # EarlyStoppingCallback earns its keep for cross-encoders: CE rerankers
212 # typically peak mid-training and then degrade, so stopping at the best
213 # eval checkpoint is load-bearing (unlike bi-encoders, which tend to
214 # plateau rather than regress).
215 trainer = CrossEncoderTrainer(
216 model=model,
217 args=args,
218 train_dataset=train_dataset,
219 eval_dataset=eval_dataset,
220 loss=loss,
221 evaluator=evaluator,
222 callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
223 )
224 if not SMOKE_TEST:
225 log_trackio_dashboard()
226 trainer.train()
227
228 logging.info("Post-training evaluation:")
229 with autocast_ctx():
230 score = evaluator(model)[evaluator.primary_metric]
231 delta = score - baseline_eval
232 verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
233 logging.info(f"VERDICT: {verdict} | score={score:.4f} | baseline={baseline_eval:.4f} | delta={delta:+.4f}")
234
235 final_dir = f"{OUTPUT_DIR}/final"
236 model.save_pretrained(final_dir)
237 logging.info(f"Saved final model to {final_dir}")
238
239 if SMOKE_TEST:
240 logging.info("SMOKE_TEST=1: skipping Hub push")
241 return
242
243 try:
244 commit_url = model.push_to_hub(RUN_NAME)
245 logging.info(f"Pushed model to {commit_url.rsplit('/commit/', 1)[0]}")
246 except Exception:
247 import traceback
248
249 logging.error(f"Hub push failed:\n{traceback.format_exc()}")
250
251
252if __name__ == "__main__":
253 main()