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

2022年01月14日 阅读数:3
这篇文章主要向大家介绍效果提高28个点!基于领域预训练和对比学习SimCSE的语义检索,主要内容包括基础应用、实用技巧、原理机制等方面,希望对大家有所帮助。

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

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

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

 

 

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

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

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

  • step-1:利用预训练模型离线构建候选语料库;框架

  • step-2:召回模块,对于在线查询Query,利用Milvus快速检索获得top1000候选集;dom

  • step-3:排序模块,对于召回的top1000,再作更精细化的排序,获得top100结果返回给用户。函数

语义检索技术框架图工具

 

 

总体概览

1.1 数据

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

1.2 代码

首先clone代码:

git clone git@github.com:PaddlePaddle/PaddleNLP.git  
cd applications/neural_search

运行环境是:

还有一些依赖包能够参考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

|—— scripts
    |—— feature_extract.sh  #提取特征向量的bash脚本
├── base_model.py # 语义索引模型基类
├── config.py  # milvus配置文件
├── data.py # 数据处理函数
├── embedding_insert.py # 插入向量
├── embedding_recall.py # 检索topK类似结果 / ANN
├── inference.py # 动态图模型向量抽取脚本
├── feature_extract.py # 批量抽取向量脚本
├── milvus_insert.py # 插入向量工具类
├── milvus_recall.py # 向量召回工具类
├── README.md
└── server_config.yml # milvus的config文件,本项目所用的配置

2.1 抽取向量

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

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

2.2 插入向量

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

embeddings=np.load('corpus_embedding.npy') 
embedding_ids = [i for i in range(embeddings.shape[0])]

client = VecToMilvus()
collection_name = 'literature_search'
partition_tag = 'partition_2'
data_size=len(embedding_ids)
batch_size=100000
for i in tqdm(range(0,data_size,batch_size)):
    cur_end=i+batch_size
    if(cur_end>data_size):
        cur_end=data_size
    batch_emb=embeddings[np.arange(i,cur_end)]
    status, ids = client.insert(collection_name=collection_name, vectors=batch_emb.tolist(), ids=embedding_ids[i:i+batch_size],partition_tag=partition_tag)

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

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

 

 

文档召回

召回阶段的目的是从海量的资源库中,快速地检索出符合Query要求的相关文档Doc。出于计算量和对线上延迟的要求,通常的召回模型都会设计成双塔形式,Doc塔离线建库,Query塔实时处理线上请求。

召回模型采用Domain-adaptive Pretraining + SimCSE + In-batch Negatives方案。

另外,若是只是想快速测试或部署,PaddleNLP也贴心地开源了训练好的模型文件,下载便可用,这里直接贴出模型连接:

  • 领域预训练ERNIE

    https://bj.bcebos.com/v1/paddlenlp/models/ernie_pretrain.zip

  • 无监督SimCSE:

    https://bj.bcebos.com/v1/paddlenlp/models/simcse_model.zip

  • 有监督In-batch Negatives:

    https://bj.bcebos.com/v1/paddlenlp/models/inbatch_model.zip

 

3.1 领域预训练

Domain-adaptive Pretraining的优点在以前文章已有具体介绍,再也不赘述。直接给代码,具体功能都标注在后面。

domain_adaptive_pretraining/
|—— scripts
    |—— run_pretrain_static.sh # 静态图与训练bash脚本
├── ernie_static_to_dynamic.py # 静态图转动态图
├── run_pretrain_static.py # ernie1.0静态图预训练
├── args.py # 预训练的参数配置文件
└── data_tools # 预训练数据处理文件目录

 

3.2 SimCSE无监督预训练

双塔模型,采用ERNIE 1.0热启,引入 SimCSE 策略。训练数据示例以下代码结构以下,各个文件的功能都有备注在后面,清晰明了。

simcse/
├── model.py # SimCSE 模型组网代码
|—— deploy
    |—— python
        |—— predict.py # PaddleInference
        ├── deploy.sh # Paddle Inference的bash脚本
|—— scripts
    ├── export_model.sh # 动态图转静态图bash脚本
    ├── predict.sh # 预测的bash脚本
    ├── evaluate.sh # 召回评估bash脚本
    ├── run_build_index.sh  # 索引的构建脚本
    ├── train.sh # 训练的bash脚本
