bert ranking pairwise demo

下面是用bert 训练pairwise rank 的 demo

python 复制代码
import torch
from torch.utils.data import DataLoader, Dataset
from transformers import BertModel, BertTokenizer
from sklearn.metrics import pairwise_distances_argmin_min

class PairwiseRankingDataset(Dataset):
    def __init__(self, sentence_pairs, tokenizer, max_length):
        self.input_ids = []
        self.attention_masks = []
        
        for pair in sentence_pairs:
            encoded_pair = tokenizer(pair, padding='max_length', truncation=True, max_length=max_length, return_tensors='pt')
            self.input_ids.append(encoded_pair['input_ids'])
            self.attention_masks.append(encoded_pair['attention_mask'])
        
        self.input_ids = torch.cat(self.input_ids, dim=0)
        self.attention_masks = torch.cat(self.attention_masks, dim=0)
        
    def __len__(self):
        return len(self.input_ids)
    
    def __getitem__(self, idx):
        input_id = self.input_ids[idx]
        attention_mask = self.attention_masks[idx]
        return input_id, attention_mask

class BERTPairwiseRankingModel(torch.nn.Module):
    def __init__(self, bert_model_name):
        super(BERTPairwiseRankingModel, self).__init__()
        self.bert = BertModel.from_pretrained(bert_model_name)
        self.dropout = torch.nn.Dropout(0.1)
        self.fc = torch.nn.Linear(self.bert.config.hidden_size, 1)
        
    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = self.dropout(outputs[1])
        logits = self.fc(pooled_output)
        return logits.squeeze()

# 初始化BERT模型和分词器
bert_model_name = 'bert-base-uncased'
tokenizer = BertTokenizer.from_pretrained(bert_model_name)

# 示例输入数据
sentence_pairs = [
    ('I like cats', 'I like dogs'),
    ('The sun is shining', 'It is raining'),
    ('Apple is a fruit', 'Car is a vehicle')
]

# 超参数
batch_size = 8
max_length = 128
learning_rate = 1e-5
num_epochs = 5

# 创建数据集和数据加载器
dataset = PairwiseRankingDataset(sentence_pairs, tokenizer, max_length)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# 初始化模型并加载预训练权重
model = BERTPairwiseRankingModel(bert_model_name)
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

# 训练模型
model.train()

for epoch in range(num_epochs):
    total_loss = 0
    
    for input_ids, attention_masks in dataloader:
        optimizer.zero_grad()
        
        logits = model(input_ids, attention_masks)
        
        # 计算损失函数(使用对比损失函数)
        pos_scores = logits[::2]  # 正样本分数
        neg_scores = logits[1::2]  # 负样本分数
        loss = torch.relu(1 - pos_scores + neg_scores).mean()
        
        total_loss += loss.item()
        
        loss.backward()
        optimizer.step()
    
    print(f"Epoch {epoch+1}/{num_epochs} - Loss: {total_loss:.4f}")

# 推断模型
model.eval()

with torch.no_grad():
    embeddings = model.bert.embeddings.word_embeddings(dataset.input_ids)
    pairwise_distances = pairwise_distances_argmin_min(embeddings.numpy())

# 输出结果
for i, pair in enumerate(sentence_pairs):
    pos_idx = pairwise_distances[0][2 * i]
    neg_idx = pairwise_distances[0][2 * i + 1]
    pos_dist = pairwise_distances[1][2 * i]
    neg_dist = pairwise_distances[1][2 * i + 1]
    
    print(f"Pair: {pair}")
    print(f"Positive example index: {pos_idx}, Distance: {pos_dist:.4f}")
    print(f"Negative example index: {neg_idx}, Distance: {neg_dist:.4f}")
    print()
相关推荐
金玉满堂@bj13 分钟前
PyCharm 中 Python 解释器的添加选项及作用
ide·python·pycharm
程序员三藏18 分钟前
如何使用Pytest进行测试?
自动化测试·软件测试·python·测试工具·职场和发展·测试用例·pytest
TY-202522 分钟前
三、神经网络——网络优化方法
人工智能·深度学习·神经网络
随心点儿1 小时前
使用python 将多个docx文件合并为一个word
开发语言·python·多个word合并为一个
不学无术の码农1 小时前
《Effective Python》第十三章 测试与调试——使用 Mock 测试具有复杂依赖的代码
开发语言·python
sleepybear11131 小时前
在Ubuntu上从零开始编译并运行Home Assistant源码并集成HACS与小米开源的Ha Xiaomi Home
python·智能家居·小米·home assistant·米家·ha xiaomi home
纪伊路上盛名在1 小时前
(鱼书)深度学习入门1:python入门
人工智能·python·深度学习
夏末蝉未鸣011 小时前
python transformers笔记(TrainingArguments类)
python·自然语言处理·transformer
德育处主任Pro2 小时前
「py数据分析」04如何将 Python 爬取的数据保存为 CSV 文件
数据库·python·数据分析
咸鱼鲸2 小时前
【PyTorch】PyTorch中数据准备工作(AI生成)
人工智能·pytorch·python