Setting the file. One moment.
Train Sentence Transformer Static Embedding Example · Train Sentence Transformers · huggingface/skills · Skills Docs
Back to the top of the page Losses Cross Encoder
scripts/ train_sentence_transformer_static_embedding_example.py
Python · 273 lines · 10 KB
the tokens of an input. No transformer, no attention. Inference is ~20x faster
17 on GPU and ~80x faster on CPU than a small encoder, with surprisingly competitive
18 quality on retrieval benchmarks when trained on >=1M contrastive pairs.
19
20 Two init paths via `WARMSTART` constant:
21 - `WARMSTART=False` (default): random init. Use with >=1M contrastive pairs.
22 Reaches a higher ceiling than warm-start when given enough data. The default
23 dataset below (GooAQ, ~3M pairs) is comfortably in this regime.
24 - `WARMSTART=True`: `StaticEmbedding.from_model2vec(...)` (distil from a
25 model2vec checkpoint). Flip to True if you swap in a smaller dataset (<1M
26 pairs). Converges faster and reaches better quality at lower data scales.
27
28 Demonstrates:
29 - MultipleNegativesRankingLoss wrapped in MatryoshkaLoss for nested embedding dims
30 - Large batch size (1024+) with a high LR (~2e-1 for random init, ~5e-2 for warm-
31 start) since the loss surface for a token-bag is much flatter than for a
32 pretrained encoder
33 - BatchSamplers.NO_DUPLICATES (load-bearing for in-batch negatives with duplicated
34 anchors)
35 - NanoBEIREvaluator at full embedding dim
36 - Auto model card + optional Hub push
37
38 Run locally (CPU works for inference, but training needs a GPU for batch=1024+):
39 pip install "sentence-transformers[train]>=5.0"
40 python train_sentence_transformer_static_embedding_example.py
41
42 Multi-GPU:
43 accelerate launch train_sentence_transformer_static_embedding_example.py
44
45 Hugging Face Jobs: paste this file's contents as the `script` in hf_jobs(...).
46
47 References:
48 - HF blog post: https://huggingface.co/blog/static-embeddings
49 - Module docs: sentence_transformers.sentence_transformer.modules.StaticEmbedding
50 """
51
52 from __future__ import annotations
53
54 import argparse
55 import logging
56 import os
57 from contextlib import nullcontext
58
59 import datasets
60 import torch
61 from datasets import load_dataset
62 from tokenizers import Tokenizer
63
64 from sentence_transformers import (
65 SentenceTransformer,
66 SentenceTransformerModelCardData,
67 SentenceTransformerTrainer,
68 SentenceTransformerTrainingArguments,
69 )
70 from sentence_transformers.base.sampler import BatchSamplers
71 from sentence_transformers.sentence_transformer.evaluation import NanoBEIREvaluator
72 from sentence_transformers.sentence_transformer.losses import (
73 MatryoshkaLoss,
74 MultipleNegativesRankingLoss,
75 )
76 from sentence_transformers.sentence_transformer.modules import StaticEmbedding
77
78
79 def autocast_ctx ():
80 """bf16/fp16 autocast for evaluator calls outside the trainer."""
81 if not torch.cuda.is_available():
82 return nullcontext()
83 dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
84 return torch.autocast( "cuda" , dtype = dtype)
85
86
87 def log_trackio_dashboard ():
88 """Surface the Trackio dashboard URL so the user can watch training live."""
89 try :
90 from huggingface_hub import whoami
91
92 hf_user = whoami().get( "name" )
93 if hf_user:
94 logging.info(
95 f "Trackio dashboard (live training progress): https://huggingface.co/spaces/ { hf_user } /trackio"
96 )
97 except Exception :
98 pass
99
100
101 TOKENIZER_NAME = "google-bert/bert-base-uncased"
102 EMBEDDING_DIM = 1024
103 MATRYOSHKA_DIMS = [ 1024 , 512 , 256 , 128 , 64 , 32 ] # ordered largest-first per MatryoshkaLoss
104
105 # False: random init (recommended for >=1M pairs, reaches a higher ceiling).
106 # True: warm-start from a model2vec checkpoint (recommended for <1M pairs).
107 # Default False because the example dataset (GooAQ, ~3M pairs) is well above the threshold.
108 WARMSTART = False
109 WARMSTART_MODEL2VEC = "minishlab/potion-base-8M"
110
111 OUTPUT_DIR = "models/static-embedding-bert-uncased"
112 RUN_NAME = "static-embedding-bert-uncased"
113 SMOKE_TEST = os.environ.get( "SMOKE_TEST" ) == "1"
114
115
116 def setup_logging ():
117 """Configure logging + TF32. Tees to logs/{RUN_NAME}.log and silences HTTP spam."""
118 os.makedirs( "logs" , exist_ok = True )
119 logging.basicConfig(
120 format = " %(asctime)s - %(message)s " ,
121 datefmt = "%Y-%m- %d %H:%M:%S" ,
122 level = logging. INFO ,
123 handlers = [logging.StreamHandler(), logging.FileHandler( f "logs/ { RUN_NAME } .log" )],
124 force = True ,
125 )
126 for noisy in ( "httpx" , "httpcore" , "huggingface_hub" , "urllib3" , "filelock" , "fsspec" ):
127 logging.getLogger(noisy).setLevel(logging. WARNING )
128 if torch.cuda.is_available():
129 torch.set_float32_matmul_precision( "high" )
130
131
132 def load_pair_dataset () -> datasets.Dataset:
133 """Load a contrastive-pair dataset for training.
134
135 StaticEmbedding starts from random initialization, so it needs *a lot* of
136 contrastive signal to converge. GooAQ alone provides ~3M (question, answer)
137 pairs, comfortably over the >=1M threshold below which a warm-start would
138 beat random init. For stronger production models, concatenate more sources
139 (NaturalQuestions, MSMARCO, MIRACL, etc.) and shuffle, in the same family of
140 sources used in `sentence-transformers/static-retrieval-mrl-en-v1`.
141 """
142 return (
143 load_dataset( "sentence-transformers/gooaq" , split = "train" )
144 .rename_columns({ "question" : "anchor" , "answer" : "positive" })
145 .select_columns([ "anchor" , "positive" ])
146 )
147
148
149 def main () -> None :
150 parser = argparse.ArgumentParser()
151 parser.add_argument(
152 "--eval-only" , type = str , default = None , help = "Skip training; load this saved model and run only the evaluator."
153 )
154 cli, _ = parser.parse_known_args()
155
156 setup_logging()
157
158 if cli.eval_only:
159 logging.info( f "Eval-only mode: loading model from { cli.eval_only } " )
160 model = SentenceTransformer(cli.eval_only)
161 evaluator = NanoBEIREvaluator( dataset_names = [ "msmarco" , "nfcorpus" , "nq" ])
162 with autocast_ctx():
163 evaluator(model)
164 return
165
166 if WARMSTART :
167 logging.info( f "Warm-starting StaticEmbedding from model2vec: { WARMSTART_MODEL2VEC } " )
168 # `StaticEmbedding.from_distillation("<bi-encoder>", vocabulary=...)` is the
169 # alternative warm-start path (distil from a stronger teacher's vectors). Pick
170 # one. model2vec is faster to load and converges quickly on smaller datasets.
171 static_embedding = StaticEmbedding.from_model2vec( WARMSTART_MODEL2VEC )
172 else :
173 logging.info( f "Random-init StaticEmbedding from { TOKENIZER_NAME } tokenizer (dim= { EMBEDDING_DIM } )" )
174 tokenizer = Tokenizer.from_pretrained( TOKENIZER_NAME )
175 static_embedding = StaticEmbedding(tokenizer, embedding_dim = EMBEDDING_DIM )
176 model = SentenceTransformer(
177 modules = [static_embedding],
178 model_card_data = SentenceTransformerModelCardData(
179 language = "en" ,
180 license = "apache-2.0" ,
181 model_name = f "Static embedding ( { EMBEDDING_DIM } d) trained on contrastive pairs" ,
182 ),
183 )
184
185 logging.info( "Loading + concatenating training datasets" )
186 full = load_pair_dataset()
187 if SMOKE_TEST :
188 logging.info( "SMOKE_TEST=1: trimmed dataset; will run max_steps=1 and skip Hub push" )
189 full = full.select( range ( min ( 200 , len (full))))
190 eval_size = 20 if SMOKE_TEST else 10_000
191 split = full.train_test_split( test_size = eval_size, seed = 12 )
192 train_dataset = split[ "train" ]
193 eval_dataset = split[ "test" ]
194 logging.info( f " train: { len (train_dataset) :,} rows | eval: { len (eval_dataset) :,} rows" )
195 logging.info( f " columns: { train_dataset.column_names } " )
196
197 inner = MultipleNegativesRankingLoss(model)
198 loss = MatryoshkaLoss(model, inner, matryoshka_dims = MATRYOSHKA_DIMS )
199
200 evaluator = NanoBEIREvaluator( dataset_names = [ "msmarco" , "nfcorpus" , "nq" ])
201 logging.info( "Baseline evaluation (random init scores near zero; warm-start scores 0.3+):" )
202 with autocast_ctx():
203 # Must run before deriving metric_key: evaluator(model) mutates primary_metric to add the name_ prefix.
204 baseline_eval = evaluator(model)[evaluator.primary_metric]
205 metric_key = f "eval_ { evaluator.primary_metric } "
206
207 args = SentenceTransformerTrainingArguments(
208 output_dir = OUTPUT_DIR ,
209 num_train_epochs = 1 ,
210 max_steps = 1 if SMOKE_TEST else - 1 ,
211 per_device_train_batch_size = 2048 ,
212 per_device_eval_batch_size = 2048 ,
213 learning_rate = 5e-2
214 if WARMSTART
215 else 2e-1 , # warm-start needs less LR. Both far higher than encoder fine-tuning
216 weight_decay = 0.0 , # weight decay on a token-bag is usually harmful
217 warmup_steps = 0.1 ,
218 lr_scheduler_type = "linear" ,
219 bf16 = True ,
220 batch_sampler = BatchSamplers. NO_DUPLICATES ,
221 eval_strategy = "steps" ,
222 eval_steps = 0.1 ,
223 save_strategy = "steps" ,
224 save_steps = 0.1 ,
225 save_total_limit = 2 ,
226 logging_steps = 0.01 ,
227 logging_first_step = True ,
228 load_best_model_at_end = True ,
229 metric_for_best_model = metric_key,
230 greater_is_better = True ,
231 report_to = "none" if SMOKE_TEST else "trackio" ,
232 run_name = RUN_NAME ,
233 seed = 12 ,
234 )
235
236 trainer = SentenceTransformerTrainer(
237 model = model,
238 args = args,
239 train_dataset = train_dataset,
240 eval_dataset = eval_dataset,
241 loss = loss,
242 evaluator = evaluator,
243 )
244 if not SMOKE_TEST :
245 log_trackio_dashboard()
246 trainer.train()
247
248 logging.info( "Post-training evaluation:" )
249 with autocast_ctx():
250 score = evaluator(model)[evaluator.primary_metric]
251 delta = score - baseline_eval
252 verdict = "WIN" if delta >= 0.005 else "MARGINAL" if delta >= 0 else "REGRESSION"
253 logging.info( f "VERDICT: { verdict } | score= { score :.4f} | baseline= { baseline_eval :.4f} | delta= { delta :+.4f} " )
254
255 final_dir = f " { OUTPUT_DIR } /final"
256 model.save_pretrained(final_dir)
257 logging.info( f "Saved final model to { final_dir } " )
258
259 if SMOKE_TEST :
260 logging.info( "SMOKE_TEST=1: skipping Hub push" )
261 return
262
263 try :
264 commit_url = model.push_to_hub( RUN_NAME )
265 logging.info( f "Pushed model to { commit_url.rsplit( '/commit/' , 1 )[ 0 ] } " )
266 except Exception :
267 import traceback
268
269 logging.error( f "Hub push failed: \n{ traceback.format_exc() } " )
270
271
272 if __name__ == "__main__" :
273 main()