Fix: remove CHROMA_DIR, HF-dataset-based RAG

config.py

@@ -1,10 +1,10 @@
-import os
-
-KB_DIR = "kb"
-VECTOR_DB_DIR = "vector_db"
+# config.py
 
+MODEL_ID = "microsoft/Phi-3-mini-4k-instruct"
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
-LLM_MODEL = "microsoft/Phi-3-mini-4k-instruct"
 
-os.makedirs(VECTOR_DB_DIR, exist_ok=True)
-os.makedirs(KB_DIR, exist_ok=True)
+# Hugging Face Dataset repo where PDFs live
+HF_DATASET_REPO = "Zubaish/HubRAG-docs"
+
+# Retrieval
+TOP_K = 3

rag.py

@@ -1,69 +1,82 @@
 # rag.py
-from typing import List, Tuple
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain_community.vectorstores import Chroma
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from config import (
-    EMBEDDING_MODEL,
-    LLM_MODEL,
-    CHROMA_DIR,
-    TOP_K,
-)
-import torch
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain.schema import Document
+from datasets import load_dataset
 
+from config import MODEL_ID, EMBEDDING_MODEL, HF_DATASET_REPO, TOP_K
+
+
+# ----------------------------
+# Load PDFs from HF Dataset
+# ----------------------------
+def load_documents():
+    ds = load_dataset(HF_DATASET_REPO, split="train")
+
+    docs = []
+    for row in ds:
+        text = row.get("text", "").strip()
+        if text:
+            docs.append(Document(page_content=text))
+
+    return docs
 
-# --- Embeddings ---
-embeddings = HuggingFaceEmbeddings(
-    model_name=EMBEDDING_MODEL
-)
 
-# --- Vector DB (safe load) ---
-try:
-    vectordb = Chroma(
-        persist_directory=CHROMA_DIR,
-        embedding_function=embeddings,
-    )
-except Exception:
-    vectordb = None
+# ----------------------------
+# Build vector store (in-memory)
+# ----------------------------
+documents = load_documents()
 
+if not documents:
+    raise RuntimeError("No documents loaded from HF Dataset")
 
-# --- LLM ---
-tokenizer = AutoTokenizer.from_pretrained(
-    LLM_MODEL,
-    trust_remote_code=True
+embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
+
+vectordb = Chroma.from_documents(
+    documents=documents,
+    embedding=embeddings
 )
 
+
+# ----------------------------
+# Load LLM (NO device_map)
+# ----------------------------
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
 model = AutoModelForCausalLM.from_pretrained(
-    LLM_MODEL,
-    trust_remote_code=True,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto",
+    MODEL_ID,
+    torch_dtype="auto"
 )
 
+llm = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=256,
+    temperature=0.2
+)
 
-def ask_rag_with_status(question: str) -> Tuple[str, List[str]]:
-    status = []
 
-    if not vectordb:
-        return (
-            "⚠️ Knowledge base is not loaded yet. Upload documents first.",
-            ["Vector DB not initialized"],
-        )
+# ----------------------------
+# Public API
+# ----------------------------
+def ask_rag_with_status(question: str):
+    status = []
 
+    status.append("Retrieving relevant documents…")
     docs = vectordb.similarity_search(question, k=TOP_K)
 
     if not docs:
-        return (
-            "⚠️ I could not find relevant information in the knowledge base.",
-            ["No documents retrieved"],
-        )
+        return {
+            "answer": "No relevant documents found.",
+            "status": status
+        }
 
     context = "\n\n".join(d.page_content for d in docs)
-    status.append(f"Retrieved {len(docs)} chunks")
 
-    prompt = f"""
-You are a helpful assistant.
-Answer ONLY using the context below.
+    prompt = f"""Use the context below to answer the question.
 
 Context:
 {context}
@@ -71,20 +84,14 @@ Context:
 Question:
 {question}
 
-Answer:
-"""
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+Answer:"""
 
-    with torch.no_grad():
-        output = model.generate(
-            **inputs,
-            max_new_tokens=256,
-            do_sample=True,
-            temperature=0.7,
-        )
+    status.append("Generating answer…")
+    result = llm(prompt)[0]["generated_text"]
 
-    answer = tokenizer.decode(output[0], skip_special_tokens=True)
-    answer = answer.split("Answer:")[-1].strip()
+    answer = result.split("Answer:")[-1].strip()
 
-    return answer, status
+    return {
+        "answer": answer,
+        "status": status
+    }