from datasets import load_dataset
from nltk.tokenize import sent_tokenize, word_tokenize

def load_wikipedia_corpus(corpus_size: int,
                          min_sentence_length: int=10,
                          max_sentence_length: int=100) -> list[str]:

	# load the Wikipedia dataset hosted on Hugging Face
    wikipedia_articles = load_dataset(
        path='wikipedia',
        name='20220301.en',
        split='train',
        streaming=True
    )

    corpus = []
    for article in wikipedia_articles:

		# tokenize the article into sentences
        sentences = sent_tokenize(article['text'])
        for sentence in sentences:

            # add sentence if it is between min and max length
            word_count = len(sentence.split())
            if word_count >= min_sentence_length and word_count <= max_sentence_length:
                corpus.append(sentence)

		# stop after corpus_size sentences
		if len(corpus) >= corpus_size:
			return corpus[:corpus_size]

from torch.nn.functional import cosine_similarity

"""
(1) Let's start by loading our corpus with 10,000 sentences from English Wikipedia.
"""
corpus = load_wikipedia_corpus(10_000)

"""
(2) The first step in our semantic search pipeline is to embed the corpus. We'll use our `generate_bert_embeddings` function to generate BERT embeddings in chunks of 100 documents at a time. Transformer models are memory-intensive due to their large number of parameters and have memory footprints that increase proportional to the input size. Processing in smaller batches will prevent out-of-memory errors.

Each batch of embeddings we generate will be a 2-dimensional tensor with shape (100, 384). The `torch.vstack` function allows us to then stack these tensors vertically to create a single 2-dimensional tensor with shape (10000, 384).
"""
batch_size = 100
corpus_embeddings = torch.vstack([
    generate_bert_embeddings(corpus[i:i+batch_size], tokenizer, model)
    for i in range(0, len(corpus), batch_size)
])

"""
(3) The next step in our pipeline is to define and embed our query. Let's write out a query and then embed it with the same routine that we used to embed the corpus. Since the `generate_bert_embeddings` function normally expects a list of documents and we only have a single document here, we'll wrap our query in a list.
"""
query = 'Ayn Rand is a Russian-American writer and phisopher'
query_embedding = generate_bert_embeddings([query], tokenizer, model)

"""
(4) Now we can compute the cosine similarity between the query embedding and each document embedding in the corpus. To do so, we'll use the handy `cosine_similarity` function from PyTorch that computes the cosine similarity between two tensors. The result is a 1-dimensional tensor with shape (10000,) containing the cosine similarity between the query and each of the 10,000 documents in the corpus.
"""
similarities = cosine_similarity(corpus_embeddings, query_embedding)

"""
(5) Finally, since greater similarity scores indicate a better match, we can sort the similarities in descending order and print the top 5 results. To do so, we use the PyTorch tensor's `argsort` method, which returns the indices of the sorted similarities.
"""
top_k = 5
top_k_indices = similarities.argsort(descending=True)[:top_k]
for index in top_k_indices:
    print(f'Index: {index}')
    print(f'Similarity: {similarities[index]:.3f}')
    print(f'Document: {corpus[index]}\n')

# select GPU if available
if torch.cuda.is_available():
    device = torch.device('cuda')

# otherwise use CPU
else:
  device = torch.device('cpu')

# move the BERT model to the GPU
model.to(device)

# set the default device type for tensors to the GPU
torch.set_default_device(device)