RAG in 85 lines of code

Learn how to implement a RAG pipeline in just 85 lines of code.

There's a lot of theory and context to think about when it comes to RAG, but let's start with a quick implementation in code to motivate what follows. The following 85 lines do the following:

load some data (a fictional dataset about 'ZenML World') as our corpus

process that text (split it into chunks and 'tokenize' it (i.e. split into words))

take a query as input and find the most relevant chunks of text from our corpus data

use OpenAI's GPT-3.5 model to answer the question based on the relevant chunks

import os

import re

import string

from openai import OpenAI

def preprocess_text(text):

text = text.lower()

text = text.translate(str.maketrans("", "", string.punctuation))

text = re.sub(r"\s+", " ", text).strip()

return text

def tokenize(text):

return preprocess_text(text).split()

def retrieve_relevant_chunks(query, corpus, top_n=2):

query_tokens = set(tokenize(query))

similarities = []

for chunk in corpus:

chunk_tokens = set(tokenize(chunk))

similarity = len(query_tokens.intersection(chunk_tokens)) / len(

query_tokens.union(chunk_tokens)

similarities.append((chunk, similarity))

similarities.sort(key=lambda x: x[1], reverse=True)

return [chunk for chunk, _ in similarities[:top_n]]

def answer_question(query, corpus, top_n=2):

relevant_chunks = retrieve_relevant_chunks(query, corpus, top_n)

if not relevant_chunks:

return "I don't have enough information to answer the question."

context = "\n".join(relevant_chunks)

client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

chat_completion = client.chat.completions.create(

messages=[

"role": "system",

"content": f"Based on the provided context, answer the following question: {query}\n\nContext:\n{context}",

},

"role": "user",

"content": query,

},

],

model="gpt-3.5-turbo",

return chat_completion.choices[0].message.content.strip()

# Sci-fi themed corpus about "ZenML World"

corpus = [