1.1 数据
数据来源于某文献检索系统,分为有监督(少量)和无监督(大量)两种。
1.2 代码
git clone git@github.com:PaddlePaddle/PaddleNLP.git
cd applications/neural_search
离线建库
从上面的语义检索技术框架图中可以看出,首先我们需要一个语义模型对输入的Query/Doc文本提取向量,这里选用基于对比学习的SimCSE,核心思想是使语义相近的句子在向量空间中临近,语义不同的互相远离。
|—— 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 抽取向量
2.2 插入向量
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)
文档召回
召回阶段的目的是从海量的资源库中,快速地检索出符合Query要求的相关文档Doc。出于计算量和对线上延迟的要求,一般的召回模型都会设计成双塔形式,Doc塔离线建库,Query塔实时处理线上请求。
3.1 领域预训练
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无监督预训练
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 模型训练、评估逻辑
3.3 有监督微调
|—— 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 语义模型效果
3.5 向量召回
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)
文档排序
不同于召回,排序阶段由于面向的打分集合相对小很多,一般只有几千级别,所以可以使用更复杂的模型,这里采用 ERNIE-Gram 预训练模型,loss选用 margin_ranking_loss。
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 单塔匹配模型预测脚本,输出文本对是相似度
是文化差异。', '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分钟即可完成。
注:作者cafedo @NewBeeNLP