bert 相似度任务训练完整版

任务

之前写了一个相似度任务的版本：bert 相似度任务训练简单版本,faiss 寻找相似 topk-CSDN博客

相似度用的是 0，1，相当于分类任务，现在我们相似度有评分，不再是 0,1 了，分数为 0-5，数字越大代表两个句子越相似，这一次的比较完整，评估，验证集，相似度模型都有了。

数据集

链接：https://pan.baidu.com/s/1B1-PKAKNoT_JwMYJx_zT1g
提取码：er1z
原始数据好几千条，我训练数据用了部分 2500 条，验证，测试 300 左右，使用 cpu 也用了好几个小时

train.py

import torch
import os
import time
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer, BertModel, AdamW, get_cosine_schedule_with_warmup
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np


# 设备选择
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = 'cpu'


# 定义文本相似度数据集类
class TextSimilarityDataset(Dataset):
    def __init__(self, file_path, tokenizer, max_len=128):
        self.data = []
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f.readlines():
                text1, text2, similarity_score = line.strip().split('\t')
                inputs1 = tokenizer(text1, padding='max_length', truncation=True, max_length=max_len)
                inputs2 = tokenizer(text2, padding='max_length', truncation=True, max_length=max_len)
                self.data.append({
                    'input_ids1': inputs1['input_ids'],
                    'attention_mask1': inputs1['attention_mask'],
                    'input_ids2': inputs2['input_ids'],
                    'attention_mask2': inputs2['attention_mask'],
                    'similarity_score': float(similarity_score),
                })

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


def cosine_similarity_torch(vec1, vec2, eps=1e-8):
    dot_product = torch.mm(vec1, vec2.t())
    norm1 = torch.norm(vec1, 2, dim=1, keepdim=True)
    norm2 = torch.norm(vec2, 2, dim=1, keepdim=True)
    similarity_scores = dot_product / (norm1 * norm2.t()).clamp(min=eps)
    return similarity_scores


# 定义模型，这里我们不仅计算两段文本的[CLS] token的点积，而是整个句向量的余弦相似度
class BertSimilarityModel(torch.nn.Module):
    def __init__(self, pretrained_model):
        super(BertSimilarityModel, self).__init__()
        self.bert = BertModel.from_pretrained(pretrained_model)
        self.dropout = torch.nn.Dropout(p=0.1)  # 引入Dropout层以防止过拟合

    def forward(self, input_ids1, attention_mask1, input_ids2, attention_mask2):
        embeddings1 = self.dropout(self.bert(input_ids=input_ids1, attention_mask=attention_mask1)['last_hidden_state'])
        embeddings2 = self.dropout(self.bert(input_ids=input_ids2, attention_mask=attention_mask2)['last_hidden_state'])

        # 计算两个文本向量的余弦相似度
        embeddings1 = torch.mean(embeddings1, dim=1)
        embeddings2 = torch.mean(embeddings2, dim=1)

        similarity_scores = cosine_similarity_torch(embeddings1, embeddings2)

        # 映射到[0, 5]评分范围
        normalized_similarities = (similarity_scores + 1) * 2.5
        return normalized_similarities.unsqueeze(1)


# 自定义损失函数，使用Smooth L1 Loss，更适合处理回归问题
class SmoothL1Loss(torch.nn.Module):
    def __init__(self):
        super(SmoothL1Loss, self).__init__()

    def forward(self, predictions, targets):
        diff = predictions - targets
        abs_diff = torch.abs(diff)
        quadratic = torch.where(abs_diff < 1, 0.5 * diff ** 2, abs_diff - 0.5)
        return torch.mean(quadratic)


def train_model(model, train_loader, val_loader, epochs=3, model_save_path='../output/bert_similarity_model.pth'):
    model.to(device)
    criterion = SmoothL1Loss()  # 使用自定义的Smooth L1 Loss
    optimizer = AdamW(model.parameters(), lr=5e-5)  # 调整初始学习率为5e-5
    num_training_steps = len(train_loader) * epochs
    scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=0.1*num_training_steps, num_training_steps=num_training_steps)  # 使用带有warmup的余弦退火学习率调度

    best_val_loss = float('inf')
    for epoch in range(epochs):
        model.train()
        for batch in train_loader:
            input_ids1 = batch['input_ids1'].to(device)
            attention_mask1 = batch['attention_mask1'].to(device)
            input_ids2 = batch['input_ids2'].to(device)
            attention_mask2 = batch['attention_mask2'].to(device)
            similarity_scores = batch['similarity_score'].to(device)

            optimizer.zero_grad()
            outputs = model(input_ids1, attention_mask1, input_ids2, attention_mask2)
            loss = criterion(outputs, similarity_scores.unsqueeze(1))
            loss.backward()
            optimizer.step()
            scheduler.step()

        # 验证阶段
        model.eval()
        with torch.no_grad():
            val_loss = 0
            total_val_samples = 0

            for batch in val_loader:
                input_ids1 = batch['input_ids1'].to(device)
                attention_mask1 = batch['attention_mask1'].to(device)
                input_ids2 = batch['input_ids2'].to(device)
                attention_mask2 = batch['attention_mask2'].to(device)
                similarity_scores = batch['similarity_score'].to(device)

                val_outputs = model(input_ids1, attention_mask1, input_ids2, attention_mask2)
                val_loss += criterion(val_outputs, similarity_scores.unsqueeze(1)).item()
                total_val_samples += len(similarity_scores)

            val_loss /= len(val_loader)
            print(f'Epoch {epoch + 1}, Validation Loss: {val_loss:.4f}')

            if val_loss < best_val_loss:
                best_val_loss = val_loss
                torch.save(model.state_dict(), model_save_path)


