Building a RAG Pipeline from Scratch with Python

from typing import List

def load_documents(paths: List[str]) -> List[str]:
    """Load plain-text documents from file paths."""
    docs = []
    for path in paths:
        with open(path, "r", encoding="utf-8") as f:
            docs.append(f.read())
    return docs


def chunk_text(text: str, chunk_size: int = 300, overlap: int = 50) -> List[str]:
    """Split text into overlapping word-level chunks."""
    words = text.split()
    chunks = []
    start = 0
    while start < len(words):
        end = start + chunk_size
        chunk = " ".join(words[start:end])
        chunks.append(chunk)
        start += chunk_size - overlap  # slide with overlap
    return chunks


def build_corpus(paths: List[str]) -> List[str]:
    docs = load_documents(paths)
    corpus = []
    for doc in docs:
        corpus.extend(chunk_text(doc))
    return corpus

from sentence_transformers import SentenceTransformer
import numpy as np

EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

def embed_corpus(corpus: List[str], model_name: str = EMBED_MODEL) -> np.ndarray:
    """
    Encode all corpus chunks into dense vectors.
    Returns array of shape (n_chunks, embedding_dim).
    """
    model = SentenceTransformer(model_name)
    print(f"Embedding {len(corpus)} chunks...")
    embeddings = model.encode(
        corpus,
        batch_size=64,
        show_progress_bar=True,
        normalize_embeddings=True  # L2-normalise for cosine similarity via dot product
    )
    return embeddings

import faiss

def build_index(embeddings: np.ndarray) -> faiss.IndexFlatIP:
    """
    Build a flat inner-product FAISS index.
    With L2-normalised vectors, inner product == cosine similarity.
    """
    dim = embeddings.shape[1]
    index = faiss.IndexFlatIP(dim)
    index.add(embeddings.astype(np.float32))
    print(f"Index built: {index.ntotal} vectors, dim={dim}")
    return index


def save_index(index: faiss.Index, path: str):
    faiss.write_index(index, path)

def load_index(path: str) -> faiss.Index:
    return faiss.read_index(path)

def retrieve(
    query: str,
    index: faiss.Index,
    corpus: List[str],
    model_name: str = EMBED_MODEL,
    top_k: int = 5
) -> List[str]:
    """
    Embed query, search FAISS index, return top-k text chunks.
    """
    model = SentenceTransformer(model_name)
    query_vec = model.encode(
        [query],
        normalize_embeddings=True
    ).astype(np.float32)

    scores, indices = index.search(query_vec, top_k)

    results = []
    for score, idx in zip(scores[0], indices[0]):
        results.append({
            "text": corpus[idx],
            "score": float(score)
        })
    return results

from transformers import pipeline

def build_prompt(query: str, context_chunks: List[dict]) -> str:
    context = "\n\n---\n\n".join(
        [f"[Passage {i+1}]\n{c['text']}" for i, c in enumerate(context_chunks)]
    )
    prompt = f"""You are a helpful assistant. Answer the question using ONLY the provided context.
If the answer is not in the context, say "I don't have enough information to answer this."

Context:
{context}

Question: {query}

Answer:"""
    return prompt


def generate_answer(prompt: str, model_name: str = "google/flan-t5-base") -> str:
    """
    Generate an answer using a local seq2seq model.
    Swap model_name for any HuggingFace-compatible model.
    """
    generator = pipeline("text2text-generation", model=model_name)
    output = generator(
        prompt,
        max_new_tokens=256,
        do_sample=False
    )
    return output[0]["generated_text"]

import os

# ── Config ──────────────────────────────────────────────────────────────────
DOCS_DIR    = "./documents"
INDEX_PATH  = "./rag_index.faiss"
CORPUS_PATH = "./corpus.txt"

# ── Build (run once) ─────────────────────────────────────────────────────────
doc_paths = [
    os.path.join(DOCS_DIR, f)
    for f in os.listdir(DOCS_DIR) if f.endswith(".txt")
]
corpus    = build_corpus(doc_paths)
embeddings = embed_corpus(corpus)
index     = build_index(embeddings)
save_index(index, INDEX_PATH)

# Save corpus for retrieval lookups
with open(CORPUS_PATH, "w") as f:
    for chunk in corpus:
        f.write(chunk + "\n<<<CHUNK_SEP>>>\n")

# ── Query (run anytime) ───────────────────────────────────────────────────────
index  = load_index(INDEX_PATH)
with open(CORPUS_PATH) as f:
    corpus = f.read().split("\n<<<CHUNK_SEP>>>\n")

query  = "What are the main challenges in retrieval-augmented generation?"
chunks = retrieve(query, index, corpus, top_k=4)
prompt = build_prompt(query, chunks)
answer = generate_answer(prompt)

print(f"\nQuery : {query}")
print(f"Answer: {answer}")
print("\nSources:")
for i, c in enumerate(chunks):
    print(f"  [{i+1}] score={c['score']:.3f} | {c['text'][:80]}...")