Embedding和Rerank模型类模板

写bug如流水

已于 2024-04-25 14:21:53 修改

阅读量1.3k

点赞数 8

CC 4.0 BY-SA版权

分类专栏：大模型文章标签： embedding

于 2024-04-25 13:53:58 首次发布

本文链接：https://blog.csdn.net/h1773655323/article/details/138187325

大模型专栏收录该内容

6 篇文章

订阅专栏

本文介绍了如何使用PyTorch和HuggingFaceTransformers库实现两种模型：基于Embedding的模型，用于生成文本句子的嵌入表示，并通过L2归一化改进任务性能；以及Rerank模型，用于处理句子对并计算相关性得分。代码示例展示了如何在GPU或CPU上初始化和使用这些模型。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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 转换为概率。