def collate_to_tensors(batch):
    '''把数据处理为模型可用的数据，不同任务可能需要修改一下，'''
    input_ids1 = torch.tensor([example['input_ids1'] for example in batch])
    attention_mask1 = torch.tensor([example['attention_mask1'] for example in batch])
    input_ids2 = torch.tensor([example['input_ids2'] for example in batch])
    attention_mask2 = torch.tensor([example['attention_mask2'] for example in batch])
    similarity_score = torch.tensor([example['similarity_score'] for example in batch])

    return {'input_ids1': input_ids1, 'attention_mask1': attention_mask1, 'input_ids2': input_ids2,
            'attention_mask2': attention_mask2, 'similarity_score': similarity_score}


# 加载数据集和预训练模型
tokenizer = BertTokenizer.from_pretrained('../bert-base-chinese')
model = BertSimilarityModel('../bert-base-chinese')

# 加载数据并创建
train_data = TextSimilarityDataset('../data/STS-B/STS-B.train - 副本.data', tokenizer)
val_data = TextSimilarityDataset('../data/STS-B/STS-B.valid - 副本.data', tokenizer)
test_data = TextSimilarityDataset('../data/STS-B/STS-B.test - 副本.data', tokenizer)

train_loader = DataLoader(train_data, batch_size=32, shuffle=True, collate_fn=collate_to_tensors)
val_loader = DataLoader(val_data, batch_size=32, collate_fn=collate_to_tensors)
test_loader = DataLoader(test_data, batch_size=32, collate_fn=collate_to_tensors)

optimizer = AdamW(model.parameters(), lr=2e-5)

# 开始训练
train_model(model, train_loader, val_loader)

# 加载最佳模型进行测试
model.load_state_dict(torch.load('../output/bert_similarity_model.pth'))
test_loss = 0
total_test_samples = 0

with torch.no_grad():
    for batch in test_loader:
        input_ids1 = batch['input_ids1'].to(device)
        attention_mask1 = batch['attention_mask1'].to(device)
        input_ids2 = batch['input_ids2'].to(device)
        attention_mask2 = batch['attention_mask2'].to(device)
        similarity_scores = batch['similarity_score'].to(device)

        test_outputs = model(input_ids1, attention_mask1, input_ids2, attention_mask2)
        test_loss += torch.nn.functional.mse_loss(test_outputs, similarity_scores.unsqueeze(1)).item()
        total_test_samples += len(similarity_scores)

test_loss /= len(test_loader)
print(f'Test Loss: {test_loss:.4f}')

predit.py

这个脚本是用来看看效果的，直接传入两个文本，使用训练好的模型来计算相似度的

import torch
from transformers import BertTokenizer, BertModel


def cosine_similarity_torch(vec1, vec2, eps=1e-8):
    dot_product = torch.mm(vec1, vec2.t())
    norm1 = torch.norm(vec1, 2, dim=1, keepdim=True)
    norm2 = torch.norm(vec2, 2, dim=1, keepdim=True)
    similarity_scores = dot_product / (norm1 * norm2.t()).clamp(min=eps)
    return similarity_scores


# 定义模型，这里我们不仅计算两段文本的[CLS] token的点积，而是整个句向量的余弦相似度
class BertSimilarityModel(torch.nn.Module):
    def __init__(self, pretrained_model):
        super(BertSimilarityModel, self).__init__()
        self.bert = BertModel.from_pretrained(pretrained_model)
        self.dropout = torch.nn.Dropout(p=0.1)  # 引入Dropout层以防止过拟合

    def forward(self, input_ids1, attention_mask1, input_ids2, attention_mask2):
        '''如果是用来预测，forward 会被禁用'''
        pass


# 加载预训练模型和分词器
tokenizer = BertTokenizer.from_pretrained('../bert-base-chinese')
model = BertSimilarityModel('../bert-base-chinese')
model.load_state_dict(torch.load('../output/bert_similarity_model.pth'))  # 请确保路径正确
model.eval()  # 设置模型为评估模式


def calculate_similarity(text1, text2):
    # 对输入文本进行编码
    inputs1 = tokenizer(text1, padding='max_length', truncation=True, max_length=128, return_tensors='pt')
    inputs2 = tokenizer(text2, padding='max_length', truncation=True, max_length=128, return_tensors='pt')

    # 计算相似度
    with torch.no_grad():
        embeddings1 = model.bert(**inputs1.to('cpu'))['last_hidden_state'][:, 0]
        embeddings2 = model.bert(**inputs2.to('cpu'))['last_hidden_state'][:, 0]
        similarity_score = cosine_similarity_torch(embeddings1, embeddings2).item()

    # 映射到[0, 5]评分范围（假设训练时有此步骤）
    normalized_similarity = (similarity_score + 1) * 2.5

    return normalized_similarity


# 示例
text1 = "瑞典驻利比亚班加西领事馆发生汽车炸弹袭击，无人员伤亡"
text2 = "汽车炸弹击中瑞典驻班加西领事馆，无人受伤。"
similarity = calculate_similarity(text1, text2)
print(f"两个句子的相似度为：{similarity:.2f}")