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

行业资讯 / 2022-01-14 09:31:37

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