Setting the file. One moment.
Mine Hard Negatives · Train Sentence Transformers · huggingface/skills · Skills Docs
Back to the top of the page Losses Cross Encoder
scripts/mine_hard_negatives.py
scripts/ mine_hard_negatives.py
Python · 199 lines · 8 KB
17 2. Pick a retriever model (can be your current base model or a stronger one).
18 3. Run this script to produce a new dataset with N mined negatives per anchor.
19 4. Train with MultipleNegativesRankingLoss or CachedMultipleNegativesRankingLoss.
20
21 Usage:
22 python mine_hard_negatives.py \\
23 --dataset sentence-transformers/gooaq \\
24 --model sentence-transformers/all-MiniLM-L6-v2 \\
25 --num-negatives 5 \\
26 --output-path data/gooaq-hard-negatives
27
28 # Mine from a separate document corpus (recommended for production):
29 python mine_hard_negatives.py \\
30 --dataset sentence-transformers/gooaq \\
31 --model sentence-transformers/all-MiniLM-L6-v2 \\
32 --corpus-dataset sentence-transformers/wikipedia-en-passages \\
33 --corpus-column text \\
34 --num-negatives 5 \\
35 --output-path data/gooaq-hn-wiki
36
37 # With a cross-encoder as an "oracle" to filter negatives by score:
38 python mine_hard_negatives.py \\
39 --dataset sentence-transformers/gooaq \\
40 --model sentence-transformers/all-MiniLM-L6-v2 \\
41 --cross-encoder cross-encoder/ms-marco-MiniLM-L-6-v2 \\
42 --num-negatives 5 \\
43 --max-score 0.9 \\
44 --relative-margin 0.05 \\
45 --output-path data/gooaq-hn-filtered
46
47 # Push the mined dataset to the Hub:
48 python mine_hard_negatives.py \\
49 --dataset sentence-transformers/gooaq --model ... --num-negatives 5 \\
50 --push-to-hub your-username/gooaq-hard-negatives
51
52 Key options:
53 --num-negatives How many hard negatives to mine per anchor (default 3).
54 --range-min/max Which retrieval-rank window to sample from (default 0..100).
55 --sampling-strategy "top" (rank-1 hardest) or "random" (within the window).
56 --relative-margin Require that negative_score < positive_score * (1 - margin).
57 --max-score Filter candidates above this score (likely false negatives).
58 --cross-encoder Use a cross-encoder to re-score candidates before filtering.
59 --corpus-dataset Mine from a separate document pool instead of the input
60 dataset's positives. Recommended for production: typical
61 retrieval corpora (Wikipedia, MSMARCO passages) are far
62 larger than your training-pair pool, giving harder negatives.
63
64 See the `mine_hard_negatives` API reference for full semantics and all flags.
65 """
66
67 from __future__ import annotations
68
69 import argparse
70 import logging
71 import sys
72
73 from datasets import load_dataset
74
75 from sentence_transformers import CrossEncoder, SentenceTransformer
76 from sentence_transformers.util import mine_hard_negatives
77
78 logging.basicConfig( format = " %(asctime)s - %(message)s " , datefmt = "%Y-%m- %d %H:%M:%S" , level = logging. INFO )
79 for _noisy in ( "httpx" , "httpcore" , "huggingface_hub" , "urllib3" , "filelock" , "fsspec" ):
80 logging.getLogger(_noisy).setLevel(logging. WARNING )
81
82
83 def build_parser () -> argparse.ArgumentParser:
84 p = argparse.ArgumentParser( description = __doc__ , formatter_class = argparse.RawDescriptionHelpFormatter)
85 p.add_argument( "--dataset" , required = True )
86 p.add_argument( "--subset" , default = None )
87 p.add_argument( "--split" , default = "train" )
88 p.add_argument( "--model" , required = True , help = "Retriever / bi-encoder used to score candidates" )
89 p.add_argument( "--cross-encoder" , default = None , help = "Optional CrossEncoder to re-score and filter" )
90 p.add_argument( "--anchor-column" , default = None )
91 p.add_argument( "--positive-column" , default = None )
92 p.add_argument(
93 "--num-negatives" ,
94 type = int ,
95 default = 3 ,
96 help = "Number of hard negatives to mine per anchor. Default 3 matches the library." ,
97 )
98 p.add_argument( "--range-min" , type = int , default = 0 )
99 p.add_argument( "--range-max" , type = int , default = 100 )
100 p.add_argument( "--sampling-strategy" , choices = [ "top" , "random" ], default = "top" )
101 p.add_argument(
102 "--max-score" , type = float , default = None , help = "Drop candidates scoring above this (likely false negatives)"
103 )
104 p.add_argument( "--min-score" , type = float , default = None )
105 p.add_argument( "--absolute-margin" , type = float , default = None )
106 p.add_argument( "--relative-margin" , type = float , default = None )
107 p.add_argument(
108 "--output-format" ,
109 choices = [ "triplet" , "n-tuple" , "labeled-pair" , "labeled-list" ],
110 default = "triplet" ,
111 )
112 p.add_argument( "--include-positives" , action = "store_true" )
113 p.add_argument( "--output-scores" , action = "store_true" )
114 p.add_argument( "--batch-size" , type = int , default = 32 )
115 p.add_argument( "--use-faiss" , action = "store_true" )
116 p.add_argument(
117 "--corpus-dataset" ,
118 default = None ,
119 help = "Optional Hub dataset id or local path for a separate document pool to mine from. "
120 "If unset, mines negatives from the input dataset's positives." ,
121 )
122 p.add_argument( "--corpus-subset" , default = None , help = "Subset of --corpus-dataset (optional)" )
123 p.add_argument( "--corpus-split" , default = "train" , help = "Split of --corpus-dataset (default 'train')" )
124 p.add_argument(
125 "--corpus-column" , default = "text" , help = "Text column to extract from --corpus-dataset (default 'text')"
126 )
127 p.add_argument( "--output-path" , default = None , help = "Local directory to save the mined dataset to" )
128 p.add_argument( "--push-to-hub" , default = None , help = "Hub repo id to push the mined dataset to (optional)" )
129 p.add_argument( "--private" , action = "store_true" , help = "Push as a private repo" )
130 return p
131
132
133 def main () -> int :
134 args = build_parser().parse_args()
135
136 dataset = (
137 load_dataset(args.dataset, args.subset, split = args.split)
138 if args.subset
139 else load_dataset(args.dataset, split = args.split)
140 )
141 print ( f "Loaded { len (dataset) :,} rows from { args.dataset } (split= { args.split } )" )
142
143 model = SentenceTransformer(args.model)
144 cross_encoder = CrossEncoder(args.cross_encoder) if args.cross_encoder else None
145
146 corpus = None
147 if args.corpus_dataset:
148 corpus_ds = (
149 load_dataset(args.corpus_dataset, args.corpus_subset, split = args.corpus_split)
150 if args.corpus_subset
151 else load_dataset(args.corpus_dataset, split = args.corpus_split)
152 )
153 if args.corpus_column not in corpus_ds.column_names:
154 raise SystemExit (
155 f "--corpus-column ' { args.corpus_column } ' not in { args.corpus_dataset } columns: "
156 f " { corpus_ds.column_names } "
157 )
158 corpus = list (corpus_ds[args.corpus_column])
159 print ( f "Loaded corpus: { len (corpus) :,} documents from { args.corpus_dataset } . { args.corpus_column } " )
160
161 mined = mine_hard_negatives(
162 dataset = dataset,
163 model = model,
164 corpus = corpus,
165 cross_encoder = cross_encoder,
166 anchor_column_name = args.anchor_column,
167 positive_column_name = args.positive_column,
168 num_negatives = args.num_negatives,
169 range_min = args.range_min,
170 range_max = args.range_max,
171 sampling_strategy = args.sampling_strategy,
172 max_score = args.max_score,
173 min_score = args.min_score,
174 absolute_margin = args.absolute_margin,
175 relative_margin = args.relative_margin,
176 output_format = args.output_format,
177 include_positives = args.include_positives,
178 output_scores = args.output_scores,
179 batch_size = args.batch_size,
180 use_faiss = args.use_faiss,
181 )
182 print ( f "Mined dataset: { len (mined) :,} rows | columns: { mined.column_names } " )
183
184 if args.output_path:
185 mined.save_to_disk(args.output_path)
186 print ( f "Saved to { args.output_path } " )
187
188 if args.push_to_hub:
189 mined.push_to_hub(args.push_to_hub, private = args.private)
190 print ( f "Pushed to https://huggingface.co/datasets/ { args.push_to_hub } " )
191
192 if not args.output_path and not args.push_to_hub:
193 print ( "No --output-path or --push-to-hub provided; nothing persisted." )
194
195 return 0
196
197
198 if __name__ == "__main__" :
199 sys.exit(main())