• Loads and parses a PDF file
• Chunks the text and creates embeddings
• Finds the most relevant chunks for any question
• Generates answers grounded in the document content
• Cites which section the answer comes from

pip install openai python-dotenv pypdf numpy

import re
from pypdf import PdfReader

def load_pdf(path: str) -> list[dict]:
    """Load a PDF and return pages as a list of {page, text} dicts."""
    reader = PdfReader(path)
    pages = []
    for i, page in enumerate(reader.pages):
        text = page.extract_text() or ""
        text = re.sub(r"\s+", " ", text).strip()
        if text:
            pages.append({"page": i + 1, "text": text})
    return pages

def chunk_text(pages: list[dict], chunk_size: int = 500, overlap: int = 50) -> list[dict]:
    """Split page text into overlapping chunks."""
    chunks = []
    for page in pages:
        words = page["text"].split()
        for i in range(0, len(words), chunk_size - overlap):
            chunk = " ".join(words[i:i + chunk_size])
            if chunk.strip():
                chunks.append({"page": page["page"], "text": chunk})
    return chunks

import numpy as np
from openai import OpenAI

client = OpenAI()

def get_embeddings(texts: list[str]) -> list[list[float]]:
    """Embed texts in batches — the API accepts up to 2048 inputs per call."""
    embeddings = []
    for i in range(0, len(texts), 2048):
        response = client.embeddings.create(
            model="text-embedding-3-small",
            input=texts[i:i + 2048]
        )
        embeddings.extend(item.embedding for item in response.data)
    return embeddings

def cosine_similarity(a: list, b: list) -> float:
    a, b = np.array(a), np.array(b)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

def find_relevant_chunks(query: str, chunks: list[dict],
                         embeddings: list[list[float]], top_k: int = 4) -> list[dict]:
    query_embedding = get_embeddings([query])[0]
    scored = [
        {"chunk": c, "score": cosine_similarity(query_embedding, e)}
        for c, e in zip(chunks, embeddings)
    ]
    scored.sort(key=lambda x: x["score"], reverse=True)
    return [s["chunk"] for s in scored[:top_k]]

def answer_question(question: str, relevant_chunks: list[dict]) -> str:
    context_parts = [
        f"[Page {c['page']}]: {c['text']}"
        for c in relevant_chunks
    ]
    context = "\n\n".join(context_parts)

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "system",
                "content": """Answer the question using only the provided context.
Always cite which page your answer comes from (e.g., "According to page 3...").
If the answer is not in the context, say: "I couldn't find that in the document." """
            },
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}
        ],
        temperature=0.2,
        max_tokens=400,
    )
    return response.choices[0].message.content

import sys

def main():
    if len(sys.argv) < 2:
        print("Usage: python qa.py document.pdf")
        sys.exit(1)

    pdf_path = sys.argv[1]
    print(f"Loading {pdf_path}...")

    pages = load_pdf(pdf_path)
    chunks = chunk_text(pages)
    print(f"Created {len(chunks)} chunks. Embedding...")

    # Embed all chunks upfront
    texts = [c["text"] for c in chunks]
    embeddings = get_embeddings(texts)
    print(f"Ready! Ask questions about the document (type 'quit' to exit)\n")

    while True:
        question = input("Question: ").strip()
        if question.lower() in ("quit", "exit"):
            break
        if not question:
            continue

        relevant = find_relevant_chunks(question, chunks, embeddings)
        answer = answer_question(question, relevant)
        print(f"\nAnswer: {answer}\n")

if __name__ == "__main__":
    main()

• [ ] PDF loads and chunks correctly
• [ ] Embeddings created for all chunks at startup
• [ ] Relevant chunks found via cosine similarity
• [ ] Answers cite page numbers
• [ ] "Not found" response when answer isn't in document