效果提升28个点!基于领域预训练和对比学习SimCSE的语义检索

语义检索相比传统基于字面关键词的检索有诸多优势,广泛应用于问答、搜索系统中。今天小编就手把手带大家完成一个基于领域预训练和对比学习SimCSE的语义检索小系统。

所谓语义检索(也称基于向量的检索),是指检索系统不再拘泥于用户Query字面本身(例如BM25检索),而是能精准捕捉到用户Query背后的真正意图并以此来搜索,从而向用户返回更准确的结果。

最终可视化demo如下,一方面可以获取文本的向量表示;另一方面可以做文本检索,即得到输入Query的top-K相关文档!

语义检索,底层技术是语义匹配,是NLP最基础常见的任务之一。从广度上看,语义匹配可以应用到QA、搜索、推荐、广告等各大方向;从技术深度上看,语义匹配需要融合各种SOTA模型、双塔和交互两种常用框架的魔改、以及样本处理的艺术和各种工程tricks。

比较有趣的是,在查相关资料的时候,发现百度飞桨PaddleNLP最近刚开源了类似的功能,可谓国货之光!之前使用过PaddleNLP,基本覆盖了NLP的各种应用和SOTA模型,调用起来也非常方便,强烈推荐大家试试!

接下来基于PaddleNLP提供的轮子一步步搭建语义检索系统。整体框架如下,由于计算量与资源的限制,一般工业界的搜索系统都会设计成多阶段级联结构,主要有召回、排序(粗排、精排、重排)等模块,各司其职。

  • step-1:利用预训练模型离线构建候选语料库;
  • step-2:召回模块,对于在线查询Query,利用Milvus快速检索得到top1000候选集;
  • step-3:排序模块,对于召回的top1000,再做更精细化的排序,得到top100结果返回给用户。

语义检索技术框架图

 整体概览 

1.1 数据

数据来源于某文献检索系统,分为有监督(少量)和无监督(大量)两种。

1.2 代码

首先clone代码:

运行环境是:

  • python3
  • paddlepaddle==2.2.1
  • paddlenlp==2.2.1

还有一些依赖包可以参考requirements.txt。

 离线建库 

从上面的语义检索技术框架图中可以看出,首先我们需要一个语义模型对输入的Query/Doc文本提取向量,这里选用基于对比学习的SimCSE,核心思想是使语义相近的句子在向量空间中临近,语义不同的互相远离。

那么,如何训练才能充分利用好模型,达到更高的精度呢?对于预训练模型,一般常用的训练范式已经从『通用预训练->领域微调』的两阶段范式变成了『通用预训练->领域预训练->领域微调』三阶段范式。

具体地,在这里我们的模型训练分为几步(代码和相应数据在下一节介绍):
1.在无监督的领域数据集上对通用ERNIE 1.0 进一步领域预训练,得到领域ERNIE;
2.以领域ERNIE为热启,在无监督的文献数据集上对 SimCSE 做预训练;
3.在有监督的文献数据集上结合In-Batch Negatives策略微调步骤2模型,得到最终的模型,用于抽取文本向量表示,即我们所需的语义模型,用于建库和召回。

由于召回模块需要从千万量级数据中快速召回候选集合,通用的做法是借助向量搜索引擎实现高效 ANN,从而实现候选集召回。这里采用Milvus开源工具,关于Milvus的搭建教程可以参考官方教程
https://milvus.io/cn/docs/v1.1.1/

Milvus是一款国产高性能检索库, 和Facebook开源的Faiss功能类似。
离线建库的代码位于PaddleNLP/applications/neural_search/recall/milvus

2.1 抽取向量

依照Milvus教程搭建完向量引擎后,就可以利用预训练语义模型提取文本向量了。运行feature_extract.py即可,注意修改需要建库的数据源路径。

运行结束会生成1000万条的文本数据,保存为corpus_embedding.npy。

2.2 插入向量

接下来,修改config.py中的Milvus ip等配置,将上一步生成的向量导入到Milvus库中。

抽取和插入向量两步,如果机器资源不是很"富裕"的话,可能会花费很长时间。这里建议可以先用一小部分数据进行功能测试,快速感知,等真实部署的阶段再进行全库的操作。

插入完成后,我们就可以通过Milvus提供的可视化工具[1]查看向量数据,分别是文档对应的ID和向量。


在线客服
在线客服