目录

任务

代码

train.py

predit.py

faiss 最相似的 topk

数

---

任务

使用 bert-base-chinese 训练相似度任务，参考：微调BERT模型实现相似性判断 - 知乎

参考他上面代码，他使用的是 BertForNextSentencePrediction 模型，BertForNextSentencePrediction 原本是设计用于下一个句子预测任务的。在BERT的原始训练中，模型会接收到一对句子，并试图预测第二个句子是否紧跟在第一个句子之后；所以使用这个模型标签(label)只能是 0,1，相当于二分类任务了

但其实在相似度任务中，我们每一条数据都是【text1\ttext2\tlabel】的形式，其中 label 代表相似度，可以给两个文本打分表示相似度，也可以映射为分类任务，0 代表不相似，1 代表相似，他这篇文章利用了这种思想，对新手还挺有用的。

现在我搞了一个招聘数据，里面有办公区域列，处理过了，每一行代表【地址1\t地址2\t相似度】

只要两文本中有一个地址相似我就作为相似，标签为 1，否则 0

利用这数据微调，没有使用验证数据集，就最后使用测试集来看看效果。

代码

train.py

import json
import torch
from transformers import BertTokenizer, BertForNextSentencePrediction
from torch.utils.data import DataLoader, Dataset


# 能用gpu就用gpu
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")

bacth_size = 32
epoch = 3
auto_save_batch = 5000
learning_rate = 2e-5


# 准备数据集
class MyDataset(Dataset):
    def __init__(self, data_file_paths):
        self.texts = []
        self.labels = []
        # 分词器用默认的
        self.tokenizer = BertTokenizer.from_pretrained('../bert-base-chinese')
        # 自己实现对数据集的解析
        with open(data_file_paths, 'r', encoding='utf-8') as f:
            for line in f:
                text1, text2, label = line.split('\t')
                self.texts.append((text1, text2))
                self.labels.append(int(label))

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

    def __getitem__(self, idx):
        text1, text2 = self.texts[idx]
        label = self.labels[idx]
        encoded_text = self.tokenizer(text1, text2, padding='max_length', truncation=True, max_length=128, return_tensors='pt')
        return encoded_text, label


# 训练数据文件路径
train_dataset = MyDataset('../data/train.txt')

# 定义模型
# num_labels=5 定义相似度评分有几个
model = BertForNextSentencePrediction.from_pretrained('../bert-base-chinese', num_labels=6)
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
train_loader = DataLoader(train_dataset, batch_size=bacth_size, shuffle=True)
trained_data = 0
batch_after_last_save = 0
total_batch = 0
total_epoch = 0

for epoch in range(epoch):
    trained_data = 0
    for batch in train_loader:
        inputs, labels = batch
        # 不知道为啥，出来的数据维度是 (batch_size, 1, 128)，需要把第二维去掉
        inputs['input_ids'] = inputs['input_ids'].squeeze(1)
        inputs['token_type_ids'] = inputs['token_type_ids'].squeeze(1)
        inputs['attention_mask'] = inputs['attention_mask'].squeeze(1)
        # 因为要用GPU，将数据传输到gpu上
        inputs = inputs.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(**inputs, labels=labels)
        loss, logits = outputs[:2]
        loss.backward()
        optimizer.step()
        trained_data += len(labels)
        trained_process = float(trained_data) / len(train_dataset)
        batch_after_last_save += 1
        total_batch += 1
        # 每训练 auto_save_batch 个 batch，保存一次模型
        if batch_after_last_save >= auto_save_batch:
            batch_after_last_save = 0
            model.save_pretrained(f'../output/cn_equal_model_{total_epoch}_{total_batch}.pth')
            print("保存模型：cn_equal_model_{}_{}.pth".format(total_epoch, total_batch))
        print("训练进度：{:.2f}%, loss={:.4f}".format(trained_process * 100, loss.item()))
    total_epoch += 1
    model.save_pretrained(f'../output/cn_equal_model_{total_epoch}_{total_batch}.pth')
    print("保存模型：cn_equal_model_{}_{}.pth".format(total_epoch, total_batch))

训练好后的文件，输出的最后一个文件夹才是效果最好的模型：

predit.py

import torch
from transformers import BertTokenizer, BertForNextSentencePrediction


tokenizer = BertTokenizer.from_pretrained('../bert-base-chinese')
model = BertForNextSentencePrediction.from_pretrained('../output/cn_equal_model_3_171.pth')

