DEV Community

Cover image for Build a RAG System in Google Colab Before Your Next AI Interview
PandeyC
PandeyC

Posted on

Build a RAG System in Google Colab Before Your Next AI Interview

You don't need to study machine learning for six months to understand RAG.

You don't need Kubernetes.

You don't need a vector database.

You don't even need an OpenAI API key.

You need:

  • Google Colab
  • Python
  • A few paragraphs of text
  • About 20 minutes

By the end of this article, you will have built the core retrieval mechanism behind a RAG system.

More importantly, you will understand what is actually happening when someone says:

"We convert documents into embeddings, store them in a vector database, retrieve relevant chunks, and send them to an LLM."

Let's build it.

What Are We Building?

Imagine that we have a small knowledge base about AWS services.

A user asks:

"Which AWS service should I use to decouple applications?"

Our system will:

  1. Break documents into chunks.
  2. Convert the chunks into embeddings.
  3. Convert the question into an embedding.
  4. Compare the question with every chunk.
  5. Retrieve the most relevant chunks.

That is the retrieval part of Retrieval-Augmented Generation.

The architecture looks like this:

Documents
    ↓
Chunking
    ↓
Embeddings
    ↓
Vector Store

User Question
    ↓
Question Embedding
    ↓
Similarity Search
    ↓
Relevant Chunks
    ↓
LLM
    ↓
Answer
Enter fullscreen mode Exit fullscreen mode

In a production application, you might use Pinecone, OpenSearch, pgvector, Qdrant, or another vector database.

For learning, we don't need any of them.

We will store our embeddings in memory and use cosine similarity.

Step 1: Open Google Colab

Create a new Google Colab notebook.

Run:

!pip install -q sentence-transformers
Enter fullscreen mode Exit fullscreen mode

We will use the sentence-transformers library to generate embeddings.

No API key is required.

Step 2: Create Our Knowledge Base

Let's create a tiny collection of documents.

documents = [
    """
    Amazon S3 is an object storage service designed for storing
    and retrieving files. It provides high durability and is
    commonly used for backups, static websites, data lakes,
    and application assets.
    """,

    """
    Amazon SQS is a managed message queue service.

    It allows applications to communicate asynchronously.

    Producers send messages to a queue and consumers process
    those messages independently.

    SQS is commonly used to decouple distributed applications.
    """,

    """
    AWS Lambda is a serverless compute service.

    Developers upload code and AWS executes the code in response
    to events.

    Lambda automatically manages servers and scales applications
    based on incoming requests.
    """,

    """
    Amazon DynamoDB is a managed NoSQL database.

    It provides low-latency access to data and automatically
    scales to handle large workloads.
    """
]
Enter fullscreen mode Exit fullscreen mode

Obviously, real RAG systems contain thousands or millions of documents.

But the mechanism is exactly the same.

Step 3: Chunk the Documents

Why do we need chunks?

Because embedding an entire book or a 200-page PDF as one vector would produce a poor representation for individual questions.

Instead, RAG systems divide documents into smaller pieces.

Let's write a very simple chunker.

def chunk_text(text, chunk_size=250):
    words = text.split()

    chunks = []

    for i in range(0, len(words), chunk_size):
        chunk = " ".join(words[i:i + chunk_size])
        chunks.append(chunk)

    return chunks


chunks = []

for document in documents:
    chunks.extend(chunk_text(document))


print("Number of chunks:", len(chunks))

for index, chunk in enumerate(chunks):
    print(f"\nCHUNK {index}")
    print(chunk)
Enter fullscreen mode Exit fullscreen mode

Production systems usually use more sophisticated strategies.

For example:

  • token-based chunking
  • overlapping chunks
  • recursive text splitting
  • semantic chunking

But our goal is to understand the mechanism first.

Step 4: Generate Embeddings

Now we need to convert text into vectors.

We will use a small embedding model.

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")
Enter fullscreen mode Exit fullscreen mode

Now generate embeddings for our chunks.

chunk_embeddings = model.encode(chunks)

print(chunk_embeddings.shape)
Enter fullscreen mode Exit fullscreen mode

You should see something similar to:

(4, 384)
Enter fullscreen mode Exit fullscreen mode

What does that mean?

We have four chunks.

Each chunk has been converted into a vector containing 384 numbers.

Something like this:

[0.023, -0.041, 0.087, ...]
Enter fullscreen mode Exit fullscreen mode

Humans cannot interpret these numbers directly.

But mathematically, texts with similar meanings tend to have vectors that are closer together.

That is what makes semantic search possible.

Step 5: Ask a Question

Let's ask:

question = "Which AWS service can help decouple applications?"
Enter fullscreen mode Exit fullscreen mode

Convert the question into an embedding.

question_embedding = model.encode([question])
Enter fullscreen mode Exit fullscreen mode

Now we have:

Documents → Vectors

Question → Vector
Enter fullscreen mode Exit fullscreen mode

We need to find which document vectors are closest to the question vector.

Step 6: Perform Similarity Search

We will use cosine similarity.

from sklearn.metrics.pairwise import cosine_similarity

scores = cosine_similarity(
    question_embedding,
    chunk_embeddings
)[0]
Enter fullscreen mode Exit fullscreen mode

Let's inspect the results.

for index, score in enumerate(scores):
    print(round(score, 3), chunks[index][:100])
Enter fullscreen mode Exit fullscreen mode

You should see the SQS document receive the highest similarity score.

Now retrieve the best chunk.

best_chunk_index = scores.argmax()

retrieved_chunk = chunks[best_chunk_index]

