This project implements an extremely simple Retrieval Augmented Generation system. Using the Wikipedia page for Interstellar as a miniature document database, Spotify Annoy for fast embeddings and local indexing, and finally Cohere for text embedding and text generation, this project showcases a minimal question-answering pipeline.
It demonstrates qualitative improvements on similarity search and reranking, as well as hallucination risk mitigation.
Question → Embedding → Annoy → Similarity Search → Rerank → LLM → AnswerAlthough this is a toy project, its architecture presents some advantages for prototypes and small scale projects:
- If the document database is small and unchanging, locally indexing embeddings is a simple yet effective form of retrieval.
- Deployment of text generation models is expensive, requiring specific hardware for proper acceleration. Furthermore, open-source models tend to perform worse than those offered by AI companies. Therefore, using an API such as Cohere's presents a fast and scalable implementation for text generation.
In this section I qualitatively compare results from 3 steps of the RAG pipeline (similarity search, reranking and retrieval augmented generation), on a small set of questions.
Example 1: factual question (expand)
Query: How much did Interstellar make?
Similarity search: It was a commercial success, grossing $681 million worldwide during its initial theatrical run, and $769 million worldwide with subsequent releases, making it the 10th-highest-grossing film of 2014
Reranking: It was a commercial success, grossing $681 million worldwide during its initial theatrical run, and $769 million worldwide with subsequent releases, making it the 10th-highest-grossing film of 2014
RAG: According to the provided excerpts, Interstellar grossed $681 million worldwide during its initial theatrical run, and $769 million worldwide with subsequent releases.
Example 2: tricky question (expand)
Query: When did Interstellar production happen?
Similarity search: Interstellar premiered at the TCL Chinese Theatre on October 26, 2014, and was released in theaters in the United States on November 5, and in the United Kingdom on November 7
Reranking: Interstellar is a 2014 epic science fiction film directed by Christopher Nolan, who co-wrote the screenplay with his brother Jonathan Nolan
RAG: According to the provided excerpts, filming for Interstellar began in late 2013. This indicates that the production of the film took place in late 2013.
Example 3: out-of-context question (expand)
Query: How many people watched Interstellar on its opening weekend in Brazil?
Similarity search: It was a commercial success, grossing $681 million worldwide during its initial theatrical run, and $769 million worldwide with subsequent releases, making it the 10th-highest-grossing film of 2014
Reranking: Interstellar premiered at the TCL Chinese Theatre on October 26, 2014, and was released in theaters in the United States on November 5, and in the United Kingdom on November 7
RAG: The provided excerpts do not contain information about the number of people who watched Interstellar on its opening weekend in Brazil. Therefore, I don't have enough information to answer the question.
This project dependency is managed by blazingly fast uv.
Other package managers can install dependencies listed at pyproject.toml.
The following command will create a Python virtual environment at .venv and install all dependencies.
uv syncTo reproduce presented results:
uv run -m simplest_rag.mainThis project is inspired by the DeepLearning.ai Course "Large Language Models with Semantic Search". I adapted the foundation for: - building and querying a local embeds index - text embedding and reranking with Cohere
As of when I took the course, it was not up-to-date with Cohere API major version change; this project is compatible with Cohere v2 API. I also slightly extended the database by using the whole Wikipedia article.