bm25s_quickstart_tokenize_index_and_retrieve_documents.py

1import bm25s
2import stemmer
3
4# Create a sample corpus and tokenize it
5corpus = [
6    "a cat is a feline",
7    "a dog is a canine",
8    "a bird is an avian",
9]
10
11# Tokenize the corpus and only keep the tokens
12corpus_tokens = bm25s.tokenize(corpus, stopwords="en")
13
14# Create the BM25 model and index the corpus
15retriever = bm25s.BM25()
16retriever.index(corpus_tokens)
17
18# Query the corpus
19query = "does the fish fly?"
20query_tokens = bm25s.tokenize(query, stopwords="en")
21
22# Get top-k results as a tuple of (indices, scores)
23results, scores = retriever.retrieve(query_tokens, corpus=corpus, k=2)
24
25for i in range(results.shape[1]):
26    print(f"Rank {i+1}: {results[0, i]} (score: {scores[0, i]:.4f})")