retriv: 一款面向人类的Python搜索引擎

retriv

引言

在当今数字时代,高效的信息检索已成为不可或缺的需求。为了满足这一需求,AmenRa开发了retriv - 一款面向人类的Python搜索引擎。retriv不仅功能强大,而且操作简单,旨在为用户提供一种直观且高效的搜索解决方案。

retriv的核心特性

retriv是一个用Python实现的搜索引擎,具有以下几个主要特点:

多种检索模式支持:
- 稀疏检索:传统的基于词法匹配的搜索,包括BM25和TF-IDF算法。
- 密集检索:基于语义相似度的搜索。
- 混合检索:结合稀疏检索和密集检索的优势。
简单易用:只需一行代码即可构建搜索引擎。
高性能:基于Numba实现高速向量运算和自动并行化。
先进技术支持:集成了PyTorch、Transformers等先进的机器学习库。
自动调优:提供自动调优功能,可以优化内部组件以提高检索效果。

技术原理

retriv的强大功能源于其先进的技术架构:

Numba:用于高速向量运算和自动并行化,提升检索效率。
PyTorch和Transformers:便于访问和使用基于Transformer的语言模型。
Faiss:用于近似最近邻搜索,提高大规模数据集的检索速度。

这些技术的结合使retriv能够在保持高效率的同时,提供精准的搜索结果。

retriv流程图

主要功能

1. 检索器

retriv提供了多种检索器,以适应不同的搜索需求:

稀疏检索器:基于词法匹配的标准搜索器,主要实现了BM25算法,同时也支持TF-IDF。它配备了多种词干提取器、分词器和停用词列表,支持多种语言。
密集检索器:执行语义搜索的检索模型,能够捕捉查询和文档之间的语义相似性。
混合检索器:结合稀疏检索和密集检索的优势,进一步提高检索效果。
高级检索器:支持过滤功能的实验性稀疏检索器。

2. 统一搜索接口

所有检索器都共享相同的搜索接口,包括:

search:标准搜索功能。
msearch:同时处理多个查询,利用自动并行化提高效率。
bsearch:批量搜索功能,可以处理大量查询并将结果动态写入磁盘。

3. 自动调优

retriv提供了多项自动调优功能:

自动调整Faiss配置,以保证10ms的响应时间。
自动优化BM25参数,只需最少的用户干预。
自动平衡混合检索器中词法和语义相关性分数的重要性。

retriv自动调优

使用示例

retriv的使用非常简单,以下是一个最小工作示例:

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)

results = se.search("witches masses")