Building a Document Q&A Chatbot with RAG

import json
import re
import urllib.request

from pypdf import PdfReader
from sentence_transformers import SentenceTransformer

URL = "http://localhost:11434/api/chat"
OPTIONS = {"temperature": 0, "seed": 42, "num_ctx": 4096, "num_predict": 150}

# 1. Ingest: split each PDF page into ~120-word chunks that remember their page number
pages = [" ".join((p.extract_text() or "").split()) for p in PdfReader("nist-ai-rmf.pdf").pages]
chunks = []
for number, text in enumerate(pages, 1):
    words = text.split()
    for start in range(0, len(words), 100):                      # 20 words of overlap
        chunks.append({"page": number, "text": " ".join(words[start:start + 120])})
embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
vectors = embedder.encode([c["text"] for c in chunks], normalize_embeddings=True)
print(f"{len(pages)} pages -> {len(chunks)} chunks")


def retrieve(query, k=3):
    scores = vectors @ embedder.encode([query], normalize_embeddings=True)[0]
    return [chunks[i] for i in scores.argsort()[::-1][:k]]


def chat(messages):
    body = {"model": "llama3.2", "messages": messages, "stream": False, "options": OPTIONS}
    req = urllib.request.Request(URL, json.dumps(body).encode(), {"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=900) as r:
        return json.load(r)["message"]["content"].strip()


def standalone(question, history):
    # 2. Rewrite a follow-up into a question that makes sense on its own
    if not history:
        return question
    recent = "\n".join(f"{role}: {text}" for role, text in history[-4:])
    return chat([{"role": "user", "content":
                  f"Conversation so far:\n{recent}\n\nRewrite this follow-up as one standalone question. "
                  f"Reply with the question only.\nFollow-up: {question}"}])


def answer(question, found, history):
    # the model sees the conversation in both runs; only the search query differs
    context = "\n\n".join(f"[page {c['page']}] {c['text']}" for c in found)
    past = [{"role": role, "content": text} for role, text in history]
    return chat([{"role": "system", "content":
                  "Answer from the excerpts only, in at most two sentences, and cite pages like (page 5). "
                  "If the excerpts do not contain the answer, reply: Not found in the document."},
                 *past,
                 {"role": "user", "content": f"Excerpts:\n{context}\n\nQuestion: {question}"}])


# 3. One conversation, run twice: retrieving with the raw question, and with the rewritten one
TURNS = [("What are the four functions of the AI RMF Core?", ["govern", "map", "measure", "manage"]),
         ("Which of them is described as cross-cutting?", ["govern"]),
         ("What characteristics of trustworthy AI does the framework list?", ["valid", "reliable", "fair"]),
         ("Which one is shown as the base for the others?", ["valid"]),
         ("What fines does the EU AI Act set for violations?", ["not found"])]

for mode in ("raw question", "rewritten question"):
    print(f"\n=== retrieval with the {mode}")
    history, right = [], 0
    for question, needed in TURNS:
        query = standalone(question, history) if mode == "rewritten question" else question
        found = retrieve(query)
        reply = answer(question, found, history)
        ok = all(term in reply.lower() for term in needed)
        right += ok
        history += [("user", question), ("assistant", reply)]
        print(f"Q: {question}")
        if query != question:
            print(f"   searched for: {query}")
        print(f"   pages {[c['page'] for c in found]} -> {' '.join(reply.split())[:170]}")
        print(f"   {'OK' if ok else 'MISSED'}")
    print(f"{right} of {len(TURNS)} answered as expected")

48 pages -> 186 chunks

=== retrieval with the raw question
Q: What are the four functions of the AI RMF Core?
   pages [25, 25, 4] -> The four functions of the AI RMF Core are GOVERN, MAP, MEASURE, and MANAGE (page 25). Governance is a cross-cutting function that informs and is infused throughout the ot
   OK
Q: Which of them is described as cross-cutting?
   pages [32, 5, 43] -> Governance is described as a cross-cutting function (page 5). It informs and is infused throughout the other three functions, namely GOVERN, MAP, MEASURE, and MANAGE.
   OK
Q: What characteristics of trustworthy AI does the framework list?
   pages [17, 5, 7] -> The framework lists the following characteristics of trustworthy AI systems: 1. Valid and reliable 2. Safe 3. Secure and resilient 4. Accountable and transparent 5. Expla
   OK
Q: Which one is shown as the base for the others?
   pages [32, 43, 26] -> The base for the others is "software or hardware systems" (page 26). It serves as a foundation for the other categories, such as data, and is also subject to similar cont
   MISSED
Q: What fines does the EU AI Act set for violations?
   pages [11, 21, 34] -> Not found in the document.
   OK
4 of 5 answered as expected

=== retrieval with the rewritten question
Q: What are the four functions of the AI RMF Core?
   pages [25, 25, 4] -> The four functions of the AI RMF Core are GOVERN, MAP, MEASURE, and MANAGE (page 25). Governance is a cross-cutting function that informs and is infused throughout the ot
   OK
Q: Which of them is described as cross-cutting?
   searched for: Is governance the cross-cutting function among GOVERN, MAP, MEASURE, and MANAGE?
   pages [5, 26, 41] -> Governance is described as a cross-cutting function (page 5). It informs and is infused throughout the other three functions.
   OK
Q: What characteristics of trustworthy AI does the framework list?
   searched for: What are the characteristics of trustworthy AI listed by the framework, specifically in relation to the four functions of the AI RMF Core (GOVERN, MAP, MEASURE, and MANAGE)?
   pages [7, 25, 17] -> The Framework lists the following characteristics of trustworthy AI systems: 1. Valid and reliable 2. Safe 3. Secure and resilient 4. Accountable and transparent 5. Expla
   OK
Q: Which one is shown as the base for the others?
   searched for: What characteristic of trustworthy AI systems is described as the base or foundation for the others, according to the framework?
   pages [5, 17, 17] -> Valid & Reliable is shown as the base for other trustworthiness characteristics (page 17).
   OK
Q: What fines does the EU AI Act set for violations?
   searched for: What are the specific fines imposed by the EU AI Act for non-compliance with its regulations?
   pages [47, 27, 7] -> Not found in the document.
   OK
5 of 5 answered as expected