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()
相关推荐
逢生博客19 分钟前
使用 Python 项目管理工具 uv 快速创建 MCP 服务(Cherry Studio、Trae 添加 MCP 服务)
python·sqlite·uv·deepseek·trae·cherry studio·mcp服务
xwz小王子22 分钟前
Nature Communications 面向形状可编程磁性软材料的数据驱动设计方法—基于随机设计探索与神经网络的协同优化框架
深度学习
堕落似梦26 分钟前
Pydantic增强SQLALchemy序列化(FastAPI直接输出SQLALchemy查询集)
python
生信碱移1 小时前
大语言模型时代,单细胞注释也需要集思广益(mLLMCelltype)
人工智能·经验分享·深度学习·语言模型·自然语言处理·数据挖掘·数据可视化
坐吃山猪1 小时前
Python-Agent调用多个Server-FastAPI版本
开发语言·python·fastapi
Bruce-li__2 小时前
使用Django REST Framework快速开发API接口
python·django·sqlite
小兜全糖(xdqt)2 小时前
python 脚本引用django中的数据库model
python·django
Arenaschi2 小时前
SQLite 是什么?
开发语言·网络·python·网络协议·tcp/ip
纪元A梦2 小时前
华为OD机试真题——推荐多样性(2025A卷:200分)Java/python/JavaScript/C++/C语言/GO六种最佳实现
java·javascript·c++·python·华为od·go·华为od机试题
硅谷秋水2 小时前
通过模仿学习实现机器人灵巧操作:综述(上)
人工智能·深度学习·机器学习·计算机视觉·语言模型·机器人