Chat with your PDF (RAG)

pip install pypdf sentence-transformers numpy

ollama pull llama3.2      # about 2 GB
ollama serve              # usually already running after install

pip install reportlab

from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer

TEXT = [
    ("Nimbus Software - Employee Handbook", "Heading1"),
    ("Section 1. Leave policy", "Heading2"),
    ("Every full-time employee receives 24 days of paid leave per calendar year. "
     "Leave accrues at 2 days per completed month of service. Unused leave carries "
     "forward to the next year up to a maximum of 12 days; anything above that lapses "
     "on 31 December. Leave must be applied for at least 3 working days in advance, "
     "except for sick leave.", "BodyText"),
    ("Section 2. Sick leave and medical certificates", "Heading2"),
    ("Sick leave is separate from paid leave and is capped at 10 days per year. "
     "A medical certificate is required for any absence of more than 2 consecutive "
     "working days. Sick leave does not carry forward.", "BodyText"),
    ("Section 3. Working hours and remote work", "Heading2"),
    ("Core hours are 11:00 to 17:00 IST, during which every employee must be "
     "reachable. Employees may work remotely up to 8 days per month with manager "
     "approval. Fully remote arrangements require approval from the Head of "
     "Engineering and are reviewed every 6 months.", "BodyText"),
    ("Section 4. Expense reimbursement", "Heading2"),
    ("Claims must be submitted within 30 days of the expense date. Meals during "
     "domestic travel are reimbursed up to Rs 1200 per day. Airfare must be economy "
     "class unless the flight exceeds 6 hours. Reimbursement is paid with the "
     "following month's salary.", "BodyText"),
    ("Section 5. Notice period", "Heading2"),
    ("The notice period is 60 days for engineers and 30 days for all other roles. "
     "The company may waive part of the notice period at its discretion. Unused "
     "paid leave is encashed at the basic salary rate on the final working day.", "BodyText"),
]

styles = getSampleStyleSheet()
doc = SimpleDocTemplate("handbook.pdf", pagesize=A4)
flow = []
for text, style in TEXT:
    flow.append(Paragraph(text, styles[style]))
    flow.append(Spacer(1, 8))
doc.build(flow)
print("wrote handbook.pdf")

import json
import re
import sys
import urllib.error
import urllib.request

import numpy as np
from pypdf import PdfReader
from sentence_transformers import SentenceTransformer

PDF = "handbook.pdf"
EMBEDDER = "sentence-transformers/all-MiniLM-L6-v2"
LLM = "llama3.2"

# The ONE line you must change for your own PDF. There is no universal heading
# detector - open your document and describe what its headings actually look like.
HEADING = re.compile(r"^(Section \d+\.|Chapter \d+|\d+\.\d+\s)")

QUESTIONS = [
    "How many paid leave days do I get?",
    "Do I need a doctor's note if I am off sick?",
    "What is the parental leave policy?",
]


def read_lines(path):
    """[(page_number, line)] for every non-empty line. The page number becomes a citation."""
    lines = []
    for number, page in enumerate(PdfReader(path).pages, start=1):
        for line in (page.extract_text() or "").split("\n"):
            line = re.sub(r"[ \t]+", " ", line).strip()
            if line:
                lines.append((number, line))
    return lines


def split_sections(lines):
    """Group lines under the heading above them. Structure first, size second."""
    sections, heading, body, page = [], "(front matter)", [], 1

    def flush():
        if body:
            sections.append({"heading": heading, "page": page, "text": " ".join(body)})

    for number, line in lines:
        if HEADING.match(line):
            flush()
            heading, body, page = line, [], number
        else:
            body.append(line)
    flush()
    return sections


def chunk(sections, size=500, overlap=100):
    """Cut only the sections that are too long, and overlap so a split fact survives."""
    out = []
    for section in sections:
        text = section["text"]
        if len(text) <= size:
            out.append(dict(section))
            continue
        start = 0
        while start < len(text):
            out.append({**section, "text": text[start:start + size]})
            start += size - overlap
    return out


def build_index(chunks, model):
    # each chunk is embedded as "heading + body" so a loose sentence keeps its context
    texts = [f"{c['heading']} {c['text']}" for c in chunks]
    return model.encode(texts, normalize_embeddings=True, show_progress_bar=False)


