简介:基于Python实现的文本匹配算法源码,涵盖PointWise单塔、DSSM与Sentence BERT双塔三类主流模型,面向自然语言处理入门者及需要完成课程设计、毕业设计的学生开发者,可快速理解文本相似度计算的建模与训练流程,既适合原理学习,也适合作为工程基线。资源共34个文件,含18个Python脚本、4个Shell启动脚本、2个TSV数据集文件及说明文档,压缩包约7.86MB,目录结构清晰,便于按模型模块对照学习。目前已有270人学习浏览。代码均在PyTorch+Transformers环境下测试通过,附带训练、推理及嵌入提取等脚本,可直接运行或按需改造;数据集与使用说明同步提供,并支持远程教学答疑;可对照源码和文档逐步复现模型效果,适合作为算法对比实验、项目初期演示或毕设基础。资源内代码经完整测试运行成功,答辩平均分达96分,整体完成度较高。
1. 文本匹配不是玄学:这包 Python 代码把单塔和双塔都跑通了
当你在做 FAQ 问答匹配、搜索召回或者论文查重时,最直接的需求就是计算两个文本之间的相似度。很多人第一步就是拿 BERT 的 CLS 向量算个 cosine,但真正业务里会发现效果飘忽不定——原因是你没区分场景就选了模型结构。这份基于 PyTorch 和 transformers 的源码包,一次把单塔 PointWise、双塔 DSSM、双塔 Sentence-BERT 三种主流文本匹配方案都实现了,还附带了好用的数据集和训练脚本。适合正在做毕设、课程设计,或者刚入 NLP 坑想快速跑通基线的工程师。拿到手不是先跑代码,而是先搞清楚“你在哪种场景下该选哪种塔”,否则改了参数也只是在错误的地基上盖楼。下面我按“选型 → 源码拆解 → 跑通训练 → 避坑 → 进阶”的顺序把这包东西吃透。
2. 选型先于编码:单塔点乘和双塔向量召回分别解决哪类相似度
2.1 单塔与双塔的本质区别
文本匹配表面上是“计算两个句子有多像”,但实际落地有两种完全不同的诉求:一种是精度优先,一种是速度优先。PointWise 单塔模型把两句话拼成一条序列喂给预训练模型,输出一个相似度分数,这类方法被称为交互式(Interaction-based)。DSSM 和 Sentence-BERT 则属于双塔模型,先把两句话分别编码成向量,再用余弦相似度或者点积算距离,这种被称为表示式(Representation-based)。
单塔的最大优势是能充分看到两个文本之间的交互信息。你在 query 和 passage 之间能拿到 token 级的 attention,模型能捕捉到“否定词在哪一侧”这种精细信号。对应的代价是推理时必须把两句话同时过一遍模型,没法预计算,线上响应时间基本等于一次完整 forward 的时间。双塔模型则相反,每个塔独立把文本映射成一个固定维度的向量,向量可以离线批量算好、存进索引里,线上只做一次向量检索,延迟能压到毫秒级,代价是两句话之间的细粒度交互被向量距离给压缩掉了。
我一般跟别人讲选型时直接给一个粗暴结论:如果你的业务是“句子对分类”,比如判断两个问题是不是同一个意思、判断评论是正面还是负面(对文本对做二分类),优先用单塔 PointWise;如果你的业务是“从海量候选里召回 top-K”,比如搜索引擎、推荐系统、知识库段落检索,优先用双塔 DSSM 或者 Sentence-BERT。用一张表来说更直观:
| 维度 | 单塔 PointWise | 双塔 DSSM / Sentence-BERT |
|---|---|---|
| 输入方式 | 两文本拼接成一个序列 | 两文本分别独立编码 |
| 交互能力 | 有 token 级 attention,捕捉细粒度关系 | 无显式交互,靠向量距离表达关系 |
| 推理效率 | 无法预计算,实时双输入 forward | 向量可离线预计算,线上只做相似度计算 |
| 精度上限 | 通常更高,尤其数据量足时 | 有交互损失,但对长尾泛化友好 |
| 典型场景 | 句子对分类、语义蕴含 | 召回、去重、embedding 检索 |
这包代码里点开 model.py 就能看到两种结构的实现分支。单塔走的是拼接后过一个 BertEncoder,双塔是两个独立的编码器输出后做 cosine。理解这个区别之后,你再去看训练脚本里的数据组织方式,就不会觉得奇怪了——单塔需要把一个 pair 压成一条样本,双塔需要把 query 和 doc 拆成两条流。
2.2 本项目数据集的文本格式与预处理流程
这套资源里的 data 目录放的是数据集。常见的文本匹配数据集格式有两种:一种是含 label 的句子对,比如 LCQMC、Quora 这种,字段就是 text1、text2、label;另一种是只有 query 和正负 doc 的三元组格式。从项目里的 utils.py 能看出,它默认读取的是句子对文件,每行是一个样本,列用分割符隔开。我刚拿到手里时第一件事不是去看 model.py,而是直接打开数据集文件确认分隔符和标签类型,因为这个位置一旦看错,后面所有训练脚本全部白搭。
预处理流程一般是这样:先用加载函数读取每一行,把 text1 和 text2 截断到 max_len,再拼上 [CLS] 和 [SEP] token。对单塔模型,两个文本会拼成长度不超过 max_len 的单条序列;对双塔模型,每个文本单独 tokenize,各截断到 max_len。分词器用的是 BertTokenizer,默认需要把模型名传进去,比如 bert-base-chinese 或者 bert-base-uncased。下面是 utils.py 里常见的构建训练样本逻辑:
import torch from torch.utils.data import Dataset from transformers import BertTokenizer def load_text_pairs(file_path, delimiter="\t", has_label=True): """加载文本对数据集,返回 (text1_list, text2_list, label_list)""" text1_list, text2_list, label_list = [], [], [] with open(file_path, "r", encoding="utf-8") as f: for line in f: parts = line.strip().split(delimiter) if has_label and len(parts) >= 3: text1_list.append(parts[0]) text2_list.append(parts[1]) label_list.append(int(parts[2])) elif not has_label and len(parts) >= 2: text1_list.append(parts[0]) text2_list.append(parts[1]) return text1_list, text2_list, label_list class TextPairDataset(Dataset): def __init__(self, text1_list, text2_list, label_list, tokenizer, max_len): self.text1_list = text1_list self.text2_list = text2_list self.label_list = label_list self.tokenizer = tokenizer self.max_len = max_len def __len__(self): return len(self.text1_list) def __getitem__(self, idx): text1 = str(self.text1_list[idx]) text2 = str(self.text2_list[idx]) # 单塔模型:直接拼接两段文本,tokenizer 内部会加 [CLS] 和 [SEP] encoded = self.tokenizer( text1, text2, max_length=self.max_len, truncation=True, padding="max_length", return_tensors="pt" ) label = torch.tensor(self.label_list[idx], dtype=torch.long) return { "input_ids": encoded["input_ids"].squeeze(0), "attention_mask": encoded["attention_mask"].squeeze(0), "token_type_ids": encoded["token_type_ids"].squeeze(0), "label": label }这段代码里最关键的是tokenizer(text1, text2, ...)这个调用。transformers 的 BertTokenizer 支持向第二个参数传入 text2,它会自动生成 [CLS] text1 [SEP] text2 [SEP] 的结构,这就是单塔模型需要的输入。如果你用的是双塔模型,这里的处理方式就不一样了,不能把两个文本塞进同一条输入,而是分别调用tokenizer(text1)和tokenizer(text2),得到两条独立的 input_ids。参数方面,max_len决定了句子的最大长度,文本匹配的句子一般不超过 128,太长训练会很慢,太短会截断关键信息;padding="max_length"是为了让 batch 内所有样本长度一致,否则 DataLoader 会报维度不匹配的错误。
这里踩过最大的坑是 token_type_ids。单塔模型里,第一句的 token_type_ids 是 0,第二句是 1,如果你在加载数据时没把第二句传进去,而是自己手动拼接两个句子然后用一个 tokenizer 编码,那么第二句的 token_type_ids 会全部是 0,模型会把两段话当成一段连续文本,语义完全错乱。
3. 三种模型实现拆解:model.py 的结构与相似度计算逻辑
3.1 model.py 里的 BertEncoder 和相似度输出层
打开 model.py 会看到,整个文件的骨架并不复杂,核心是先把预训练模型加载进来,然后在上面接入不同的分类头或者向量映射层。对文本匹配来说,单塔和双塔的差异就体现在 forward 函数怎么组织输入和输出。单塔的 forward 接收拼接后的 input_ids、attention_mask、token_type_ids,把它们一起喂给 BertEncoder,拿到最后一层的隐藏状态,取 [CLS] 位置对应的向量作为整句话对的特征,再过一个线性层和 softmax,输出二分类概率。双塔模型则是分别编码 query 和 doc,各自取 [CLS] 向量,过一层全连接映射得到固定维度的向量,然后计算 cosine 相似度。
我一般会把 model.py 改造成一个基类,再用三个子类分别实现 PointWise、DSSM、Sentence-BERT。这样一个文件里就能同时看到三种结构的异同,方便对比改参数。核心代码如下:
import torch.nn as nn from transformers import BertModel class TextMatchingModel(nn.Module): def __init__(self, model_name, num_labels=2, embedding_size=768): super().__init__() self.encoder = BertModel.from_pretrained(model_name) self.embedding_size = embedding_size self.num_labels = num_labels def get_cls_vector(self, input_ids, attention_mask, token_type_ids=None): outputs = self.encoder( input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids ) # 取 [CLS] 位置向量,transformer 输出为 (batch, seq_len, hidden) return outputs.last_hidden_state[:, 0, :] class PointWiseModel(TextMatchingModel): def __init__(self, model_name, num_labels=2): super().__init__(model_name, num_labels) self.classifier = nn.Linear(self.embedding_size, num_labels) def forward(self, input_ids, attention_mask, token_type_ids): cls_vec = self.get_cls_vector(input_ids, attention_mask, token_type_ids) logits = self.classifier(cls_vec) return logits class DSSMModel(TextMatchingModel): def __init__(self, model_name, emb_size=128): super().__init__(model_name) self.query_fc = nn.Linear(self.embedding_size, emb_size) self.doc_fc = nn.Linear(self.embedding_size, emb_size) def forward(self, query_input_ids, query_mask, doc_input_ids, doc_mask): query_vec = self.get_cls_vector(query_input_ids, query_mask) doc_vec = self.get_cls_vector(doc_input_ids, doc_mask) # 映射到低维空间后计算 cosine 相似度 query_vec = self.query_fc(query_vec) doc_vec = self.doc_fc(doc_vec) return nn.functional.cosine_similarity(query_vec, doc_vec, dim=-1)这里有个值得注意的参数:embedding_size=768。这是 Bert base 模型的隐藏层维度,如果你换成 RoBERTa-large 或者别的模型,这个数值要跟着改。DSSM 里我用emb_size=128,把 768 维压缩到 128 维是为了后面建索引检索时内存更友好,但压缩太狠会损失信息,建议在 128 到 512 之间试。PointWise 直接输出 logits,训练时配合 CrossEntropyLoss 就行,而 DSSM 输出的是 cosine 相似度,它是一个连续值,训练目标通常是让正样本对的相似度接近 1、负样本对接近 0,损失函数可以定义成 MarginRankingLoss 或者带温度的对比损失。
3.2 train_pointwise.py 的训练流程与 loss 设计
PointWise 的训练脚本是最容易上手的,因为它就是一个标准的句子对分类任务。train_pointwise.py 的逻辑跟常规分类训练没有本质区别:加载数据集、tokenize、构造 DataLoader、循环训练。区别主要在 loss 的设计上。如果是二分类任务,正样本和负样本是 1:1 的时候,直接用 CrossEntropyLoss 没问题;但如果正负样本比例失衡,就得考虑给正样本加权重,或者改成 Focal Loss。
从代码结构上看,train_pointwise.py 会先实例化 PointWiseModel,然后定义优化器和学习率调度器。学习率这块,预训练模型微调一般用 2e-5 到 5e-5。transformer 类的模型对学习率特别敏感,超过 5e-5 经常出现 loss 震荡上升,这也是为什么这个项目里用了 warmup 的原因。下面是我认为这个脚本里最值得抄的一段:
from transformers import AdamW, get_linear_schedule_with_warmup def build_optimizer(model, learning_rate=2e-5, warmup_ratio=0.1): # 只对需要梯度的参数做优化 no_decay = ["bias", "LayerNorm.weight"] optimizer_grouped_parameters = [ { "params": [ p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay) ], "weight_decay": 0.01, }, { "params": [ p for n, p in model.named_parameters() if any(nd in n for nd in no_decay) ], "weight_decay": 0.0, }, ] optimizer = AdamW(optimizer_grouped_parameters, lr=learning_rate) total_steps = len(train_loader) * epochs scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=int(total_steps * warmup_ratio), num_training_steps=total_steps ) return optimizer, scheduler这里有两个坑。第一是 weight_decay 对 bias 和 LayerNorm 不生效,这是 transformers 官方 tokenizer 的常见做法,不做的话权重衰减会把这些参数拉偏。第二是 warmup 的比例,我一般设 0.1,也就是前 10% 的步数学习率线性上升,后面线性降到 0。这个做法能让模型在初期不因为突然吃到大量数据而把预训练权重冲坏。
3.3 train_dssm.py 和 train_sentence_transformer.py 的双塔训练细节
双塔模型跟单塔最大的不同在于数据的 batch 组织方式。单塔里一组样本是 (text1, text2, label),双塔里一组样本是 (query, pos_doc, neg_doc) 或者 (query, doc, label)。train_dssm.py 里用的是这种结构:query 和 doc 分别过各自的编码器,得到向量后算相似度,然后用对比损失拉近正样本对、推远负样本对。常见做法是采用 in-batch negative 策略——在一个 batch 里,把每一条 query 自己的 doc 当作正样本,其他所有 batch 里的 doc 当作负样本,这样不用额外构造负样本文件,batch 越大负样本越丰富。
Sentence-BERT 的训练逻辑跟 DSSM 很相似,区别在于损失函数更丰富。train_sentence_transformer.py 里你可能会看到 MultipleNegativesRankingLoss 或者 SoftmaxLoss。MultipleNegativesRankingLoss 本质上就是 in-batch negative 的对比损失,把每行的第一列当 anchor,第二列当 positive,其他行的 positive 当 negative。这个损失函数很适合句子向量训练,因为它在 batch 内自动构造了足够多的负样本。
import torch.nn.functional as F def contrastive_loss(query_vec, doc_vec, batch_size, temperature=0.05): # 归一化后算相似度矩阵 query_vec = F.normalize(query_vec, dim=-1) doc_vec = F.normalize(doc_vec, dim=-1) sim_matrix = torch.matmul(query_vec, doc_vec.T) / temperature # 对角线是正样本对,目标是让对角线分数远高于非对角线 labels = torch.arange(batch_size, device=query_vec.device) loss = F.cross_entropy(sim_matrix, labels) return loss温度系数temperature是训练对比模型最敏感的参数。设成 0.05 会让相似度分布更尖锐,模型能更好地区分正负样本,但太小会导致梯度爆炸;常见取值范围是 0.05 到 0.1。如果你发现训练时 loss 降到 0 但下游评测指标很差,八成是温度太低模型过拟合到训练数据上去了。
4. 从零跑通训练与推理:命令、日志监控和结果验证
4.1 环境安装与数据集准备
这个项目基于 pytorch + transformers,先把依赖装上。在项目根目录执行 requirements.txt 安装,三件套缺一不可:torch、transformers、datasets(如果有的话)。注意 transformers 版本最好锁死,用太新或者太旧的版本都会出现 API 不兼容——老版本里BertModel.from_pretrained的返回值结构就有差异。我在一个新环境下跑通整个流程的命令大概是这样:
# 进入项目根目录 cd text_match-master # 安装依赖(用豆瓣源会快很多) pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple # 需要 GPU 的话,确认 cuda 可用 python -c "import torch; print(torch.cuda.is_available())"装依赖这一条我没少翻车。最常见的错误是 transformers 和 torch 版本不匹配,导致加载模型时直接报key not found。我建议你装完依赖后先随手跑一下import transformers,能看到版本号之后再开始。数据集方面,data 目录里已经准备好了可供训练的文本对文件,你可以先看两行手工确认分隔符。如果数据集是 JSON 格式,那就得改 utils.py 里的 load 函数,别硬套 tab 分隔的逻辑。
4.2 三种训练脚本的启动方式与超参
项目中提供了三个 shell 脚本:train_pointwise.sh、train_dssm.sh、train_sentence_transformer.sh。直接用 sh 跑就行。每个脚本里面都写了模型名称、学习率、batch size、epochs 这些超参数。train_pointwise.sh 的典型内容如下:
# 单塔 PointWise 训练脚本 export CUDA_VISIBLE_DEVICES=0 python train_pointwise.py \ --model_name bert-base-chinese \ --train_data ./data/train_pointwise.tsv \ --dev_data ./data/dev_pointwise.tsv \ --batch_size 16 \ --learning_rate 2e-5 \ --epochs 3 \ --max_len 128 \ --output_dir ./checkpoints/pointwise这里--model_name bert-base-chinese是中文场景的默认选择,如果是英文文本就换成bert-base-uncased。--batch_size 16在显存只有 6G 的显卡上跑 BERT-base 是安全的;如果你的卡显存小,改成 8 再加上梯度累积。--max_len 128对短文本对话场景足够了,但如果你做的是长文档匹配,可能要拉到 256,代价是训练时间翻倍。train_dssm.sh 和 train_sentence_transformer.sh 里的参数结构差不多,只是脚本名和模型结构不同。
4.3 推理脚本输出与效果核对
训练完成后,你会有三个 checkpoint 目录。推理脚本分别是 inference_pointwise.py、inference_dssm.py、inference_sentence_transformer.py。它们做的事情本质一样:加载训练好的 checkpoint,把新的文本对读进来,输出相似度分数或者分类概率。下面是我最常见的推理用法:
# inference_pointwise.py 的核心推理逻辑 from model import PointWiseModel from transformers import BertTokenizer import torch model = PointWiseModel.from_pretrained("./checkpoints/pointwise", model_name="bert-base-chinese") tokenizer = BertTokenizer.from_pretrained("bert-base-chinese") def predict(text1, text2, max_len=128): encoded = tokenizer( text1, text2, max_length=max_len, truncation=True, padding="max_length", return_tensors="pt" ) model.eval() with torch.no_grad(): logits = model(encoded["input_ids"], encoded["attention_mask"], encoded["token_type_ids"]) prob = torch.softmax(logits, dim=-1).squeeze(0) return prob.tolist() print(predict("今天天气怎么样", "明天会下雨吗"))推理脚本里最容易忽略的是model.eval()和torch.no_grad()。缺了这两个,模型的 dropout 还在生效,推理结果每次都会变,而且还会累计计算图导致内存爆炸。我在之前的项目里没加no_grad(),跑了几百条样本之后显存直接满掉,翻车翻得特别彻底。
4.4 用 iTrainingLogger 看训练曲线
这个项目里有个 iTrainingLogger.py 文件,实现了轻量级的日志记录。它比直接 print 强的地方是,能把 loss 和 eval 指标写进本地日志文件,方便训练结束后画曲线。用法也很简单,在训练循环里每若干步调用一次写日志。这个文件不是强制依赖,但你如果想判断“模型到底收敛了没有”,光看最后一轮 loss 是不够的。我一般会每 500 步记录一次 loss,每 1 个 epoch 在 dev 集上算一次准确率。下面是训练循环里常见的日志记录片段:
# 训练循环内,每 100 步写一条 loss if step % 100 == 0: logger.log_metrics({ "loss": loss.item(), "epoch": epoch, "step": step, })你要注意logger.log_metrics这个方法来自 iTrainingLogger.py,它接收一个 dict,内部会按时间戳追加到日志文件。训练结束后你想看曲线,就把日志文件里的 loss 序列拉出来 plot 一下,如果 loss 在 1 个 epoch 后就一路走平甚至微升,说明学习率大了,赶紧停下来降温。别等 3 个 epoch 全部跑完再回来看,那是浪费时间。
5. 避坑指南:五个踩过就回头的坑
5.1 现象:loss 直接变成 nan,或者训练到一半 loss 飙升到几十
我第一次跑 train_pointwise.py 的时候,刚跑了 200 步 loss 就变成 nan。查了日志,前面几步 loss 还在正常下降,突然就爆上去。原因是我把学习率设成了 0.01,这对 BERT 这种 transformer 模型来说完全是灾难。预训练模型微调的学习率基本不能超过 5e-5,常见做法是 2e-5 到 3e-5。解决方法是把学习率降到 2e-5,并且加上 warmup,让模型前 10% 的步数慢慢进入状态。如果你已经降了学习率还是 nan,就检查输入文本里有没有空值或者乱码,tokenizer 遇到异常字符会产出全零的 input_ids,直接把梯度打穿。
5.2 现象:双塔模型推理时向量维度对不上,算 cosine 报维度错误
这个问题典型出现在你训练好了 DSSM,想用 sentence embedding 做检索。训练脚本里 query_fc 把 768 维映射到 128 维,但推理脚本里你直接拿编码器的原始输出当向量,维度是 768。两边一算 cosine,torch 直接报 batch 维度不匹配。解决方法是推理时必须复用训练时那个线性映射层,不用重新初始化一个。我在 inference_dssm.py 里加了一行代码,把编码器输出过一遍 query_fc 再保存,后面所有下游任务都不会踩这个坑。
5.3 现象:训练日志里一直看不到 eval 指标,或者 dev 集准确率纹丝不动
这种情况多半是你把 dev 集和训练集混在一起了。项目里 data 目录下既有训练集也有验证集,但有些同学为了图方便,直接把训练集当验证集用,结果每个 epoch 的准确率都在 99% 上下浮动,自己还以为模型很牛。等你真正换到线上数据,立刻打回原形。解决方法是把数据集按一定比例切分,比如 8:1:1 分成 train、dev、test,每个集合互不重叠。另外,如果 dev 集太少(比如只有 50 条),准确率的方差会很大,这时候改用 F1 或者 AUC 指标会更稳。
5.4 现象:中文文本输入后,[CLS] 向量全是同一个值
这是个非常隐蔽的 bug。当你的 tokenizer 用的是bert-base-uncased,但文本是中文时,tokenizer 会把中文字符全拆成[UNK],导致输入全部是同一个 token。模型吃进去之后,[CLS] 向量自然固定不变,loss 永远降不下去。解决方法是换成bert-base-chinese或者chinese-roberta-wwm-ext。判断方法很简单:tokenizer 一下你的文本,打印出 input_ids,如果里面[UNK]对应的 id(100)占了一大半,那就是 tokenizer 选错了。
5.5 现象:多卡训练时产出结果不一致,或者运行时提示 CUDA 版本不对
这个项目默认单卡训练,但有人会改CUDA_VISIBLE_DEVICES=0,1想用多卡。如果没有改 DataLoader 的sampler,两个卡上会出现数据重复,导致 loss 波动。还有一次是在没装对应 CUDA 版本的 torch 环境里直接跑,报错说torch.cuda.is_available()是 False。解决方法是先检查nvidia-smi看驱动支持的 CUDA 版本,再安装匹配的 torch。常见做法是直接用官网的 pip 命令安装对应 cuda 版本的 torch,别图方便装 CPU 版本。多卡场景下,我建议你保持单卡训练,等调通基线再考虑分布式。
6. 进阶:把训练好的模型包装成一个稳定可用的文本相似度接口
模型跑通只是第一步,真正能用起来还得把它封装成可调用的服务。我一般会做这样几个动作:先用验证集找到合适的相似度阈值,再用 Flask 写一个极简的 HTTP 接口,最后把 checkpoint 里的参数和 tokenizer 一起固化下来。这里有个值得记住的技巧——阈值不是拍脑袋定的,而是通过在 dev 集上画 PR 曲线,找到你业务能接受的召回和精度的平衡点。拿单塔模型举例,你可以把 dev 集里所有样本的相似度分数和真实标签拉出来,试一堆阈值,看哪个阈值上 F1 最高。这一步做完,你的模型才真正具备上线价值。
from flask import Flask, request, jsonify app = Flask(__name__) # 加载训练好的模型和 tokenizer(单塔和双塔都适用) model = PointWiseModel.from_pretrained("./checkpoints/pointwise", model_name="bert-base-chinese") tokenizer = BertTokenizer.from_pretrained("bert-base-chinese") model.eval() def compute_similarity(text1, text2): encoded = tokenizer( text1, text2, max_length=128, truncation=True, padding="max_length", return_tensors="pt" ) with torch.no_grad(): logits = model(encoded["input_ids"], encoded["attention_mask"], encoded["token_type_ids"]) prob = torch.softmax(logits, dim=-1).squeeze(0) # 返回正样本概率作为相似度 return prob[1].item() @app.route("/match", methods=["POST"]) def match(): data = request.get_json() text1 = data.get("text1", "") text2 = data.get("text2", "") score = compute_similarity(text1, text2) return jsonify({"score": score, "labels": "相似" if score > 0.6 else "不相似"}) if __name__ == "__main__": app.run(host="0.0.0.0", port=8890)这个接口本身不复杂,但我在上线以后才意识到一个问题:每次请求都即时 forward,单塔模型吞吐量很低。如果你们公司的业务是百万级文本互相比,就不要用这种接口,而是应该把双塔模型离线把所有文本向量化,然后接 faiss 做 ANN 检索。我在这个项目上吃过这个亏,当时用单塔模型做了一个问答匹配接口,线上 QPS 不到 20,后来才回头把双塔模型重新训练了一版,向量存进 faiss,QPS 直接拉到几百。从那以后我每次做文本匹配项目,都强制自己先想清楚“要的是精度还是吞吐”,再决定用哪个模型封装;如果是毕设或者 demo,单塔就够,如果是生产,必须双塔加索引。希望帮到你。
本文还有配套的精品资源,点击获取