1. 模型介绍
介绍见链接:Embedding和Rerank模型介绍
2. Embedding模型模板
from transformers import AutoModel, AutoTokenizer
import torch
class EmbeddingModel:
def __init__(self, model_name, device='cuda'):
"""
Initializes the embedding model with the specified model name and device.
"""
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModel.from_pretrained(model_name)
self.device = device
self.model.to(self.device)
def get_embeddings(self, sentences):
"""
Processes a list of sentences to produce their embeddings.
"""
inputs = self.tokenizer(sentences, padding=True, truncation=True, max_length=512, return_tensors="pt")
inputs_on_device = {k: v.to(self.device) for k, v in inputs.items()}
outputs = self.model(**inputs_on_device, return_dict=True)
embeddings = outputs.last_hidden_state[:, 0] # cls token
normalized_embeddings = embeddings / embeddings.norm(dim=1, keepdim=True)
return normalized_embeddings
# Usage example:
model_name = 'maidalun1020/bce-embedding-base_v1'
sentences = ['sentence_0', 'sentence_1']
embedding_model = EmbeddingModel(model_name, device='cuda' if torch.cuda.is_available() else 'cpu')
embeddings = embedding_model.get_embeddings(sentences)
print(embeddings)
说明
- 初始化:在构造函数中初始化分词器和模型,并将模型转移到指定的计算设备(如 GPU 或 CPU)。
- 获取嵌入:定义了一个方法
get_embeddings
,它接受一系列句子,将它们转换为嵌入表示。 - 标准化:最后,嵌入通过其L2范数进行标准化,以改善某些下游任务的性能。
3. Rerank模型模板
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
class RerankerModel:
def __init__(self, model_name, device='cuda'):
"""
Initializes the reranker model with the specified model name and device.
"""
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
self.device = device if torch.cuda.is_available() else 'cpu'
self.model.to(self.device)
def score_pairs(self, sentence_pairs):
"""
Processes a list of sentence pairs to produce their scores.
"""
inputs = self.tokenizer(sentence_pairs, padding=True, truncation=True, max_length=512, return_tensors="pt")
inputs_on_device = {k: v.to(self.device) for k, v in inputs.items()}
logits = self.model(**inputs_on_device, return_dict=True).logits.view(-1,).float()
scores = torch.sigmoid(logits)
return scores
# Usage example:
model_name = 'maidalun1020/bce-reranker-base_v1'
sentence_pairs = [
("The weather is nice today.", "It's a beautiful day."),
("He won the race.", "He came first in the competition.")
]
reranker_model = RerankerModel(model_name, device='cuda')
scores = reranker_model.score_pairs(sentence_pairs)
print(scores)
说明
- 初始化:在构造函数中,加载了分词器和模型,并将模型转移到指定的设备上(GPU 或 CPU)。
- 得分计算:定义了
score_pairs
方法来接受句子对,将其转换为模型可接受的格式,并计算得分。这里使用sigmoid
函数是因为通常在二元分类任务中用来将 logits 转换为概率。