ratatouille-intro.py

from ragatouille import RAGPretrainedModel
from ragatouille.utils import get_wikipedia_page


topics = ["Barack_Obama", "Nelson_Mandela", "king"]

RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")

my_documents = []

# get documents for each topic
for topic in topics:
    my_documents.append(get_wikipedia_page(topic))

# create index
index_path = RAG.index(
    index_name="leaders",
    collection=my_documents,
    document_ids=topics,
    max_document_length=180, 
    split_documents=True
)

query = "Who was president of the united states of america in 2010 ?"

#retrieve relevant chunks from index
results = RAG.search(query, k=3)

print(results)