Embedding和Rerank模型类模板

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)

说明

1. 初始化：在构造函数中初始化分词器和模型，并将模型转移到指定的计算设备（如 GPU 或 CPU）。
2. 获取嵌入：定义了一个方法 get_embeddings，它接受一系列句子，将它们转换为嵌入表示。
3. 标准化：最后，嵌入通过其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)

说明

1. 初始化：在构造函数中，加载了分词器和模型，并将模型转移到指定的设备上（GPU 或 CPU）。
2. 得分计算：定义了 score_pairs 方法来接受句子对，将其转换为模型可接受的格式，并计算得分。这里使用 sigmoid 函数是因为通常在二元分类任务中用来将 logits 转换为概率。