with torch.no_grad():
    with open('../data/test.txt', 'r', encoding='utf8') as f:
        lines = f.readlines()
        correct = 0
        for i, line in enumerate(lines):
            text1, text2, label = line.split('\t')
            encoded_text = tokenizer(text1, text2, padding='max_length', truncation=True, max_length=128, return_tensors='pt')
            outputs = model(**encoded_text)
            res = torch.argmax(outputs.logits, dim=1).item()
            print(text1, text2, label, res)
            if str(res) == label.strip('\n'):
                correct += 1
            print(f'{i + 1}/{len(lines)}')
        print(f'acc:{correct / len(lines)}')

可以看到还是较好的学习了我数据特征：只要两文本中有一个地址相似我就作为相似，标签为 1，否则 0

faiss 最相似的 topk

使用 faiss 寻找 topk 相似的，从结果上看最相似的基本都还是找到排到较为靠前的位置

import torch
import faiss
import pandas as pd
import numpy as np
from transformers import BertTokenizer, BertModel


# 假设有一个数据集df，其中包含'index'列和'text'列
df = pd.read_csv('../data/DataAnalyst.csv', encoding='gbk')  # 根据实际情况加载数据集
df = df.dropna().drop_duplicates().reset_index()
df['index'] = df.index
df = df[['index', '公司所在商区']]  # 保留所需列
df['公司所在商区'] = df['公司所在商区'].map(lambda row: ','.join(eval(row)))

# device = torch.device('gpu' if torch.cuda.is_available() else 'cpu')
device = torch.device('cpu')

# 加载微调好的模型和tokenizer
tokenizer = BertTokenizer.from_pretrained('../bert-base-chinese')
model = BertModel.from_pretrained('../output/cn_equal_model_3_171.pth')
model.eval()


# 将数据集转化为模型所需的格式并计算所有样本的向量表示
def encode_texts(df):
    text_vectors = []
    for index, row in df.iterrows():
        text = row['公司所在商区']
        inputs = tokenizer(text, padding='max_length', truncation=True, max_length=128, return_tensors='pt')
        with torch.no_grad():
            embeddings = model(**inputs.to(device))['last_hidden_state'][:, 0]
            text_vectors.append(embeddings.cpu().numpy())
        print(f'{index + 1}/{len(df)}')
    return np.vstack(text_vectors)

# 加载数据集并计算所有样本的向量
print('enbedding all data...')
all_embeddings = encode_texts(df)

# 初始化Faiss索引
print('init faiss all embedding...')
index = faiss.IndexFlatIP(all_embeddings.shape[1])  # 使用内积空间，适用于余弦相似度
index.add(all_embeddings)
print('init faiss all embedding finish~~~')


# 定义查找最相似样本的函数
def find_top_k_similar(query_text, k=100):
    print('当前 query_text embedding.')
    query_embedding = encode_single_text(query_text)
    print('begin to search topk....')
    D, I = index.search(query_embedding, k)  # 返回距离和索引
    top_k_indices = df.iloc[I[0]].index.tolist()  # 将索引转换为原始数据集的索引
    return top_k_indices


# 编码单个文本的函数
def encode_single_text(text):
    inputs = tokenizer(text, padding='max_length', truncation=True, max_length=128, return_tensors='pt')
    with torch.no_grad():
        embedding = model(**inputs.to(device))['last_hidden_state'][:, 0].cpu().numpy()
    print('当前 query_text embedding finish!')
    return embedding


# 示例：找一个query_text的top10相似样本
query_text = "左家庄,国展,西坝河"
top10_indices = find_top_k_similar(query_text)
# 获取与查询文本最相似的前10条原始文本
top10_texts = [df.loc[index, '公司所在商区'] for index in top10_indices]

print(f"与'{query_text}'最相似的前100条样本及其文本：")
for i, (idx, text) in enumerate(zip(top10_indices, top10_texts)):
    print(f"{i+1}. 索引：{idx}，文本：{text}")

数据

链接：https://pan.baidu.com/s/1Cpr-ZD9Neakt73naGdsVTw 
提取码：eryw 
链接：https://pan.baidu.com/s/1qHYjXC7UCeUsXVnYTQIPCg 
提取码：o8py 
链接：https://pan.baidu.com/s/1CTntG1Z6AIhiPt6i8Ad97Q 
提取码：x6sz