retriv

🔥 新闻

[2023年8月23日] retriv 0.2.2 发布！
此版本新增了对多字段文档和过滤器的实验性支持。请参阅高级检索器文档。
[2023年2月18日] retriv 0.2.0 发布！
此版本新增了对密集检索和混合检索的支持。密集检索利用查询和文档向量表示的语义相似性，这些向量可以直接由 retriv 计算或从其他来源导入。混合检索将传统检索（非正式称为稀疏检索）和密集检索的结果结合，进一步提高检索效果。由于库几乎完全重写，不再支持之前版本构建的索引。

⚡️ 简介

retriv 是一个用户友好且高效的搜索引擎，使用 Python 实现，支持稀疏检索（基于 BM25、TF-IDF 的传统搜索）、密集检索（语义搜索）和混合检索（稀疏检索和密集检索的结合）。它允许您通过__单行代码__构建搜索引擎。

retriv 基于 Numba 实现高速向量运算和自动并行化，利用 PyTorch 和 Transformers 实现对基于 Transformer 的语言模型的便捷访问和使用，并使用 Faiss 进行近似最近邻搜索。此外，它还提供自动调优功能，让您能以最少的干预调整其内部组件。

✨ 主要特性

检索器

稀疏检索器：基于词法匹配的标准搜索器。 retriv 将 BM25 作为其主要检索模型。为了教育目的，也支持 TF-IDF。稀疏检索器配备了多种词干提取器、分词器和停用词列表，支持多种语言。点击此处了解更多。
密集检索器：执行语义搜索的检索模型。点击此处了解更多。
混合检索器：基于稀疏检索器和密集检索器构建的检索模型。点击此处了解更多。
高级检索器：支持过滤器的高级稀疏检索器。这是一项实验性功能。点击此处了解更多。

统一搜索接口

所有支持的检索器共享相同的搜索接口：

search：标准搜索功能，即您对搜索引擎的期望。
msearch：一次性计算多个查询的结果。它尽可能利用自动并行化。
bsearch：与msearch类似，但自动生成要评估的查询批次，并允许以JSONl格式动态将搜索结果写入磁盘。bsearch适用于计算数十万甚至数百万个查询的结果，而不会占用大量内存。预计算的结果可用于信息检索中神经网络模型训练期间的负采样。

自动调优

retriv通过利用AutoFaiss自动调优Faiss配置，以根据您可用的硬件保证10毫秒的响应时间。此外，它还为BM25参数提供自动调优功能，只需最少的用户干预。在后台，retriv利用Optuna（一个超参数优化框架）和ranx（一个信息检索评估库）来测试BM25的多种参数配置并选择最佳方案。最后，它可以自动平衡混合检索器计算的词法和语义相关性得分的重要性，以最大化检索效果。

📚 文档

🔌 要求

python>=3.8

💾 安装

pip install retriv

💡 最小工作示例

# 注意：SearchEngine是SparseRetriever的别名
from retriv import SearchEngine

collection = [
  {"id": "doc_1", "text": "Generals gathered in their masses"},
  {"id": "doc_2", "text": "Just like witches at black masses"},
  {"id": "doc_3", "text": "Evil minds that plot destruction"},
  {"id": "doc_4", "text": "Sorcerer of death's construction"},
]

se = SearchEngine("new-index").index(collection)

se.search("witches masses")

输出：

[
  {
    "id": "doc_2",
    "text": "Just like witches at black masses",
    "score": 1.7536403
  },
  {
    "id": "doc_1",
    "text": "Generals gathered in their masses",
    "score": 0.6931472
  }
]