目录
基于某检索平台开源的数据集构造生成了面向语义索引的召回库。
使用 Milvus 搭建召回系统，然后使用训练好的语义索引模型，抽取向量，插入到 Milvus 中，然后进行检索。
环境依赖
- python >= 3.6.2
- paddlepaddle >= 2.2
- paddlenlp >= 2.2
- milvus >= 2.1.0
- pymilvus >= 2.1.0
|—— scripts
|—— feature_extract.sh 提取特征向量的bash脚本
|—— search.sh 插入向量和向量检索bash脚本
├── base_model.py # 语义索引模型基类
├── config.py # milvus配置文件
├── data.py # 数据处理函数
├── milvus_ann_search.py # 向量插入和检索的脚本
├── inference.py # 动态图模型向量抽取脚本
├── feature_extract.py # 批量抽取向量脚本
├── milvus_util.py # milvus的工具类
└── README.md
数据集的样例如下，有两种，第一种是 title+keywords 进行拼接；第二种是一句话。
煤矸石-污泥基活性炭介导强化污水厌氧消化煤矸石,污泥,复合基活性炭,厌氧消化,直接种间电子传递
睡眠障碍与常见神经系统疾病的关系睡眠觉醒障碍,神经系统疾病,睡眠,快速眼运动,细胞增殖,阿尔茨海默病
城市道路交通流中观仿真研究智能运输系统;城市交通管理;计算机仿真;城市道路;交通流;路径选择
....
├── milvus # milvus建库数据集
├── milvus_data.csv. # 构建召回库的数据
├── recall # 召回（语义索引）数据集
├── corpus.csv # 用于测试的召回库
├── dev.csv # 召回验证集
├── test.csv # 召回测试集
├── train.csv # 召回训练集
├── train_unsupervised.csv # 无监督训练集
├── sort # 排序数据集
├── test_pairwise.csv # 排序测试集
├── dev_pairwise.csv # 排序验证集
└── train_pairwise.csv # 排序训练集
数据准备结束以后，我们开始搭建 Milvus 的语义检索引擎，用于语义向量的快速检索，我们使用Milvus开源工具进行召回，Milvus 的搭建教程请参考官方教程 Milvus官方安装教程本案例使用的是 Milvus 的2.1版本，建议使用官方的 Docker 安装方式，简单快捷。
Milvus 搭建完系统以后就可以插入和检索向量了，首先生成 embedding 向量，每个样本生成256维度的向量，使用的是32GB的V100的卡进行的提取：
CUDA_VISIBLE_DEVICES=0 python feature_extract.py \
--model_dir=./output \
--model_name_or_path rocketqa-zh-base-query-encoder \
--corpus_file "data/milvus_data.csv"
其中 output 目录下存放的是召回的 Paddle Inference 静态图模型。
| 数据量 | 时间 |
|---|---|
| 1000万条 | 3hour40min39s |
运行结束后会生成 corpus_embedding.npy
生成了向量后，需要把数据插入到 Milvus 库中，首先修改配置：
修改 config.py 的配置 ip 和端口，本项目使用的是8530端口，而 Milvus 默认的是19530，需要根据情况进行修改：
MILVUS_HOST='your milvus ip'
MILVUS_PORT = 8530
然后运行下面的命令把向量插入到Milvus库中：
python milvus_ann_search.py --data_path milvus/milvus_data.csv \
--embedding_path corpus_embedding.npy \
--batch_size 100000 \
--insert
参数含义说明
data_path: 数据的路径embedding_path: 数据对应向量的路径index: 选择检索向量的索引，用于向量检索insert: 是否插入向量search: 是否检索向量batch_size: 表示的是一次性插入的向量的数量
| 数据量 | 时间 |
|---|---|
| 1000万条 | 21min12s |
另外，Milvus提供了可视化的管理界面，可以很方便的查看数据，安装地址为Attu.
运行召回脚本：
python milvus_ann_search.py --data_path milvus/milvus_data.csv \
--embedding_path corpus_embedding.npy \
--batch_size 100000 \
--index 18 \
--search
运行以后的结果的输出为：
hit: (distance: 0.0, id: 18), text field: 吉林铁合金集团资产管理现状分析及对策资产管理;资金控制;应收帐款风险;造价控制;集中化财务控制
hit: (distance: 0.45325806736946106, id: 7611689), text field: 哈药集团应收账款分析应收账款,流动资产,财务报告
hit: (distance: 0.5440893769264221, id: 4297885), text field: 宝钢集团负债经营风险控制策略研究钢铁行业;负债经营;风险控制
hit: (distance: 0.5455711483955383, id: 5661135), text field: 浅谈电网企业固定资产风险管理大数据,固定资产,风险管理
...
返回的是向量的距离，向量的id，以及对应的文本。
也可以一键执行上述的过程：
sh scripts/search.sh
首先修改代码的模型路径和样本：
params_path='checkpoints/model_40/model_state.pdparams'
id2corpus={0:'国有企业引入非国有资本对创新绩效的影响——基于制造业国有上市公司的经验证据'}
运行命令
python3 inference.py
运行的输出为，分别是抽取的向量和召回的结果：
[1, 256]
Tensor(shape=[1, 256], dtype=float32, place=Place(gpu:0), stop_gradient=True,
[[ 0.07830613, -0.14036864, 0.03433795, -0.14967985, -0.03386058,
0.06630671, 0.01357946, 0.03531205, 0.02411086, 0.02000865,
0.05724005, -0.08119474, 0.06286906, 0.06509133, 0.07193415,
....
hit: (distance: 0.40141725540161133, id: 2742485), text field: 完善国有企业技术创新投入机制的探讨--基于经济责任审计实践国有企业,技术创新,投
入机制
hit: (distance: 0.40258315205574036, id: 1472893), text field: 企业技术创新与组织冗余--基于国有企业与非国有企业的情境研究
hit: (distance: 0.4121206998825073, id: 51831), text field: 企业创新影响对外直接投资决策—基于中国制造业上市公司的研究企业创新;对外直接投资;
制造业;上市公司
hit: (distance: 0.42234909534454346, id: 8682312), text field: 政治关联对企业创新绩效的影响——国有企业与民营企业的对比政治关联,创新绩效,国有
企业,民营企业,双重差分
hit: (distance: 0.46187296509742737, id: 9324797), text field: 财务杠杆、股权激励与企业创新——基于中国A股制造业经验数据制造业;上市公司;股权激
励;财务杠杆;企业创新
....
使用的是近似索引，详情请参考Milvus官方文档，索引创建机制