|—— ann_util.py # Ann 建索引库相关函数
├── data.py # 无监督语义匹配训练数据、测试数据的读取逻辑
├── export_model.py # 动态图转静态图
├── predict.py # 基于训练好的无监督语义匹配模型计算文本 Pair 类似度
├── evaluate.py # 根据召回结果和评估集计算评估指标
|—— inference.py # 动态图抽取向量
|—— recall.py # 基于训练好的语义索引模型,从召回库中召回给定文本的类似文本
└── train.py # SimCSE 模型训练、评估逻辑

对于训练、评估和预测分别运行scripts目录下对应的脚本便可。训练获得模型,咱们一方面能够用于提取文本的语义向量表示,另外一方面也能够用于计算文本对的语义类似度,只须要调整下数据输入格式便可。

3.3 有监督微调

对上一步的模型进行有监督数据微调,训练数据示例以下,每行由一对语义类似的文本对组成,tab分割,负样原本源于引入In-batch Negatives采样策略。

 

关于In-batch Negatives 的细节,能够参考文章:

大规模搜索+预训练,百度是如何落地的?

https://mp.weixin.qq.com/s/MyVK6iKTiI-VpP1LKf4LIA

总体代码结构以下:

|—— data.py # 数据读取、数据转换等预处理逻辑
|—— base_model.py # 语义索引模型基类
|—— train_batch_neg.py # In-batch Negatives 策略的训练主脚本
|—— batch_negative
    |—— model.py # In-batch Negatives 策略核心网络结构
|—— ann_util.py # Ann 建索引库相关函数
|—— recall.py # 基于训练好的语义索引模型,从召回库中召回给定文本的类似文本
|—— evaluate.py # 根据召回结果和评估集计算评估指标
|—— predict.py # 给定输入文件,计算文本 pair 的类似度
|—— export_model.py # 动态图转换成静态图
|—— scripts
    |—— export_model.sh  # 动态图转换成静态图脚本
    |—— predict.sh  # 预测bash版本
    |—— evaluate.sh # 评估bash版本
    |—— run_build_index.sh # 构建索引bash版本
    |—— train_batch_neg.sh  # 训练bash版本
|—— deploy
    |—— python
        |—— predict.py # PaddleInference
        |—— deploy.sh # Paddle Inference部署脚本
|—— inference.py # 动态图抽取向量

训练、评估、预测的步骤和上一步无监督的相似,聪明的你确定一看就懂了!

3.4 语义模型效果

前面说了那么多,来看看几个模型的效果到底怎么样?对于匹配或者检索模型,经常使用的评价指标是Recall@K,即前TOP-K个结果检索出的正确结果数与全库中全部正确结果数的比值。

 

对比能够发现,首先利用ERNIE 1.0作Domain-adaptive Pretraining,而后把训练好的模型加载到SimCSE上进行无监督训练,最后利用In-batch Negatives 在有监督数据上进行训练能得到最佳的性能。

3.5 向量召回

终于到了召回,回顾一下,在这以前咱们已经训练好了语义模型、搭建完了召回库,接下来只须要去库中检索便可。代码位于

PaddleNLP/applications/neural_search/recall/milvus/inference.py

def search_in_milvus(text_embedding):
    collection_name = 'literature_search'  # 以前搭建好的Milvus库
    partition_tag = 'partition_2'
    client = RecallByMilvus()
    status, results = client.search(collection_name=collection_name, vectors=text_embedding.tolist(),
                                    partition_tag=partition_tag)

    corpus_file = "../../data/milvus/milvus_data.csv"
    id2corpus = gen_id2corpus(corpus_file)
    for line in results:
        for item in line:
            idx = item.id
            distance = item.distance
            text = id2corpus[idx]
            print(idx, text, distance)

 

以输入 国有企业引入非国有资本对创新绩效的影响——基于制造业国有上市公司的经验证据 为例,检索返回效果以下

 

返回结果的最后一列为类似度,Milvus默认使用的是欧式距离,若是想换成余弦类似度,能够在Milvus的配置文件中修改。

 

