BGE的技术亮点：
- 高效预训练和大规模文本微调；
- 在两个大规模语料集上采用了RetroMAE预训练算法，进一步增强了模型的语义表征能力；
- 通过负采样和难负样例挖掘，增强了语义向量的判别力；
- 借鉴Instruction Tuning的策略，增强了在多任务场景下的通用能力。
数据集的构成：
RetroMAE预训练
主要思想是：encoder用小一点的mask rate得到sentence embedding，然后decoder用大一点的mask rate结合encoder得到的sentence embedding进行重构
此外，为了使得每个token使用的context信息不同，RetroMAE还使用了增强解码的方法
- 解码的时候每一行都带，上下文信息和位置信息
对应代码
import logging
import os
import torch
from torch import nn
from transformers import BertForMaskedLM, AutoModelForMaskedLM
from transformers.modeling_outputs import MaskedLMOutput
from .arguments import ModelArguments
from .enhancedDecoder import BertLayerForDecoder
logger = logging.getLogger(__name__)
class RetroMAEForPretraining(nn.Module):
def __init__(
self,
bert: BertForMaskedLM,
model_args: ModelArguments,
):
super(RetroMAEForPretraining, self).__init__()
self.lm = bert
if hasattr(self.lm, 'bert'):
self.decoder_embeddings = self.lm.bert.embeddings
elif hasattr(self.lm, 'roberta'):
self.decoder_embeddings = self.lm.roberta.embeddings
else:
self.decoder_embeddings = self.lm.bert.embeddings
self.c_head = BertLayerForDecoder(bert.config)
self.c_head.apply(self.lm._init_weights)
self.cross_entropy = nn.CrossEntropyLoss()
self.model_args = model_args
def gradient_checkpointing_enable(self, **kwargs):
self.lm.gradient_checkpointing_enable(**kwargs)
def forward(self,
encoder_input_ids, encoder_attention_mask, encoder_labels,
decoder_input_ids, decoder_attention_mask, decoder_labels):
lm_out: MaskedLMOutput = self.lm(
encoder_input_ids, encoder_attention_mask,
labels=encoder_labels,
output_hidden_states=True,
return_dict=True
)
cls_hiddens = lm_out.hidden_states[-1][:, :1] # B 1 D
decoder_embedding_output = self.decoder_embeddings(input_ids=decoder_input_ids)
hiddens = torch.cat([cls_hiddens, decoder_embedding_output[:, 1:]], dim=1)
decoder_position_ids = self.lm.bert.embeddings.position_ids[:, :decoder_input_ids.size(1)]
decoder_position_embeddings = self.lm.bert.embeddings.position_embeddings(decoder_position_ids) # B L D
query = decoder_position_embeddings + cls_hiddens
# cls_hiddens = cls_hiddens.expand(hiddens.size(0), hiddens.size(1), hiddens.size(2))
# query = self.decoder_embeddings(inputs_embeds=cls_hiddens)
matrix_attention_mask = self.lm.get_extended_attention_mask(
decoder_attention_mask,
decoder_attention_mask.shape,
decoder_attention_mask.device
)
hiddens = self.c_head(query=query,
key=hiddens,
value=hiddens,
attention_mask=matrix_attention_mask)[0]
pred_scores, loss = self.mlm_loss(hiddens, decoder_labels)
return (loss + lm_out.loss,)
def mlm_loss(self, hiddens, labels):
if hasattr(self.lm, 'cls'):
pred_scores = self.lm.cls(hiddens)
elif hasattr(self.lm, 'lm_head'):
pred_scores = self.lm.lm_head(hiddens)
else:
raise NotImplementedError
masked_lm_loss = self.cross_entropy(
pred_scores.view(-1, self.lm.config.vocab_size),
labels.view(-1)
)
return pred_scores, masked_lm_loss
def save_pretrained(self, output_dir: str):
self.lm.save_pretrained(os.path.join(output_dir, "encoder_model"))
torch.save(self.state_dict(), os.path.join(output_dir, 'pytorch_model.bin'))
@classmethod
def from_pretrained(
cls, model_args: ModelArguments,
*args, **kwargs
):
hf_model = AutoModelForMaskedLM.from_pretrained(*args, **kwargs)
model = cls(hf_model, model_args)
return model
训练策略
- 主要通过对比学习和Instruction Tuning的思想
对比学习是一种训练模型的方法，通过比较正例和反例来学习数据的表示。
-
输入数据的格式：模型接受三元组格式的数据作为输入，包括一个查询（query），一个正例（positive），和一个反例（negative）。
-
in-batch negatives 策略：除了上述三元组中的反例外，他们还采用了“in-batch negatives”策略，意思是在同一个批次的数据中，使用其他数据作为额外的反例。
-
cross-device negatives sharing method：这是一种在不同的GPU之间共享反例的方法，目的是大大增加反例的数量。
-
训练硬件和参数：使用了48个A100(40G)的GPU进行训练。批次大小为32,768，因此每个查询在批次中有65,535个反例。使用了AdamW优化器，学习率为1e-5。对比损失的温度为0.01。
-
在训练中为检索任务的查询添加了instruction。 对于英语，指令是Represent this sentence for searching relevant passages: ; 对于中文，指令是为这个句子生成表示以用于检索相关文章：. 在评测中，针对段落检索任务的任务需要在查询中添加指令，但不需要为段落文档添加指令。
Embedding
Todo：代码阅读
- https://github.com/FlagOpen/FlagEmbedding/blob/master/FlagEmbedding/flag_models.py
Embedding微调代码
数据格式
{"query": str, "pos": List[str], "neg":List[str]}
query是查询，pos是肯定文本列表，neg是否定文本列表。
如果查询没有否定文本，您可以从整个语料库中随机采样一些文本作为否定文本。
demo.json
{"query": "Five women walk along a beach wearing flip-flops.", "pos": ["Some women with flip-flops on, are walking along the beach"], "neg": ["The 4 women are sitting on the beach.", "There was a reform in 1996.", "She's not going to court to clear her record.", "The man is talking about hawaii.", "A woman is standing outside.", "The battle was over. ", "A group of people plays volleyball."]}
{"query": "A woman standing on a high cliff on one leg looking over a river.", "pos": ["A woman is standing on a cliff."], "neg": ["A woman sits on a chair.", "George Bush told the Republicans there was no way he would let them even consider this foolish idea, against his top advisors advice.", "The family was falling apart.", "no one showed up to the meeting", "A boy is sitting outside playing in the sand.", "Ended as soon as I received the wire.", "A child is reading in her bedroom."]}
{"query": "Two woman are playing instruments; one a clarinet, the other a violin.", "pos": ["Some people are playing a tune."], "neg": ["Two women are playing a guitar and drums.", "A man is skiing down a mountain.", "The fatal dose was not taken when the murderer thought it would be.", "Person on bike", "The girl is standing, leaning against the archway.", "A group of women watch soap operas.", "No matter how old people get they never forget. "]}
{"query": "A girl with a blue tank top sitting watching three dogs.", "pos": ["A girl is wearing blue."], "neg": ["A girl is with three cats.", "The people are watching a funeral procession.", "The child is wearing black.", "Financing is an issue for us in public schools.", "Kids at a pool.", "It is calming to be assaulted.", "I face a serious problem at eighteen years old. "]}
{"query": "A yellow dog running along a forest path.", "pos": ["a dog is running"], "neg": ["a cat is running", "Steele did not keep her original story.", "The rule discourages people to pay their child support.", "A man in a vest sits in a car.", "Person in black clothing, with white bandanna and sunglasses waits at a bus stop.", "Neither the Globe or Mail had comments on the current state of Canada's road system. ", "The Spring Creek facility is old and outdated."]}
{"query": "It sets out essential activities in each phase along with critical factors related to those activities.", "pos": ["Critical factors for essential activities are set out."], "neg": ["It lays out critical activities but makes no provision for critical factors related to those activities.", "People are assembled in protest.", "The state would prefer for you to do that.", "A girl sits beside a boy.", "Two males are performing.", "Nobody is jumping", "Conrad was being plotted against, to be hit on the head."]}
{"query": "A man giving a speech in a restaurant.", "pos": ["A person gives a speech."], "neg": ["The man sits at the table and eats food.", "This is definitely not an endorsement.", "They sold their home because they were retiring and not because of the loan.", "The seal of Missouri is perfect.", "Someone is raising their hand.", "An athlete is competing in the 1500 meter swimming competition.", "Two men watching a magic show."]}
{"query": "Indians having a gathering with coats and food and drinks.", "pos": ["A group of Indians are having a gathering with food and drinks"], "neg": ["A group of Indians are having a funeral", "It is only staged on Winter afternoons in Palma's large bullring.", "Right information can empower the legal service practices and the justice system. ", "Meanwhile, the mainland was empty of population.", "Two children is sleeping.", "a fisherman is trying to catch a monkey", "the people are in a train"]}
{"query": "A woman with violet hair rides her bicycle outside.", "pos": ["A woman is riding her bike."], "neg": ["A woman is jogging in the park.", "The street was lined with white-painted houses.", "A group watches a movie inside.", "man at picnics cut steak", "Several chefs are sitting down and talking about food.", "The Commission notes that no significant alternatives were considered.", "We ran out of firewood and had to use pine needles for the fire."]}
{"query": "A man pulls two women down a city street in a rickshaw.", "pos": ["A man is in a city."], "neg": ["A man is a pilot of an airplane.", "It is boring and mundane.", "The morning sunlight was shining brightly and it was warm. ", "Two people jumped off the dock.", "People watching a spaceship launch.", "Mother Teresa is an easy choice.", "It's worth being able to go at a pace you prefer."]}
启动脚本
难负样本
难负样本是一种广泛使用的提高句子嵌入质量的方法。您可以按照以下命令微调
python -m FlagEmbedding.baai_general_embedding.finetune.hn_mine \
--model_name_or_path BAAI/bge-base-en-v1.5 \
--input_file toy_finetune_data.jsonl \
--output_file toy_finetune_data_minedHN.jsonl \
--range_for_sampling 2-200 \
--negative_number 15 \
--use_gpu_for_searching
- input_file：用于微调的 json 数据。该脚本将为每个查询检索 top-k 文档，并从 top-k 文档中随机抽取负样本（不包括正文档）。
- output_file：保存带有挖掘的硬底片以进行微调的 JSON 数据的路径
- negative_number：采样负数的数量
- range_for_sampling：在哪里采样负数。例如，2-100表示negative_number从 top2-top200 文档中采样负数。您可以设置更大的值来降低负数的难度（例如，设置60-300为从top60-300段落中采样负数）
- candidate_pool：要检索的池。默认值为 None，此脚本将从 中所有neg内容的组合中检索input_file。该文件的格式与预训练数据相同。如果输入候选池，该脚本将从该文件中检索底片。
- use_gpu_for_searching：是否使用 faiss-gpu 检索底片。
train
torchrun --nproc_per_node {number of gpus} \
-m FlagEmbedding.baai_general_embedding.finetune.run \
--output_dir {path to save model} \
--model_name_or_path BAAI/bge-large-zh-v1.5 \
--train_data ./toy_finetune_data.jsonl \
--learning_rate 1e-5 \
--fp16 \
--num_train_epochs 5 \
--per_device_train_batch_size {large batch size; set 1 for toy data} \
--dataloader_drop_last True \
--normlized True \
--temperature 0.02 \
--query_max_len 64 \
--passage_max_len 256 \
--train_group_size 2 \
--negatives_cross_device \
--logging_steps 10 \
--query_instruction_for_retrieval ""
- per_device_train_batch_size：训练中的批量大小。在大多数情况下，更大的批量大小会带来更强的性能。您可以通过启用–fp16、–deepspeed ./df_config.json（df_config.json 可以参考ds_config.json）–gradient_checkpointing等来扩展它。
- train_group_size：训练中查询的正数和负数的数量。总是有一个正数，因此该参数将控制负数的数量 (#negatives=train_group_size-1)。请注意，负数的数量不应大于 data 中负数的数量"neg":List[str]。除了该组中的底片之外，批量内的底片也将用于微调。
- negatives_cross_device：共享所有 GPU 的负面影响。这个论点将扩大否定的数量。
- learning_rate：选择适合您的型号的。对于大型/基础/小型，推荐 1e-5/2e-5/3e-5。
- temperature：会影响相似度分数的分布。
- query_max_len：查询的最大长度。请根据您数据中查询的平均长度进行设置。
- passage_max_len：通道的最大长度。请根据您的数据中段落的平均长度进行设置。
- query_instruction_for_retrieval：查询指令，会添加到每个查询中。您还可以将其设置""为不添加任何查询。
- use_inbatch_neg：使用同一批中的段落作为底片。默认值为 True。
模型合并
微调基础 BGE 模型可以提高其在目标任务上的性能，但可能会导致目标域之外的模型一般功能的严重退化（例如，在 c-mteb 任务上的性能较低）。通过合并微调模型和基础模型，LM-Cocktail 可以显着增强下游任务的性能，同时保持其他不相关任务的性能。
from LM_Cocktail import mix_models, mix_models_with_data
example_data = [
{"query": "How does one become an actor in the Telugu Film Industry?", "pos": [" How do I become an actor in Telugu film industry?"], "neg": [" What is the story of Moses and Ramesses?", " Does caste system affect economic growth of India?"]},
{"query": "Why do some computer programmers develop amazing software or new concepts, while some are stuck with basic programming work?", "pos": [" Why do some computer programmers develops amazing softwares or new concepts, while some are stuck with basics programming works?"], "neg": [" When visiting a friend, do you ever think about what would happen if you did something wildly inappropriate like punch them or destroy their furniture?", " What is the difference between a compliment and flirting?"]}
]
model = mix_models_with_data(
model_names_or_paths=["BAAI/bge-base-en-v1.5", "Shitao/bge-hotpotqa", "Shitao/bge-quora"],
model_type='encoder',
example_ata=example_data,
temperature=5.0,
max_input_length=512,
neg_number=2)
评估
-
语料库
-
query验证集
python -m FlagEmbedding.baai_general_embedding.finetune.eval_msmarco \
--encoder BAAI/bge-base-en-v1.5 \
--fp16 \
--add_instruction \
--k 100
Rerank
Todo：代码阅读
微调代码
reranker使用问题和文档作为输入，直接输出相似度而不是embedding。您可以通过向重新排序器输入查询和段落来获得相关性分数。
-
重排序器是基于交叉熵损失进行优化的，因此相关性得分不受特定范围的限制。
-
数据格式和Embedding 一样，
query,pos,neg，Hard Negative
train
torchrun --nproc_per_node {number of gpus} \
-m FlagEmbedding.reranker.run \
--output_dir {path to save model} \
--model_name_or_path BAAI/bge-reranker-base \
--train_data ./toy_finetune_data.jsonl \
--learning_rate 6e-5 \
--fp16 \
--num_train_epochs 5 \
--per_device_train_batch_size {batch size; set 1 for toy data} \
--gradient_accumulation_steps 4 \
--dataloader_drop_last True \
--train_group_size 16 \
--max_len 512 \
--weight_decay 0.01 \
--logging_steps 10
merge
from LM_Cocktail import mix_models, mix_models_with_data
# Mix fine-tuned model and base model; then save it to output_path: ./mixed_model_1
model = mix_models(
model_names_or_paths=["BAAI/bge-reranker-base", "your_fine-tuned_model"],
model_type='encoder',
weights=[0.5, 0.5], # you can change the weights to get a better trade-off.
output_path='./mixed_model_1')
推理
from FlagEmbedding import FlagReranker
reranker = FlagReranker('BAAI/bge-reranker-large', use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation
score = reranker.compute_score(['query', 'passage'])
print(score)
scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']])
print(scores)
or
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, BatchEncoding, PreTrainedTokenizerFast
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-base')
model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-base')
model.eval()
pairs = [['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']]
with torch.no_grad():
inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
print(scores)
FlagEmbedding 代码
import logging
import torch
from torch import nn
from transformers import AutoModelForSequenceClassification, PreTrainedModel, TrainingArguments
from transformers.modeling_outputs import SequenceClassifierOutput
from .arguments import ModelArguments, DataArguments
logger = logging.getLogger(__name__)
class CrossEncoder(nn.Module):
def __init__(self, hf_model: PreTrainedModel, model_args: ModelArguments, data_args: DataArguments,
train_args: TrainingArguments):
super().__init__()
self.hf_model = hf_model
self.model_args = model_args
self.train_args = train_args
self.data_args = data_args
self.config = self.hf_model.config
self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')
self.register_buffer(
'target_label',
torch.zeros(self.train_args.per_device_train_batch_size, dtype=torch.long)
)
def gradient_checkpointing_enable(self, **kwargs):
self.hf_model.gradient_checkpointing_enable(**kwargs)
def forward(self, batch):
ranker_out: SequenceClassifierOutput = self.hf_model(**batch, return_dict=True)
logits = ranker_out.logits
if self.training:
scores = logits.view(
self.train_args.per_device_train_batch_size,
self.data_args.train_group_size
)
loss = self.cross_entropy(scores, self.target_label)
return SequenceClassifierOutput(
loss=loss,
**ranker_out,
)
else:
return ranker_out
@classmethod
def from_pretrained(
cls, model_args: ModelArguments, data_args: DataArguments, train_args: TrainingArguments,
*args, **kwargs
):
hf_model = AutoModelForSequenceClassification.from_pretrained(*args, **kwargs)
reranker = cls(hf_model, model_args, data_args, train_args)
return reranker
def save_pretrained(self, output_dir: str):
state_dict = self.hf_model.state_dict()
state_dict = type(state_dict)(
{k: v.clone().cpu()
for k,
v in state_dict.items()})
self.hf_model.save_pretrained(output_dir, state_dict=state_dict)
FlagModel
from typing import cast, List, Union, Tuple
import numpy as np
import torch
from tqdm import tqdm
from transformers import AutoModel, AutoTokenizer, AutoModelForSequenceClassification
class FlagModel:
def __init__(
self,
model_name_or_path: str = None,
pooling_method: str = 'cls',
normalize_embeddings: bool = True,
query_instruction_for_retrieval: str = None,
use_fp16: bool = True
) -> None:
self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
self.model = AutoModel.from_pretrained(model_name_or_path)
self.query_instruction_for_retrieval = query_instruction_for_retrieval
self.normalize_embeddings = normalize_embeddings
self.pooling_method = pooling_method
if torch.cuda.is_available():
self.device = torch.device("cuda")
elif torch.backends.mps.is_available():
self.device = torch.device("mps")
else:
self.device = torch.device("cpu")
use_fp16 = False
if use_fp16: self.model.half()
self.model = self.model.to(self.device)
self.num_gpus = torch.cuda.device_count()
if self.num_gpus > 1:
print(f"----------using {self.num_gpus}*GPUs----------")
self.model = torch.nn.DataParallel(self.model)
def encode_queries(self, queries: Union[List[str], str],
batch_size: int = 256,
max_length: int = 512,
convert_to_numpy: bool = True) -> np.ndarray:
'''
This function will be used for retrieval task
if there is a instruction for queries, we will add it to the query text
'''
if self.query_instruction_for_retrieval is not None:
if isinstance(queries, str):
input_texts = self.query_instruction_for_retrieval + queries
else:
input_texts = ['{}{}'.format(self.query_instruction_for_retrieval, q) for q in queries]
else:
input_texts = queries
return self.encode(input_texts, batch_size=batch_size, max_length=max_length, convert_to_numpy=convert_to_numpy)
def encode_corpus(self,
corpus: Union[List[str], str],
batch_size: int = 256,
max_length: int = 512,
convert_to_numpy: bool = True) -> np.ndarray:
'''
This function will be used for retrieval task
encode corpus for retrieval task
'''
return self.encode(corpus, batch_size=batch_size, max_length=max_length, convert_to_numpy=convert_to_numpy)
@torch.no_grad()
def encode(self,
sentences: Union[List[str], str],
batch_size: int = 256,
max_length: int = 512,
convert_to_numpy: bool = True) -> np.ndarray:
if self.num_gpus > 0:
batch_size = batch_size * self.num_gpus
self.model.eval()
input_was_string = False
if isinstance(sentences, str):
sentences = [sentences]
input_was_string = True
all_embeddings = []
for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
disable=len(sentences) < 256):
sentences_batch = sentences[start_index:start_index + batch_size]
inputs = self.tokenizer(
sentences_batch,
padding=True,
truncation=True,
return_tensors='pt',
max_length=max_length,
).to(self.device)
# 输出取的是最后一层的隐变量
last_hidden_state = self.model(**inputs, return_dict=True).last_hidden_state
# embedding 是池化后的
embeddings = self.pooling(last_hidden_state, inputs['attention_mask'])
if self.normalize_embeddings:
embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
embeddings = cast(torch.Tensor, embeddings)
if convert_to_numpy:
embeddings = embeddings.cpu().numpy()
all_embeddings.append(embeddings)
if convert_to_numpy:
all_embeddings = np.concatenate(all_embeddings, axis=0)
else:
all_embeddings = torch.stack(all_embeddings)
if input_was_string:
return all_embeddings[0]
return all_embeddings
def pooling(self,
last_hidden_state: torch.Tensor,
attention_mask: torch.Tensor = None):
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
class FlagReranker:
def __init__(
self,
model_name_or_path: str = None,
use_fp16: bool = False
) -> None:
self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
self.model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path)
if torch.cuda.is_available():
self.device = torch.device('cuda')
elif torch.backends.mps.is_available():
self.device = torch.device('mps')
else:
self.device = torch.device('cpu')
use_fp16 = False
if use_fp16:
self.model.half()
self.model = self.model.to(self.device)
self.model.eval()
self.num_gpus = torch.cuda.device_count()
if self.num_gpus > 1:
print(f"----------using {self.num_gpus}*GPUs----------")
self.model = torch.nn.DataParallel(self.model)
@torch.no_grad()
def compute_score(self, sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]], batch_size: int = 256,
max_length: int = 512) -> List[float]:
if self.num_gpus > 0:
batch_size = batch_size * self.num_gpus
assert isinstance(sentence_pairs, list)
if isinstance(sentence_pairs[0], str):
sentence_pairs = [sentence_pairs]
all_scores = []
for start_index in tqdm(range(0, len(sentence_pairs), batch_size), desc="Compute Scores",
disable=len(sentence_pairs) < 128):
sentences_batch = sentence_pairs[start_index:start_index + batch_size]
inputs = self.tokenizer(
sentences_batch,
padding=True,
truncation=True,
return_tensors='pt',
max_length=max_length,
).to(self.device)
scores = self.model(**inputs, return_dict=True).logits.view(-1, ).float()
all_scores.extend(scores.cpu().numpy().tolist())
if len(all_scores) == 1:
return all_scores[0]
return all_scores
class LLMEmbedder:
instructions = {
"qa": {
"query": "Represent this query for retrieving relevant documents: ",
"key": "Represent this document for retrieval: ",
},
"convsearch": {
"query": "Encode this query and context for searching relevant passages: ",
"key": "Encode this passage for retrieval: ",
},
"chat": {
"query": "Embed this dialogue to find useful historical dialogues: ",
"key": "Embed this historical dialogue for retrieval: ",
},
"lrlm": {
"query": "Embed this text chunk for finding useful historical chunks: ",
"key": "Embed this historical text chunk for retrieval: ",
},
"icl": {
"query": "Convert this example into vector to look for useful examples: ",
"key": "Convert this example into vector for retrieval: ",
},
"tool": {
"query": "Transform this user request for fetching helpful tool descriptions: ",
"key": "Transform this tool description for retrieval: "
},
}
def __init__(
self,
model_name_or_path: str = None,
pooling_method: str = 'cls',
normalize_embeddings: bool = True,
use_fp16: bool = True
) -> None:
self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
self.model = AutoModel.from_pretrained(model_name_or_path)
self.normalize_embeddings = normalize_embeddings
self.pooling_method = pooling_method
if torch.cuda.is_available():
self.device = torch.device("cuda")
elif torch.backends.mps.is_available():
self.device = torch.device("mps")
else:
self.device = torch.device("cpu")
use_fp16 = False
if use_fp16: self.model.half()
self.model = self.model.to(self.device)
self.num_gpus = torch.cuda.device_count()
if self.num_gpus > 1:
print(f"----------using {self.num_gpus}*GPUs----------")
self.model = torch.nn.DataParallel(self.model)
def encode_queries(self, queries: Union[List[str], str],
batch_size: int = 256,
max_length: int = 256,
task: str = 'qa') -> np.ndarray:
'''
Encode queries into dense vectors.
Automatically add instructions according to given task.
'''
instruction = self.instructions[task]["query"]
if isinstance(queries, str):
input_texts = instruction + queries
else:
input_texts = [instruction + q for q in queries]
return self._encode(input_texts, batch_size=batch_size, max_length=max_length)
def encode_keys(self, keys: Union[List[str], str],
batch_size: int = 256,
max_length: int = 512,
task: str = 'qa') -> np.ndarray:
'''
Encode keys into dense vectors.
Automatically add instructions according to given task.
'''
instruction = self.instructions[task]["key"]
if isinstance(keys, str):
input_texts = instruction + keys
else:
input_texts = [instruction + k for k in keys]
return self._encode(input_texts, batch_size=batch_size, max_length=max_length)
@torch.no_grad()
def _encode(self, sentences: Union[List[str], str], batch_size: int = 256, max_length: int = 512) -> np.ndarray:
if self.num_gpus > 0:
batch_size = batch_size * self.num_gpus
self.model.eval()
input_was_string = False
if isinstance(sentences, str):
sentences = [sentences]
input_was_string = True
all_embeddings = []
for start_index in tqdm(range(0, len(sentences), batch_size), desc="Inference Embeddings",
disable=len(sentences) < 256):
sentences_batch = sentences[start_index:start_index + batch_size]
inputs = self.tokenizer(
sentences_batch,
padding=True,
truncation=True,
return_tensors='pt',
max_length=max_length,
).to(self.device)
last_hidden_state = self.model(**inputs, return_dict=True).last_hidden_state
embeddings = self.pooling(last_hidden_state, inputs['attention_mask'])
if self.normalize_embeddings:
embeddings = torch.nn.functional.normalize(embeddings, dim=-1)
embeddings = cast(torch.Tensor, embeddings)
all_embeddings.append(embeddings.cpu().numpy())
all_embeddings = np.concatenate(all_embeddings, axis=0)
if input_was_string:
return all_embeddings[0]
return all_embeddings
def pooling(self,
last_hidden_state: torch.Tensor,
attention_mask: torch.Tensor = None):
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!")
1881
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
