目录

任务

代码

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 = 0for epoch in range(epoch):trained_data = 0for 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 += 1total_batch += 1# 每训练 auto_save_batch 个 batch，保存一次模型if batch_after_last_save >= auto_save_batch:batch_after_last_save = 0model.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 += 1model.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, BertForNextSentencePredictiontokenizer = 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 = 0for 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 += 1print(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