Fix: HF dataset PDF loading + stable RAG

rag.py

@@ -1,89 +1,133 @@
 # rag.py
 
-from datasets import load_dataset
+import os
+from typing import List, Tuple
+
+from huggingface_hub import hf_hub_download, list_repo_files
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.schema import Document
 from transformers import pipeline
 
-from config import HF_DATASET_REPO, EMBEDDING_MODEL, LLM_MODEL
+from config import (
+    HF_DATASET_REPO,
+    EMBEDDING_MODEL,
+    LLM_MODEL,
+    CHROMA_DIR,
+    CHUNK_SIZE,
+    CHUNK_OVERLAP,
+)
 
-# ------------------------
-# Load documents from HF Dataset
-# ------------------------
+# -----------------------------
+# Load PDFs from HF Dataset repo
+# -----------------------------
 def load_documents():
-    ds = load_dataset(HF_DATASET_REPO, split="train")
-
     docs = []
-    for row in ds:
-        text = row.get("text") or row.get("content")
-        if text and text.strip():
-            docs.append(Document(page_content=text))
+
+    try:
+        files = list_repo_files(
+            repo_id=HF_DATASET_REPO,
+            repo_type="dataset"
+        )
+    except Exception as e:
+        print("❌ Could not access dataset:", e)
+        return []
+
+    pdf_files = [f for f in files if f.lower().endswith(".pdf")]
+
+    if not pdf_files:
+        print("⚠️ No PDFs found in dataset")
+        return []
+
+    os.makedirs("kb", exist_ok=True)
+
+    for pdf in pdf_files:
+        local_path = hf_hub_download(
+            repo_id=HF_DATASET_REPO,
+            filename=pdf,
+            repo_type="dataset"
+        )
+
+        loader = PyPDFLoader(local_path)
+        docs.extend(loader.load())
 
     return docs
 
 
-# ------------------------
-# Build Vector DB (ONCE)
-# ------------------------
-documents = load_documents()
+# -----------------------------
+# Build vector DB (safe)
+# -----------------------------
+def build_vectorstore():
+    documents = load_documents()
 
-if not documents:
-    raise RuntimeError("No documents loaded from HF Dataset")
+    if not documents:
+        print("⚠️ No documents loaded, vector DB will be empty")
+        return None
 
-splitter = RecursiveCharacterTextSplitter(
-    chunk_size=500,
-    chunk_overlap=50,
-)
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=CHUNK_SIZE,
+        chunk_overlap=CHUNK_OVERLAP,
+    )
 
-chunks = splitter.split_documents(documents)
+    splits = splitter.split_documents(documents)
 
-embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
+    embeddings = HuggingFaceEmbeddings(
+        model_name=EMBEDDING_MODEL
+    )
+
+    vectordb = Chroma.from_documents(
+        documents=splits,
+        embedding=embeddings,
+        persist_directory=CHROMA_DIR
+    )
+
+    return vectordb
 
-vectordb = Chroma.from_documents(
-    documents=chunks,
-    embedding=embeddings,
-)
 
-retriever = vectordb.as_retriever(search_kwargs={"k": 3})
+# Build once at startup
+VECTOR_DB = build_vectorstore()
 
-# ------------------------
-# LLM (CPU SAFE)
-# ------------------------
-llm = pipeline(
+# -----------------------------
+# LLM (CPU-safe)
+# -----------------------------
+qa_pipeline = pipeline(
     "text2text-generation",
     model=LLM_MODEL,
-    max_new_tokens=256,
+    max_new_tokens=256
 )
 
-# ------------------------
-# RAG Query
-# ------------------------
-def ask_rag_with_status(question: str):
+
+# -----------------------------
+# Public API
+# -----------------------------
+def ask_rag_with_status(question: str) -> Tuple[str, List[str]]:
     status = []
 
-    status.append("🔎 Retrieving documents")
+    if VECTOR_DB is None:
+        return "No documents available.", ["Vector DB not initialized"]
+
+    retriever = VECTOR_DB.as_retriever(search_kwargs={"k": 3})
     docs = retriever.get_relevant_documents(question)
 
     if not docs:
-        return "No relevant documents found.", status
+        return "No relevant information found.", ["No matching chunks"]
 
     context = "\n\n".join(d.page_content for d in docs)
 
     prompt = f"""
-Answer the question using the context below.
+Answer the question using ONLY the context below.
 
 Context:
 {context}
 
 Question:
 {question}
-
-Answer:
 """
 
-    status.append("🧠 Generating answer")
-    result = llm(prompt)[0]["generated_text"]
+    result = qa_pipeline(prompt)[0]["generated_text"]
+
+    status.append(f"Retrieved {len(docs)} chunks")
+    status.append("Answer generated")
 
     return result.strip(), status