文档排序

不一样于召回,排序阶段因为面向的打分集合相对小不少,通常只有几千级别,因此可使用更复杂的模型,这里采用 ERNIE-Gram 预训练模型,loss选用 margin_ranking_loss。

训练数据示例以下,三列,分别为(query,title,neg_title),tab分割。对于真实搜索场景,训练数据一般来源业务线上的点击日志,构造出正样本和强负样本。

 

代码结构以下

ernie_matching/
├── deply # 部署
    └── python
        ├── deploy.sh # 预测部署bash脚本
        └── predict.py # python 预测部署示例
|—— scripts
    ├── export_model.sh # 动态图参数导出静态图参数的bash文件
    ├── train_pairwise.sh # Pair-wise 单塔匹配模型训练的bash文件
    ├── evaluate.sh # 评估验证文件bash脚本
    ├── predict_pairwise.sh # Pair-wise 单塔匹配模型预测脚本的bash文件
├── export_model.py # 动态图参数导出静态图参数脚本
├── model.py #  Pair-wise 匹配模型组网
├── data.py #  Pair-wise 训练样本的转换逻辑 、Pair-wise 生成随机负例的逻辑
├── train_pairwise.py # Pair-wise 单塔匹配模型训练脚本
├── evaluate.py # 评估验证文件
├── predict_pairwise.py # Pair-wise 单塔匹配模型预测脚本,输出文本对是类似度

 

训练运行sh scripts/train_pairwise.sh便可。

一样,PaddleNLP也开源了排序模型:

https://bj.bcebos.com/v1/paddlenlp/models/ernie_gram_sort.zip

对于预测,准备数据为每行一个文本对,最终预测返回文本对的语义类似度。

是文化差别。', 'pred_prob': 0.85112214}
{'query': '中西方语言与文化的差别', 'title': '跨文化视角下中国文化对外传播路径琐谈跨文化,中国文化,传播,翻译', 'pred_prob': 0.78629625}
{'query': '中西方语言与文化的差别', 'title': '从中西方民族文化心理的差别看英汉翻译语言,文化,民族文化心理,思惟方式,翻译', 'pred_prob': 0.91767526}
{'query': '中西方语言与文化的差别', 'title': '中英文化差别对翻译的影响中英文化,差别,翻译的影响', 'pred_prob': 0.8601749}
{'query': '中西方语言与文化的差别', 'title': '浅谈文化与语言习得文化,语言,文化与语言的关系,文化与语言习得意识,跨文化交际', 'pred_prob': 0.8944413}

 

总结

本文基于PaddleNLP提供的Neural Search功能本身快速搭建了一套语义检索系统。相对于本身从零开始,PaddleNLP很是好地提供了一套轮子。若是直接下载PaddleNLP开源训练好的模型文件,对于语义类似度任务,调用现成的脚本几分钟便可搞定!对于语义检索任务,须要将全量数据导入Milvus构建索引,除训练和建库时间外,整个流程预计30-50分钟便可完成。

在训练的间隙还研究了下,发现GitHub上的文档也很清晰详细啊,对于小白入门同窗,作到了一键运行,不至于被繁杂的流程步骤困住而逐渐失去兴趣;模型所有开源,拿来即用;对于想要深刻研究的同窗,PaddleNLP也开源了数据和代码,能够进一步学习,赞!照着跑下来,发现PaddleNLP太香了!赶忙Star收藏一下,持续跟进最新能力吧,也表示对开源社区的一点支持~

https://github.com/PaddlePaddle/PaddleNLP

另外咱们还能够基于这些功能进行本身额外的开发,譬如开篇的动图,搭建一个更直观的语义向量生成和检索服务。Have Fun!

在跑代码过程当中也遇到一些问题,很是感谢飞桨同窗的耐心解答。而且得知针对这个项目还有一节视频课程已经公开,点击连接便可观看课程:

https://aistudio.baidu.com/aistudio/course/introduce/24902

最后附上本次实践项目的代码:

https://github.com/PaddlePaddle/PaddleNLP/tree/develop/applications/neural_search

点击进入得到更多技术信息~~