AmenRa / retriv

A Python Search Engine for Humans 🥸
MIT License
174 stars 20 forks source link
bm25 dense-retrieval hybrid-retrieval information-retrieval numba search search-engine search-engine-optimization semantic-search sparse-retrieval tf-idf

PyPI version License: MIT

🔥 News

⚡️ Introduction

retriv is a user-friendly and efficient search engine implemented in Python supporting Sparse (traditional search with BM25, TF-IDF), Dense (semantic search) and Hybrid retrieval (a mix of Sparse and Dense Retrieval). It allows you to build a search engine in a single line of code.

retriv is built upon Numba for high-speed vector operations and automatic parallelization, PyTorch and Transformers for easy access and usage of Transformer-based Language Models, and Faiss for approximate nearest neighbor search. In addition, it provides automatic tuning functionalities to allow you to tune its internal components with minimal intervention.

✨ Main Features

Retrievers

Unified Search Interface

All the supported retrievers share the same search interface:

AutoTune

retriv automatically tunes Faiss configuration for approximate nearest neighbors search by leveraging AutoFaiss to guarantee 10ms response time based on your available hardware. Moreover, it offers an automatic tuning functionality for BM25's parameters, which require minimal user intervention. Under the hood, retriv leverages Optuna, a hyperparameter optimization framework, and ranx, an Information Retrieval evaluation library, to test several parameter configurations for BM25 and choose the best one. Finally, it can automatically balance the importance of lexical and semantic relevance scores computed by the Hybrid Retriever to maximize retrieval effectiveness.

📚 Documentation

🔌 Requirements

python>=3.8

💾 Installation

pip install retriv

💡 Minimal Working Example

# Note: SearchEngine is an alias for the 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")

Output:

[
  {
    "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
  }
]

🎁 Feature Requests

Would you like to see other features implemented? Please, open a feature request.

🤘 Want to contribute?

Would you like to contribute? Please, drop me an e-mail.

📄 License

retriv is an open-sourced software licensed under the MIT license.