def search(question, chunks, matrix, model, k=3):
    q = model.encode([question], normalize_embeddings=True)[0]
    scores = matrix @ q                     # both sides are unit length, so this is cosine
    return [(float(scores[i]), chunks[i]) for i in np.argsort(-scores)[:k]]


def build_prompt(question, hits):
    context = "\n\n".join(
        f"[page {c['page']} | {c['heading']}]\n{c['text']}" for _, c in hits
    )
    return (
        "You answer questions about a company handbook.\n"
        "Use ONLY the context below. If the answer is not in the context, reply "
        'exactly: "That is not in the handbook." Cite the page number you used.\n\n'
        f"CONTEXT:\n{context}\n\nQUESTION: {question}\nANSWER:"
    )


def ask_llm(prompt):
    payload = {"model": LLM, "prompt": prompt, "stream": False,
               "options": {"temperature": 0.0}}   # 0 keeps it as repeatable as it can be
    request = urllib.request.Request(
        "http://localhost:11434/api/generate",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(request, timeout=180) as response:
            return json.loads(response.read())["response"].strip()
    except urllib.error.HTTPError as exc:
        # a 404 here means the server is up but the model was never pulled
        return f"[ollama error {exc.code}: run `ollama pull {LLM}` first]"
    except urllib.error.URLError as exc:
        return f"[no ollama on port 11434: {exc.reason}]"


def main():
    retrieval_only = "--retrieval-only" in sys.argv

    model = SentenceTransformer(EMBEDDER)
    chunks = chunk(split_sections(read_lines(PDF)))
    matrix = build_index(chunks, model)
    print(f"{len(chunks)} chunks indexed, {matrix.shape[1]} numbers per chunk")

    for question in QUESTIONS:
        hits = search(question, chunks, matrix, model)
        print(f"\nQ: {question}")
        for score, c in hits:
            print(f"   {score:.2f}  page {c['page']}  {c['heading']}")
        if not retrieval_only:
            print("   A:", ask_llm(build_prompt(question, hits)))


if __name__ == "__main__":
    main()

python rag.py --retrieval-only

python rag.py

import hashlib, json
from pathlib import Path
import numpy as np
import rag
from sentence_transformers import SentenceTransformer

CACHE = Path("index.npz")

def load_or_build(pdf, model):
    stamp = hashlib.sha256(Path(pdf).read_bytes()).hexdigest()[:16]
    if CACHE.exists():
        blob = np.load(CACHE, allow_pickle=False)
        if str(blob["stamp"]) == stamp:              # same file, reuse the vectors
            return json.loads(str(blob["chunks"])), blob["matrix"]
    chunks = rag.chunk(rag.split_sections(rag.read_lines(pdf)))
    matrix = rag.build_index(chunks, model)
    np.savez(CACHE, stamp=np.array(stamp),
             chunks=np.array(json.dumps(chunks)), matrix=matrix)
    return chunks, matrix

model = SentenceTransformer(rag.EMBEDDER)
chunks, matrix = load_or_build("handbook.pdf", model)
print("chunks:", len(chunks), "  vectors:", matrix.shape)

ollama pull nomic-embed-text

import json, urllib.request
import numpy as np

def embed(texts):
    """Same job as SentenceTransformer.encode, using a model Ollama already holds."""
    payload = {"model": "nomic-embed-text", "input": texts}
    req = urllib.request.Request("http://localhost:11434/api/embed",
                                 data=json.dumps(payload).encode(),
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=300) as r:
        vectors = np.array(json.loads(r.read())["embeddings"], dtype="float32")
    return vectors / np.linalg.norm(vectors, axis=1, keepdims=True)

v = embed(["How many paid leave days do I get?",
           "Every full-time employee receives 24 days of paid leave per calendar year.",
           "Core hours are 11:00 to 17:00 IST."])
print("shape:", v.shape)
print("question vs leave sentence :", round(float(v[0] @ v[1]), 3))
print("question vs hours sentence :", round(float(v[0] @ v[2]), 3))

def recall_at_k(pairs, chunks, matrix, model, k=3):
    """pairs = [(question, heading_that_answers_it), ...]"""
    hits = 0
    for question, want in pairs:
        found = [c["heading"] for _, c in search(question, chunks, matrix, model, k)]
        hits += want in found
    return hits / len(pairs)