Summary: This tutorial provides implementation guidance for the BGE (BAAI General Embedding) series of embedding models, specifically helping with text embedding and retrieval tasks. It covers code examples for different BGE variants including base models, BGE v1.5, BGE M3 (multi-functional/lingual/granularity), BGE Multilingual Gemma2, and BGE ICL (in-context learning). Key functionalities include dense/sparse retrieval, multilingual support (100+ languages), long text processing (up to 8192 tokens), and few-shot learning capabilities. The tutorial demonstrates essential implementation patterns using the FlagEmbedding library, including model initialization, query/corpus encoding, and parameter configuration for optimization.

# BGE Series

In this Part, we will walk through the BGE series and introduce how to use the BGE embedding models.

## 1. BAAI General Embedding

BGE stands for BAAI General Embedding, it's a series of embeddings models developed and published by Beijing Academy of Artificial Intelligence (BAAI).

A full support of APIs and related usages of BGE is maintained in [FlagEmbedding](https://github.com/FlagOpen/FlagEmbedding) on GitHub.

Run the following cell to install FlagEmbedding in your environment.


```python
%%capture
%pip install -U FlagEmbedding
```


```python
import os 
os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = 'true'
# single GPU is better for small tasks
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
```

The collection of BGE models can be found in [Huggingface collection](https://huggingface.co/collections/BAAI/bge-66797a74476eb1f085c7446d).

## 2. BGE Series Models

### 2.1 BGE

The very first version of BGE has 6 models, with 'large', 'base', and 'small' for English and Chinese. 

| Model  | Language |   Parameters   |   Model Size   |    Description    |   Base Model     |
|:-------|:--------:|:--------------:|:--------------:|:-----------------:|:----------------:|
| [BAAI/bge-large-en](https://huggingface.co/BAAI/bge-large-en)   | English |    500M    |    1.34 GB   |              Embedding Model which map text into vector                            |  BERT  |
| [BAAI/bge-base-en](https://huggingface.co/BAAI/bge-base-en)     | English |    109M    |    438 MB    |          a base-scale model but with similar ability to `bge-large-en`  |  BERT  |
| [BAAI/bge-small-en](https://huggingface.co/BAAI/bge-small-en)   | English |    33.4M   |    133 MB    |          a small-scale model but with competitive performance                    |  BERT  |
| [BAAI/bge-large-zh](https://huggingface.co/BAAI/bge-large-zh)   | Chinese |    326M    |    1.3 GB    |              Embedding Model which map text into vector                            |  BERT  |
| [BAAI/bge-base-zh](https://huggingface.co/BAAI/bge-base-zh)     | Chinese |    102M    |    409 MB    |           a base-scale model but with similar ability to `bge-large-zh`           |  BERT  |
| [BAAI/bge-small-zh](https://huggingface.co/BAAI/bge-small-zh)   | Chinese |    24M     |    95.8 MB   |           a small-scale model but with competitive performance                    |  BERT  |

For inference, simply import FlagModel from FlagEmbedding and initialize the model.


```python
from FlagEmbedding import FlagModel

# Load BGE model
model = FlagModel(
    'BAAI/bge-base-en',
    query_instruction_for_retrieval="Represent this sentence for searching relevant passages:",
    query_instruction_format='{}{}',
)

queries = ["query 1", "query 2"]
corpus = ["passage 1", "passage 2"]

# encode the queries and corpus
q_embeddings = model.encode_queries(queries)
p_embeddings = model.encode_corpus(corpus)

# compute the similarity scores
scores = q_embeddings @ p_embeddings.T
print(scores)
```

For general encoding, use either `encode()`:
```python
FlagModel.encode(sentences, batch_size=256, max_length=512, convert_to_numpy=True)
```
or `encode_corpus()` that directly calls `encode()`:
```python
FlagModel.encode_corpus(corpus, batch_size=256, max_length=512, convert_to_numpy=True)
```
The *encode_queries()* function concatenate the `query_instruction_for_retrieval` with each of the input query to form the new sentences and then feed them to `encode()`.
```python
FlagModel.encode_queries(queries, batch_size=256, max_length=512, convert_to_numpy=True)
```

### 2.2 BGE v1.5

BGE 1.5 alleviate the issue of the similarity distribution, and enhance retrieval ability without instruction.

| Model  | Language |   Parameters   |   Model Size   |    Description    |   Base Model     |
|:-------|:--------:|:--------------:|:--------------:|:-----------------:|:----------------:|
| [BAAI/bge-large-en-v1.5](https://huggingface.co/BAAI/bge-large-en-v1.5)   | English |    335M    |    1.34 GB   |     version 1.5 with more reasonable similarity distribution      |   BERT   |
| [BAAI/bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5)     | English |    109M    |    438 MB    |     version 1.5 with more reasonable similarity distribution      |   BERT   |
| [BAAI/bge-small-en-v1.5](https://huggingface.co/BAAI/bge-small-en-v1.5)   | English |    33.4M   |    133 MB    |     version 1.5 with more reasonable similarity distribution      |   BERT   |
| [BAAI/bge-large-zh-v1.5](https://huggingface.co/BAAI/bge-large-zh-v1.5)   | Chinese |    326M    |    1.3 GB    |     version 1.5 with more reasonable similarity distribution      |   BERT   |
| [BAAI/bge-base-zh-v1.5](https://huggingface.co/BAAI/bge-base-zh-v1.5)     | Chinese |    102M    |    409 MB    |     version 1.5 with more reasonable similarity distribution      |   BERT   |
| [BAAI/bge-small-zh-v1.5](https://huggingface.co/BAAI/bge-small-zh-v1.5)   | Chinese |    24M     |    95.8 MB   |     version 1.5 with more reasonable similarity distribution      |   BERT   |

You can use BGE 1.5 models exactly same to BGE v1 models.


```python
model = FlagModel(
    'BAAI/bge-base-en-v1.5',
    query_instruction_for_retrieval="Represent this sentence for searching relevant passages:",
    query_instruction_format='{}{}'
)

queries = ["query 1", "query 2"]
corpus = ["passage 1", "passage 2"]

# encode the queries and corpus
q_embeddings = model.encode_queries(queries)
p_embeddings = model.encode_corpus(corpus)

# compute the similarity scores
scores = q_embeddings @ p_embeddings.T
print(scores)
```

### 2.3 BGE M3

BGE-M3 is the new version of BGE models that is distinguished for its versatility in:
- Multi-Functionality: Simultaneously perform the three common retrieval functionalities of embedding model: dense retrieval, multi-vector retrieval, and sparse retrieval.
- Multi-Linguality: Supports more than 100 working languages.
- Multi-Granularity: Can proces inputs with different granularityies, spanning from short sentences to long documents of up to 8192 tokens.

For more details, feel free to check out the [paper](https://arxiv.org/pdf/2402.03216).

| Model  | Language |   Parameters   |   Model Size   |    Description    |   Base Model     |
|:-------|:--------:|:--------------:|:--------------:|:-----------------:|:----------------:|
| [BAAI/bge-m3](https://huggingface.co/BAAI/bge-m3)                   |    Multilingual     |   568M   |  2.27 GB  |  Multi-Functionality(dense retrieval, sparse retrieval, multi-vector(colbert)), Multi-Linguality, and Multi-Granularity(8192 tokens) | XLM-RoBERTa |


```python
from FlagEmbedding import BGEM3FlagModel

model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)

sentences = ["What is BGE M3?", "Defination of BM25"]
```

```python
BGEM3FlagModel.encode(
    sentences, 
    batch_size=12, 
    max_length=8192, 
    return_dense=True, 
    return_sparse=False, 
    return_colbert_vecs=False
)
```
It returns a dictionary like:
```python
{
    'dense_vecs':       # array of dense embeddings of inputs if return_dense=True, otherwise None,
    'lexical_weights':  # array of dictionaries with keys and values are ids of tokens and their corresponding weights if return_sparse=True, otherwise None,
    'colbert_vecs':     # array of multi-vector embeddings of inputs if return_cobert_vecs=True, otherwise None,'
}
```


```python
# If you don't need such a long length of 8192 input tokens, you can set max_length to a smaller value to speed up encoding.
embeddings = model.encode(
    sentences, 
    max_length=10,
    return_dense=True, 
    return_sparse=True, 
    return_colbert_vecs=True
)
```


```python
print(f"dense embedding:\n{embeddings['dense_vecs']}")
print(f"sparse embedding:\n{embeddings['lexical_weights']}")
print(f"multi-vector:\n{embeddings['colbert_vecs']}")
```

### 2.4 BGE Multilingual Gemma2

BGE Multilingual Gemma2 is a LLM-based Multi-Lingual embedding model.

| Model  | Language |   Parameters   |   Model Size   |    Description    |   Base Model     |
|:-------|:--------:|:--------------:|:--------------:|:-----------------:|:----------------:|
| [BAAI/bge-multilingual-gemma2](https://huggingface.co/BAAI/bge-multilingual-gemma2)                   |    Multilingual     |   9.24B   |  37 GB  |  LLM-based multilingual embedding model with SOTA results on multilingual benchmarks | Gemma2-9B |


```python
from FlagEmbedding import FlagLLMModel

queries = ["how much protein should a female eat", "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."
]

model = FlagLLMModel('BAAI/bge-multilingual-gemma2', 
                     query_instruction_for_retrieval="Given a web search query, retrieve relevant passages that answer the query.",
                     use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation

embeddings_1 = model.encode_queries(queries)
embeddings_2 = model.encode_corpus(documents)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

```

### 2.4 BGE ICL

BGE ICL stands for in-context learning. By providing few-shot examples in the query, it can significantly enhance the model's ability to handle new tasks.

| Model  | Language |   Parameters   |   Model Size   |    Description    |   Base Model     |
|:-------|:--------:|:--------------:|:--------------:|:-----------------:|:----------------:|
| [BAAI/bge-en-icl](https://huggingface.co/BAAI/bge-en-icl)                   |    English     |   7.11B   |  28.5 GB  |  LLM-based English embedding model with excellent in-context learning ability. | Mistral-7B |


```python
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."
]

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."
    }
]

queries = ["how much protein should a female eat", "summit define"]
```


```python
from FlagEmbedding import FlagICLModel
import os

model = FlagICLModel('BAAI/bge-en-icl', 
                     examples_for_task=examples,  # set `examples_for_task=None` to use model without examples
                    #  examples_instruction_format="<instruct>{}\n<query>{}\n<response>{}" # specify the format to use examples_for_task
                     )

embeddings_1 = model.encode_queries(queries)
embeddings_2 = model.encode_corpus(documents)
similarity = embeddings_1 @ embeddings_2.T

print(similarity)
```
