BAAI
/

Shitao commited on
Commit
44482ec
1 Parent(s): 74c4854

Upload README.md with huggingface_hub

Browse files
Files changed (1) hide show
  1. README.md +369 -1
README.md CHANGED
@@ -1,3 +1,371 @@
1
  ---
2
- license: mit
 
 
 
 
 
 
3
  ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
+ license: apache-2.0
3
+ pipeline_tag: text-classification
4
+ tags:
5
+ - transformers
6
+ - sentence-transformers
7
+ language:
8
+ - multilingual
9
  ---
10
+
11
+ # Reranker
12
+
13
+ **More details please refer to our Github: [FlagEmbedding](https://github.com/FlagOpen/FlagEmbedding/tree/master).**
14
+
15
+ - [Model List](#model-list)
16
+ - [Usage](#usage)
17
+ - [Fine-tuning](#fine-tune)
18
+ - [Evaluation](#evaluation)
19
+ - [Citation](#citation)
20
+
21
+ Different from embedding model, reranker uses question and document as input and directly output similarity instead of embedding.
22
+ You can get a relevance score by inputting query and passage to the reranker.
23
+ And the score can be mapped to a float value in [0,1] by sigmoid function.
24
+
25
+
26
+ ## Model List
27
+
28
+ | Model | Base model | Language | layerwise | feature |
29
+ |:--------------------------------------------------------------------------|:--------:|:-----------------------------------------------------------------------------------------------------------------------------------:|:----------------------------------------------------------------------------------------------:|:----------------------------------------------------------------------------------------------:|
30
+ | [BAAI/bge-reranker-base](https://huggingface.co/BAAI/bge-reranker-base) | [xlm-roberta-base](https://huggingface.co/xlm-roberta-base) | Chinese and English | - | Lightweight reranker model, easy to deploy, with fast inference. |
31
+ | [BAAI/bge-reranker-large](https://huggingface.co/BAAI/bge-reranker-large) | [xlm-roberta-large](https://huggingface.co/FacebookAI/xlm-roberta-large) | Chinese and English | - | Lightweight reranker model, easy to deploy, with fast inference. |
32
+ | [BAAI/bge-reranker-v2-m3](https://huggingface.co/BAAI/bge-reranker-v2-m3) | [bge-m3](https://huggingface.co/BAAI/bge-m3) | Multilingual | - | Lightweight reranker model, possesses strong multilingual capabilities, easy to deploy, with fast inference. |
33
+ | [BAAI/bge-reranker-v2-gemma](https://huggingface.co/BAAI/bge-reranker-v2-gemma) | [google/gemma-2b](https://huggingface.co/google/gemma-2b) | Multilingual | - | Suitable for multilingual contexts, performs well in both English proficiency and multilingual capabilities. |
34
+ | [BAAI/bge-reranker-v2-minicpm-layerwise](https://huggingface.co/BAAI/bge-reranker-v2-minicpm-layerwise) | [openbmb/MiniCPM-2B-dpo-fp16](https://huggingface.co/openbmb/MiniCPM-2B-dpo-fp16/tree/main) | Multilingual | 8-40 | Suitable for multilingual contexts, performs well in both English and Chinese proficiency, allows freedom to select layers for output, facilitating accelerated inference. |
35
+
36
+
37
+ You can select the model according your senario and resource.
38
+ - For **multilingual**, utilize [BAAI/bge-reranker-v2-m3](https://huggingface.co/BAAI/bge-reranker-v2-m3) and [BAAI/bge-reranker-v2-gemma](https://huggingface.co/BAAI/bge-reranker-v2-gemma)
39
+
40
+ - For **Chinese or English**, utilize [BAAI/bge-reranker-v2-m3](https://huggingface.co/BAAI/bge-reranker-v2-m3) and [BAAI/bge-reranker-v2-minicpm-layerwise](https://huggingface.co/BAAI/bge-reranker-v2-minicpm-layerwise).
41
+
42
+ - For **efficiency**, utilize [BAAI/bge-reranker-v2-m3](https://huggingface.co/BAAI/bge-reranker-v2-m3) and the low layer of [BAAI/bge-reranker-v2-minicpm-layerwise](https://huggingface.co/BAAI/bge-reranker-v2-minicpm-layerwise).
43
+
44
+ - For better performance, recommand [BAAI/bge-reranker-v2-minicpm-layerwise](https://huggingface.co/BAAI/bge-reranker-v2-minicpm-layerwise) and [BAAI/bge-reranker-v2-gemma](https://huggingface.co/BAAI/bge-reranker-v2-gemma)
45
+
46
+ ## Usage
47
+ ### Using FlagEmbedding
48
+
49
+ ```
50
+ pip install -U FlagEmbedding
51
+ ```
52
+
53
+ #### For normal reranker (bge-reranker-base / bge-reranker-large / bge-reranker-v2-m3 )
54
+
55
+ Get relevance scores (higher scores indicate more relevance):
56
+
57
+ ```python
58
+ from FlagEmbedding import FlagReranker
59
+ reranker = FlagReranker('BAAI/bge-reranker-v2-m3', use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation
60
+
61
+ score = reranker.compute_score(['query', 'passage'])
62
+ print(score) # -5.65234375
63
+
64
+ # You can map the scores into 0-1 by set "normalize=True", which will apply sigmoid function to the score
65
+ score = reranker.compute_score(['query', 'passage'], normalize=True)
66
+ print(score) # 0.003497010252573502
67
+
68
+ scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']])
69
+ print(scores) # [-8.1875, 5.26171875]
70
+
71
+ # You can map the scores into 0-1 by set "normalize=True", which will apply sigmoid function to the score
72
+ scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']], normalize=True)
73
+ print(scores) # [0.00027803096387751553, 0.9948403768236574]
74
+ ```
75
+
76
+ #### For LLM-based reranker
77
+
78
+ ```python
79
+ from FlagEmbedding import FlagLLMReranker
80
+ reranker = FlagLLMReranker('BAAI/bge-reranker-v2-gemma', use_bf16=True) # Setting use_bf16 to True speeds up computation with a slight performance degradation
81
+
82
+ score = reranker.compute_score(['query', 'passage'])
83
+ print(score) # 2.15625
84
+
85
+ scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']])
86
+ print(scores) # [-0.84765625, 10.625]
87
+ ```
88
+
89
+ #### For LLM-based layerwise reranker
90
+
91
+ ```python
92
+ from FlagEmbedding import LayerWiseFlagLLMReranker
93
+ reranker = LayerWiseFlagLLMReranker('BAAI/bge-reranker-v2-minicpm-layerwise', use_bf16=True) # Setting use_bf16 to True speeds up computation with a slight performance degradation
94
+
95
+ score = reranker.compute_score(['query', 'passage'], cutoff_layers=[28]) # Adjusting 'cutoff_layers' to pick which layers are used for computing the score.
96
+ print(score) # -7.03125
97
+
98
+ scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']], cutoff_layers=[28])
99
+ print(scores) # [-10.0, 1.8203125]
100
+ ```
101
+
102
+ ### Using Huggingface transformers
103
+
104
+ #### For normal reranker (bge-reranker-base / bge-reranker-large / bge-reranker-v2-m3 )
105
+
106
+ Get relevance scores (higher scores indicate more relevance):
107
+
108
+ ```python
109
+ import torch
110
+ from transformers import AutoModelForSequenceClassification, AutoTokenizer
111
+
112
+ tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-m3')
113
+ model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-v2-m3')
114
+ model.eval()
115
+
116
+ pairs = [['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']]
117
+ with torch.no_grad():
118
+ inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
119
+ scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
120
+ print(scores)
121
+ ```
122
+
123
+ #### For LLM-based reranker
124
+
125
+ ```python
126
+ import torch
127
+ from transformers import AutoModelForCausalLM, AutoTokenizer
128
+
129
+ def get_inputs(pairs, tokenizer, prompt=None, max_length=1024):
130
+ if prompt is None:
131
+ prompt = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'."
132
+ sep = "\n"
133
+ prompt_inputs = tokenizer(prompt,
134
+ return_tensors=None,
135
+ add_special_tokens=False)['input_ids']
136
+ sep_inputs = tokenizer(sep,
137
+ return_tensors=None,
138
+ add_special_tokens=False)['input_ids']
139
+ inputs = []
140
+ for query, passage in pairs:
141
+ query_inputs = tokenizer(f'A: {query}',
142
+ return_tensors=None,
143
+ add_special_tokens=False,
144
+ max_length=max_length * 3 // 4,
145
+ truncation=True)
146
+ passage_inputs = tokenizer(f'B: {passage}',
147
+ return_tensors=None,
148
+ add_special_tokens=False,
149
+ max_length=max_length,
150
+ truncation=True)
151
+ item = tokenizer.prepare_for_model(
152
+ [tokenizer.bos_token_id] + query_inputs['input_ids'],
153
+ sep_inputs + passage_inputs['input_ids'],
154
+ truncation='only_second',
155
+ max_length=max_length,
156
+ padding=False,
157
+ return_attention_mask=False,
158
+ return_token_type_ids=False,
159
+ add_special_tokens=False
160
+ )
161
+ item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
162
+ item['attention_mask'] = [1] * len(item['input_ids'])
163
+ inputs.append(item)
164
+ return tokenizer.pad(
165
+ inputs,
166
+ padding=True,
167
+ max_length=max_length + len(sep_inputs) + len(prompt_inputs),
168
+ pad_to_multiple_of=8,
169
+ return_tensors='pt',
170
+ )
171
+
172
+ tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-gemma')
173
+ model = AutoModelForCausalLM.from_pretrained('BAAI/bge-reranker-v2-gemma')
174
+ yes_loc = tokenizer('Yes', add_special_tokens=False)['input_ids'][0]
175
+ model.eval()
176
+
177
+ pairs = [['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']]
178
+ with torch.no_grad():
179
+ inputs = get_inputs(pairs, tokenizer)
180
+ scores = model(**inputs, return_dict=True).logits[:, -1, yes_loc].view(-1, ).float()
181
+ print(scores)
182
+ ```
183
+
184
+ #### For LLM-based layerwise reranker
185
+
186
+ ```python
187
+ import torch
188
+ from transformers import AutoModelForCausalLM, AutoTokenizer
189
+
190
+ def get_inputs(pairs, tokenizer, prompt=None, max_length=1024):
191
+ if prompt is None:
192
+ prompt = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'."
193
+ sep = "\n"
194
+ prompt_inputs = tokenizer(prompt,
195
+ return_tensors=None,
196
+ add_special_tokens=False)['input_ids']
197
+ sep_inputs = tokenizer(sep,
198
+ return_tensors=None,
199
+ add_special_tokens=False)['input_ids']
200
+ inputs = []
201
+ for query, passage in pairs:
202
+ query_inputs = tokenizer(f'A: {query}',
203
+ return_tensors=None,
204
+ add_special_tokens=False,
205
+ max_length=max_length * 3 // 4,
206
+ truncation=True)
207
+ passage_inputs = tokenizer(f'B: {passage}',
208
+ return_tensors=None,
209
+ add_special_tokens=False,
210
+ max_length=max_length,
211
+ truncation=True)
212
+ item = tokenizer.prepare_for_model(
213
+ [tokenizer.bos_token_id] + query_inputs['input_ids'],
214
+ sep_inputs + passage_inputs['input_ids'],
215
+ truncation='only_second',
216
+ max_length=max_length,
217
+ padding=False,
218
+ return_attention_mask=False,
219
+ return_token_type_ids=False,
220
+ add_special_tokens=False
221
+ )
222
+ item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
223
+ item['attention_mask'] = [1] * len(item['input_ids'])
224
+ inputs.append(item)
225
+ return tokenizer.pad(
226
+ inputs,
227
+ padding=True,
228
+ max_length=max_length + len(sep_inputs) + len(prompt_inputs),
229
+ pad_to_multiple_of=8,
230
+ return_tensors='pt',
231
+ )
232
+
233
+ tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-minicpm-layerwise', trust_remote_code=True, torch_dtype=torch.bfloat16)
234
+ model = AutoModelForCausalLM.from_pretrained('BAAI/bge-reranker-v2-minicpm-layerwise', trust_remote_code=True, torch_dtype=torch.bfloat16)
235
+ model = model.to('cuda')
236
+ model.eval()
237
+
238
+ pairs = [['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']]
239
+ with torch.no_grad():
240
+ inputs = get_inputs(pairs, tokenizer).to(model.device)
241
+ all_scores = model(**inputs, return_dict=True, cutoff_layers=[28])
242
+ all_scores = [scores[:, -1].view(-1, ).float() for scores in all_scores[0]]
243
+ print(all_scores)
244
+ ```
245
+
246
+ ## Fine-tune
247
+
248
+ You can fine-tune the reranker with the following code:
249
+
250
+ **For llm-based reranker**
251
+
252
+ ```shell
253
+ torchrun --nproc_per_node {number of gpus} \
254
+ -m FlagEmbedding.llm_reranker.finetune_for_instruction.run \
255
+ --output_dir {path to save model} \
256
+ --model_name_or_path BAAI/bge-reranker-v2-gemma \
257
+ --train_data ./toy_finetune_data.jsonl \
258
+ --learning_rate 2e-4 \
259
+ --num_train_epochs 1 \
260
+ --per_device_train_batch_size 1 \
261
+ --gradient_accumulation_steps 16 \
262
+ --dataloader_drop_last True \
263
+ --query_max_len 512 \
264
+ --passage_max_len 512 \
265
+ --train_group_size 16 \
266
+ --logging_steps 1 \
267
+ --save_steps 2000 \
268
+ --save_total_limit 50 \
269
+ --ddp_find_unused_parameters False \
270
+ --gradient_checkpointing \
271
+ --deepspeed stage1.json \
272
+ --warmup_ratio 0.1 \
273
+ --bf16 \
274
+ --use_lora True \
275
+ --lora_rank 32 \
276
+ --lora_alpha 64 \
277
+ --use_flash_attn True \
278
+ --target_modules q_proj k_proj v_proj o_proj
279
+ ```
280
+
281
+ **For llm-based layerwise reranker**
282
+
283
+ ```shell
284
+ torchrun --nproc_per_node {number of gpus} \
285
+ -m FlagEmbedding.llm_reranker.finetune_for_layerwise.run \
286
+ --output_dir {path to save model} \
287
+ --model_name_or_path BAAI/bge-reranker-v2-minicpm-layerwise \
288
+ --train_data ./toy_finetune_data.jsonl \
289
+ --learning_rate 2e-4 \
290
+ --num_train_epochs 1 \
291
+ --per_device_train_batch_size 1 \
292
+ --gradient_accumulation_steps 16 \
293
+ --dataloader_drop_last True \
294
+ --query_max_len 512 \
295
+ --passage_max_len 512 \
296
+ --train_group_size 16 \
297
+ --logging_steps 1 \
298
+ --save_steps 2000 \
299
+ --save_total_limit 50 \
300
+ --ddp_find_unused_parameters False \
301
+ --gradient_checkpointing \
302
+ --deepspeed stage1.json \
303
+ --warmup_ratio 0.1 \
304
+ --bf16 \
305
+ --use_lora True \
306
+ --lora_rank 32 \
307
+ --lora_alpha 64 \
308
+ --use_flash_attn True \
309
+ --target_modules q_proj k_proj v_proj o_proj \
310
+ --start_layer 8 \
311
+ --head_multi True \
312
+ --head_type simple
313
+ ```
314
+
315
+ Our rerankers are initialized from [google/gemma-2b](https://huggingface.co/google/gemma-2b) (for llm-based reranker) and [openbmb/MiniCPM-2B-dpo-fp16](https://huggingface.co/openbmb/MiniCPM-2B-dpo-fp16/tree/main) (for llm-based layerwise reranker), and we train it on a mixture of multilingual datasets:
316
+
317
+ - [bge-m3-data](https://huggingface.co/datasets/Shitao/bge-m3-data)
318
+ - [quora train data](https://huggingface.co/datasets/quora)
319
+ - [fever train data](https://fever.ai/dataset/fever.html)
320
+
321
+ ## Evaluation
322
+
323
+ - llama-index.
324
+
325
+ ![image-20240317193909373](./assets/llama-index.png)
326
+
327
+
328
+ - BEIR.
329
+
330
+ rereank the top 100 results from bge-en-v1.5 large.
331
+
332
+ ![image-20240317174633333](./assets/BEIR-bge-en-v1.5.png)
333
+
334
+ rereank the top 100 results from e5 mistral 7b instruct.
335
+
336
+ ![image-20240317172949713](./assets/BEIR-e5-mistral.png)
337
+
338
+ - CMTEB-retrieval.
339
+ It rereank the top 100 results from bge-zh-v1.5 large.
340
+
341
+ ![image-20240317173026235](./assets/CMTEB-retrieval-bge-zh-v1.5.png)
342
+
343
+ - miracl (multi-language).
344
+ It rereank the top 100 results from bge-m3.
345
+
346
+ ![image-20240317173117639](./assets/miracl-bge-m3.png)
347
+
348
+
349
+
350
+ ## Citation
351
+
352
+ If you find this repository useful, please consider giving a star and citation
353
+
354
+ ```bibtex
355
+ @misc{li2023making,
356
+ title={Making Large Language Models A Better Foundation For Dense Retrieval},
357
+ author={Chaofan Li and Zheng Liu and Shitao Xiao and Yingxia Shao},
358
+ year={2023},
359
+ eprint={2312.15503},
360
+ archivePrefix={arXiv},
361
+ primaryClass={cs.CL}
362
+ }
363
+ @misc{chen2024bge,
364
+ title={BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation},
365
+ author={Jianlv Chen and Shitao Xiao and Peitian Zhang and Kun Luo and Defu Lian and Zheng Liu},
366
+ year={2024},
367
+ eprint={2402.03216},
368
+ archivePrefix={arXiv},
369
+ primaryClass={cs.CL}
370
+ }
371
+ ```