print(retrieved_chunk)
Enter fullscreen mode Exit fullscreen mode

Congratulations.

You just built semantic retrieval.

This is the foundation of a RAG system.

Let's Make It Reusable

Let's wrap everything into a function.

def search(question, top_k=2):

    question_embedding = model.encode([question])

    scores = cosine_similarity(
        question_embedding,
        chunk_embeddings
    )[0]

    top_indices = scores.argsort()[::-1][:top_k]

    results = []

    for index in top_indices:
        results.append({
            "score": float(scores[index]),
            "text": chunks[index]
        })

    return results
Enter fullscreen mode Exit fullscreen mode

Now try asking different questions.

questions = [
    "Where should I store application files?",
    "How can I run code without managing servers?",
    "Which database provides low latency access?",
    "How can microservices communicate asynchronously?"
]

for question in questions:

    print("\nQUESTION:", question)

    results = search(question)

    for result in results:
        print(
            round(result["score"], 3),
            result["text"][:150]
        )
Enter fullscreen mode Exit fullscreen mode

You now have a tiny semantic search engine.

But Where Is the LLM?

This is where many developers misunderstand RAG.

The vector database does not answer the question.

The embedding model does not answer the question.

The retrieval system finds relevant information.

The retrieved information is then placed into the prompt sent to the LLM.

Conceptually:

def build_prompt(question, retrieved_chunks):

    context = "\n\n".join(retrieved_chunks)

    prompt = f"""
    Answer the question using only the context below.

    CONTEXT:

    {context}

    QUESTION:

    {question}
    """

    return prompt
Enter fullscreen mode Exit fullscreen mode

Let's generate the prompt.

question = "Which AWS service should I use to decouple applications?"

results = search(question)

retrieved_chunks = [
    result["text"]
    for result in results
]

prompt = build_prompt(
    question,
    retrieved_chunks
)

print(prompt)
Enter fullscreen mode Exit fullscreen mode

This prompt could now be sent to an LLM.

That completes the RAG pipeline:

Documents
    ↓
Chunks
    ↓
Embeddings
    ↓
Vector Search
    ↓
Relevant Context
    ↓
Prompt
    ↓
LLM
    ↓
Answer
Enter fullscreen mode Exit fullscreen mode

Five Experiments You Should Try

Don't stop after running the notebook.

Change it.

First, add documents that use similar terminology.

For example:

SQS decouples applications.

EventBridge connects applications using events.

SNS distributes messages to multiple subscribers.
Enter fullscreen mode Exit fullscreen mode

Does the correct document still rank first?

Second, change the chunk size.

Try:

chunk_size = 50
Enter fullscreen mode Exit fullscreen mode

Then:

chunk_size = 500
Enter fullscreen mode Exit fullscreen mode

What happens to retrieval quality?

Third, retrieve more documents.

Change:

top_k = 1
Enter fullscreen mode Exit fullscreen mode

to:

top_k = 3
Enter fullscreen mode Exit fullscreen mode

Would sending more context always produce a better answer?

Fourth, add irrelevant documents.

Does retrieval quality change?

Fifth, ask ambiguous questions.

For example:

Which service should I use for messaging?
Enter fullscreen mode Exit fullscreen mode

Now you are starting to encounter the problems that real RAG systems must solve.

The Interview Questions Hidden Inside This Project

If you understand the notebook above, you should be able to discuss several common AI interview questions.

Why do RAG systems chunk documents?

Because retrieval usually needs to identify specific passages rather than entire documents.

What is an embedding?

A numerical representation of data that allows semantic relationships to be compared mathematically.

Why use cosine similarity?

Because it measures the direction between vectors and is commonly used to compare embedding similarity.

Does a vector database generate answers?

No.

It stores vectors and helps retrieve relevant information.

The LLM generates the answer.

Does retrieving more chunks always improve the answer?

No.

More context can increase cost, introduce irrelevant information, and make it harder for the model to identify the useful evidence.

What happens when the correct document isn't retrieved?

The LLM never receives the information required to answer correctly.

This is why evaluating retrieval quality is one of the most important parts of building a production RAG system.

The 20-Line Demo Is the Easy Part

You can build a RAG demo in twenty minutes.

Production RAG systems are harder.

You have to make decisions about:

  • chunk size and overlap
  • embedding models
  • metadata filtering
  • top-k retrieval
  • similarity thresholds
  • hybrid search
  • reranking
  • context-window management
  • hallucination control
  • retrieval evaluation
  • answer evaluation

And interviewers increasingly ask about these trade-offs.

Not because they expect you to memorize definitions.

They want to know whether you understand how the system behaves.

Want to Go Deeper?

If you ran this notebook and found yourself asking:

How should I choose chunk size?

When should I use a vector database?

How do I know whether retrieval is actually working?

What is the difference between semantic search, hybrid search, and reranking?

How would I design this system for thousands or millions of documents?

Those are exactly the questions worth exploring next.

I'm covering embeddings, vector search, RAG architecture, and the engineering decisions behind production AI systems on ConfidentPrep.

You can explore the deeper RAG learning material here:

https://confidentprep.com/courses/ai-ml-for-interview/3-embeddings-and-rag/

And if you'd rather learn by discussing these concepts live, asking questions, and working through interview scenarios with other developers, join one of the upcoming live sessions:

https://confidentprep.com/live/

Build the small version first.

Understand why it works.

Then learn how to explain the trade-offs.

That's a much better way to prepare for an AI interview than memorizing another list of 100 RAG questions.

Top comments (0)