<jupyter_start><jupyter_text>Question-Answering using Simple Wikipedia IndexThis examples demonstrates the setup for Query / Question-Answer-Retrieval.You can input a query or a question. The script then uses semantic searchto find relevant passages in Simple English Wikipedia (as it is smaller and fits better in RAM).For semantic search, we use SentenceTransformer('msmarco-distilbert-base-v2') and retrieve100 potentially passages that answer the input query.Next, we use a more powerful CrossEncoder (cross_encoder = CrossEncoder('cross-encoder/ms-marco-TinyBERT-L-6')) thatscores the query and all retrieved passages for their relevancy. The cross-encoder is neccessary to filter out certain noisethat might be retrieved from the semantic search step.<jupyter_code>!pip install -U sentence-transformers rank_bm25
import json
from sentence_transformers import SentenceTransformer, CrossEncoder, util
import time
import gzip
import os
import torch

if not torch.cuda.is_available():
    print("Warning: No GPU found. Please add GPU to your notebook")

# We use the Bi-Encoder to encode all passages, so that we can use it with sematic search
model_name = 'msmarco-MiniLM-L-6-v3'
bi_encoder = SentenceTransformer(model_name)
top_k = 100  # Number of passages we want to retrieve with the bi-encoder

# The bi-encoder will retrieve 100 documents. We use a cross-encoder, to re-rank the results list to improve the quality
cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

# As dataset, we use Simple English Wikipedia. Compared to the full English wikipedia, it has only
# about 170k articles. We split these articles into paragraphs and encode them with the bi-encoder
wikipedia_filepath = 'simplewiki-2020-11-01.jsonl.gz'

if not os.path.exists(wikipedia_filepath):
    util.http_get('http://sbert.net/datasets/simplewiki-2020-11-01.jsonl.gz', wikipedia_filepath)

passages = []
with gzip.open(wikipedia_filepath, 'rt', encoding='utf8') as fIn:
    for line in fIn:
        data = json.loads(line.strip())
        passages.extend(data['paragraphs'])

print("Passages:", len(passages))

# To speed things up, pre-computed embeddings are downloaded.
# The provided file encoded the passages with the model 'msmarco-distilbert-base-v2'
if model_name == 'msmarco-MiniLM-L-6-v3':
    embeddings_filepath = f'simplewiki-2020-11-01-{model_name}.pt'
    if not os.path.exists(embeddings_filepath):
        util.http_get(f'http://sbert.net/datasets/simplewiki-2020-11-01-{model_name}.pt', embeddings_filepath)

    corpus_embeddings = torch.load(embeddings_filepath)
    if torch.cuda.is_available():
        corpus_embeddings = corpus_embeddings.to('cuda')
else:  # Here, we compute the corpus_embeddings from scratch (which can take a while depending on the GPU)
    corpus_embeddings = bi_encoder.encode(passages, convert_to_tensor=True, show_progress_bar=True)

# We also compare the results to lexical search (keyword search). Here, we use 
# the BM25 algorithm which is implemented in the rank_bm25 package.

from rank_bm25 import BM25Okapi
from sklearn.feature_extraction import stop_words
import string
from tqdm.autonotebook import tqdm
import numpy as np


# We lower case our text and remove stop-words from indexing
def bm25_tokenizer(text):
    tokenized_doc = []
    for token in text.lower().split():
        token = token.strip(string.punctuation)

        if len(token) > 0 and token not in stop_words.ENGLISH_STOP_WORDS:
            tokenized_doc.append(token)
    return tokenized_doc


tokenized_corpus = []
for passage in tqdm(passages):
    tokenized_corpus.append(bm25_tokenizer(passage))

bm25 = BM25Okapi(tokenized_corpus)

# This function will search all wikipedia articles for passages that
# answer the query
def search(query):
    print("Input question:", query)

    ##### BM25 search (lexical search) #####
    bm25_scores = bm25.get_scores(bm25_tokenizer(query))
    top_n = np.argpartition(bm25_scores, -5)[-5:]
    bm25_hits = [{'corpus_id': idx, 'score': bm25_scores[idx]} for idx in top_n]
    bm25_hits = sorted(bm25_hits, key=lambda x: x['score'], reverse=True)
    
    print("Top-3 lexical search (BM25) hits")
    for hit in bm25_hits[0:3]:
        print("\t{:.3f}\t{}".format(hit['score'], passages[hit['corpus_id']].replace("\n", " ")))

    ##### Sematic Search #####
    # Encode the query using the bi-encoder and find potentially relevant passages
    question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
    question_embedding = question_embedding.cuda()
    hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
    hits = hits[0]  # Get the hits for the first query

    ##### Re-Ranking #####
    # Now, score all retrieved passages with the cross_encoder
    cross_inp = [[query, passages[hit['corpus_id']]] for hit in hits]
    cross_scores = cross_encoder.predict(cross_inp)

    # Sort results by the cross-encoder scores
    for idx in range(len(cross_scores)):
        hits[idx]['cross-score'] = cross_scores[idx]

    # Output of top-5 hits from bi-encoder
    print("\n-------------------------\n")
    print("Top-3 Bi-Encoder Retrieval hits")
    hits = sorted(hits, key=lambda x: x['score'], reverse=True)
    for hit in hits[0:3]:
        print("\t{:.3f}\t{}".format(hit['score'], passages[hit['corpus_id']].replace("\n", " ")))

    # Output of top-5 hits from re-ranker
    print("\n-------------------------\n")
    print("Top-3 Cross-Encoder Re-ranker hits")
    hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
    for hit in hits[0:3]:
        print("\t{:.3f}\t{}".format(hit['cross-score'], passages[hit['corpus_id']].replace("\n", " ")))

search(query = "What is the capital of the United States?")
search(query = "What is the best orchestra in the world?")
search(query = "Number countries Europe")
search(query = "When did the cold war end?")
search(query = "How long do cats live?")
search(query = "How many people live in Toronto?")
search(query = "Oldest US president")
search(query = "Coldest place earth")
search(query = "Elon Musk year birth")
search(query = "Paris eiffel tower")
search(query = "Which US president was killed?")
search(query="When is Chinese New Year")<jupyter_output>Input question: When is Chinese New Year
Top-3 lexical search (BM25) hits
	18.606	Today in China the Gregorian calendar is used for most activities. At the same time, the Chinese calendar is still used for traditional Chinese holidays like Chinese New Year or Lunar New Year.
	18.151	Before that, the holiday was usually just called the "NewYear". Because the traditional Chinese calendar is mostly based on the changes in the moon, the Chinese New Year is also known in English as the "Lunar New Year" or "Chinese Lunar New Year". This name comes from "Luna", an old Latin name for the moon. The Indonesian name for the holiday is Imlek, which comes from the Hokkien word for the old Chinese calendar and is therefore also like saying "Lunar New Year".
	18.011	Spring Festival is the Chinese New Year.

-------------------------

Top-3 Bi-Encoder Retrieval hits
	0.739	Chinese New Year, known in China as the SpringFestival and in Singapore as the LunarNewYear, is a holiday on and around the new mo[...]
{'fraction_non_alphanumeric': 0.11098840419657648, 'fraction_numerical': 0.018636112644947543, 'mean_word_length': 4.219740634005763, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '148220', 'license': 'Apache-2.0', 'n_tokens_mistral': 2296, 'n_tokens_neox': 2162, 'n_words': 761, 'path': 'examples/applications/retrieve_rerank/retrieve_rerank_simple_wikipedia.ipynb', 'repo_name': 'mhmohona/sentence-transformers', 'size': '7244'}