chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,506 @@
|
||||
# Embedder
|
||||
|
||||
- [Model List](#model-list)
|
||||
- [Usage](#usage)
|
||||
- [Citation](#citation)
|
||||
|
||||
An embedder can encode text into embeddings.
|
||||
|
||||
When provided with a query and a passage, the embedder encodes both separately, and then uses the similarity between their embeddings as the similarity score.
|
||||
|
||||
For more detailed using, you can look [embedder-encoder only](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/inference/embedder/encoder_only) or [embedder-decoder only](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/inference/embedder/decoder_only)
|
||||
|
||||
|
||||
## Model List
|
||||
|
||||
`bge` is short for `BAAI general embedding`.
|
||||
|
||||
| Model | Language | Description | query instruction for retrieval |
|
||||
| :----------------------------------------------------------- | :-----------------: | :----------------------------------------------------------: | :----------------------------------------------------------: |
|
||||
| [BAAI/bge-en-icl](https://huggingface.co/BAAI/bge-en-icl) | English | A LLM-based embedding model with in-context learning capabilities, which can fully leverage the model's potential based on a few shot examples | Provide instructions and few-shot examples freely based on the given task. |
|
||||
| [BAAI/bge-multilingual-gemma2](https://huggingface.co/BAAI/bge-multilingual-gemma2) | Multilingual | A LLM-based multilingual embedding model, trained on a diverse range of languages and tasks. | Provide instructions based on the given task. |
|
||||
| [BAAI/bge-m3](https://huggingface.co/BAAI/bge-m3) | Multilingual | Multi-Functionality(dense retrieval, sparse retrieval, multi-vector(colbert)), Multi-Linguality, and Multi-Granularity(8192 tokens) | |
|
||||
| [BAAI/bge-large-en-v1.5](https://huggingface.co/BAAI/bge-large-en-v1.5) | English | version 1.5 with more reasonable similarity distribution | `Represent this sentence for searching relevant passages: ` |
|
||||
| [BAAI/bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5) | English | version 1.5 with more reasonable similarity distribution | `Represent this sentence for searching relevant passages: ` |
|
||||
| [BAAI/bge-small-en-v1.5](https://huggingface.co/BAAI/bge-small-en-v1.5) | English | version 1.5 with more reasonable similarity distribution | `Represent this sentence for searching relevant passages: ` |
|
||||
| [BAAI/bge-large-zh-v1.5](https://huggingface.co/BAAI/bge-large-zh-v1.5) | Chinese | version 1.5 with more reasonable similarity distribution | `为这个句子生成表示以用于检索相关文章：` |
|
||||
| [BAAI/bge-base-zh-v1.5](https://huggingface.co/BAAI/bge-base-zh-v1.5) | Chinese | version 1.5 with more reasonable similarity distribution | `为这个句子生成表示以用于检索相关文章：` |
|
||||
| [BAAI/bge-small-zh-v1.5](https://huggingface.co/BAAI/bge-small-zh-v1.5) | Chinese | version 1.5 with more reasonable similarity distribution | `为这个句子生成表示以用于检索相关文章：` |
|
||||
| [BAAI/bge-large-en](https://huggingface.co/BAAI/bge-large-en) | English | Embedding Model which map text into vector | `Represent this sentence for searching relevant passages: ` |
|
||||
| [BAAI/bge-base-en](https://huggingface.co/BAAI/bge-base-en) | English | a base-scale model but with similar ability to `bge-large-en` | `Represent this sentence for searching relevant passages: ` |
|
||||
| [BAAI/bge-small-en](https://huggingface.co/BAAI/bge-small-en) | English | a small-scale model but with competitive performance | `Represent this sentence for searching relevant passages: ` |
|
||||
| [BAAI/bge-large-zh](https://huggingface.co/BAAI/bge-large-zh) | Chinese | Embedding Model which map text into vector | `为这个句子生成表示以用于检索相关文章：` |
|
||||
| [BAAI/bge-base-zh](https://huggingface.co/BAAI/bge-base-zh) | Chinese | a base-scale model but with similar ability to `bge-large-zh` | `为这个句子生成表示以用于检索相关文章：` |
|
||||
| [BAAI/bge-small-zh](https://huggingface.co/BAAI/bge-small-zh) | Chinese | a small-scale model but with competitive performance | `为这个句子生成表示以用于检索相关文章：` |
|
||||
|
||||
## Usage
|
||||
|
||||
### Using FlagEmbedding
|
||||
|
||||
#### 1. Auto Model
|
||||
|
||||
You can use `FlagAutoModel` to load the model. For the **custom model** (not included in [`AUTO_EMBEDDER_MAPPING`](https://github.com/FlagOpen/FlagEmbedding/blob/master/FlagEmbedding/inference/embedder/model_mapping.py#L39)), you must specify the `model_class` parameter. You can also submit a pull request to add your **released model** to the [`AUTO_EMBEDDER_MAPPING`](https://github.com/FlagOpen/FlagEmbedding/blob/master/FlagEmbedding/inference/embedder/model_mapping.py#L39) dictionary. If need, you can create a new `<model>.py` file in [here](https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/inference/embedder/encoder_only) or [here](https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/inference/embedder/decoder_only).
|
||||
|
||||
```python
|
||||
from FlagEmbedding import FlagAutoModel
|
||||
sentences_1 = ["样例数据-1", "样例数据-2"]
|
||||
sentences_2 = ["样例数据-3", "样例数据-4"]
|
||||
model = FlagAutoModel.from_finetuned('BAAI/bge-large-zh-v1.5',
|
||||
query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：",
|
||||
use_fp16=True,
|
||||
devices=['cuda:1']) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
embeddings_1 = model.encode(sentences_1)
|
||||
embeddings_2 = model.encode(sentences_2)
|
||||
similarity = embeddings_1 @ embeddings_2.T
|
||||
print(similarity)
|
||||
|
||||
# for s2p(short query to long passage) retrieval task, suggest to use encode_queries() which will automatically add the instruction to each query
|
||||
# corpus in retrieval task can still use encode_corpus(), since they don't need instruction
|
||||
queries = ['query_1', 'query_2']
|
||||
passages = ["样例文档-1", "样例文档-2"]
|
||||
q_embeddings = model.encode_queries(queries)
|
||||
p_embeddings = model.encode_corpus(passages)
|
||||
scores = q_embeddings @ p_embeddings.T
|
||||
print(scores)
|
||||
```
|
||||
|
||||
For your **custom model** (assume the model is finetuned from `BAAI/bge-large-zh-v1.5`, then the model class is `encoder-only-base`), you can use the following code:
|
||||
|
||||
```python
|
||||
from FlagEmbedding import FlagAutoModel
|
||||
sentences_1 = ["样例数据-1", "样例数据-2"]
|
||||
sentences_2 = ["样例数据-3", "样例数据-4"]
|
||||
model = FlagAutoModel.from_finetuned('your_model_name_or_path',
|
||||
model_class='encoder-only-base', # specify the model class
|
||||
query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：",
|
||||
pooling_method='cls', # specify the pooling method
|
||||
use_fp16=True,
|
||||
devices=['cuda:1']) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
embeddings_1 = model.encode(sentences_1)
|
||||
embeddings_2 = model.encode(sentences_2)
|
||||
similarity = embeddings_1 @ embeddings_2.T
|
||||
print(similarity)
|
||||
|
||||
# for s2p(short query to long passage) retrieval task, suggest to use encode_queries() which will automatically add the instruction to each query
|
||||
# corpus in retrieval task can still use encode_corpus(), since they don't need instruction
|
||||
queries = ['query_1', 'query_2']
|
||||
passages = ["样例文档-1", "样例文档-2"]
|
||||
q_embeddings = model.encode_queries(queries)
|
||||
p_embeddings = model.encode_corpus(passages)
|
||||
scores = q_embeddings @ p_embeddings.T
|
||||
print(scores)
|
||||
```
|
||||
|
||||
The `model_class` parameter currently includes the following options:
|
||||
- `encoder-only-base`: for encoder-only normal model, such as `BAAI/bge-large-en-v1.5`
|
||||
- `encoder-only-m3`: for encoder-only M3 model, such as `BAAI/bge-m3`
|
||||
- `decoder-only-base`: for decoder-only normal model, such as `BAAI/bge-multilingual-gemma2`
|
||||
- `decoder-only-icl`: for decoder-only ICL model, such as `BAAI/bge-en-icl`
|
||||
|
||||
#### 2. Normal Model
|
||||
|
||||
For `FlagModel`, it supports `BAAI/bge-large-en-v1.5`, `BAAI/bge-base-en-v1.5`, `BAAI/bge-small-en-v1.5`, `BAAI/bge-large-zh-v1.5`, `BAAI/bge-base-zh-v1.5`, `BAAI/bge-small-zh-v1.5`, `BAAI/bge-large-en`, `BAAI/bge-base-en`, `BAAI/bge-small-en`, `BAAI/bge-large-zh`, `BAAI/bge-base-zh`, `BAAI/bge-small-zh'`:
|
||||
|
||||
```python
|
||||
from FlagEmbedding import FlagModel
|
||||
sentences_1 = ["样例数据-1", "样例数据-2"]
|
||||
sentences_2 = ["样例数据-3", "样例数据-4"]
|
||||
model = FlagModel('BAAI/bge-large-zh-v1.5',
|
||||
query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：",
|
||||
use_fp16=True,
|
||||
devices=['cuda:1']) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
embeddings_1 = model.encode(sentences_1)
|
||||
embeddings_2 = model.encode(sentences_2)
|
||||
similarity = embeddings_1 @ embeddings_2.T
|
||||
print(similarity)
|
||||
|
||||
# for s2p(short query to long passage) retrieval task, suggest to use encode_queries() which will automatically add the instruction to each query
|
||||
# corpus in retrieval task can still use encode_corpus(), since they don't need instruction
|
||||
queries = ['query_1', 'query_2']
|
||||
passages = ["样例文档-1", "样例文档-2"]
|
||||
q_embeddings = model.encode_queries(queries)
|
||||
p_embeddings = model.encode_corpus(passages)
|
||||
scores = q_embeddings @ p_embeddings.T
|
||||
print(scores)
|
||||
```
|
||||
|
||||
#### 3. M3 Model
|
||||
|
||||
For `BGEM3FlagModel`, it supports `BAAI/bge-m3`:
|
||||
|
||||
```python
|
||||
from FlagEmbedding import BGEM3FlagModel
|
||||
sentences_1 = ["样例数据-1", "样例数据-2"]
|
||||
sentences_2 = ["样例数据-3", "样例数据-4"]
|
||||
model = BGEM3FlagModel('BAAI/bge-m3',
|
||||
use_fp16=True,
|
||||
pooling_method='cls',
|
||||
devices=['cuda:1']) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
embeddings_1 = model.encode(
|
||||
sentences_1,
|
||||
return_dense=True,
|
||||
return_sparse=True,
|
||||
return_colbert_vecs=False,
|
||||
)
|
||||
embeddings_2 = model.encode(
|
||||
sentences_2,
|
||||
return_dense=True,
|
||||
return_sparse=True,
|
||||
return_colbert_vecs=False,
|
||||
)
|
||||
dense_similarity = embeddings_1["dense_vecs"] @ embeddings_2["dense_vecs"].T
|
||||
print('dense similarity:', dense_similarity)
|
||||
sparse_similarity = model.compute_lexical_matching_score(
|
||||
embeddings_1["lexical_weights"],
|
||||
embeddings_2["lexical_weights"],
|
||||
)
|
||||
print('sparse similarity:', sparse_similarity)
|
||||
|
||||
queries = ['query_1', 'query_2']
|
||||
passages = ["样例文档-1", "样例文档-2"]
|
||||
q_embeddings = model.encode_queries(
|
||||
queries,
|
||||
return_dense=True,
|
||||
return_sparse=True,
|
||||
return_colbert_vecs=False,
|
||||
)
|
||||
p_embeddings = model.encode_corpus(
|
||||
passages,
|
||||
return_dense=True,
|
||||
return_sparse=True,
|
||||
return_colbert_vecs=False,
|
||||
)
|
||||
dense_scores = embeddings_1["dense_vecs"] @ embeddings_2["dense_vecs"].T
|
||||
print('dense scores:', dense_scores)
|
||||
sparse_scores = model.compute_lexical_matching_score(
|
||||
embeddings_1["lexical_weights"],
|
||||
embeddings_2["lexical_weights"],
|
||||
)
|
||||
print('sparse similarity:', sparse_scores)
|
||||
```
|
||||
|
||||
#### 4. LLM-based Model
|
||||
|
||||
For `FlagLLMModel`, it supports `BAAI/bge-multilingual-gemma2`, `Alibaba-NLP/gte-Qwen2-7B-instruct`, `intfloat/e5-mistral-7b-instruct`, .etc:
|
||||
|
||||
```python
|
||||
from FlagEmbedding import FlagLLMModel
|
||||
sentences_1 = ["样例数据-1", "样例数据-2"]
|
||||
sentences_2 = ["样例数据-3", "样例数据-4"]
|
||||
model = FlagLLMModel('BAAI/bge-multilingual-gemma2',
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
query_instruction_format="<instruct>{}\n<query>{}",
|
||||
use_fp16=True,
|
||||
devices=['cuda:1']) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
queries = ['query_1', 'query_2']
|
||||
passages = ["样例文档-1", "样例文档-2"]
|
||||
q_embeddings = model.encode_queries(queries)
|
||||
p_embeddings = model.encode_corpus(passages)
|
||||
scores = q_embeddings @ p_embeddings.T
|
||||
print(scores)
|
||||
```
|
||||
|
||||
#### 5. LLM-based ICL Model
|
||||
|
||||
For `FlagICLModel`, it supports `BAAI/bge-en-icl`:
|
||||
|
||||
```python
|
||||
from FlagEmbedding import FlagICLModel
|
||||
|
||||
examples = [
|
||||
{
|
||||
'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',
|
||||
'query': 'what is a virtual interface',
|
||||
'response': "A virtual interface is a software-defined abstraction that mimics the behavior and characteristics of a physical network interface. It allows multiple logical network connections to share the same physical network interface, enabling efficient utilization of network resources. Virtual interfaces are commonly used in virtualization technologies such as virtual machines and containers to provide network connectivity without requiring dedicated hardware. They facilitate flexible network configurations and help in isolating network traffic for security and management purposes."
|
||||
},
|
||||
{
|
||||
'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',
|
||||
'query': 'causes of back pain in female for a week',
|
||||
'response': "Back pain in females lasting a week can stem from various factors. Common causes include muscle strain due to lifting heavy objects or improper posture, spinal issues like herniated discs or osteoporosis, menstrual cramps causing referred pain, urinary tract infections, or pelvic inflammatory disease. Pregnancy-related changes can also contribute. Stress and lack of physical activity may exacerbate symptoms. Proper diagnosis by a healthcare professional is crucial for effective treatment and management."
|
||||
}
|
||||
]
|
||||
model = FlagICLModel(
|
||||
'BAAI/bge-en-icl',
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
query_instruction_format="<instruct>{}\n<query>{}",
|
||||
examples_for_task=examples,
|
||||
examples_instruction_format="<instruct>{}\n<query>{}\n<response>{}",
|
||||
use_fp16=True,
|
||||
devices=['cuda:1']
|
||||
) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define"
|
||||
]
|
||||
passages = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
|
||||
]
|
||||
q_embeddings = model.encode_queries(queries)
|
||||
p_embeddings = model.encode_corpus(passages)
|
||||
scores = q_embeddings @ p_embeddings.T
|
||||
print(scores)
|
||||
```
|
||||
|
||||
### Using HuggingFace Transformers
|
||||
|
||||
#### 1. Normal Model
|
||||
|
||||
It supports `BAAI/bge-large-en-v1.5`, `BAAI/bge-base-en-v1.5`, `BAAI/bge-small-en-v1.5`, `BAAI/bge-large-zh-v1.5`, `BAAI/bge-base-zh-v1.5`, `BAAI/bge-small-zh-v1.5`, `BAAI/bge-large-en`, `BAAI/bge-base-en`, `BAAI/bge-small-en`, `BAAI/bge-large-zh`, `BAAI/bge-base-zh`, `BAAI/bge-small-zh'`, the **dense method** of `BAAI/bge-m3`:
|
||||
|
||||
```python
|
||||
import torch
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-large-zh-v1.5')
|
||||
model = AutoModel.from_pretrained('BAAI/bge-large-zh-v1.5')
|
||||
model.eval()
|
||||
|
||||
sentences_1 = ["样例数据-1", "样例数据-2"]
|
||||
sentences_2 = ["样例数据-3", "样例数据-4"]
|
||||
with torch.no_grad():
|
||||
encoded_input_1 = tokenizer(sentences_1, padding=True, truncation=True, return_tensors='pt')
|
||||
encoded_input_2 = tokenizer(sentences_2, padding=True, truncation=True, return_tensors='pt')
|
||||
model_output_1 = model(**encoded_input_1)
|
||||
model_output_2 = model(**encoded_input_2)
|
||||
embeddings_1 = model_output_1[0][:, 0]
|
||||
embeddings_2 = model_output_2[0][:, 0]
|
||||
similarity = embeddings_1 @ embeddings_2.T
|
||||
print(similarity)
|
||||
```
|
||||
|
||||
#### 2. M3 Model
|
||||
|
||||
It only supports the **dense method** of `BAAI/bge-m3`, you can refer to the above code.
|
||||
|
||||
#### 3. LLM-based Model
|
||||
|
||||
It supports `BAAI/bge-multilingual-gemma2`:
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from torch import Tensor
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
|
||||
|
||||
def last_token_pool(last_hidden_states: Tensor,
|
||||
attention_mask: Tensor) -> Tensor:
|
||||
left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
|
||||
if left_padding:
|
||||
return last_hidden_states[:, -1]
|
||||
else:
|
||||
sequence_lengths = attention_mask.sum(dim=1) - 1
|
||||
batch_size = last_hidden_states.shape[0]
|
||||
return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
|
||||
|
||||
|
||||
def get_detailed_instruct(task_description: str, query: str) -> str:
|
||||
return f'<instruct>{task_description}\n<query>{query}'
|
||||
|
||||
|
||||
task = 'Given a web search query, retrieve relevant passages that answer the query.'
|
||||
queries = [
|
||||
get_detailed_instruct(task, 'how much protein should a female eat'),
|
||||
get_detailed_instruct(task, 'summit define')
|
||||
]
|
||||
# No need to add instructions for documents
|
||||
documents = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
|
||||
]
|
||||
input_texts = queries + documents
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-multilingual-gemma2')
|
||||
model = AutoModel.from_pretrained('BAAI/bge-multilingual-gemma2')
|
||||
model.eval()
|
||||
|
||||
max_length = 4096
|
||||
# Tokenize the input texts
|
||||
batch_dict = tokenizer(input_texts, max_length=max_length, padding=True, truncation=True, return_tensors='pt', pad_to_multiple_of=8)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = model(**batch_dict)
|
||||
embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
|
||||
|
||||
# normalize embeddings
|
||||
embeddings = F.normalize(embeddings, p=2, dim=1)
|
||||
scores = (embeddings[:2] @ embeddings[2:].T) * 100
|
||||
print(scores.tolist())
|
||||
# [[55.92064666748047, 1.6549524068832397], [-0.2698777914047241, 49.95653533935547]]
|
||||
```
|
||||
|
||||
#### 4. LLM-based ICL Model
|
||||
|
||||
It supports `BAAI/bge-en-icl`:
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from torch import Tensor
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
|
||||
|
||||
def last_token_pool(last_hidden_states: Tensor,
|
||||
attention_mask: Tensor) -> Tensor:
|
||||
left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
|
||||
if left_padding:
|
||||
return last_hidden_states[:, -1]
|
||||
else:
|
||||
sequence_lengths = attention_mask.sum(dim=1) - 1
|
||||
batch_size = last_hidden_states.shape[0]
|
||||
return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
|
||||
|
||||
|
||||
def get_detailed_instruct(task_description: str, query: str) -> str:
|
||||
return f'<instruct>{task_description}\n<query>{query}'
|
||||
|
||||
def get_detailed_example(task_description: str, query: str, response: str) -> str:
|
||||
return f'<instruct>{task_description}\n<query>{query}\n<response>{response}'
|
||||
|
||||
def get_new_queries(queries, query_max_len, examples_prefix, tokenizer):
|
||||
inputs = tokenizer(
|
||||
queries,
|
||||
max_length=query_max_len - len(tokenizer('<s>', add_special_tokens=False)['input_ids']) - len(
|
||||
tokenizer('\n<response></s>', add_special_tokens=False)['input_ids']),
|
||||
return_token_type_ids=False,
|
||||
truncation=True,
|
||||
return_tensors=None,
|
||||
add_special_tokens=False
|
||||
)
|
||||
prefix_ids = tokenizer(examples_prefix, add_special_tokens=False)['input_ids']
|
||||
suffix_ids = tokenizer('\n<response>', add_special_tokens=False)['input_ids']
|
||||
new_max_length = (len(prefix_ids) + len(suffix_ids) + query_max_len + 8) // 8 * 8 + 8
|
||||
new_queries = tokenizer.batch_decode(inputs['input_ids'])
|
||||
for i in range(len(new_queries)):
|
||||
new_queries[i] = examples_prefix + new_queries[i] + '\n<response>'
|
||||
return new_max_length, new_queries
|
||||
|
||||
task = 'Given a web search query, retrieve relevant passages that answer the query.'
|
||||
examples = [
|
||||
{'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',
|
||||
'query': 'what is a virtual interface',
|
||||
'response': "A virtual interface is a software-defined abstraction that mimics the behavior and characteristics of a physical network interface. It allows multiple logical network connections to share the same physical network interface, enabling efficient utilization of network resources. Virtual interfaces are commonly used in virtualization technologies such as virtual machines and containers to provide network connectivity without requiring dedicated hardware. They facilitate flexible network configurations and help in isolating network traffic for security and management purposes."},
|
||||
{'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',
|
||||
'query': 'causes of back pain in female for a week',
|
||||
'response': "Back pain in females lasting a week can stem from various factors. Common causes include muscle strain due to lifting heavy objects or improper posture, spinal issues like herniated discs or osteoporosis, menstrual cramps causing referred pain, urinary tract infections, or pelvic inflammatory disease. Pregnancy-related changes can also contribute. Stress and lack of physical activity may exacerbate symptoms. Proper diagnosis by a healthcare professional is crucial for effective treatment and management."}
|
||||
]
|
||||
examples = [get_detailed_example(e['instruct'], e['query'], e['response']) for e in examples]
|
||||
examples_prefix = '\n\n'.join(examples) + '\n\n' # if there not exists any examples, just set examples_prefix = ''
|
||||
queries = [
|
||||
get_detailed_instruct(task, 'how much protein should a female eat'),
|
||||
get_detailed_instruct(task, 'summit define')
|
||||
]
|
||||
documents = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
|
||||
]
|
||||
query_max_len, doc_max_len = 512, 512
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-en-icl')
|
||||
model = AutoModel.from_pretrained('BAAI/bge-en-icl')
|
||||
model.eval()
|
||||
|
||||
new_query_max_len, new_queries = get_new_queries(queries, query_max_len, examples_prefix, tokenizer)
|
||||
|
||||
query_batch_dict = tokenizer(new_queries, max_length=new_query_max_len, padding=True, truncation=True, return_tensors='pt')
|
||||
doc_batch_dict = tokenizer(documents, max_length=doc_max_len, padding=True, truncation=True, return_tensors='pt')
|
||||
|
||||
with torch.no_grad():
|
||||
query_outputs = model(**query_batch_dict)
|
||||
query_embeddings = last_token_pool(query_outputs.last_hidden_state, query_batch_dict['attention_mask'])
|
||||
doc_outputs = model(**doc_batch_dict)
|
||||
doc_embeddings = last_token_pool(doc_outputs.last_hidden_state, doc_batch_dict['attention_mask'])
|
||||
|
||||
# normalize embeddings
|
||||
query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
|
||||
doc_embeddings = F.normalize(doc_embeddings, p=2, dim=1)
|
||||
scores = (query_embeddings @ doc_embeddings.T) * 100
|
||||
print(scores.tolist())
|
||||
```
|
||||
|
||||
### Using Sentence-Transformers
|
||||
|
||||
You can also use the `bge` models with [sentence-transformers](https://www.sbert.net/). It currently supports `BAAI/bge-large-en-v1.5`, `BAAI/bge-base-en-v1.5`, `BAAI/bge-small-en-v1.5`, `BAAI/bge-large-zh-v1.5`, `BAAI/bge-base-zh-v1.5`, `BAAI/bge-small-zh-v1.5`, `BAAI/bge-large-en`, `BAAI/bge-base-en`, `BAAI/bge-small-en`, `BAAI/bge-large-zh`, `BAAI/bge-base-zh`, `BAAI/bge-small-zh'`, the **dense method** of `BAAI/bge-m3`, `BAAI/bge-multilingual-gemma2`:
|
||||
|
||||
```
|
||||
pip install -U sentence-transformers
|
||||
```
|
||||
|
||||
```shell
|
||||
from sentence_transformers import SentenceTransformer
|
||||
sentences_1 = ["样例数据-1", "样例数据-2"]
|
||||
sentences_2 = ["样例数据-3", "样例数据-4"]
|
||||
model = SentenceTransformer('BAAI/bge-large-zh-v1.5')
|
||||
embeddings_1 = model.encode(sentences_1, normalize_embeddings=True)
|
||||
embeddings_2 = model.encode(sentences_2, normalize_embeddings=True)
|
||||
similarity = embeddings_1 @ embeddings_2.T
|
||||
print(similarity)
|
||||
```
|
||||
|
||||
For s2p(short query to long passage) retrieval task, each short query should start with an instruction (instructions see [Model List](https://github.com/FlagOpen/FlagEmbedding/tree/master#model-list)). But the instruction is not needed for passages.
|
||||
|
||||
```shell
|
||||
from sentence_transformers import SentenceTransformer
|
||||
queries = ['query_1', 'query_2']
|
||||
passages = ["样例文档-1", "样例文档-2"]
|
||||
instruction = "为这个句子生成表示以用于检索相关文章："
|
||||
|
||||
model = SentenceTransformer('BAAI/bge-large-zh-v1.5')
|
||||
q_embeddings = model.encode([instruction+q for q in queries], normalize_embeddings=True)
|
||||
p_embeddings = model.encode(passages, normalize_embeddings=True)
|
||||
scores = q_embeddings @ p_embeddings.T
|
||||
```
|
||||
|
||||
### Using Langchain
|
||||
|
||||
You can use `bge` in langchain like this:
|
||||
|
||||
```python
|
||||
from langchain.embeddings import HuggingFaceBgeEmbeddings
|
||||
model_name = "BAAI/bge-large-en-v1.5"
|
||||
model_kwargs = {'device': 'cuda'}
|
||||
encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity
|
||||
model = HuggingFaceBgeEmbeddings(
|
||||
model_name=model_name,
|
||||
model_kwargs=model_kwargs,
|
||||
encode_kwargs=encode_kwargs,
|
||||
query_instruction="为这个句子生成表示以用于检索相关文章："
|
||||
)
|
||||
model.query_instruction = "为这个句子生成表示以用于检索相关文章："
|
||||
```
|
||||
|
||||
## Citation
|
||||
|
||||
If you find this repository useful, please consider giving a star :star: and citation
|
||||
|
||||
```
|
||||
@misc{bge_embedding,
|
||||
title={C-Pack: Packaged Resources To Advance General Chinese Embedding},
|
||||
author={Shitao Xiao and Zheng Liu and Peitian Zhang and Niklas Muennighoff},
|
||||
year={2023},
|
||||
eprint={2309.07597},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CL}
|
||||
}
|
||||
@misc{bge-m3,
|
||||
title={BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation},
|
||||
author={Jianlv Chen and Shitao Xiao and Peitian Zhang and Kun Luo and Defu Lian and Zheng Liu},
|
||||
year={2024},
|
||||
eprint={2402.03216},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CL}
|
||||
}
|
||||
@misc{li2024makingtextembeddersfewshot,
|
||||
title={Making Text Embedders Few-Shot Learners},
|
||||
author={Chaofan Li and MingHao Qin and Shitao Xiao and Jianlyu Chen and Kun Luo and Yingxia Shao and Defu Lian and Zheng Liu},
|
||||
year={2024},
|
||||
eprint={2409.15700},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.IR},
|
||||
url={https://arxiv.org/abs/2409.15700},
|
||||
}
|
||||
```
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoModel
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagAutoModel.from_finetuned(
|
||||
'BAAI/bge-multilingual-gemma2',
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
devices=["cuda:0", "cuda:1"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define"
|
||||
] * 100
|
||||
passages = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.558 0.02113 ]\n [0.01643 0.526 ]]")
|
||||
@@ -0,0 +1,34 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoModel
|
||||
|
||||
|
||||
def test_base_single_device():
|
||||
model = FlagAutoModel.from_finetuned(
|
||||
'BAAI/bge-multilingual-gemma2',
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
devices="cuda:0", # if you don't have a GPU, you can use "cpu"
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define"
|
||||
] * 100
|
||||
passages = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_single_device()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.558 0.0212 ]\n [0.01651 0.526 ]]")
|
||||
@@ -0,0 +1,48 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoModel
|
||||
|
||||
|
||||
def test_icl_multi_devices():
|
||||
examples = [
|
||||
{
|
||||
'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',
|
||||
'query': 'what is a virtual interface',
|
||||
'response': "A virtual interface is a software-defined abstraction that mimics the behavior and characteristics of a physical network interface. It allows multiple logical network connections to share the same physical network interface, enabling efficient utilization of network resources. Virtual interfaces are commonly used in virtualization technologies such as virtual machines and containers to provide network connectivity without requiring dedicated hardware. They facilitate flexible network configurations and help in isolating network traffic for security and management purposes."
|
||||
},
|
||||
{
|
||||
'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',
|
||||
'query': 'causes of back pain in female for a week',
|
||||
'response': "Back pain in females lasting a week can stem from various factors. Common causes include muscle strain due to lifting heavy objects or improper posture, spinal issues like herniated discs or osteoporosis, menstrual cramps causing referred pain, urinary tract infections, or pelvic inflammatory disease. Pregnancy-related changes can also contribute. Stress and lack of physical activity may exacerbate symptoms. Proper diagnosis by a healthcare professional is crucial for effective treatment and management."
|
||||
}
|
||||
]
|
||||
model = FlagAutoModel.from_finetuned(
|
||||
'BAAI/bge-en-icl',
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
examples_for_task=examples,
|
||||
examples_instruction_format="<instruct>{}\n<query>{}\n<response>{}",
|
||||
devices=["cuda:0", "cuda:1"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define"
|
||||
] * 100
|
||||
passages = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_icl_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.579 0.2776]\n [0.2249 0.5146]]")
|
||||
@@ -0,0 +1,48 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoModel
|
||||
|
||||
|
||||
def test_icl_single_device():
|
||||
examples = [
|
||||
{
|
||||
'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',
|
||||
'query': 'what is a virtual interface',
|
||||
'response': "A virtual interface is a software-defined abstraction that mimics the behavior and characteristics of a physical network interface. It allows multiple logical network connections to share the same physical network interface, enabling efficient utilization of network resources. Virtual interfaces are commonly used in virtualization technologies such as virtual machines and containers to provide network connectivity without requiring dedicated hardware. They facilitate flexible network configurations and help in isolating network traffic for security and management purposes."
|
||||
},
|
||||
{
|
||||
'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',
|
||||
'query': 'causes of back pain in female for a week',
|
||||
'response': "Back pain in females lasting a week can stem from various factors. Common causes include muscle strain due to lifting heavy objects or improper posture, spinal issues like herniated discs or osteoporosis, menstrual cramps causing referred pain, urinary tract infections, or pelvic inflammatory disease. Pregnancy-related changes can also contribute. Stress and lack of physical activity may exacerbate symptoms. Proper diagnosis by a healthcare professional is crucial for effective treatment and management."
|
||||
}
|
||||
]
|
||||
model = FlagAutoModel.from_finetuned(
|
||||
'BAAI/bge-en-icl',
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
examples_for_task=examples,
|
||||
examples_instruction_format="<instruct>{}\n<query>{}\n<response>{}",
|
||||
devices="cuda:0", # if you don't have a GPU, you can use "cpu"
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define"
|
||||
] * 100
|
||||
passages = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_icl_single_device()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.579 0.2776]\n [0.2249 0.5146]]")
|
||||
@@ -0,0 +1,42 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoModel
|
||||
|
||||
|
||||
def test_auto_pseudo_moe_multi_devices():
|
||||
model_name_or_path = "geevec-ai/geevec-embeddings-1.0-lite"
|
||||
|
||||
model = FlagAutoModel.from_finetuned(
|
||||
model_name_or_path,
|
||||
model_class="decoder-only-pseudo_moe",
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
query_instruction_format="Instruct: {}\nQuery: {}",
|
||||
domain_for_pseudo_moe="coding",
|
||||
use_fp16=False,
|
||||
use_bf16=True,
|
||||
trust_remote_code=True,
|
||||
devices=["cuda:0", "cuda:1"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv("HF_HUB_CACHE", None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define",
|
||||
] * 100
|
||||
passages = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day.",
|
||||
"Definition of summit for English Language Learners: the highest point of a mountain; the highest level; a meeting between leaders.",
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_auto_pseudo_moe_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.700 0.246]\n [0.158 0.654]]")
|
||||
@@ -0,0 +1,42 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoModel
|
||||
|
||||
|
||||
def test_auto_pseudo_moe_single_device():
|
||||
model_name_or_path = "geevec-ai/geevec-embeddings-1.0-lite"
|
||||
|
||||
model = FlagAutoModel.from_finetuned(
|
||||
model_name_or_path,
|
||||
model_class="decoder-only-pseudo_moe",
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
query_instruction_format="Instruct: {}\nQuery: {}",
|
||||
domain_for_pseudo_moe="reasoning",
|
||||
use_fp16=False,
|
||||
use_bf16=True,
|
||||
trust_remote_code=True,
|
||||
devices="cuda:0", # if you don't have a GPU, you can use "cpu"
|
||||
cache_dir=os.getenv("HF_HUB_CACHE", None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define",
|
||||
] * 10
|
||||
passages = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day.",
|
||||
"Definition of summit for English Language Learners: the highest point of a mountain; the highest level; a meeting between leaders.",
|
||||
] * 10
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_auto_pseudo_moe_single_device()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.844 0.466 ]\n [0.395 0.684 ]]")
|
||||
@@ -0,0 +1,35 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagLLMModel
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagLLMModel(
|
||||
'BAAI/bge-multilingual-gemma2',
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
query_instruction_format="<instruct>{}\n<query>{}",
|
||||
devices=["cuda:0", "cuda:1"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define"
|
||||
] * 100
|
||||
passages = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.558 0.02113 ]\n [0.01643 0.526 ]]")
|
||||
@@ -0,0 +1,35 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagLLMModel
|
||||
|
||||
|
||||
def test_base_single_device():
|
||||
model = FlagLLMModel(
|
||||
'BAAI/bge-multilingual-gemma2',
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
query_instruction_format="<instruct>{}\n<query>{}",
|
||||
devices="cuda:0", # if you don't have a GPU, you can use "cpu"
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define"
|
||||
] * 100
|
||||
passages = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_single_device()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.558 0.0212 ]\n [0.01651 0.526 ]]")
|
||||
@@ -0,0 +1,49 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagICLModel
|
||||
|
||||
|
||||
def test_icl_multi_devices():
|
||||
examples = [
|
||||
{
|
||||
'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',
|
||||
'query': 'what is a virtual interface',
|
||||
'response': "A virtual interface is a software-defined abstraction that mimics the behavior and characteristics of a physical network interface. It allows multiple logical network connections to share the same physical network interface, enabling efficient utilization of network resources. Virtual interfaces are commonly used in virtualization technologies such as virtual machines and containers to provide network connectivity without requiring dedicated hardware. They facilitate flexible network configurations and help in isolating network traffic for security and management purposes."
|
||||
},
|
||||
{
|
||||
'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',
|
||||
'query': 'causes of back pain in female for a week',
|
||||
'response': "Back pain in females lasting a week can stem from various factors. Common causes include muscle strain due to lifting heavy objects or improper posture, spinal issues like herniated discs or osteoporosis, menstrual cramps causing referred pain, urinary tract infections, or pelvic inflammatory disease. Pregnancy-related changes can also contribute. Stress and lack of physical activity may exacerbate symptoms. Proper diagnosis by a healthcare professional is crucial for effective treatment and management."
|
||||
}
|
||||
]
|
||||
model = FlagICLModel(
|
||||
'BAAI/bge-en-icl',
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
query_instruction_format="<instruct>{}\n<query>{}",
|
||||
examples_for_task=examples,
|
||||
examples_instruction_format="<instruct>{}\n<query>{}\n<response>{}",
|
||||
devices=["cuda:0", "cuda:1"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define"
|
||||
] * 100
|
||||
passages = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_icl_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.579 0.2776]\n [0.2249 0.5146]]")
|
||||
@@ -0,0 +1,49 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagICLModel
|
||||
|
||||
|
||||
def test_icl_single_device():
|
||||
examples = [
|
||||
{
|
||||
'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',
|
||||
'query': 'what is a virtual interface',
|
||||
'response': "A virtual interface is a software-defined abstraction that mimics the behavior and characteristics of a physical network interface. It allows multiple logical network connections to share the same physical network interface, enabling efficient utilization of network resources. Virtual interfaces are commonly used in virtualization technologies such as virtual machines and containers to provide network connectivity without requiring dedicated hardware. They facilitate flexible network configurations and help in isolating network traffic for security and management purposes."
|
||||
},
|
||||
{
|
||||
'instruct': 'Given a web search query, retrieve relevant passages that answer the query.',
|
||||
'query': 'causes of back pain in female for a week',
|
||||
'response': "Back pain in females lasting a week can stem from various factors. Common causes include muscle strain due to lifting heavy objects or improper posture, spinal issues like herniated discs or osteoporosis, menstrual cramps causing referred pain, urinary tract infections, or pelvic inflammatory disease. Pregnancy-related changes can also contribute. Stress and lack of physical activity may exacerbate symptoms. Proper diagnosis by a healthcare professional is crucial for effective treatment and management."
|
||||
}
|
||||
]
|
||||
model = FlagICLModel(
|
||||
'BAAI/bge-en-icl',
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
query_instruction_format="<instruct>{}\n<query>{}",
|
||||
examples_for_task=examples,
|
||||
examples_instruction_format="<instruct>{}\n<query>{}\n<response>{}",
|
||||
devices="cuda:0", # if you don't have a GPU, you can use "cpu"
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define"
|
||||
] * 100
|
||||
passages = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
|
||||
"Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_icl_single_device()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.579 0.2776]\n [0.2249 0.5146]]")
|
||||
@@ -0,0 +1,41 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagPseudoMoEModel
|
||||
|
||||
|
||||
def test_pseudo_moe_multi_devices():
|
||||
model_name_or_path = "geevec-ai/geevec-embeddings-1.0-lite"
|
||||
|
||||
model = FlagPseudoMoEModel(
|
||||
model_name_or_path,
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
query_instruction_format="Instruct: {}\nQuery: {}",
|
||||
domain_for_pseudo_moe="reasoning",
|
||||
use_fp16=False,
|
||||
use_bf16=True,
|
||||
trust_remote_code=True,
|
||||
devices=["cuda:0", "cuda:1"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv("HF_HUB_CACHE", None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define",
|
||||
] * 100
|
||||
passages = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day.",
|
||||
"Definition of summit for English Language Learners: the highest point of a mountain; the highest level; a meeting between leaders.",
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_pseudo_moe_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.844 0.466 ]\n [0.395 0.684 ]]")
|
||||
@@ -0,0 +1,41 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagPseudoMoEModel
|
||||
|
||||
|
||||
def test_pseudo_moe_single_device():
|
||||
model_name_or_path = "geevec-ai/geevec-embeddings-1.0-lite"
|
||||
|
||||
model = FlagPseudoMoEModel(
|
||||
model_name_or_path,
|
||||
query_instruction_for_retrieval="Given a question, retrieve passages that answer the question.",
|
||||
query_instruction_format="Instruct: {}\nQuery: {}",
|
||||
domain_for_pseudo_moe="coding",
|
||||
use_fp16=False,
|
||||
use_bf16=True,
|
||||
trust_remote_code=True,
|
||||
devices="cuda:0", # if you don't have a GPU, you can use "cpu"
|
||||
cache_dir=os.getenv("HF_HUB_CACHE", None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"how much protein should a female eat",
|
||||
"summit define",
|
||||
] * 10
|
||||
passages = [
|
||||
"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day.",
|
||||
"Definition of summit for English Language Learners: the highest point of a mountain; the highest level; a meeting between leaders.",
|
||||
] * 10
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_pseudo_moe_single_device()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.700 0.246]\n [0.158 0.654]]")
|
||||
@@ -0,0 +1,34 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoModel
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagAutoModel.from_finetuned(
|
||||
'BAAI/bge-small-en-v1.5',
|
||||
query_instruction_for_retrieval="Represent this sentence for searching relevant passages: ",
|
||||
devices=["cuda:0", "cuda:1"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"What is the capital of France?",
|
||||
"What is the population of China?",
|
||||
] * 100
|
||||
passages = [
|
||||
"Paris is the capital of France.",
|
||||
"The population of China is over 1.4 billion people."
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.7944 0.4492]\n [0.5806 0.801 ]]")
|
||||
@@ -0,0 +1,34 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoModel
|
||||
|
||||
|
||||
def test_base_single_device():
|
||||
model = FlagAutoModel.from_finetuned(
|
||||
'BAAI/bge-small-en-v1.5',
|
||||
query_instruction_for_retrieval="Represent this sentence for searching relevant passages: ",
|
||||
devices="cuda:0", # if you don't have a GPU, you can use "cpu"
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"What is the capital of France?",
|
||||
"What is the population of China?",
|
||||
] * 100
|
||||
passages = [
|
||||
"Paris is the capital of France.",
|
||||
"The population of China is over 1.4 billion people."
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_single_device()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.7944 0.4492]\n [0.58 0.801 ]]")
|
||||
@@ -0,0 +1,52 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoModel
|
||||
|
||||
|
||||
def test_m3_multi_devices():
|
||||
model = FlagAutoModel.from_finetuned(
|
||||
'BAAI/bge-m3',
|
||||
devices=["cuda:0", "cuda:1"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"What is BGE M3?",
|
||||
"Defination of BM25"
|
||||
] * 100
|
||||
passages = [
|
||||
"BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
|
||||
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(
|
||||
queries,
|
||||
return_dense=True,
|
||||
return_sparse=True,
|
||||
return_colbert_vecs=False,
|
||||
)
|
||||
passages_embeddings = model.encode_corpus(
|
||||
passages,
|
||||
return_dense=True,
|
||||
return_sparse=True,
|
||||
return_colbert_vecs=False,
|
||||
)
|
||||
|
||||
dense_scores = queries_embeddings["dense_vecs"] @ passages_embeddings["dense_vecs"].T
|
||||
sparse_scores = model.compute_lexical_matching_score(
|
||||
queries_embeddings["lexical_weights"],
|
||||
passages_embeddings["lexical_weights"],
|
||||
)
|
||||
|
||||
print("Dense score:\n", dense_scores[:2, :2])
|
||||
print("Sparse score:\n", sparse_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_m3_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("Dense score:")
|
||||
print(" [[0.626 0.3477]\n [0.3499 0.678 ]]")
|
||||
print("Sparse score:")
|
||||
print(" [[0.19561768 0.00878906]\n [0. 0.18030453]]")
|
||||
@@ -0,0 +1,52 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoModel
|
||||
|
||||
|
||||
def test_m3_single_device():
|
||||
model = FlagAutoModel.from_finetuned(
|
||||
'BAAI/bge-m3',
|
||||
devices="cuda:0", # if you don't have a GPU, you can use "cpu"
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"What is BGE M3?",
|
||||
"Defination of BM25"
|
||||
] * 100
|
||||
passages = [
|
||||
"BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
|
||||
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(
|
||||
queries,
|
||||
return_dense=True,
|
||||
return_sparse=True,
|
||||
return_colbert_vecs=False,
|
||||
)
|
||||
passages_embeddings = model.encode_corpus(
|
||||
passages,
|
||||
return_dense=True,
|
||||
return_sparse=True,
|
||||
return_colbert_vecs=False,
|
||||
)
|
||||
|
||||
dense_scores = queries_embeddings["dense_vecs"] @ passages_embeddings["dense_vecs"].T
|
||||
sparse_scores = model.compute_lexical_matching_score(
|
||||
queries_embeddings["lexical_weights"],
|
||||
passages_embeddings["lexical_weights"],
|
||||
)
|
||||
|
||||
print("Dense score:\n", dense_scores[:2, :2])
|
||||
print("Sparse score:\n", sparse_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_m3_single_device()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("Dense score:")
|
||||
print(" [[0.626 0.3477]\n [0.3496 0.678 ]]")
|
||||
print("Sparse score:")
|
||||
print(" [[0.19554901 0.00880432]\n [0. 0.18036556]]")
|
||||
@@ -0,0 +1,36 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagModel
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagModel(
|
||||
'BAAI/bge-small-en-v1.5',
|
||||
query_instruction_for_retrieval="Represent this sentence for searching relevant passages: ",
|
||||
query_instruction_format="{}{}",
|
||||
devices=["cuda:0", "cuda:1"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
pooling_method='cls',
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"What is the capital of France?",
|
||||
"What is the population of China?",
|
||||
] * 100
|
||||
passages = [
|
||||
"Paris is the capital of France.",
|
||||
"The population of China is over 1.4 billion people."
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.7944 0.4492]\n [0.5806 0.801 ]]")
|
||||
@@ -0,0 +1,36 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagModel
|
||||
|
||||
|
||||
def test_base_single_device():
|
||||
model = FlagModel(
|
||||
'BAAI/bge-small-en-v1.5',
|
||||
query_instruction_for_retrieval="Represent this sentence for searching relevant passages: ",
|
||||
query_instruction_format="{}{}",
|
||||
devices="cuda:0", # if you don't have a GPU, you can use "cpu"
|
||||
pooling_method='cls',
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"What is the capital of France?",
|
||||
"What is the population of China?",
|
||||
] * 100
|
||||
passages = [
|
||||
"Paris is the capital of France.",
|
||||
"The population of China is over 1.4 billion people."
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(queries)
|
||||
passages_embeddings = model.encode_corpus(passages)
|
||||
|
||||
cos_scores = queries_embeddings @ passages_embeddings.T
|
||||
print(cos_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_single_device()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[[0.7944 0.4492]\n [0.58 0.801 ]]")
|
||||
@@ -0,0 +1,53 @@
|
||||
import os
|
||||
from FlagEmbedding import BGEM3FlagModel
|
||||
|
||||
|
||||
def test_m3_multi_devices():
|
||||
model = BGEM3FlagModel(
|
||||
'BAAI/bge-m3',
|
||||
devices=["cuda:0", "cuda:1"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
pooling_method='cls',
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"What is BGE M3?",
|
||||
"Defination of BM25"
|
||||
] * 100
|
||||
passages = [
|
||||
"BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
|
||||
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(
|
||||
queries,
|
||||
return_dense=True,
|
||||
return_sparse=True,
|
||||
return_colbert_vecs=False,
|
||||
)
|
||||
passages_embeddings = model.encode_corpus(
|
||||
passages,
|
||||
return_dense=True,
|
||||
return_sparse=True,
|
||||
return_colbert_vecs=False,
|
||||
)
|
||||
|
||||
dense_scores = queries_embeddings["dense_vecs"] @ passages_embeddings["dense_vecs"].T
|
||||
sparse_scores = model.compute_lexical_matching_score(
|
||||
queries_embeddings["lexical_weights"],
|
||||
passages_embeddings["lexical_weights"],
|
||||
)
|
||||
|
||||
print("Dense score:\n", dense_scores[:2, :2])
|
||||
print("Sparse score:\n", sparse_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_m3_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("Dense score:")
|
||||
print(" [[0.626 0.3477]\n [0.3499 0.678 ]]")
|
||||
print("Sparse score:")
|
||||
print(" [[0.19561768 0.00878906]\n [0. 0.18030453]]")
|
||||
@@ -0,0 +1,55 @@
|
||||
import os
|
||||
from FlagEmbedding import BGEM3FlagModel
|
||||
|
||||
|
||||
def test_m3_multi_devices():
|
||||
model = BGEM3FlagModel(
|
||||
'BAAI/bge-m3',
|
||||
devices=["cuda:0", "cuda:1"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
pooling_method='cls',
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"What is BGE M3?",
|
||||
"Defination of BM25"
|
||||
] * 100
|
||||
passages = [
|
||||
"BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
|
||||
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"
|
||||
] * 100
|
||||
|
||||
sentence_pairs = list(zip(queries, passages))
|
||||
scores_dict = model.compute_score(
|
||||
sentence_pairs,
|
||||
weights_for_different_modes=[1., 0.3, 1.]
|
||||
)
|
||||
|
||||
queries.reverse()
|
||||
sentence_pairs = list(zip(queries, passages))
|
||||
|
||||
scores_dict_reverse = model.compute_score(
|
||||
sentence_pairs,
|
||||
weights_for_different_modes=[1., 0.3, 1.]
|
||||
)
|
||||
|
||||
scores_dict = {
|
||||
key: value[:2]
|
||||
for key, value in scores_dict.items()
|
||||
}
|
||||
scores_dict_reverse = {
|
||||
key: value[:2]
|
||||
for key, value in scores_dict_reverse.items()
|
||||
}
|
||||
|
||||
print(scores_dict)
|
||||
print(scores_dict_reverse)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_m3_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("{'colbert': [0.7798609733581543, 0.7897368669509888], 'sparse': [0.1956787109375, 0.1802978515625], 'dense': [0.6259765625, 0.67822265625], 'sparse+dense': [0.5266770720481873, 0.5633169412612915], 'colbert+sparse+dense': [0.6367570757865906, 0.6617604494094849]}")
|
||||
print("{'colbert': [0.4524071514606476, 0.4619773030281067], 'sparse': [0.0, 0.0087890625], 'dense': [0.349853515625, 0.34765625], 'sparse+dense': [0.2691181004047394, 0.269456148147583], 'colbert+sparse+dense': [0.34880897402763367, 0.3531610071659088]}")
|
||||
@@ -0,0 +1,53 @@
|
||||
import os
|
||||
from FlagEmbedding import BGEM3FlagModel
|
||||
|
||||
|
||||
def test_m3_single_device():
|
||||
model = BGEM3FlagModel(
|
||||
'BAAI/bge-m3',
|
||||
devices="cuda:0", # if you don't have a GPU, you can use "cpu"
|
||||
pooling_method='cls',
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"What is BGE M3?",
|
||||
"Defination of BM25"
|
||||
] * 100
|
||||
passages = [
|
||||
"BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
|
||||
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"
|
||||
] * 100
|
||||
|
||||
queries_embeddings = model.encode_queries(
|
||||
queries,
|
||||
return_dense=True,
|
||||
return_sparse=True,
|
||||
return_colbert_vecs=False,
|
||||
)
|
||||
passages_embeddings = model.encode_corpus(
|
||||
passages,
|
||||
return_dense=True,
|
||||
return_sparse=True,
|
||||
return_colbert_vecs=False,
|
||||
)
|
||||
|
||||
dense_scores = queries_embeddings["dense_vecs"] @ passages_embeddings["dense_vecs"].T
|
||||
sparse_scores = model.compute_lexical_matching_score(
|
||||
queries_embeddings["lexical_weights"],
|
||||
passages_embeddings["lexical_weights"],
|
||||
)
|
||||
|
||||
print("Dense score:\n", dense_scores[:2, :2])
|
||||
print("Sparse score:\n", sparse_scores[:2, :2])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_m3_single_device()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("Dense score:")
|
||||
print(" [[0.626 0.3477]\n [0.3496 0.678 ]]")
|
||||
print("Sparse score:")
|
||||
print(" [[0.19554901 0.00880432]\n [0. 0.18036556]]")
|
||||
@@ -0,0 +1,55 @@
|
||||
import os
|
||||
from FlagEmbedding import BGEM3FlagModel
|
||||
|
||||
|
||||
def test_m3_single_device():
|
||||
model = BGEM3FlagModel(
|
||||
'BAAI/bge-m3',
|
||||
devices="cuda:0", # if you don't have a GPU, you can use "cpu"
|
||||
pooling_method='cls',
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
queries = [
|
||||
"What is BGE M3?",
|
||||
"Defination of BM25"
|
||||
] * 100
|
||||
passages = [
|
||||
"BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
|
||||
"BM25 is a bag-of-words retrieval function that ranks a set of documents based on the query terms appearing in each document"
|
||||
] * 100
|
||||
|
||||
sentence_pairs = list(zip(queries, passages))
|
||||
scores_dict = model.compute_score(
|
||||
sentence_pairs,
|
||||
weights_for_different_modes=[1., 0.3, 1.]
|
||||
)
|
||||
|
||||
queries.reverse()
|
||||
sentence_pairs = list(zip(queries, passages))
|
||||
|
||||
scores_dict_reverse = model.compute_score(
|
||||
sentence_pairs,
|
||||
weights_for_different_modes=[1., 0.3, 1.]
|
||||
)
|
||||
|
||||
scores_dict = {
|
||||
key: value[:2]
|
||||
for key, value in scores_dict.items()
|
||||
}
|
||||
scores_dict_reverse = {
|
||||
key: value[:2]
|
||||
for key, value in scores_dict_reverse.items()
|
||||
}
|
||||
|
||||
print(scores_dict)
|
||||
print(scores_dict_reverse)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_m3_single_device()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("{'colbert': [0.7798250317573547, 0.7899274826049805], 'sparse': [0.195556640625, 0.180419921875], 'dense': [0.6259765625, 0.67822265625], 'sparse+dense': [0.5266488790512085, 0.5633450746536255], 'colbert+sparse+dense': [0.6367254853248596, 0.6618592143058777]}")
|
||||
print("{'colbert': [0.4524373412132263, 0.46213820576667786], 'sparse': [0.0, 0.0088043212890625], 'dense': [0.349609375, 0.34765625], 'sparse+dense': [0.2689302861690521, 0.26945966482162476], 'colbert+sparse+dense': [0.34871599078178406, 0.3532329499721527]}")
|
||||
@@ -0,0 +1,472 @@
|
||||
# Reranker
|
||||
|
||||
- [Model List](#model-list)
|
||||
- [Usage](#usage)
|
||||
- [Citation](#citation)
|
||||
|
||||
Different from embedding model, reranker uses question and document as input and directly output similarity instead of embedding.
|
||||
You can get a relevance score by inputting query and passage to the reranker.
|
||||
And the score can be mapped to a float value in [0,1] by sigmoid function.
|
||||
|
||||
For more detailed using, you can look [reranker-encoder only](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/inference/reranker/encoder_only) or [reranker-decoder only](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/inference/reranker/decoder_only)
|
||||
|
||||
## Model List
|
||||
|
||||
| Model | Base model | Language | layerwise | feature |
|
||||
|:--------------------------------------------------------------------------|:--------:|:-----------------------------------------------------------------------------------------------------------------------------------:|:----------------------------------------------------------------------------------------------:|:----------------------------------------------------------------------------------------------:|
|
||||
| [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. |
|
||||
| [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. |
|
||||
| [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. |
|
||||
| [BAAI/bge-reranker-v2-gemma](https://huggingface.co/BAAI/bge-reranker-v2-gemma) | [gemma-2b](https://huggingface.co/google/gemma-2b) | Multilingual | - | Suitable for multilingual contexts, performs well in both English proficiency and multilingual capabilities. |
|
||||
| [BAAI/bge-reranker-v2-minicpm-layerwise](https://huggingface.co/BAAI/bge-reranker-v2-minicpm-layerwise) | [MiniCPM-2B-dpo-bf16](https://huggingface.co/openbmb/MiniCPM-2B-dpo-bf16) | 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. |
|
||||
|
||||
|
||||
You can select the model according your senario and resource.
|
||||
- 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)
|
||||
|
||||
- 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).
|
||||
|
||||
- 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).
|
||||
|
||||
- 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)
|
||||
|
||||
## Usage
|
||||
### Using FlagEmbedding
|
||||
|
||||
#### 1. Auto Reranker
|
||||
|
||||
You can use `FlagAutoReranker` to load the model. For the **custom model** (not included in [`AUTO_RERANKER_MAPPING`](https://github.com/FlagOpen/FlagEmbedding/blob/master/FlagEmbedding/inference/reranker/model_mapping.py#L31)), you must specify the `model_class` parameter. You can also submit a pull request to add your **released model** to the [`AUTO_RERANKER_MAPPING`](https://github.com/FlagOpen/FlagEmbedding/blob/master/FlagEmbedding/inference/reranker/model_mapping.py#L31) dictionary. If need, you can create a new `<model>.py` file in [here](https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/inference/reranker/encoder_only) or [here](https://github.com/FlagOpen/FlagEmbedding/tree/master/FlagEmbedding/inference/reranker/decoder_only).
|
||||
|
||||
```python
|
||||
from FlagEmbedding import FlagAutoReranker
|
||||
reranker = FlagAutoReranker.from_finetuned('BAAI/bge-reranker-large',
|
||||
query_max_length=256,
|
||||
passage_max_length=512,
|
||||
use_fp16=True,
|
||||
devices=['cuda:1']) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
score = reranker.compute_score(['query', 'passage'])
|
||||
print(score) # -1.5263671875
|
||||
|
||||
# You can map the scores into 0-1 by set "normalize=True", which will apply sigmoid function to the score
|
||||
score = reranker.compute_score(['query', 'passage'], normalize=True)
|
||||
print(score) # 0.1785258315203034
|
||||
|
||||
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.']])
|
||||
print(scores) # [-5.60546875, 5.76171875]
|
||||
|
||||
# You can map the scores into 0-1 by set "normalize=True", which will apply sigmoid function to the score
|
||||
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)
|
||||
print(scores) # [0.0036642203307843528, 0.9968641641227171]
|
||||
```
|
||||
|
||||
For your **custom model** (assume the model is finetuned from `BAAI/bge-reranker-large`, then the model class is `encoder-only-base`), you can use the following code:
|
||||
|
||||
```python
|
||||
from FlagEmbedding import FlagAutoReranker
|
||||
reranker = FlagAutoReranker.from_finetuned('your_model_name_or_path',
|
||||
model_class='encoder-only-base',
|
||||
query_max_length=256,
|
||||
passage_max_length=512,
|
||||
use_fp16=True,
|
||||
devices=['cuda:1']) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
score = reranker.compute_score(['query', 'passage'])
|
||||
print(score)
|
||||
```
|
||||
|
||||
The `model_class` parameter currently includes the following options:
|
||||
- `encoder-only-base`: for encoder-only reranker model, such as `BAAI/bge-reranker-large`
|
||||
- `decoder-only-base`: for decoder-only reranker model, such as `BAAI/bge-reranker-v2-gemma`
|
||||
- `decoder-only-layerwise`: for decoder-only layerwise reranker model, such as `BAAI/bge-reranker-v2-minicpm-layerwise`
|
||||
- `decoder-only-lightweight`: for decoder-only lightweight reranker model, such as `BAAI/bge-reranker-v2.5-gemma2-lightweight`
|
||||
|
||||
#### 2. Normal Reranker
|
||||
|
||||
For `FlagReranker`, it supports `BAAI/bge-reranker-base`, `BAAI/bge-reranker-large`, `BAAI/bge-reranker-v2-m3`:
|
||||
|
||||
```python
|
||||
from FlagEmbedding import FlagReranker
|
||||
reranker = FlagReranker(
|
||||
'BAAI/bge-reranker-v2-m3',
|
||||
query_max_length=256,
|
||||
passage_max_length=512,
|
||||
use_fp16=True,
|
||||
devices=['cuda:1']
|
||||
) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
|
||||
score = reranker.compute_score(['query', 'passage'])
|
||||
print(score) # -5.65234375
|
||||
|
||||
# You can map the scores into 0-1 by set "normalize=True", which will apply sigmoid function to the score
|
||||
score = reranker.compute_score(['query', 'passage'], normalize=True)
|
||||
print(score) # 0.003497010252573502
|
||||
|
||||
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.']])
|
||||
print(scores) # [-8.1875, 5.26171875]
|
||||
|
||||
# You can map the scores into 0-1 by set "normalize=True", which will apply sigmoid function to the score
|
||||
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)
|
||||
print(scores) # [0.00027803096387751553, 0.9948403768236574]
|
||||
```
|
||||
|
||||
#### 3. LLM-based Reranker
|
||||
|
||||
For `FlagLLMReranker`, it supports `BAAI/bge-reranker-v2-gemma`:
|
||||
|
||||
```python
|
||||
from FlagEmbedding import FlagLLMReranker
|
||||
reranker = FlagLLMReranker(
|
||||
'BAAI/bge-reranker-v2-gemma',
|
||||
query_max_length=256,
|
||||
passage_max_length=512,
|
||||
use_fp16=True,
|
||||
devices=['cuda:1']
|
||||
) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
|
||||
score = reranker.compute_score(['query', 'passage'])
|
||||
print(score)
|
||||
|
||||
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.']])
|
||||
print(scores)
|
||||
```
|
||||
|
||||
#### 4. LLM-based Layerwise Reranker
|
||||
|
||||
For `LayerWiseFlagLLMReranker`, it supports `BAAI/bge-reranker-v2-minicpm-layerwise`:
|
||||
|
||||
```python
|
||||
from FlagEmbedding import LayerWiseFlagLLMReranker
|
||||
reranker = LayerWiseFlagLLMReranker(
|
||||
'BAAI/bge-reranker-v2-minicpm-layerwise',
|
||||
query_max_length=256,
|
||||
passage_max_length=512,
|
||||
use_fp16=True,
|
||||
devices=['cuda:1']
|
||||
) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
|
||||
score = reranker.compute_score(['query', 'passage'], cutoff_layers=[28]) # Adjusting 'cutoff_layers' to pick which layers are used for computing the score.
|
||||
print(score)
|
||||
|
||||
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])
|
||||
print(scores)
|
||||
```
|
||||
|
||||
#### 5. LLM-based lightweight Reranker
|
||||
|
||||
For `LightWeightFlagLLMReranker`, it supports `BAAI/bge-reranker-v2.5-gemma2-lightweight`:
|
||||
|
||||
```python
|
||||
from FlagEmbedding import LightWeightFlagLLMReranker
|
||||
reranker = LightWeightFlagLLMReranker(
|
||||
'BAAI/bge-reranker-v2.5-gemma2-lightweight',
|
||||
query_max_length=256,
|
||||
passage_max_length=512,
|
||||
use_fp16=True,
|
||||
devices=['cuda:1']
|
||||
) # Setting use_fp16 to True speeds up computation with a slight performance degradation
|
||||
|
||||
score = reranker.compute_score(['query', 'passage'], cutoff_layers=[28], compress_ratio=2, compress_layers=[24, 40]) # Adjusting 'cutoff_layers' to pick which layers are used for computing the score.
|
||||
print(score)
|
||||
|
||||
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], compress_ratio=2, compress_layers=[24, 40])
|
||||
print(scores)
|
||||
```
|
||||
|
||||
### Using Huggingface transformers
|
||||
|
||||
#### 1. Normal Reranker
|
||||
|
||||
It supports `BAAI/bge-reranker-base`, `BAAI/bge-reranker-large`, `BAAI/bge-reranker-v2-m3`:
|
||||
|
||||
```python
|
||||
import torch
|
||||
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-m3')
|
||||
model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-v2-m3')
|
||||
model.eval()
|
||||
|
||||
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.']]
|
||||
with torch.no_grad():
|
||||
inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
|
||||
scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
|
||||
print(scores)
|
||||
```
|
||||
|
||||
#### 2. LLM-based reranker
|
||||
|
||||
It supports `BAAI/bge-reranker-v2-gemma`:
|
||||
|
||||
```python
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
def get_inputs(pairs, tokenizer, prompt=None, max_length=1024):
|
||||
if prompt is None:
|
||||
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'."
|
||||
sep = "\n"
|
||||
prompt_inputs = tokenizer(prompt,
|
||||
return_tensors=None,
|
||||
add_special_tokens=False)['input_ids']
|
||||
sep_inputs = tokenizer(sep,
|
||||
return_tensors=None,
|
||||
add_special_tokens=False)['input_ids']
|
||||
inputs = []
|
||||
for query, passage in pairs:
|
||||
query_inputs = tokenizer(f'A: {query}',
|
||||
return_tensors=None,
|
||||
add_special_tokens=False,
|
||||
max_length=max_length * 3 // 4,
|
||||
truncation=True)
|
||||
passage_inputs = tokenizer(f'B: {passage}',
|
||||
return_tensors=None,
|
||||
add_special_tokens=False,
|
||||
max_length=max_length,
|
||||
truncation=True)
|
||||
item = tokenizer.prepare_for_model(
|
||||
[tokenizer.bos_token_id] + query_inputs['input_ids'],
|
||||
sep_inputs + passage_inputs['input_ids'],
|
||||
truncation='only_second',
|
||||
max_length=max_length,
|
||||
padding=False,
|
||||
return_attention_mask=False,
|
||||
return_token_type_ids=False,
|
||||
add_special_tokens=False
|
||||
)
|
||||
item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
|
||||
item['attention_mask'] = [1] * len(item['input_ids'])
|
||||
inputs.append(item)
|
||||
return tokenizer.pad(
|
||||
inputs,
|
||||
padding=True,
|
||||
max_length=max_length + len(sep_inputs) + len(prompt_inputs),
|
||||
pad_to_multiple_of=8,
|
||||
return_tensors='pt',
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-gemma')
|
||||
model = AutoModelForCausalLM.from_pretrained('BAAI/bge-reranker-v2-gemma')
|
||||
yes_loc = tokenizer('Yes', add_special_tokens=False)['input_ids'][0]
|
||||
model.eval()
|
||||
|
||||
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.']]
|
||||
with torch.no_grad():
|
||||
inputs = get_inputs(pairs, tokenizer)
|
||||
scores = model(**inputs, return_dict=True).logits[:, -1, yes_loc].view(-1, ).float()
|
||||
print(scores)
|
||||
```
|
||||
|
||||
#### 3. LLM-based layerwise reranker
|
||||
|
||||
It supports `BAAI/bge-reranker-v2-minicpm-layerwise`:
|
||||
|
||||
```python
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
def get_inputs(pairs, tokenizer, prompt=None, max_length=1024):
|
||||
if prompt is None:
|
||||
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'."
|
||||
sep = "\n"
|
||||
prompt_inputs = tokenizer(prompt,
|
||||
return_tensors=None,
|
||||
add_special_tokens=False)['input_ids']
|
||||
sep_inputs = tokenizer(sep,
|
||||
return_tensors=None,
|
||||
add_special_tokens=False)['input_ids']
|
||||
inputs = []
|
||||
for query, passage in pairs:
|
||||
query_inputs = tokenizer(f'A: {query}',
|
||||
return_tensors=None,
|
||||
add_special_tokens=False,
|
||||
max_length=max_length * 3 // 4,
|
||||
truncation=True)
|
||||
passage_inputs = tokenizer(f'B: {passage}',
|
||||
return_tensors=None,
|
||||
add_special_tokens=False,
|
||||
max_length=max_length,
|
||||
truncation=True)
|
||||
item = tokenizer.prepare_for_model(
|
||||
[tokenizer.bos_token_id] + query_inputs['input_ids'],
|
||||
sep_inputs + passage_inputs['input_ids'],
|
||||
truncation='only_second',
|
||||
max_length=max_length,
|
||||
padding=False,
|
||||
return_attention_mask=False,
|
||||
return_token_type_ids=False,
|
||||
add_special_tokens=False
|
||||
)
|
||||
item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
|
||||
item['attention_mask'] = [1] * len(item['input_ids'])
|
||||
inputs.append(item)
|
||||
return tokenizer.pad(
|
||||
inputs,
|
||||
padding=True,
|
||||
max_length=max_length + len(sep_inputs) + len(prompt_inputs),
|
||||
pad_to_multiple_of=8,
|
||||
return_tensors='pt',
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-minicpm-layerwise', trust_remote_code=True)
|
||||
model = AutoModelForCausalLM.from_pretrained('BAAI/bge-reranker-v2-minicpm-layerwise', trust_remote_code=True, torch_dtype=torch.bfloat16)
|
||||
model = model.to('cuda')
|
||||
model.eval()
|
||||
|
||||
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.']]
|
||||
with torch.no_grad():
|
||||
inputs = get_inputs(pairs, tokenizer).to(model.device)
|
||||
all_scores = model(**inputs, return_dict=True, cutoff_layers=[28])
|
||||
all_scores = [scores[:, -1].view(-1, ).float() for scores in all_scores[0]]
|
||||
print(all_scores)
|
||||
```
|
||||
|
||||
#### 4. LLM-based lightweight reranker
|
||||
|
||||
It supports `BAAI/bge-reranker-v2.5-gemma2-lightweight`:
|
||||
|
||||
```python
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
def last_logit_pool(logits: torch.Tensor,
|
||||
attention_mask: torch.Tensor) -> torch.Tensor:
|
||||
left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
|
||||
if left_padding:
|
||||
return logits[:, -1]
|
||||
else:
|
||||
sequence_lengths = attention_mask.sum(dim=1) - 1
|
||||
batch_size = logits.shape[0]
|
||||
return torch.stack([logits[i, sequence_lengths[i]] for i in range(batch_size)], dim=0)
|
||||
|
||||
def get_inputs(pairs, tokenizer, prompt=None, max_length=1024):
|
||||
if prompt is None:
|
||||
prompt = "Predict whether passage B contains an answer to query A."
|
||||
sep = "\n"
|
||||
prompt_inputs = tokenizer(prompt,
|
||||
return_tensors=None,
|
||||
add_special_tokens=False)['input_ids']
|
||||
sep_inputs = tokenizer(sep,
|
||||
return_tensors=None,
|
||||
add_special_tokens=False)['input_ids']
|
||||
inputs = []
|
||||
query_lengths = []
|
||||
prompt_lengths = []
|
||||
for query, passage in pairs:
|
||||
query_inputs = tokenizer(f'A: {query}',
|
||||
return_tensors=None,
|
||||
add_special_tokens=False,
|
||||
max_length=max_length * 3 // 4,
|
||||
truncation=True)
|
||||
passage_inputs = tokenizer(f'B: {passage}',
|
||||
return_tensors=None,
|
||||
add_special_tokens=False,
|
||||
max_length=max_length,
|
||||
truncation=True)
|
||||
item = tokenizer.prepare_for_model(
|
||||
[tokenizer.bos_token_id] + query_inputs['input_ids'],
|
||||
sep_inputs + passage_inputs['input_ids'],
|
||||
truncation='only_second',
|
||||
max_length=max_length,
|
||||
padding=False,
|
||||
return_attention_mask=False,
|
||||
return_token_type_ids=False,
|
||||
add_special_tokens=False
|
||||
)
|
||||
item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
|
||||
item['attention_mask'] = [1] * len(item['input_ids'])
|
||||
inputs.append(item)
|
||||
query_lengths.append(len([tokenizer.bos_token_id] + query_inputs['input_ids'] + sep_inputs))
|
||||
prompt_lengths.append(len(sep_inputs + prompt_inputs))
|
||||
|
||||
return tokenizer.pad(
|
||||
inputs,
|
||||
padding=True,
|
||||
max_length=max_length + len(sep_inputs) + len(prompt_inputs),
|
||||
pad_to_multiple_of=8,
|
||||
return_tensors='pt',
|
||||
), query_lengths, prompt_lengths
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2.5-gemma2-lightweight', trust_remote_code=True)
|
||||
tokenizer.padding_side = 'right'
|
||||
model = AutoModelForCausalLM.from_pretrained('BAAI/bge-reranker-v2.5-gemma2-lightweight', trust_remote_code=True)
|
||||
model = model.to('cuda')
|
||||
model.eval()
|
||||
|
||||
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.']]
|
||||
with torch.no_grad():
|
||||
inputs, query_lengths, prompt_lengths = get_inputs(pairs, tokenizer)
|
||||
inputs = inputs.to(model.device)
|
||||
outputs = model(**inputs,
|
||||
return_dict=True,
|
||||
cutoff_layers=[28],
|
||||
compress_ratio=2,
|
||||
compress_layer=[24, 40],
|
||||
query_lengths=query_lengths,
|
||||
prompt_lengths=prompt_lengths)
|
||||
scores = []
|
||||
for i in range(len(outputs.logits)):
|
||||
logits = last_logit_pool(outputs.logits[i], outputs.attention_masks[i])
|
||||
scores.append(logits.cpu().float().tolist())
|
||||
print(scores)
|
||||
```
|
||||
|
||||
## Load model in local
|
||||
|
||||
### Load llm-based layerwise reranker in local
|
||||
|
||||
If you download reranker-v2-minicpm-layerwise, you can load it with the following method:
|
||||
|
||||
1. make sure `configuration_minicpm_reranker.py` and `modeling_minicpm_reranker.py` from [BAAI/bge-reranker-v2-minicpm-layerwise](https://huggingface.co/BAAI/bge-reranker-v2-minicpm-layerwise) in your local path.
|
||||
2. modify the following part of `config.json`:
|
||||
|
||||
```
|
||||
"auto_map": {
|
||||
"AutoConfig": "configuration_minicpm_reranker.LayerWiseMiniCPMConfig",
|
||||
"AutoModel": "modeling_minicpm_reranker.LayerWiseMiniCPMModel",
|
||||
"AutoModelForCausalLM": "modeling_minicpm_reranker.LayerWiseMiniCPMForCausalLM"
|
||||
},
|
||||
```
|
||||
|
||||
### Load llm-based lightweight reranker in local
|
||||
|
||||
1. make sure `gemma_config.py` and `gemma_model.py` from [BAAI/bge-reranker-v2.5-gemma2-lightweight](https://huggingface.co/BAAI/bge-reranker-v2.5-gemma2-lightweight/tree/main) in your local path.
|
||||
2. modify the following part of config.json:
|
||||
|
||||
```
|
||||
"auto_map": {
|
||||
"AutoConfig": "gemma_config.CostWiseGemmaConfig",
|
||||
"AutoModel": "gemma_model.CostWiseGemmaModel",
|
||||
"AutoModelForCausalLM": "gemma_model.CostWiseGemmaForCausalLM"
|
||||
},
|
||||
```
|
||||
|
||||
## Citation
|
||||
|
||||
If you find this repository useful, please consider giving a star :star: and citation
|
||||
|
||||
```
|
||||
@misc{li2023making,
|
||||
title={Making Large Language Models A Better Foundation For Dense Retrieval},
|
||||
author={Chaofan Li and Zheng Liu and Shitao Xiao and Yingxia Shao},
|
||||
year={2023},
|
||||
eprint={2312.15503},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CL}
|
||||
}
|
||||
@misc{chen2024bge,
|
||||
title={BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation},
|
||||
author={Jianlv Chen and Shitao Xiao and Peitian Zhang and Kun Luo and Defu Lian and Zheng Liu},
|
||||
year={2024},
|
||||
eprint={2402.03216},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CL}
|
||||
}
|
||||
@misc{li2024makingtextembeddersfewshot,
|
||||
title={Making Text Embedders Few-Shot Learners},
|
||||
author={Chaofan Li and MingHao Qin and Shitao Xiao and Jianlyu Chen and Kun Luo and Yingxia Shao and Defu Lian and Zheng Liu},
|
||||
year={2024},
|
||||
eprint={2409.15700},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.IR},
|
||||
url={https://arxiv.org/abs/2409.15700},
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,32 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagAutoReranker.from_finetuned(
|
||||
'BAAI/bge-reranker-v2-gemma',
|
||||
use_fp16=True,
|
||||
query_instruction_for_rerank="A: ",
|
||||
passage_instruction_for_rerank="B: ",
|
||||
devices=["cuda:3", "cuda:4"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs)
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[ 9.1484375 -4.50390625 -5.53125 10.21875 ]")
|
||||
@@ -0,0 +1,32 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagAutoReranker.from_finetuned(
|
||||
'BAAI/bge-reranker-v2-gemma',
|
||||
use_fp16=True,
|
||||
query_instruction_for_rerank="A: ",
|
||||
passage_instruction_for_rerank="B: ",
|
||||
devices=["cuda:3"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs)
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[9.171875, -4.49609375, -5.5234375, 10.2109375]")
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagAutoReranker.from_finetuned(
|
||||
'BAAI/bge-reranker-v2-minicpm-layerwise',
|
||||
use_fp16=True,
|
||||
query_instruction_for_rerank="A: ",
|
||||
passage_instruction_for_rerank="B: ",
|
||||
trust_remote_code=True,
|
||||
devices=["cuda:3", "cuda:4"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs, cutoff_layers=[28])
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[1.939453125, -12.71875, -11.78125, 2.189453125]")
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagAutoReranker.from_finetuned(
|
||||
'BAAI/bge-reranker-v2-minicpm-layerwise',
|
||||
use_fp16=True,
|
||||
query_instruction_for_rerank="A: ",
|
||||
passage_instruction_for_rerank="B: ",
|
||||
trust_remote_code=True,
|
||||
devices=["cuda:3"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs, cutoff_layers=[28])
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[1.939453125, -12.71875, -11.78125, 2.189453125]")
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagAutoReranker.from_finetuned(
|
||||
'BAAI/bge-reranker-v2.5-gemma2-lightweight',
|
||||
use_fp16=True,
|
||||
query_instruction_for_rerank="A: ",
|
||||
passage_instruction_for_rerank="B: ",
|
||||
trust_remote_code=True,
|
||||
devices=["cuda:3", "cuda:4"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs, cutoff_layers=[28], compress_ratio=2, compress_layers=[24, 40])
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[25.375, 8.734375, 9.8359375, 26.15625]")
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagAutoReranker.from_finetuned(
|
||||
'BAAI/bge-reranker-v2.5-gemma2-lightweight',
|
||||
use_fp16=True,
|
||||
query_instruction_for_rerank="A: ",
|
||||
passage_instruction_for_rerank="B: ",
|
||||
trust_remote_code=True,
|
||||
devices=["cuda:3"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs, cutoff_layers=[28], compress_ratio=2, compress_layers=[24, 40])
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[25.375, 8.734375, 9.8359375, 26.15625]")
|
||||
@@ -0,0 +1,32 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagLLMReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagLLMReranker(
|
||||
'BAAI/bge-reranker-v2-gemma',
|
||||
use_fp16=True,
|
||||
query_instruction_for_rerank="A: ",
|
||||
passage_instruction_for_rerank="B: ",
|
||||
devices=["cuda:3", "cuda:4"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs)
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[ 9.1484375 -4.50390625 -5.53125 10.21875 ]")
|
||||
@@ -0,0 +1,32 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagLLMReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagLLMReranker(
|
||||
'BAAI/bge-reranker-v2-gemma',
|
||||
use_fp16=True,
|
||||
query_instruction_for_rerank="A: ",
|
||||
passage_instruction_for_rerank="B: ",
|
||||
devices=["cuda:3"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs)
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[9.171875, -4.49609375, -5.5234375, 10.2109375]")
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
from FlagEmbedding import LayerWiseFlagLLMReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = LayerWiseFlagLLMReranker(
|
||||
'BAAI/bge-reranker-v2-minicpm-layerwise',
|
||||
use_fp16=True,
|
||||
query_instruction_for_rerank="A: ",
|
||||
passage_instruction_for_rerank="B: ",
|
||||
trust_remote_code=True,
|
||||
devices=["cuda:3"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs, cutoff_layers=[28])
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[1.939453125, -12.71875, -11.78125, 2.189453125]")
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
from FlagEmbedding import LayerWiseFlagLLMReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = LayerWiseFlagLLMReranker(
|
||||
'BAAI/bge-reranker-v2-minicpm-layerwise',
|
||||
use_fp16=True,
|
||||
query_instruction_for_rerank="A: ",
|
||||
passage_instruction_for_rerank="B: ",
|
||||
trust_remote_code=True,
|
||||
devices=["cuda:3"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs, cutoff_layers=[28])
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[1.939453125, -12.71875, -11.78125, 2.189453125]")
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
from FlagEmbedding import LightWeightFlagLLMReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = LightWeightFlagLLMReranker(
|
||||
'BAAI/bge-reranker-v2.5-gemma2-lightweight',
|
||||
use_fp16=True,
|
||||
query_instruction_for_rerank="A: ",
|
||||
passage_instruction_for_rerank="B: ",
|
||||
trust_remote_code=True,
|
||||
devices=["cuda:3", "cuda:4"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs, cutoff_layers=[28], compress_ratio=2, compress_layers=[24, 40])
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[25.375, 8.734375, 9.8359375, 26.15625]")
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
from FlagEmbedding import LightWeightFlagLLMReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = LightWeightFlagLLMReranker(
|
||||
'BAAI/bge-reranker-v2.5-gemma2-lightweight',
|
||||
use_fp16=True,
|
||||
query_instruction_for_rerank="A: ",
|
||||
passage_instruction_for_rerank="B: ",
|
||||
trust_remote_code=True,
|
||||
devices=["cuda:3"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs, cutoff_layers=[28], compress_ratio=2, compress_layers=[24, 40])
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[25.375, 8.734375, 9.8359375, 26.15625]")
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagAutoReranker.from_finetuned(
|
||||
'BAAI/bge-reranker-large',
|
||||
use_fp16=True,
|
||||
batch_size=128,
|
||||
query_max_length=256,
|
||||
max_length=512,
|
||||
devices=["cuda:3", "cuda:4"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs)
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[ 7.97265625 -6.8515625 -7.15625 5.45703125]")
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagAutoReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagAutoReranker.from_finetuned(
|
||||
'BAAI/bge-reranker-large',
|
||||
use_fp16=True,
|
||||
batch_size=128,
|
||||
query_max_length=256,
|
||||
max_length=512,
|
||||
devices=["cuda:3"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs)
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[7.9765625, -6.84375, -7.15625, 5.453125]")
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagReranker(
|
||||
'BAAI/bge-reranker-large',
|
||||
use_fp16=True,
|
||||
batch_size=128,
|
||||
query_max_length=256,
|
||||
max_length=512,
|
||||
devices=["cuda:3", "cuda:4"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs)
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[ 7.97265625 -6.8515625 -7.15625 5.45703125]")
|
||||
@@ -0,0 +1,33 @@
|
||||
import os
|
||||
from FlagEmbedding import FlagReranker
|
||||
|
||||
|
||||
def test_base_multi_devices():
|
||||
model = FlagReranker(
|
||||
'BAAI/bge-reranker-large',
|
||||
use_fp16=True,
|
||||
batch_size=128,
|
||||
query_max_length=256,
|
||||
max_length=512,
|
||||
devices=["cuda:3"], # if you don't have GPUs, you can use ["cpu", "cpu"]
|
||||
cache_dir=os.getenv('HF_HUB_CACHE', None),
|
||||
)
|
||||
|
||||
pairs = [
|
||||
["What is the capital of France?", "Paris is the capital of France."],
|
||||
["What is the capital of France?", "The population of China is over 1.4 billion people."],
|
||||
["What is the population of China?", "Paris is the capital of France."],
|
||||
["What is the population of China?", "The population of China is over 1.4 billion people."]
|
||||
] * 100
|
||||
|
||||
scores = model.compute_score(pairs)
|
||||
|
||||
print(scores[:4])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_base_multi_devices()
|
||||
|
||||
print("--------------------------------")
|
||||
print("Expected Output:")
|
||||
print("[7.9765625, -6.84375, -7.15625, 5.453125]")
|
||||
Reference in New Issue
Block a user