File size: 1,684 Bytes
74b1bac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
"""
# Example: Retrieve from pre-built index of Natural Questions
This shows how to load an index built with BM25.index and saved with BM25.save, and retrieve
the top-k results for a set of queries from the Natural Questions dataset, via BEIR library.
To run this example, you need to install the following dependencies:
```bash
pip install beir bm25s[full]
```
To build an index, please refer to the `examples/index_nq.py` script. You
can run this script with:
```bash
python examples/index_nq.py
```
Then, run this script with:
```bash
python examples/retrieve_nq.py
```
"""
import beir.util
from beir.datasets.data_loader import GenericDataLoader
import Stemmer
import bm25s
from bm25s.utils.beir import BASE_URL
def main(index_dir="bm25s_indices/nq", data_dir="datasets", dataset="nq", mmap=True):
if mmap:
print("Using memory-mapped index (mmap) to reduce memory usage.")
# Load the queries from BEIR
data_path = beir.util.download_and_unzip(BASE_URL.format(dataset), data_dir)
loader = GenericDataLoader(data_folder=data_path)
loader._load_queries()
queries_lst = list(loader.queries.values())[:1000]
# Tokenize the queries
stemmer = Stemmer.Stemmer("english")
queries_tokenized = bm25s.tokenize(queries_lst, stemmer=stemmer)
# Load the BM25 index and retrieve the top-k results
retriever = bm25s.BM25.load(index_dir, mmap=mmap, load_corpus=True)
results = retriever.retrieve(queries_tokenized, k=20)
first_result = results.documents[0]
print(f"First score (# 1 result):{results.scores[0, 0]}")
print(f"First result (# 1 result):\n{first_result[0]}")
if __name__ == "__main__":
main() |