
近年来,基于预训练的语言模型的密集检索器取得了显着的进步。为了促进更多开发人员使用最先进的技术,该存储库提供了一个易于使用的工具包,用于运行和微调最先进的密集检索器,即Rocketqa 。该工具包具有以下优点:

我们提供两种安装方法: Python安装程序包和Docker环境
首先,安装桨板。
# GPU version:
$ pip install paddlepaddle-gpu
# CPU version:
$ pip install paddlepaddle第二,安装RocketQa软件包(最新版本:1.1.0):
$ pip install rocketqa注意:此工具包必须在Python3.6+上运行,并使用PaddlePaddle 2.0+运行。
docker pull rocketqa/rocketqa
docker run -it docker.io/rocketqa/rocketqa bash请参阅下面的示例,您可以使用几行代码来构建和运行自己的搜索引擎。我们还提供一个带有Jupyternotebook的操场。在浏览器中立即尝试Rocketqa!
Jina是一个云本地的神经搜索框架,可在几分钟内构建SOTA和可扩展的深度学习搜索应用程序。这是一个简单的示例,用于构建基于Jina和Rocketqa的搜索引擎。
cd examples/jina_example
pip3 install -r requirements.txt
# Generate vector representations and build a libray for your Documents
# JINA will automaticlly start a web service for you
python3 app.py index toy_data/test.tsv
# Try some questions related to the indexed Documents
python3 app.py query_cli请查看Jina示例以了解更多。
我们还提供了一个基于faiss的简单示例。
cd examples/faiss_example/
pip3 install -r requirements.txt
# Generate vector representations and build a libray for your Documents
python3 index.py zh ../data/dureader.para test_index
# Start a web service on http://localhost:8888/rocketqa
python3 rocketqa_service.py zh ../data/dureader.para test_index
# Try some questions related to the indexed Documents
python3 query.py您还可以轻松地将RocketQA集成到自己的任务中。我们提供两种类型的模型:基于Ernie的双重编码器,用于答案检索和基于Ernie的交叉编码器,以重新排名。对于运行我们的模型,您可以使用以下功能。
rocketqa.available_models()返回可用RocketQA型号的名称。要了解有关可用模型的更多信息,请参阅代码注释。
rocketqa.load_model(model, use_cuda=False, device_id=0, batch_size=1)返回输入参数指定的模型。它可以初始化双重编码器和交叉编码器。通过设置输入参数,您可以加载由“可用_models()”返回的RocketQA型号或您自己的检查点。
由“ load_model()”返回的双编码器支持以下功能:
model.encode_query(query: List[str])给定查询列表,返回按模型编码的表示向量。
model.encode_para(para: List[str], title: List[str])给定段落及其相应标题(可选)的列表,返回由模型编码的表示向量。
model.matching(query: List[str], para: List[str], title: List[str])给定查询和段落(和标题)的列表,返回其匹配分数(两个表示向量之间的点产品)。
model.train(train_set: str, epoch: int, save_model_path: str, args)鉴于HyperParameters train_set , epoch和save_model_path ,您可以训练自己的双编码器模型或Finetune我们的型号。其他设置(例如save_steps and learning_rate也可以在args中设置。请参阅示例/示例。
由“ load_model()”返回的交叉编码器支持以下功能:
model.matching(query: List[str], para: List[str], title: List[str])给定查询和段落(和标题)的列表,返回其匹配分数(段落是查询正确的答案的概率)。
model.train(train_set: str, epoch: int, save_model_path: str, args)给定超参数train_set , epoch和save_model_path ,您可以训练自己的交叉编码器模型或我们的型号Finetune。其他设置(例如save_steps and learning_rate也可以在args中设置。请参阅示例/示例。
按照以下示例,您可以检索文档的向量表示,并将RocketQA连接到您自己的任务。
要运行RocketQA模型,您应该在“ load_model()”中设置参数model ,其中RocketQa模型名称由'opaine_models()返回。
import rocketqa
query_list = [ "trigeminal definition" ]
para_list = [
"Definition of TRIGEMINAL. : of or relating to the trigeminal nerve.ADVERTISEMENT. of or relating to the trigeminal nerve. ADVERTISEMENT." ]
# init dual encoder
dual_encoder = rocketqa . load_model ( model = "v1_marco_de" , use_cuda = True , device_id = 0 , batch_size = 16 )
# encode query & para
q_embs = dual_encoder . encode_query ( query = query_list )
p_embs = dual_encoder . encode_para ( para = para_list )
# compute dot product of query representation and para representation
dot_products = dual_encoder . matching ( query = query_list , para = para_list )要训练自己的模型,您可以将train()功能与数据集和参数一起使用。培训数据包含4列:查询,标题,Para,标签(0或1),由“ t”隔开。有关参数和数据集的详细信息,请参阅'./ examples/example.py'
import rocketqa
# init cross encoder, and set device and batch_size
cross_encoder = rocketqa . load_model ( model = "zh_dureader_ce" , use_cuda = True , device_id = 0 , batch_size = 32 )
# finetune cross encoder based on "zh_dureader_ce_v2"
cross_encoder . train ( './examples/data/cross.train.tsv' , 2 , 'ce_models' , save_steps = 1000 , learning_rate = 1e-5 , log_folder = 'log_ce' )要运行自己的模型,您应该使用JSON配置文件中的“ Load_Model()”中设置参数model 。
import rocketqa
# init cross encoder
cross_encoder = rocketqa . load_model ( model = "./examples/ce_models/config.json" , use_cuda = True , device_id = 0 , batch_size = 16 )
# compute relevance of query and para
relevance = cross_encoder . matching ( query = query_list , para = para_list )配置是这样的JSON文件
{
"model_type": "cross_encoder",
"max_seq_len": 384,
"model_conf_path": "zh_config.json",
"model_vocab_path": "zh_vocab.txt",
"model_checkpoint_path": ${YOUR_MODEL},
"for_cn": true,
"share_parameter": 0
}
文件夹examples提供了更多详细信息。
如果您发现Rocketqa V1模型有所帮助,请随时引用我们的出版物Rocketqa:一种优化的培训方法,可用于张开通道检索,以回答
@inproceedings{rocketqa_v1,
title="RocketQA: An Optimized Training Approach to Dense Passage Retrieval for Open-Domain Question Answering",
author="Yingqi Qu, Yuchen Ding, Jing Liu, Kai Liu, Ruiyang Ren, Wayne Xin Zhao, Daxiang Dong, Hua Wu and Haifeng Wang",
year="2021",
booktitle = "In Proceedings of NAACL"
}
如果您发现配对模型有帮助,请随时引用我们的出版物对:利用以段落为中心的相似性关系以改善密集的通道检索
@inproceedings{rocketqa_pair,
title="PAIR: Leveraging Passage-Centric Similarity Relation for Improving Dense Passage Retrieval",
author="Ruiyang Ren, Shangwen Lv, Yingqi Qu, Jing Liu, Wayne Xin Zhao, Qiaoqiao She, Hua Wu, Haifeng Wang and Ji-Rong Wen",
year="2021",
booktitle = "In Proceedings of ACL Findings"
}
如果您发现Rocketqa V2型号有用
@inproceedings{rocketqa_v2,
title="RocketQAv2: A Joint Training Method for Dense Passage Retrieval and Passage Re-ranking",
author="Ruiyang Ren, Yingqi Qu, Jing Liu, Wayne Xin Zhao, Qiaoqiao She, Hua Wu, Haifeng Wang and Ji-Rong Wen",
year="2021",
booktitle = "In Proceedings of EMNLP"
}
如果您发现Dureader检索数据集有用
@inproceedings{DuReader_retrieval,
title="DuReader_retrieval: A Large-scale Chinese Benchmark for Passage Retrieval from Web Search Engine",
author="Yifu Qiu, Hongyu Li, Yingqi Qu, Ying Chen, Qiaoqiao She, Jing Liu, Hua Wu and Haifeng Wang",
booktitle = "In Proceedings of EMNLP"
year="2022"
}
如果您发现我们的调查对您的工作有用,请在基于验证的语言模型中引用以下纸张密集文本检索:一项调查
@article{DRSurvey,
title={Dense Text Retrieval based on Pretrained Language Models: A Survey},
author={Wayne Xin Zhao, Jing Liu, Ruiyang Ren, Ji-Rong Wen},
year={2022},
journal={arXiv preprint arXiv:2211.14876}
}
该存储库是在Apache-2.0许可下提供的。
有关使用RocketQA的帮助或问题,请提交GitHub问题。
有关其他沟通或合作,请联系Jing Liu([email protected])或扫描以下QR码。
