wehub/flagopen--flagembedding
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:39:21 +08:00
commit bc34f6df14
1149 changed files with 328099 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
FlagEmbedding/inference/embedder/__init__.py
+12
View File
@@ -0,0 +1,12 @@
from .encoder_only import FlagModel, BGEM3FlagModel
from .decoder_only import FlagICLModel, FlagLLMModel, FlagPseudoMoEModel
from .model_mapping import EmbedderModelClass
__all__ = [
"FlagModel",
"BGEM3FlagModel",
"FlagICLModel",
"FlagLLMModel",
"FlagPseudoMoEModel",
"EmbedderModelClass",
]
FlagEmbedding/inference/embedder/decoder_only/__init__.py
+9
View File
@@ -0,0 +1,9 @@
from .base import BaseLLMEmbedder as FlagLLMModel
from .icl import ICLLLMEmbedder as FlagICLModel
from .pseudo_moe import PseudoMoELLMEmbedder as FlagPseudoMoEModel
__all__ = [
"FlagLLMModel",
"FlagICLModel",
"FlagPseudoMoEModel",
]
FlagEmbedding/inference/embedder/decoder_only/base.py
+301
View File
@@ -0,0 +1,301 @@
from tqdm import tqdm, trange
from typing import cast, Any, List, Union, Optional
import torch
import numpy as np
from transformers import AutoModel, AutoTokenizer
from FlagEmbedding.abc.inference import AbsEmbedder
# Pooling function for LLM-based embedding models
def last_token_pool(last_hidden_states: torch.Tensor,
attention_mask: torch.Tensor) -> torch.Tensor:
"""Last token pooling method.
Args:
last_hidden_state (torch.Tensor): The last hidden state of the model.
attention_mask (torch.Tensor): Attention mask. Defaults to :data:`None`.
Returns:
torch.Tensor: The embedding vectors after pooling.
"""
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]
class BaseLLMEmbedder(AbsEmbedder):
"""Base embedder class for LLM like decoder only models.
Args:
model_name_or_path (str): If it's a path to a local model, it loads the model from the path. Otherwise tries to download and
load a model from HuggingFace Hub with the name.
normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to :data:`True`.
use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance
degradation. Defaults to :data:`True`.
query_instruction_for_retrieval (Optional[str], optional): Query instruction for retrieval tasks, which will be used with
with :attr:`query_instruction_format`. Defaults to :data:`None`.
query_instruction_format (str, optional): The template for :attr:`query_instruction_for_retrieval`. Defaults to :data:`"Instruct: {}\nQuery: {}"`.
devices (Optional[Union[str, int, List[str], List[int]]], optional): Devices to use for model inference. Defaults to :data:`None`.
trust_remote_code (bool, optional): trust_remote_code for HF datasets or models. Defaults to :data:`False`.
cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
batch_size (int, optional): Batch size for inference. Defaults to :data:`256`.
query_max_length (int, optional): Maximum length for query. Defaults to :data:`512`.
passage_max_length (int, optional): Maximum length for passage. Defaults to :data:`512`.
convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will be a Torch Tensor.
Defaults to :data:`True`.
Attributes:
DEFAULT_POOLING_METHOD: The default pooling method when running the model.
"""
DEFAULT_POOLING_METHOD = "last_token"
def __init__(
self,
model_name_or_path: str,
normalize_embeddings: bool = True,
use_fp16: bool = True,
use_bf16: bool = False,
query_instruction_for_retrieval: Optional[str] = None,
query_instruction_format: str = "Instruct: {}\nQuery: {}", # specify the format of query_instruction_for_retrieval
devices: Optional[Union[str, List[str]]] = None, # specify devices, such as "cuda:0" or ["cuda:0", "cuda:1"]
# Additional parameters for BaseLLMEmbedder
trust_remote_code: bool = False,
cache_dir: Optional[str] = None,
# inference
batch_size: int = 256,
query_max_length: int = 512,
passage_max_length: int = 512,
convert_to_numpy: bool = True,
truncate_dim: Optional[int] = None,
**kwargs: Any,
):
super().__init__(
model_name_or_path,
normalize_embeddings=normalize_embeddings,
use_fp16=use_fp16,
use_bf16=use_bf16,
query_instruction_for_retrieval=query_instruction_for_retrieval,
query_instruction_format=query_instruction_format,
devices=devices,
batch_size=batch_size,
query_max_length=query_max_length,
passage_max_length=passage_max_length,
convert_to_numpy=convert_to_numpy,
truncate_dim=truncate_dim,
**kwargs
)
self.tokenizer = AutoTokenizer.from_pretrained(
model_name_or_path,
trust_remote_code=trust_remote_code,
cache_dir=cache_dir
)
self.model = AutoModel.from_pretrained(
model_name_or_path,
trust_remote_code=trust_remote_code,
cache_dir=cache_dir,
dtype=self.get_model_torch_dtype(),
)
if self.kwargs.get("pooling_method", "last_token") != "last_token":
raise ValueError("Pooling method must be 'last_token' for LLM-based models.")
def encode_queries(
self,
queries: Union[List[str], str],
batch_size: Optional[int] = None,
max_length: Optional[int] = None,
convert_to_numpy: Optional[bool] = None,
**kwargs: Any
) -> Union[np.ndarray, torch.Tensor]:
"""Encode the queries.
Args:
queries (Union[List[str], str]): Input queries to encode.
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will
be a Torch Tensor. Defaults to :data:`None`.
Returns:
Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
"""
return super().encode_queries(
queries,
batch_size=batch_size,
max_length=max_length,
convert_to_numpy=convert_to_numpy,
**kwargs
)
def encode_corpus(
self,
corpus: Union[List[str], str],
batch_size: Optional[int] = None,
max_length: Optional[int] = None,
convert_to_numpy: Optional[bool] = None,
**kwargs: Any
) -> Union[np.ndarray, torch.Tensor]:
"""Encode the corpus.
Args:
corpus (Union[List[str], str]): Input corpus to encode.
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will
be a Torch Tensor. Defaults to :data:`None`.
Returns:
Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
"""
return super().encode_corpus(
corpus,
batch_size=batch_size,
max_length=max_length,
convert_to_numpy=convert_to_numpy,
**kwargs
)
def encode(
self,
sentences: Union[List[str], str],
batch_size: Optional[int] = None,
max_length: Optional[int] = None,
convert_to_numpy: Optional[bool] = None,
**kwargs: Any
) -> Union[np.ndarray, torch.Tensor]:
"""Encode the input sentences with the embedding model.
Args:
sentences (Union[List[str], str]): Input sentences to encode.
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will
be a Torch Tensor. Defaults to :data:`None`.
Returns:
Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
"""
return super().encode(
sentences,
batch_size=batch_size,
max_length=max_length,
convert_to_numpy=convert_to_numpy,
**kwargs
)
@torch.no_grad()
def encode_single_device(
self,
sentences: Union[List[str], str],
batch_size: int = 256,
max_length: int = 512,
convert_to_numpy: bool = True,
device: Optional[str] = None,
**kwargs: Any # add `pad_to_multiple_of=8` for bge-multilingual-gemmma2
):
"""Encode input sentences by a single device.
Args:
sentences (Union[List[str], str]): Input sentences to encode.
batch_size (int, optional): Number of sentences for each iter. Defaults to :data:`256`.
max_length (int, optional): Maximum length of tokens. Defaults to :data:`512`.
convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will
be a Torch Tensor. Defaults to :data:`True`.
device (Optional[str], optional): Device to use for encoding. Defaults to None.
Returns:
Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
"""
if device is None:
device = self.target_devices[0]
if device == "cpu":
self.model.float()
self.model.to(device)
self.model.eval()
input_was_string = False
if isinstance(sentences, str):
sentences = [sentences]
input_was_string = True
# tokenize without padding to get the correct length
all_inputs = []
for start_index in trange(0, len(sentences), batch_size, desc='pre tokenize',
disable=len(sentences) < batch_size):
sentences_batch = sentences[start_index:start_index + batch_size]
inputs_batch = self.tokenizer(
sentences_batch,
truncation=True,
max_length=max_length,
**kwargs
)
inputs_batch = [{
k: inputs_batch[k][i] for k in inputs_batch.keys()
} for i in range(len(sentences_batch))]
all_inputs.extend(inputs_batch)
# sort by length for less padding
length_sorted_idx = np.argsort([-len(x['input_ids']) for x in all_inputs])
all_inputs_sorted = [all_inputs[i] for i in length_sorted_idx]
# adjust batch size
flag = False
while flag is False:
try:
inputs_batch = self.tokenizer.pad(
all_inputs_sorted[: batch_size],
padding=True,
return_tensors='pt',
**kwargs
).to(device)
last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
flag = True
except RuntimeError as e:
batch_size = batch_size * 3 // 4
except torch.cuda.OutOfMemoryError as e:
batch_size = batch_size * 3 // 4
# encode
all_embeddings = []
for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
disable=len(sentences) < batch_size):
inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
inputs_batch = self.tokenizer.pad(
inputs_batch,
padding=True,
return_tensors='pt',
**kwargs
).to(device)
last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
embeddings = self._truncate_embeddings(embeddings)
if self.normalize_embeddings:
embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
embeddings = cast(torch.Tensor, embeddings)
if convert_to_numpy:
embeddings = self._convert_to_numpy(embeddings, device=device)
all_embeddings.append(embeddings)
if convert_to_numpy:
all_embeddings = np.concatenate(all_embeddings, axis=0)
else:
all_embeddings = torch.cat(all_embeddings, dim=0)
# adjust the order of embeddings
all_embeddings = all_embeddings[np.argsort(length_sorted_idx)]
# return the embeddings
if input_was_string:
return all_embeddings[0]
return all_embeddings
FlagEmbedding/inference/embedder/decoder_only/icl.py
+567
View File
@@ -0,0 +1,567 @@
from tqdm import tqdm, trange
from typing import cast, Any, List, Union, Optional
import queue
from multiprocessing import Queue
import gc
import torch
import numpy as np
from transformers import AutoModel, AutoTokenizer
from FlagEmbedding.abc.inference import AbsEmbedder
# Pooling function for LLM-based embedding models
def last_token_pool(last_hidden_states: torch.Tensor,
attention_mask: torch.Tensor) -> torch.Tensor:
"""Last token pooling method.
Args:
last_hidden_state (torch.Tensor): The last hidden state of the model.
attention_mask (torch.Tensor): Attention mask. Defaults to :data:`None`.
Returns:
torch.Tensor: The embedding vectors after pooling.
"""
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]
class ICLLLMEmbedder(AbsEmbedder):
"""
Embedder class for BGE-EN-icl.
Args:
model_name_or_path (str): If it's a path to a local model, it loads the model from the path. Otherwise tries to download and
load a model from HuggingFace Hub with the name.
normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to :data:`True`.
use_fp16 (bool, optional) If true, use half-precision floating-point to speed up computation with a slight performance
degradation. Defaults to :data:`True`.
query_instruction_for_retrieval (Optional[str], optional): Query instruction for retrieval tasks, which will be used with
with :attr:`query_instruction_format`. Defaults to :data:`None`.
query_instruction_format (str, optional): The template for :attr:`query_instruction_for_retrieval`. Defaults to :data:`"{}{}"`.
devices (Optional[Union[str, int, List[str], List[int]]], optional): Devices to use for model inference. Defaults to :data:`None`.
examples_for_task (Optional[List[dict]], optional): Few-shot examples for the model to enhance model's ability.
Defaults to :data:`None`.
examples_instruction_format (str, optional): Example format when using :attr:`examples_for_task`.
trust_remote_code (bool, optional): trust_remote_code for HF datasets or models. Defaults to :data:`False`.
cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
batch_size (int, optional): Batch size for inference. Defaults to :data:`256`.
query_max_length (int, optional): Maximum length for query. Defaults to :data:`512`.
passage_max_length (int, optional): Maximum length for passage. Defaults to :data:`512`.
convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will be a Torch Tensor.
Defaults to :data:`True`.
Attributes:
DEFAULT_POOLING_METHOD: The default pooling method when running the model.
"""
DEFAULT_POOLING_METHOD = "last_token"
def __init__(
self,
model_name_or_path: str,
normalize_embeddings: bool = True,
use_fp16: bool = True,
use_bf16: bool = False,
query_instruction_for_retrieval: Optional[str] = None,
query_instruction_format: str = "<instruct>{}\n<query>{}", # specify the format of query_instruction_for_retrieval
suffix: str = '\n<response>',
devices: Optional[Union[str, List[str]]] = None, # specify devices, such as "cuda:0" or ["cuda:0", "cuda:1"]
# Additional parameters for ICLLLMEmbedder
examples_for_task: Optional[List[dict]] = None,
examples_instruction_format: str = "<instruct>{}\n<query>{}\n<response>{}", # specify the format of examples_for_task
trust_remote_code: bool = False,
cache_dir: Optional[str] = None,
# inference
batch_size: int = 256,
query_max_length: int = 512,
passage_max_length: int = 512,
convert_to_numpy: bool = True,
truncate_dim: Optional[int] = None,
**kwargs: Any,
):
query_instruction_format = query_instruction_format.replace('\\n', '\n')
examples_instruction_format = examples_instruction_format.replace('\\n', '\n')
super().__init__(
model_name_or_path,
normalize_embeddings=normalize_embeddings,
use_fp16=use_fp16,
use_bf16=use_bf16,
query_instruction_for_retrieval=query_instruction_for_retrieval,
query_instruction_format=query_instruction_format,
devices=devices,
batch_size=batch_size,
query_max_length=query_max_length,
passage_max_length=passage_max_length,
convert_to_numpy=convert_to_numpy,
truncate_dim=truncate_dim,
**kwargs
)
self.tokenizer = AutoTokenizer.from_pretrained(
model_name_or_path,
trust_remote_code=trust_remote_code,
cache_dir=cache_dir
)
self.model = AutoModel.from_pretrained(
model_name_or_path,
trust_remote_code=trust_remote_code,
cache_dir=cache_dir,
torch_dtype=self.get_model_torch_dtype(),
)
self.examples_for_task = examples_for_task
self.examples_instruction_format = examples_instruction_format
if self.kwargs.get("pooling_method", "last_token") != "last_token":
raise ValueError("Pooling method must be 'last_token' for LLM-based models.")
self.set_examples()
self.suffix = suffix
self.query_pool = None
def __del__(self):
self.stop_self_pool()
self.stop_self_query_pool()
def set_examples(self, examples_for_task: Optional[List[dict]] = None):
"""Set the prefix to the provided examples.
Args:
examples_for_task (Optional[List[dict]], optional): Few-shot examples for the model to enhance model's ability.
Defaults to :data:`None`.
"""
if examples_for_task is None and self.examples_for_task is None:
self.prefix = ''
elif examples_for_task is not None:
eg_paris = []
for i in range(len(examples_for_task)):
eg_paris.append(
self.get_detailed_example(
self.examples_instruction_format,
examples_for_task[i].get('instruct', self.query_instruction_for_retrieval),
examples_for_task[i].get('query', ''),
examples_for_task[i].get('response', '')
)
)
self.prefix = '\n\n'.join(eg_paris) + '\n\n'
else:
eg_paris = []
for i in range(len(self.examples_for_task)):
eg_paris.append(
self.get_detailed_example(
self.examples_instruction_format,
self.examples_for_task[i].get('instruct', self.query_instruction_for_retrieval),
self.examples_for_task[i].get('query', ''),
self.examples_for_task[i].get('response', '')
)
)
self.prefix = '\n\n'.join(eg_paris) + '\n\n'
@staticmethod
def get_detailed_example(instruction_format: str, instruction: str, query: str, response: str):
"""Combine the instruction and sentence along with the instruction format.
Args:
instruction_format (str): Format for instruction.
instruction (str): The text of instruction.
query (str): The text of example query.
response (str): The text of example response.
Returns:
str: The complete example following the given format.
"""
if "\\n" in instruction_format:
instruction_format = instruction_format.replace("\\n", "\n")
return instruction_format.format(instruction, query, response)
def stop_self_query_pool(self):
if self.query_pool is not None:
self.stop_multi_process_pool(self.query_pool)
self.query_pool = None
try:
self.model.to('cpu')
torch.cuda.empty_cache()
except:
pass
gc.collect()
def encode_queries(
self,
queries: Union[List[str], str],
batch_size: Optional[int] = None,
max_length: Optional[int] = None,
convert_to_numpy: Optional[bool] = None,
**kwargs: Any
) -> Union[np.ndarray, torch.Tensor]:
"""Encode the queries.
Args:
queries (Union[List[str], str]): Input queries to encode.
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will
be a Torch Tensor. Defaults to :data:`None`.
Returns:
Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
"""
if batch_size is None: batch_size = self.batch_size
if max_length is None: max_length = self.query_max_length
if convert_to_numpy is None: convert_to_numpy = self.convert_to_numpy
if isinstance(queries, str) or len(self.target_devices) == 1:
return self.encode_queries_single_device(
queries,
batch_size=batch_size,
max_length=max_length,
convert_to_numpy=convert_to_numpy,
device=self.target_devices[0],
**kwargs
)
self.stop_self_pool()
if self.query_pool is None:
self.query_pool = self.start_multi_process_pool(ICLLLMEmbedder._encode_queries_multi_process_worker)
embeddings = self.encode_multi_process(
queries,
self.query_pool,
batch_size=batch_size,
max_length=max_length,
convert_to_numpy=convert_to_numpy,
**kwargs
)
return embeddings
def encode_corpus(
self,
corpus: Union[List[str], str],
batch_size: Optional[int] = None,
max_length: Optional[int] = None,
convert_to_numpy: Optional[bool] = None,
**kwargs: Any
) -> Union[np.ndarray, torch.Tensor]:
"""Encode the corpus.
Args:
corpus (Union[List[str], str]): Input corpus to encode.
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will
be a Torch Tensor. Defaults to :data:`None`.
Returns:
Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
"""
self.stop_self_query_pool()
return super().encode_corpus(
corpus,
batch_size=batch_size,
max_length=max_length,
convert_to_numpy=convert_to_numpy,
**kwargs
)
def encode(
self,
sentences: Union[List[str], str],
batch_size: Optional[int] = None,
max_length: Optional[int] = None,
convert_to_numpy: Optional[bool] = None,
**kwargs: Any
) -> Union[np.ndarray, torch.Tensor]:
"""Encode the input sentences with the embedding model.
Args:
sentences (Union[List[str], str]): Input sentences to encode.
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will
be a Torch Tensor. Defaults to :data:`None`.
Returns:
Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
"""
return super().encode(
sentences,
batch_size=batch_size,
max_length=max_length,
convert_to_numpy=convert_to_numpy,
**kwargs
)
# adapted from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L976
@staticmethod
def _encode_queries_multi_process_worker(
target_device: str, model: 'ICLLLMEmbedder', input_queue: Queue, results_queue: Queue
) -> None:
"""
Internal working process to encode sentences in multi-process setup
"""
while True:
try:
chunk_id, sentences, kwargs = (
input_queue.get()
)
embeddings = model.encode_queries_single_device(
sentences,
device=target_device,
**kwargs
)
results_queue.put([chunk_id, embeddings])
except queue.Empty:
break
@torch.no_grad()
def encode_queries_single_device(
self,
queries: Union[List[str], str],
batch_size: int = 256,
max_length: int = 512,
convert_to_numpy: bool = True,
device: Optional[str] = None,
**kwargs: Any
):
"""Encode queries by a single device.
Args:
queries (Union[List[str], str]): Input queries to encode.
batch_size (int, optional): Number of queries for each iter. Defaults to :data:`256`.
max_length (int, optional): Maximum length of tokens. Defaults to :data:`512`.
convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will
be a Torch Tensor. Defaults to :data:`True`.
device (Optional[str], optional): Device to use for encoding. Defaults to None.
Returns:
Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
"""
if device is None:
device = self.target_devices[0]
if device == "cpu":
self.model.float()
self.model.to(device)
self.model.eval()
input_was_string = False
if isinstance(queries, str):
queries = [queries]
input_was_string = True
if self.query_instruction_for_retrieval is not None:
if isinstance(queries, str):
input_texts = self.get_detailed_instruct(self.query_instruction_format, self.query_instruction_for_retrieval, queries)
else:
input_texts = [self.get_detailed_instruct(self.query_instruction_format, self.query_instruction_for_retrieval, query) for query in queries]
else:
input_texts = queries
prefix_ids = self.tokenizer(self.prefix, add_special_tokens=False)['input_ids']
suffix_ids = self.tokenizer(self.suffix, add_special_tokens=False)['input_ids']
_len_1 = len(self.tokenizer('<s>', add_special_tokens=False)['input_ids'])
_len_2 = len(self.tokenizer(f'{self.suffix}</s>', add_special_tokens=False)['input_ids'])
new_max_length = (len(prefix_ids) + len(suffix_ids) + max_length + 8) // 8 * 8 + 8
# tokenize without padding to get the correct length
all_inputs = []
for start_index in trange(0, len(input_texts), batch_size, desc='pre tokenize',
disable=len(input_texts) < batch_size):
sentences_batch = input_texts[start_index:start_index + batch_size]
inputs_batch = self.tokenizer(
sentences_batch,
truncation=True,
max_length=max_length - _len_1 - _len_2,
add_special_tokens=False,
**kwargs
)
sentences_batch = self.tokenizer.batch_decode(inputs_batch['input_ids'])
for i in range(len(sentences_batch)):
sentences_batch[i] = self.prefix + sentences_batch[i] + self.suffix
inputs_batch = self.tokenizer(
sentences_batch,
truncation=True,
max_length=new_max_length,
**kwargs
)
inputs_batch = [{
k: inputs_batch[k][i] for k in inputs_batch.keys()
} for i in range(len(sentences_batch))]
all_inputs.extend(inputs_batch)
# sort by length for less padding
length_sorted_idx = np.argsort([-len(x['input_ids']) for x in all_inputs])
all_inputs_sorted = [all_inputs[i] for i in length_sorted_idx]
sentences_sorted = [input_texts[i] for i in length_sorted_idx]
# adjust batch size
flag = False
while flag is False:
try:
inputs_batch = self.tokenizer.pad(
all_inputs_sorted[: batch_size],
padding=True,
return_tensors='pt',
**kwargs
).to(device)
last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
flag = True
except RuntimeError as e:
batch_size = batch_size * 3 // 4
except torch.cuda.OutOfMemoryError as e:
batch_size = batch_size * 3 // 4
# encode
all_embeddings = []
for start_index in tqdm(range(0, len(sentences_sorted), batch_size), desc="Inference Embeddings",
disable=len(sentences_sorted) < batch_size):
inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
inputs_batch = self.tokenizer.pad(
inputs_batch,
padding=True,
return_tensors='pt',
**kwargs
).to(device)
last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
embeddings = self._truncate_embeddings(embeddings)
if self.normalize_embeddings:
embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
embeddings = cast(torch.Tensor, embeddings)
if convert_to_numpy:
embeddings = self._convert_to_numpy(embeddings, device=device)
all_embeddings.append(embeddings)
if convert_to_numpy:
all_embeddings = np.concatenate(all_embeddings, axis=0)
else:
all_embeddings = torch.cat(all_embeddings, dim=0)
# adjust the order of embeddings
all_embeddings = all_embeddings[np.argsort(length_sorted_idx)]
# return the embeddings
if input_was_string:
return all_embeddings[0]
return all_embeddings
@torch.no_grad()
def encode_single_device(
self,
sentences: Union[List[str], str],
batch_size: int = 256,
max_length: int = 512,
convert_to_numpy: bool = True,
device: Optional[str] = None,
**kwargs: Any
):
"""Encode input sentences by a single device.
Args:
sentences (Union[List[str], str]): Input sentences to encode.
batch_size (int, optional): Number of sentences for each iter. Defaults to :data:`256`.
max_length (int, optional): Maximum length of tokens. Defaults to :data:`512`.
convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will
be a Torch Tensor. Defaults to :data:`True`.
device (Optional[str], optional): Device to use for encoding. Defaults to None.
Returns:
Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
"""
if device is None:
device = self.target_devices[0]
if device == "cpu":
self.model.float()
self.model.to(device)
self.model.eval()
input_was_string = False
if isinstance(sentences, str):
sentences = [sentences]
input_was_string = True
# tokenize without padding to get the correct length
all_inputs = []
for start_index in trange(0, len(sentences), batch_size, desc='pre tokenize',
disable=len(sentences) < batch_size):
sentences_batch = sentences[start_index:start_index + batch_size]
inputs_batch = self.tokenizer(
sentences_batch,
truncation=True,
max_length=max_length,
**kwargs
)
inputs_batch = [{
k: inputs_batch[k][i] for k in inputs_batch.keys()
} for i in range(len(sentences_batch))]
all_inputs.extend(inputs_batch)
# sort by length for less padding
length_sorted_idx = np.argsort([-len(x['input_ids']) for x in all_inputs])
all_inputs_sorted = [all_inputs[i] for i in length_sorted_idx]
# adjust batch size
flag = False
while flag is False:
try:
inputs_batch = self.tokenizer.pad(
all_inputs_sorted[: batch_size],
padding=True,
return_tensors='pt',
**kwargs
).to(device)
last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
flag = True
except RuntimeError as e:
batch_size = batch_size * 3 // 4
except torch.cuda.OutOfMemoryError as e:
batch_size = batch_size * 3 // 4
# encode
all_embeddings = []
for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
disable=len(sentences) < batch_size):
inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
inputs_batch = self.tokenizer.pad(
inputs_batch,
padding=True,
return_tensors='pt',
**kwargs
).to(device)
last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
embeddings = self._truncate_embeddings(embeddings)
if self.normalize_embeddings:
embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
embeddings = cast(torch.Tensor, embeddings)
if convert_to_numpy:
embeddings = self._convert_to_numpy(embeddings, device=device)
all_embeddings.append(embeddings)
if convert_to_numpy:
all_embeddings = np.concatenate(all_embeddings, axis=0)
else:
all_embeddings = torch.cat(all_embeddings, dim=0)
# adjust the order of embeddings
all_embeddings = all_embeddings[np.argsort(length_sorted_idx)]
# return the embeddings
if input_was_string:
return all_embeddings[0]
return all_embeddings
FlagEmbedding/inference/embedder/decoder_only/pseudo_moe.py
+194
View File
@@ -0,0 +1,194 @@
from typing import cast, Any, List, Union, Optional
import torch
import numpy as np
from .base import BaseLLMEmbedder, last_token_pool
class PseudoMoELLMEmbedder(BaseLLMEmbedder):
"""Decoder-only embedder for pseudo MoE checkpoints.
This class follows the same behavior as :class:`BaseLLMEmbedder`, but supports
selecting an active domain (e.g. ``general``, ``coding``, ``reasoning``) during
inference when the underlying model implements domain routing.
Args:
model_name_or_path (str): If it's a path to a local model, it loads the model from the path. Otherwise tries to download and
load a model from HuggingFace Hub with the name.
normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to :data:`True`.
use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance
degradation. Defaults to :data:`True`.
query_instruction_for_retrieval (Optional[str], optional): Query instruction for retrieval tasks, which will be used with
with :attr:`query_instruction_format`. Defaults to :data:`None`.
query_instruction_format (str, optional): The template for :attr:`query_instruction_for_retrieval`. Defaults to :data:`"Instruct: {}\nQuery: {}"`.
devices (Optional[Union[str, int, List[str], List[int]]], optional): Devices to use for model inference. Defaults to :data:`None`.
trust_remote_code (bool, optional): trust_remote_code for HF datasets or models. Defaults to :data:`False`.
cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
batch_size (int, optional): Batch size for inference. Defaults to :data:`256`.
query_max_length (int, optional): Maximum length for query. Defaults to :data:`512`.
passage_max_length (int, optional): Maximum length for passage. Defaults to :data:`512`.
convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will be a Torch Tensor.
Defaults to :data:`True`.
domain_for_pseudo_moe (str, optional): Specifies the active domain for the decoder-only pseudo-MoE model (e.g., "general", "coding", or "reasoning").
Defaults to "general".
Attributes:
DEFAULT_POOLING_METHOD: The default pooling method when running the model.
"""
DEFAULT_POOLING_METHOD = "last_token"
def __init__(
self,
model_name_or_path: str,
normalize_embeddings: bool = True,
use_fp16: bool = False,
use_bf16: bool = True,
query_instruction_for_retrieval: Optional[str] = None,
query_instruction_format: str = "Instruct: {}\nQuery: {}",
devices: Optional[Union[str, List[str]]] = None,
trust_remote_code: bool = True,
cache_dir: Optional[str] = None,
batch_size: int = 256,
query_max_length: int = 512,
passage_max_length: int = 512,
convert_to_numpy: bool = True,
truncate_dim: Optional[int] = None,
domain_for_pseudo_moe: Optional[str] = None,
**kwargs: Any,
):
self.domain_for_pseudo_moe = domain_for_pseudo_moe
super().__init__(
model_name_or_path=model_name_or_path,
normalize_embeddings=normalize_embeddings,
use_fp16=use_fp16,
use_bf16=use_bf16,
query_instruction_for_retrieval=query_instruction_for_retrieval,
query_instruction_format=query_instruction_format,
devices=devices,
trust_remote_code=trust_remote_code,
cache_dir=cache_dir,
batch_size=batch_size,
query_max_length=query_max_length,
passage_max_length=passage_max_length,
convert_to_numpy=convert_to_numpy,
truncate_dim=truncate_dim,
**kwargs,
)
def _resolve_domain(self, kwargs: Any) -> Optional[str]:
domain = kwargs.pop("domain_for_pseudo_moe", None)
if domain is None:
domain = kwargs.pop("domain", None)
if domain is None:
domain = self.domain_for_pseudo_moe
return domain
@torch.no_grad()
def encode_single_device(
self,
sentences: Union[List[str], str],
batch_size: int = 256,
max_length: int = 512,
convert_to_numpy: bool = True,
device: Optional[str] = None,
**kwargs: Any
):
if device is None:
device = self.target_devices[0]
if device == "cpu":
self.model.float()
self.model.to(device)
self.model.eval()
input_was_string = False
if isinstance(sentences, str):
sentences = [sentences]
input_was_string = True
domain = self._resolve_domain(kwargs)
if domain is not None and hasattr(self.model, "set_domain"):
self.model.set_domain(domain)
model_forward_kwargs = {"return_dict": True}
if domain is not None:
model_forward_kwargs["domain"] = domain
# tokenize without padding to get the correct length
all_inputs = []
for start_index in range(0, len(sentences), batch_size):
sentences_batch = sentences[start_index:start_index + batch_size]
inputs_batch = self.tokenizer(
sentences_batch,
truncation=True,
max_length=max_length,
**kwargs
)
inputs_batch = [{
k: inputs_batch[k][i] for k in inputs_batch.keys()
} for i in range(len(sentences_batch))]
all_inputs.extend(inputs_batch)
# sort by length for less padding
length_sorted_idx = np.argsort([-len(x['input_ids']) for x in all_inputs])
all_inputs_sorted = [all_inputs[i] for i in length_sorted_idx]
# adjust batch size
flag = False
while flag is False:
try:
inputs_batch = self.tokenizer.pad(
all_inputs_sorted[: batch_size],
padding=True,
return_tensors='pt',
**kwargs
).to(device)
try:
last_hidden_state = self.model(**inputs_batch, **model_forward_kwargs).last_hidden_state
except TypeError:
last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
_ = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
flag = True
except RuntimeError:
batch_size = batch_size * 3 // 4
except torch.cuda.OutOfMemoryError:
batch_size = batch_size * 3 // 4
# encode
all_embeddings = []
for start_index in range(0, len(sentences), batch_size):
inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
inputs_batch = self.tokenizer.pad(
inputs_batch,
padding=True,
return_tensors='pt',
**kwargs
).to(device)
try:
last_hidden_state = self.model(**inputs_batch, **model_forward_kwargs).last_hidden_state
except TypeError:
last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
embeddings = last_token_pool(last_hidden_state, inputs_batch['attention_mask'])
embeddings = self._truncate_embeddings(embeddings)
embeddings = torch.nan_to_num(embeddings, nan=0.0, posinf=1e4, neginf=-1e4)
if self.normalize_embeddings:
embeddings = torch.nn.functional.normalize(embeddings.float(), dim=-1)
embeddings = cast(torch.Tensor, embeddings)
if convert_to_numpy:
embeddings = self._convert_to_numpy(embeddings, device=device)
all_embeddings.append(embeddings)
if convert_to_numpy:
all_embeddings = np.concatenate(all_embeddings, axis=0)
else:
all_embeddings = torch.cat(all_embeddings, dim=0)
# adjust the order of embeddings
all_embeddings = all_embeddings[np.argsort(length_sorted_idx)]
if input_was_string:
return all_embeddings[0]
return all_embeddings
FlagEmbedding/inference/embedder/encoder_only/__init__.py
+7
View File
@@ -0,0 +1,7 @@
from .base import BaseEmbedder as FlagModel
from .m3 import M3Embedder as BGEM3FlagModel
__all__ = [
"FlagModel",
"BGEM3FlagModel",
]
FlagEmbedding/inference/embedder/encoder_only/base.py
+308
View File
@@ -0,0 +1,308 @@
from tqdm import tqdm, trange
from typing import cast, Any, List, Union, Optional
import torch
import numpy as np
from transformers import AutoModel, AutoTokenizer
from FlagEmbedding.abc.inference import AbsEmbedder
class BaseEmbedder(AbsEmbedder):
"""
Base embedder for encoder only models.
Args:
model_name_or_path (str): If it's a path to a local model, it loads the model from the path. Otherwise tries to download and
load a model from HuggingFace Hub with the name.
normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to :data:`True`.
use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance
degradation. Defaults to :data:`True`.
query_instruction_for_retrieval (Optional[str], optional): Query instruction for retrieval tasks, which will be used with
with :attr:`query_instruction_format`. Defaults to :data:`None`.
query_instruction_format (str, optional): The template for :attr:`query_instruction_for_retrieval`. Defaults to :data:`"{}{}"`.
devices (Optional[Union[str, int, List[str], List[int]]], optional): Devices to use for model inference. Defaults to :data:`None`.
pooling_method (str, optional): Pooling method to get embedding vector from the last hidden state. Defaults to :data:`"cls"`.
trust_remote_code (bool, optional): trust_remote_code for HF datasets or models. Defaults to :data:`False`.
cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
batch_size (int, optional): Batch size for inference. Defaults to :data:`256`.
query_max_length (int, optional): Maximum length for query. Defaults to :data:`512`.
passage_max_length (int, optional): Maximum length for passage. Defaults to :data:`512`.
convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will be a Torch Tensor.
Defaults to :data:`True`.
Attributes:
DEFAULT_POOLING_METHOD: The default pooling method when running the model.
"""
DEFAULT_POOLING_METHOD = "cls"
def __init__(
self,
model_name_or_path: str,
normalize_embeddings: bool = True,
use_fp16: bool = True,
use_bf16: bool = False,
query_instruction_for_retrieval: Optional[str] = None,
query_instruction_format: str = "{}{}", # specify the format of query_instruction_for_retrieval
devices: Optional[Union[str, List[str]]] = None, # specify devices, such as "cuda:0" or ["cuda:0", "cuda:1"]
# Additional parameters for BaseEmbedder
pooling_method: str = "cls",
trust_remote_code: bool = False,
cache_dir: Optional[str] = None,
# inference
batch_size: int = 256,
query_max_length: int = 512,
passage_max_length: int = 512,
convert_to_numpy: bool = True,
truncate_dim: Optional[int] = None,
**kwargs: Any,
):
super().__init__(
model_name_or_path,
normalize_embeddings=normalize_embeddings,
use_fp16=use_fp16,
use_bf16=use_bf16,
query_instruction_for_retrieval=query_instruction_for_retrieval,
query_instruction_format=query_instruction_format,
devices=devices,
batch_size=batch_size,
query_max_length=query_max_length,
passage_max_length=passage_max_length,
convert_to_numpy=convert_to_numpy,
truncate_dim=truncate_dim,
**kwargs
)
self.pooling_method = pooling_method
self.tokenizer = AutoTokenizer.from_pretrained(
model_name_or_path,
trust_remote_code=trust_remote_code,
cache_dir=cache_dir
)
self.model = AutoModel.from_pretrained(
model_name_or_path,
trust_remote_code=trust_remote_code,
cache_dir=cache_dir,
dtype=self.get_model_torch_dtype(),
)
def encode_queries(
self,
queries: Union[List[str], str],
batch_size: Optional[int] = None,
max_length: Optional[int] = None,
convert_to_numpy: Optional[bool] = None,
**kwargs: Any
) -> Union[np.ndarray, torch.Tensor]:
"""Encode the queries.
Args:
queries (Union[List[str], str]): Input queries to encode.
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will
be a Torch Tensor. Defaults to :data:`None`.
Returns:
Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
"""
return super().encode_queries(
queries,
batch_size=batch_size,
max_length=max_length,
convert_to_numpy=convert_to_numpy,
**kwargs
)
def encode_corpus(
self,
corpus: Union[List[str], str],
batch_size: Optional[int] = None,
max_length: Optional[int] = None,
convert_to_numpy: Optional[bool] = None,
**kwargs: Any
) -> Union[np.ndarray, torch.Tensor]:
"""Encode the corpus using the instruction if provided.
Args:
corpus (Union[List[str], str]): Input corpus to encode.
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will
be a Torch Tensor. Defaults to :data:`None`.
Returns:
Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
"""
return super().encode_corpus(
corpus,
batch_size=batch_size,
max_length=max_length,
convert_to_numpy=convert_to_numpy,
**kwargs
)
def encode(
self,
sentences: Union[List[str], str],
batch_size: Optional[int] = None,
max_length: Optional[int] = None,
convert_to_numpy: Optional[bool] = None,
**kwargs: Any
) -> Union[np.ndarray, torch.Tensor]:
"""Encode the input sentences with the embedding model.
Args:
sentences (Union[List[str], str]): Input sentences to encode.
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will
be a Torch Tensor. Defaults to :data:`None`.
Returns:
Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
"""
return super().encode(
sentences,
batch_size=batch_size,
max_length=max_length,
convert_to_numpy=convert_to_numpy,
**kwargs
)
@torch.no_grad()
def encode_single_device(
self,
sentences: Union[List[str], str],
batch_size: int = 256,
max_length: int = 512,
convert_to_numpy: bool = True,
device: Optional[str] = None,
**kwargs: Any
):
"""Encode input sentences by a single device.
Args:
sentences (Union[List[str], str]): Input sentences to encode.
batch_size (int, optional): Number of sentences for each iter. Defaults to :data:`256`.
max_length (int, optional): Maximum length of tokens. Defaults to :data:`512`.
convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will
be a Torch Tensor. Defaults to :data:`True`.
device (Optional[str], optional): Device to use for encoding. Defaults to None.
Returns:
Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
"""
if device is None:
device = self.target_devices[0]
if device == "cpu":
self.model.float()
self.model.to(device)
self.model.eval()
input_was_string = False
if isinstance(sentences, str):
sentences = [sentences]
input_was_string = True
# tokenize without padding to get the correct length
all_inputs = []
for start_index in trange(0, len(sentences), batch_size, desc='pre tokenize',
disable=len(sentences) < batch_size):
sentences_batch = sentences[start_index:start_index + batch_size]
inputs_batch = self.tokenizer(
sentences_batch,
truncation=True,
max_length=max_length,
**kwargs
)
inputs_batch = [{
k: inputs_batch[k][i] for k in inputs_batch.keys()
} for i in range(len(sentences_batch))]
all_inputs.extend(inputs_batch)
# sort by length for less padding
length_sorted_idx = np.argsort([-len(x['input_ids']) for x in all_inputs])
all_inputs_sorted = [all_inputs[i] for i in length_sorted_idx]
# adjust batch size
flag = False
while flag is False:
try:
inputs_batch = self.tokenizer.pad(
all_inputs_sorted[: batch_size],
padding=True,
return_tensors='pt',
**kwargs
).to(device)
last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
embeddings = self.pooling(last_hidden_state, inputs_batch['attention_mask'])
flag = True
except RuntimeError as e:
batch_size = batch_size * 3 // 4
except torch.cuda.OutOfMemoryError as e:
batch_size = batch_size * 3 // 4
# encode
all_embeddings = []
for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
disable=len(sentences) < batch_size):
inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
inputs_batch = self.tokenizer.pad(
inputs_batch,
padding=True,
return_tensors='pt',
**kwargs
).to(device)
last_hidden_state = self.model(**inputs_batch, return_dict=True).last_hidden_state
embeddings = self.pooling(last_hidden_state, inputs_batch['attention_mask'])
embeddings = self._truncate_embeddings(embeddings)
if self.normalize_embeddings:
embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
embeddings = cast(torch.Tensor, embeddings)
if convert_to_numpy:
embeddings = self._convert_to_numpy(embeddings, device=device)
all_embeddings.append(embeddings)
if convert_to_numpy:
all_embeddings = np.concatenate(all_embeddings, axis=0)
else:
all_embeddings = torch.cat(all_embeddings, dim=0)
# adjust the order of embeddings
all_embeddings = all_embeddings[np.argsort(length_sorted_idx)]
# return the embeddings
if input_was_string:
return all_embeddings[0]
return all_embeddings
def pooling(
self,
last_hidden_state: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None
):
"""The pooling function.
Args:
last_hidden_state (torch.Tensor): The last hidden state of the model.
attention_mask (Optional[torch.Tensor], optional): Attention mask. Defaults to :data:`None`.
Raises:
NotImplementedError: pooling method not implemented.
Returns:
torch.Tensor: The embedding vectors after pooling.
"""
if self.pooling_method == 'cls':
return last_hidden_state[:, 0]
elif self.pooling_method == 'mean':
s = torch.sum(last_hidden_state * attention_mask.unsqueeze(-1).float(), dim=1)
d = attention_mask.sum(dim=1, keepdim=True).float()
return s / d
else:
raise NotImplementedError(f"pooling method {self.pooling_method} not implemented")
FlagEmbedding/inference/embedder/encoder_only/m3.py
+792
View File
@@ -0,0 +1,792 @@
import math
import torch
import queue
import logging
import numpy as np
from tqdm import tqdm, trange
from multiprocessing import Queue
from collections import defaultdict
from transformers import AutoTokenizer
from typing import Any, List, Union, Dict, Literal, Tuple, Optional
from FlagEmbedding.abc.inference import AbsEmbedder
from FlagEmbedding.finetune.embedder.encoder_only.m3 import (
EncoderOnlyEmbedderM3ModelForInference, EncoderOnlyEmbedderM3Runner
)
logger = logging.getLogger(__name__)
class M3Embedder(AbsEmbedder):
"""
Embedder class for BGE-M3.
Args:
model_name_or_path (str): If it's a path to a local model, it loads the model from the path. Otherwise tries to download and
load a model from HuggingFace Hub with the name.
normalize_embeddings (bool, optional): If True, normalize the dense embedding vector. Defaults to :data:`True`.
use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance
degradation. Defaults to :data:`True`.
query_instruction_for_retrieval: (Optional[str], optional): Query instruction for retrieval tasks, which will be used with
with :attr:`query_instruction_format`. Defaults to :data:`None`.
query_instruction_format: (str, optional): The template for :attr:`query_instruction_for_retrieval`. Defaults to :data:`"{}{}"`.
devices (Optional[Union[str, int, List[str], List[int]]], optional): Devices to use for model inference. Defaults to :data:`None`.
pooling_method (str, optional): Pooling method to get embedding vector from the last hidden state. Defaults to :data:`"cls"`.
trust_remote_code (bool, optional): trust_remote_code for HF datasets or models. Defaults to :data:`False`.
cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
cobert_dim (int, optional): Dimension of colbert linear. Return the hidden_size if -1. Defaults to :data:`-1`.
batch_size (int, optional): Batch size for inference. Defaults to :data:`256`.
query_max_length (int, optional): Maximum length for query. Defaults to :data:`512`.
passage_max_length (int, optional): Maximum length for passage. Defaults to :data:`512`.
return_dense (bool, optional): If true, will return the dense embedding. Defaults to :data:`True`.
return_sparse (bool, optional): If true, will return the sparce embedding. Defaults to :data:`False`.
return_colbert_vecs (bool, optional): If true, will return the colbert vectors. Defaults to :data:`False`.
Attributes:
DEFAULT_POOLING_METHOD: The default pooling method when running the model.
"""
DEFAULT_POOLING_METHOD = "cls"
def __init__(
self,
model_name_or_path: str,
normalize_embeddings: bool = True,
use_fp16: bool = True,
use_bf16: bool = False,
query_instruction_for_retrieval: Optional[str] = None,
query_instruction_format: str = "{}{}", # specify the format of query_instruction_for_retrieval
devices: Optional[Union[str, List[str]]] = None, # specify devices, such as "cuda:0" or ["cuda:0", "cuda:1"]
# Additional parameters for M3Embedder
pooling_method: str = "cls",
trust_remote_code: bool = False,
cache_dir: Optional[str] = None,
colbert_dim: int = -1,
# inference
batch_size: int = 256,
query_max_length: int = 512,
passage_max_length: int = 512,
return_dense: bool = True,
return_sparse: bool = False,
return_colbert_vecs: bool = False,
truncate_dim: Optional[int] = None,
**kwargs: Any,
):
super().__init__(
model_name_or_path,
normalize_embeddings=normalize_embeddings,
use_fp16=use_fp16,
use_bf16=use_bf16,
query_instruction_for_retrieval=query_instruction_for_retrieval,
query_instruction_format=query_instruction_format,
devices=devices,
batch_size=batch_size,
query_max_length=query_max_length,
passage_max_length=passage_max_length,
return_dense=return_dense,
return_sparse=return_sparse,
return_colbert_vecs=return_colbert_vecs,
truncate_dim=truncate_dim,
**kwargs
)
self.pooling_method = pooling_method
self.tokenizer = AutoTokenizer.from_pretrained(
model_name_or_path,
trust_remote_code=trust_remote_code,
cache_dir=cache_dir
)
self.model = EncoderOnlyEmbedderM3ModelForInference(
EncoderOnlyEmbedderM3Runner.get_model(
model_name_or_path,
trust_remote_code=trust_remote_code,
colbert_dim=colbert_dim,
cache_dir=cache_dir,
torch_dtype=self.get_model_torch_dtype(),
),
tokenizer=self.tokenizer,
sentence_pooling_method=pooling_method,
normalize_embeddings=normalize_embeddings
)
def convert_id_to_token(self, lexical_weights: List[Dict]):
"""Convert the ids back to tokens.
Args:
lexical_weights (List[Dict]): A list of dictionaries of id & weights.
Returns:
List[Dict]: A list of dictionaries of tokens & weights.
"""
if isinstance(lexical_weights, dict):
lexical_weights = [lexical_weights]
new_lexical_weights = []
for item in lexical_weights:
new_item = {}
for id, weight in item.items():
token = self.tokenizer.decode([int(id)])
new_item[token] = weight
new_lexical_weights.append(new_item)
if len(new_lexical_weights) == 1:
new_lexical_weights = new_lexical_weights[0]
return new_lexical_weights
def compute_lexical_matching_score(
self,
lexical_weights_1: Union[Dict[str, float], List[Dict[str, float]]],
lexical_weights_2: Union[Dict[str, float], List[Dict[str, float]]]
) -> Union[np.ndarray, float]:
"""Compute the laxical matching score of two given lexical weights.
Args:
lexical_weights_1 (Union[Dict[str, float], List[Dict[str, float]]]): First array of lexical weights.
lexical_weights_2 (Union[Dict[str, float], List[Dict[str, float]]]): Second array of lexical weights.
Returns:
Union[np.ndarray, float]: The computed lexical weights across the two arries of lexical weights.
"""
def _compute_single_lexical_matching_score(lw1: Dict[str, float], lw2: Dict[str, float]):
scores = 0
for token, weight in lw1.items():
if token in lw2:
scores += weight * lw2[token]
return scores
if isinstance(lexical_weights_1, dict) and isinstance(lexical_weights_2, dict):
return _compute_single_lexical_matching_score(lexical_weights_1, lexical_weights_2)
elif isinstance(lexical_weights_1, list) and isinstance(lexical_weights_2, list):
scores_array = []
for lw1 in lexical_weights_1:
scores_array.append([
_compute_single_lexical_matching_score(lw1, lw2)
for lw2 in lexical_weights_2
])
return np.array(scores_array)
else:
raise ValueError("The input format of lexical_weights is not correct.")
def colbert_score(self, q_reps, p_reps):
"""Compute colbert scores of input queries and passages.
Args:
q_reps (np.ndarray): Multi-vector embeddings for queries.
p_reps (np.ndarray): Multi-vector embeddings for passages/corpus.
Returns:
torch.Tensor: Computed colbert scores.
"""
q_reps, p_reps = torch.from_numpy(q_reps), torch.from_numpy(p_reps)
token_scores = torch.einsum('in,jn->ij', q_reps, p_reps)
scores, _ = token_scores.max(-1)
scores = torch.sum(scores) / q_reps.size(0)
return scores
def encode_queries(
self,
queries: Union[List[str], str],
batch_size: Optional[int] = None,
max_length: Optional[int] = None,
return_dense: Optional[bool] = None,
return_sparse: Optional[bool] = None,
return_colbert_vecs: Optional[bool] = None,
**kwargs: Any
) -> Dict[
Literal["dense_vecs", "lexical_weights", "colbert_vecs"],
Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
]:
"""Encode the queries using the specified way.
Args:
queries (Union[List[str], str]): The input queries to encode.
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
return_dense (Optional[bool], optional): If True, compute and return dense embedding. Defaults to :data:`None`.
return_sparse (Optional[bool], optional): If True, compute and return sparce embedding. Defaults to :data:`None`.
return_colbert_vecs (Optional[bool], optional): If True, compute and return cobert vectors. Defaults to :data:`None`.
Returns:
Dict[Literal["dense_vecs", "lexical_weights", "colbert_vecs"], Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
"""
if batch_size is None: batch_size = self.batch_size
if max_length is None: max_length = self.query_max_length
if return_dense is None: return_dense = self.return_dense
if return_sparse is None: return_sparse = self.return_sparse
if return_colbert_vecs is None: return_colbert_vecs = self.return_colbert_vecs
return super().encode_queries(
queries,
batch_size=batch_size,
max_length=max_length,
return_dense=return_dense,
return_sparse=return_sparse,
return_colbert_vecs=return_colbert_vecs,
**kwargs
)
def encode_corpus(
self,
corpus: Union[List[str], str],
batch_size: Optional[int] = None,
max_length: Optional[int] = None,
return_dense: Optional[bool] = None,
return_sparse: Optional[bool] = None,
return_colbert_vecs: Optional[bool] = None,
**kwargs: Any
) -> Dict[
Literal["dense_vecs", "lexical_weights", "colbert_vecs"],
Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
]:
"""Encode the corpus using the specified way.
Args:
corpus (Union[List[str], str]): The input corpus to encode.
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
return_dense (Optional[bool], optional): If True, compute and return dense embedding. Defaults to :data:`None`.
return_sparse (Optional[bool], optional): If True, compute and return sparce embedding. Defaults to :data:`None`.
return_colbert_vecs (Optional[bool], optional): If True, compute and return cobert vectors. Defaults to :data:`None`.
Returns:
Dict[Literal["dense_vecs", "lexical_weights", "colbert_vecs"], Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
"""
if batch_size is None: batch_size = self.batch_size
if max_length is None: max_length = self.passage_max_length
if return_dense is None: return_dense = self.return_dense
if return_sparse is None: return_sparse = self.return_sparse
if return_colbert_vecs is None: return_colbert_vecs = self.return_colbert_vecs
return super().encode_corpus(
corpus,
batch_size=batch_size,
max_length=max_length,
return_dense=return_dense,
return_sparse=return_sparse,
return_colbert_vecs=return_colbert_vecs,
**kwargs
)
def encode(
self,
sentences: Union[List[str], str],
batch_size: Optional[int] = None,
max_length: Optional[int] = None,
return_dense: Optional[bool] = None,
return_sparse: Optional[bool] = None,
return_colbert_vecs: Optional[bool] = None,
**kwargs: Any
) -> Dict[
Literal["dense_vecs", "lexical_weights", "colbert_vecs"],
Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
]:
"""Encode the sentences using the specified way.
Args:
sentences (Union[List[str], str]): The input sentences to encode.
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
return_dense (Optional[bool], optional): If True, compute and return dense embedding. Defaults to :data:`None`.
return_sparse (Optional[bool], optional): If True, compute and return sparce embedding. Defaults to :data:`None`.
return_colbert_vecs (Optional[bool], optional): If True, compute and return cobert vectors. Defaults to :data:`None`.
Returns:
Dict[Literal["dense_vecs", "lexical_weights", "colbert_vecs"], Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
"""
if batch_size is None: batch_size = self.batch_size
if max_length is None: max_length = self.passage_max_length
if return_dense is None: return_dense = self.return_dense
if return_sparse is None: return_sparse = self.return_sparse
if return_colbert_vecs is None: return_colbert_vecs = self.return_colbert_vecs
return super().encode(
sentences,
batch_size=batch_size,
max_length=max_length,
return_dense=return_dense,
return_sparse=return_sparse,
return_colbert_vecs=return_colbert_vecs,
**kwargs
)
@torch.no_grad()
def encode_single_device(
self,
sentences: Union[List[str], str],
batch_size: int = 256,
max_length: int = 512,
return_dense: bool = True,
return_sparse: bool = False,
return_colbert_vecs: bool = False,
device: Optional[str] = None,
**kwargs: Any
):
"""Using single device to encode the input sentences.
Args:
sentences (Union[List[str], str]): The input sentences to encode.
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`256`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`512`.
return_dense (Optional[bool], optional): If True, compute and return dense embedding. Defaults to :data:`True`.
return_sparse (Optional[bool], optional): If True, compute and return sparce embedding. Defaults to :data:`False`.
return_colbert_vecs (Optional[bool], optional): If True, compute and return cobert vectors. Defaults to :data:`False`.
device (Optional[str], optional): _description_. Defaults to :data:`None`.
Returns:
Dict[Literal["dense_vecs", "lexical_weights", "colbert_vecs"], Union[np.ndarray, List[Dict[str, float]], List[np.ndarray]]
"""
# pop convert_to_numpy from kwargs
kwargs.pop("convert_to_numpy", None)
if device is None:
device = self.target_devices[0]
if device == "cpu":
self.model.float()
self.model.to(device)
self.model.eval()
input_was_string = False
if isinstance(sentences, str):
sentences = [sentences]
input_was_string = True
def _process_token_weights(token_weights: np.ndarray, input_ids: list):
# conver to dict
result = defaultdict(int)
unused_tokens = set()
for _token in ['cls_token', 'eos_token', 'pad_token', 'unk_token']:
if _token in self.tokenizer.special_tokens_map:
_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.special_tokens_map[_token])
unused_tokens.add(_token_id)
# token_weights = np.ceil(token_weights * 100)
for w, idx in zip(token_weights, input_ids):
if idx not in unused_tokens and w > 0:
idx = str(idx)
# w = int(w)
if w > result[idx]:
result[idx] = w
return result
def _process_colbert_vecs(colbert_vecs: np.ndarray, attention_mask: list):
# delte the vectors of padding tokens
tokens_num = np.sum(attention_mask)
return colbert_vecs[:tokens_num - 1] # we don't use the embedding of cls, so select tokens_num-1
# tokenize without padding to get the correct length
all_inputs = []
for start_index in trange(0, len(sentences), batch_size, desc='pre tokenize',
disable=len(sentences) < batch_size):
sentences_batch = sentences[start_index:start_index + batch_size]
inputs_batch = self.tokenizer(
sentences_batch,
truncation=True,
max_length=max_length,
**kwargs
)
inputs_batch = [{
k: inputs_batch[k][i] for k in inputs_batch.keys()
} for i in range(len(sentences_batch))]
all_inputs.extend(inputs_batch)
# sort by length for less padding
length_sorted_idx = np.argsort([-len(x['input_ids']) for x in all_inputs])
all_inputs_sorted = [all_inputs[i] for i in length_sorted_idx]
# adjust batch size
flag = False
while flag is False:
try:
inputs_batch = self.tokenizer.pad(
all_inputs_sorted[: batch_size],
padding=True,
return_tensors='pt',
**kwargs
).to(device)
outputs = self.model(
inputs_batch,
return_dense=return_dense,
return_sparse=return_sparse,
return_colbert_vecs=return_colbert_vecs
)
flag = True
except RuntimeError as e:
batch_size = batch_size * 3 // 4
except torch.cuda.OutOfMemoryError as e:
batch_size = batch_size * 3 // 4
# encode
all_dense_embeddings, all_lexical_weights, all_colbert_vecs = [], [], []
for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
disable=len(sentences) < batch_size):
inputs_batch = all_inputs_sorted[start_index:start_index + batch_size]
inputs_batch = self.tokenizer.pad(
inputs_batch,
padding=True,
return_tensors='pt',
**kwargs
).to(device)
outputs = self.model(
inputs_batch,
return_dense=return_dense,
return_sparse=return_sparse,
return_colbert_vecs=return_colbert_vecs,
truncate_dim=self.truncate_dim
)
if return_dense:
all_dense_embeddings.append(self._convert_to_numpy(outputs['dense_vecs'], device=device))
if return_sparse:
token_weights = outputs['sparse_vecs'].squeeze(-1)
all_lexical_weights.extend(
list(map(
_process_token_weights,
self._convert_to_numpy(token_weights, device=device),
self._convert_to_numpy(inputs_batch['input_ids'], device=device).tolist()
)))
if return_colbert_vecs:
all_colbert_vecs.extend(
list(map(
_process_colbert_vecs,
self._convert_to_numpy(outputs['colbert_vecs'], device=device),
self._convert_to_numpy(inputs_batch['attention_mask'], device=device)
)))
if return_dense:
all_dense_embeddings = np.concatenate(all_dense_embeddings, axis=0)
# adjust the order of embeddings
all_dense_embeddings = all_dense_embeddings[np.argsort(length_sorted_idx)]
if input_was_string:
all_dense_embeddings = all_dense_embeddings[0]
else:
all_dense_embeddings = None
if return_sparse:
# adjust the order of lexical weights
all_lexical_weights = [all_lexical_weights[i] for i in np.argsort(length_sorted_idx)]
if input_was_string:
all_lexical_weights = all_lexical_weights[0]
else:
all_lexical_weights = None
if return_colbert_vecs:
# adjust the order of embeddings
all_colbert_vecs = [all_colbert_vecs[i] for i in np.argsort(length_sorted_idx)]
if input_was_string:
all_colbert_vecs = all_colbert_vecs[0]
else:
all_colbert_vecs = None
# return the embeddings
return {
"dense_vecs": all_dense_embeddings,
"lexical_weights": all_lexical_weights,
"colbert_vecs": all_colbert_vecs
}
def compute_score(
self,
sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
batch_size: Optional[int] = None,
max_query_length: Optional[int] = None,
max_passage_length: Optional[int] = None,
weights_for_different_modes: Optional[List[float]] = None,
**kwargs: Any
) -> Dict[
Literal["colbert", "sparse", "dense", "sparse+dense", "colbert+sparse+dense"],
List[float]
]:
"""Compute the relevance score of different attributes.
Args:
sentence_pairs (Union[List[Tuple[str, str]], Tuple[str, str]]): _description_
batch_size (Optional[int], optional): _description_. Defaults to None.
max_query_length (Optional[int], optional): _description_. Defaults to None.
max_passage_length (Optional[int], optional): _description_. Defaults to None.
weights_for_different_modes (Optional[List[float]], optional): _description_. Defaults to None.
Returns:
Dict[Literal["colbert", "sparse", "dense", "sparse+dense", "colbert+sparse+dense"], List[float]]
"""
if batch_size is None: batch_size = self.batch_size
if max_query_length is None: max_query_length = self.query_max_length
if max_passage_length is None: max_passage_length = self.passage_max_length
if len(self.target_devices) == 1:
return self.compute_score_single_device(
sentence_pairs,
batch_size=batch_size,
max_query_length=max_query_length,
max_passage_length=max_passage_length,
weights_for_different_modes=weights_for_different_modes,
device=self.target_devices[0],
**kwargs
)
pool = self.start_multi_process_pool(M3Embedder._compute_score_multi_process_worker)
embeddings = self.compute_score_multi_process(
sentence_pairs,
pool,
batch_size=batch_size,
max_query_length=max_query_length,
max_passage_length=max_passage_length,
weights_for_different_modes=weights_for_different_modes,
**kwargs
)
self.stop_multi_process_pool(pool)
return embeddings
# adapted from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L877
def compute_score_multi_process(
self,
sentence_pairs: List[Tuple[str, str]],
pool: Dict[Literal["input", "output", "processes"], Any],
**kwargs
):
chunk_size = math.ceil(len(sentence_pairs) / len(pool["processes"]))
input_queue = pool["input"]
last_chunk_id = 0
chunk = []
for sentence_pair in sentence_pairs:
chunk.append(sentence_pair)
if len(chunk) >= chunk_size:
input_queue.put(
[last_chunk_id, chunk, kwargs]
)
last_chunk_id += 1
chunk = []
if len(chunk) > 0:
input_queue.put([last_chunk_id, chunk, kwargs])
last_chunk_id += 1
output_queue = pool["output"]
results_list = sorted(
[output_queue.get() for _ in trange(last_chunk_id, desc="Chunks")],
key=lambda x: x[0],
)
scores_dict = self._concatenate_compute_score_results_from_multi_process([result[1] for result in results_list])
return scores_dict
# adapted from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L976
@staticmethod
def _compute_score_multi_process_worker(
target_device: str, model: 'M3Embedder', input_queue: Queue, results_queue: Queue
) -> None:
"""
Internal working process to encode sentences in multi-process setup
"""
while True:
try:
chunk_id, sentences, kwargs = (
input_queue.get()
)
embeddings = model.compute_score_single_device(
sentences,
device=target_device,
**kwargs
)
results_queue.put([chunk_id, embeddings])
except queue.Empty:
break
@torch.no_grad()
def compute_score_single_device(
self,
sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
batch_size: int = 256,
max_query_length: int = 512,
max_passage_length: int = 512,
weights_for_different_modes: Optional[List[float]] = None,
device: Optional[str] = None,
**kwargs: Any
) -> Dict[
Literal["colbert", "sparse", "dense", "sparse+dense", "colbert+sparse+dense"],
List[float]
]:
"""Compute the relevance score of different attributes.
Args:
sentence_pairs (Union[List[Tuple[str, str]], Tuple[str, str]]): Pairs of sentences to compute the score.
batch_size (Optional[int], optional): _description_. Defaults to :data:`None`.
max_query_length (Optional[int], optional): _description_. Defaults to :data:`None`.
max_passage_length (Optional[int], optional): _description_. Defaults to :data:`None`.
weights_for_different_modes (Optional[List[float]], optional): The weights for different methods. Defaults to :data:`None`.
device (Optional[str], optional): The device to use. Defaults to :data:`None`.
Returns:
Dict[Literal["colbert", "sparse", "dense", "sparse+dense", "colbert+sparse+dense"], List[float]]
"""
def _tokenize(texts: list, max_length: int):
return self.tokenizer(
texts,
max_length=max_length,
padding=True,
return_token_type_ids=False,
truncation=True,
return_tensors='pt',
**kwargs
)
if device is None:
device = self.target_devices[0]
if device == "cpu":
self.model.float()
self.model.to(device)
self.model.eval()
if isinstance(sentence_pairs, list) and len(sentence_pairs) == 0:
return []
if isinstance(sentence_pairs[0], str):
one_input_pair = True
sentence_pairs = [sentence_pairs]
else:
one_input_pair = False
all_scores = {
'colbert': [],
'sparse': [],
'dense': [],
'sparse+dense': [],
'colbert+sparse+dense': []
}
for start_index in tqdm(range(0, len(sentence_pairs), batch_size), desc="Compute Scores",
disable=len(sentence_pairs) < batch_size):
sentences_batch = sentence_pairs[start_index:start_index + batch_size]
queries_batch = [pair[0] for pair in sentences_batch]
corpus_batch = [pair[1] for pair in sentences_batch]
queries_inputs = _tokenize(queries_batch, max_length=max_query_length).to(device)
corpus_inputs = _tokenize(corpus_batch, max_length=max_passage_length).to(device)
queries_output = self.model(
queries_inputs,
return_dense=True, return_sparse=True, return_colbert_vecs=True,
return_sparse_embedding=True
)
corpus_output = self.model(
corpus_inputs,
return_dense=True, return_sparse=True, return_colbert_vecs=True,
return_sparse_embedding=True
)
q_dense_vecs, q_sparse_vecs, q_colbert_vecs = queries_output['dense_vecs'], queries_output['sparse_vecs'], \
queries_output['colbert_vecs']
p_dense_vecs, p_sparse_vecs, p_colbert_vecs = corpus_output['dense_vecs'], corpus_output['sparse_vecs'], \
corpus_output['colbert_vecs']
dense_scores = self.model.compute_dense_score(q_dense_vecs, p_dense_vecs)
sparse_scores = self.model.compute_sparse_score(q_sparse_vecs, p_sparse_vecs)
colbert_scores = self.model.compute_colbert_score(
q_colbert_vecs, p_colbert_vecs,
q_mask=queries_inputs['attention_mask']
)
if weights_for_different_modes is None:
weights_for_different_modes = [1., 1., 1.]
weight_sum = 3
logger.info("default weights for dense, sparse, colbert are [1.0, 1.0, 1.0] ")
else:
assert len(weights_for_different_modes) == 3
weight_sum = sum(weights_for_different_modes)
inx = torch.arange(0, len(sentences_batch))
dense_scores, sparse_scores, colbert_scores = dense_scores[inx, inx].float(), sparse_scores[
inx, inx].float(), colbert_scores[inx, inx].float()
all_scores['colbert'].extend(
self._convert_to_numpy(colbert_scores, device=device).tolist()
)
all_scores['sparse'].extend(
self._convert_to_numpy(sparse_scores, device=device).tolist()
)
all_scores['dense'].extend(
self._convert_to_numpy(dense_scores, device=device).tolist()
)
all_scores['sparse+dense'].extend(
self._convert_to_numpy(
(sparse_scores * weights_for_different_modes[1] + dense_scores * weights_for_different_modes[0])
/ (weights_for_different_modes[1] + weights_for_different_modes[0]),
device=device,
).tolist()
)
all_scores['colbert+sparse+dense'].extend(
self._convert_to_numpy(
(colbert_scores * weights_for_different_modes[2]
+ sparse_scores * weights_for_different_modes[1]
+ dense_scores * weights_for_different_modes[0]) / weight_sum,
device=device,
).tolist()
)
if one_input_pair:
return {k: v[0] for k, v in all_scores.items()}
return all_scores
def _concatenate_results_from_multi_process(
self,
results_list: List[Dict[Literal["dense_vecs", "lexical_weights", "colbert_vecs"], Any]]
):
"""Concatenate and return the results from all the processes.
Args:
results_list (List[Dict[Literal[&quot;dense_vecs&quot;, &quot;lexical_weights&quot;, &quot;colbert_vecs&quot;], Any]]):
A list of results from all the processes.
Returns:
Dict: The merged encoding results from the multi processes.
"""
merged_results = {
"dense_vecs": [],
"lexical_weights": [],
"colbert_vecs": []
}
for key in merged_results.keys():
for results in results_list:
if results[key] is None:
merged_results[key] = None
break
else:
if key == "dense_vecs":
merged_results[key].append(results[key])
else:
merged_results[key].extend(results[key])
if merged_results["dense_vecs"] is not None:
merged_results["dense_vecs"] = np.concatenate(merged_results["dense_vecs"], axis=0)
return merged_results
def _concatenate_compute_score_results_from_multi_process(
self,
results_list: List[Dict[Literal["colbert", "sparse", "dense", "sparse+dense", "colbert+sparse+dense"], List[float]]]
):
"""Concatenate and return the results from all the processes.
Args:
results_list (List[Dict[Literal[&quot;colbert&quot;, &quot;sparse&quot;, &quot;dense&quot;, &quot;sparse):
A list of computed scores.
Returns:
Dict: The merged computed scores from the multi processes.
"""
merged_results = {
"colbert": [],
"sparse": [],
"dense": [],
"sparse+dense": [],
"colbert+sparse+dense": []
}
for key in merged_results.keys():
for results in results_list:
merged_results[key].extend(results[key])
return merged_results
FlagEmbedding/inference/embedder/model_mapping.py
+274
View File
@@ -0,0 +1,274 @@
from enum import Enum
from typing import Type, List
from dataclasses import dataclass
from collections import OrderedDict
from FlagEmbedding.abc.inference import AbsEmbedder
from FlagEmbedding.inference.embedder import FlagModel, BGEM3FlagModel, FlagLLMModel, FlagICLModel, FlagPseudoMoEModel
class EmbedderModelClass(Enum):
ENCODER_ONLY_BASE = "encoder-only-base"
ENCODER_ONLY_M3 = "encoder-only-m3"
DECODER_ONLY_BASE = "decoder-only-base"
DECODER_ONLY_ICL = "decoder-only-icl"
DECODER_ONLY_PSEUDO_MOE = "decoder-only-pseudo_moe"
EMBEDDER_CLASS_MAPPING = OrderedDict([
(EmbedderModelClass.ENCODER_ONLY_BASE, FlagModel),
(EmbedderModelClass.ENCODER_ONLY_M3, BGEM3FlagModel),
(EmbedderModelClass.DECODER_ONLY_BASE, FlagLLMModel),
(EmbedderModelClass.DECODER_ONLY_ICL, FlagICLModel),
(EmbedderModelClass.DECODER_ONLY_PSEUDO_MOE, FlagPseudoMoEModel)
])
class PoolingMethod(Enum):
LAST_TOKEN = "last_token"
CLS = "cls"
MEAN = "mean"
@dataclass
class EmbedderConfig:
model_class: Type[AbsEmbedder]
pooling_method: PoolingMethod
trust_remote_code: bool = False
query_instruction_format: str = "{}{}"
# BGE models mapping
BGE_MAPPING = OrderedDict([
(
"bge-reasoner-embed-qwen3-8b-0923",
EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
),
(
"bge-code-v1",
EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, trust_remote_code=True, query_instruction_format="<instruct>{}\n<query>{}")
),
(
"bge-en-icl",
EmbedderConfig(FlagICLModel, PoolingMethod.LAST_TOKEN, query_instruction_format="<instruct>{}\n<query>{}")
),
(
"bge-multilingual-gemma2",
EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="<instruct>{}\n<query>{}")
),
(
"bge-m3",
EmbedderConfig(BGEM3FlagModel, PoolingMethod.CLS)
),
(
"bge-large-en-v1.5",
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
(
"bge-base-en-v1.5",
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
(
"bge-small-en-v1.5",
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
(
"bge-large-zh-v1.5",
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
(
"bge-base-zh-v1.5",
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
(
"bge-small-zh-v1.5",
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
(
"bge-large-en",
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
(
"bge-base-en",
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
(
"bge-small-en",
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
(
"bge-large-zh",
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
(
"bge-base-zh",
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
(
"bge-small-zh",
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
])
# Qwen3-Embedding models mapping
QWEN3_EMBEDDING_MAPPING = OrderedDict([
(
"Qwen3-Embedding-0.6B",
EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery:{}")
),
(
"Qwen3-Embedding-4B",
EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery:{}")
),
(
"Qwen3-Embedding-8B",
EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery:{}")
),
])
# E5 models mapping
E5_MAPPING = OrderedDict([
(
"e5-mistral-7b-instruct",
EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
),
(
"e5-large-v2",
EmbedderConfig(FlagModel, PoolingMethod.MEAN)
),
(
"e5-base-v2",
EmbedderConfig(FlagModel, PoolingMethod.MEAN)
),
(
"e5-small-v2",
EmbedderConfig(FlagModel, PoolingMethod.MEAN)
),
(
"multilingual-e5-large-instruct",
EmbedderConfig(FlagModel, PoolingMethod.MEAN, query_instruction_format="Instruct: {}\nQuery: {}")
),
(
"multilingual-e5-large",
EmbedderConfig(FlagModel, PoolingMethod.MEAN)
),
(
"multilingual-e5-base",
EmbedderConfig(FlagModel, PoolingMethod.MEAN)
),
(
"multilingual-e5-small",
EmbedderConfig(FlagModel, PoolingMethod.MEAN)
),
(
"e5-large",
EmbedderConfig(FlagModel, PoolingMethod.MEAN)
),
(
"e5-base",
EmbedderConfig(FlagModel, PoolingMethod.MEAN)
),
(
"e5-small",
EmbedderConfig(FlagModel, PoolingMethod.MEAN)
),
])
# GTE models mapping
GTE_MAPPING = OrderedDict([
(
"gte-Qwen2-7B-instruct",
EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, trust_remote_code=True, query_instruction_format="Instruct: {}\nQuery: {}")
),
(
"gte-Qwen2-1.5B-instruct",
EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, trust_remote_code=True, query_instruction_format="Instruct: {}\nQuery: {}")
),
(
"gte-Qwen1.5-7B-instruct",
EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, trust_remote_code=True, query_instruction_format="Instruct: {}\nQuery: {}")
),
(
"gte-multilingual-base",
EmbedderConfig(FlagModel, PoolingMethod.CLS, trust_remote_code=True)
),
(
"gte-large-en-v1.5",
EmbedderConfig(FlagModel, PoolingMethod.CLS, trust_remote_code=True)
),
(
"gte-base-en-v1.5",
EmbedderConfig(FlagModel, PoolingMethod.CLS, True)
),
(
'gte-large',
EmbedderConfig(FlagModel, PoolingMethod.MEAN)
),
(
'gte-base',
EmbedderConfig(FlagModel, PoolingMethod.MEAN)
),
(
'gte-small',
EmbedderConfig(FlagModel, PoolingMethod.MEAN)
),
(
'gte-large-zh',
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
(
'gte-base-zh',
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
(
'gte-small-zh',
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
])
# SFR models mapping
SFR_MAPPING = OrderedDict([
(
'SFR-Embedding-2_R',
EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
),
(
'SFR-Embedding-Mistral',
EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
),
])
# Linq models mapping
LINQ_MAPPING = OrderedDict([
(
'Linq-Embed-Mistral',
EmbedderConfig(FlagLLMModel, PoolingMethod.LAST_TOKEN, query_instruction_format="Instruct: {}\nQuery: {}")
),
])
# BCE models mapping
BCE_MAPPING = OrderedDict([
(
'bce-embedding-base_v1',
EmbedderConfig(FlagModel, PoolingMethod.CLS)
),
])
# Combine all mappings
AUTO_EMBEDDER_MAPPING = OrderedDict()
AUTO_EMBEDDER_MAPPING.update(BGE_MAPPING)
AUTO_EMBEDDER_MAPPING.update(QWEN3_EMBEDDING_MAPPING)
AUTO_EMBEDDER_MAPPING.update(E5_MAPPING)
AUTO_EMBEDDER_MAPPING.update(GTE_MAPPING)
AUTO_EMBEDDER_MAPPING.update(SFR_MAPPING)
AUTO_EMBEDDER_MAPPING.update(LINQ_MAPPING)
AUTO_EMBEDDER_MAPPING.update(BCE_MAPPING)
# TODO: Add more models, such as Jina, Stella_v5, NV-Embed, etc.
def support_native_bge_model_list()->List[str]:
return list(BGE_MAPPING.keys())
def support_model_list()->List[str]:
return list(AUTO_EMBEDDER_MAPPING.keys())