Fix: use existing HF dataset hubrag-kb

rag.py

@@ -1,82 +1,57 @@
-# rag.py
-
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from datasets import load_dataset
+from langchain.schema import Document
 from langchain_community.vectorstores import Chroma
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain.schema import Document
-from datasets import load_dataset
+from transformers import pipeline
 
-from config import MODEL_ID, EMBEDDING_MODEL, HF_DATASET_REPO, TOP_K
+HF_DATASET_REPO = "Zubaish/hubrag-kb"
+EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+CHROMA_DIR = "./chroma"
 
-
-# ----------------------------
-# Load PDFs from HF Dataset
-# ----------------------------
 def load_documents():
+    docs = []
     ds = load_dataset(HF_DATASET_REPO, split="train")
 
-    docs = []
     for row in ds:
-        text = row.get("text", "").strip()
-        if text:
+        text = row.get("text")
+        if text and text.strip():
             docs.append(Document(page_content=text))
 
     return docs
 
-
-# ----------------------------
-# Build vector store (in-memory)
-# ----------------------------
 documents = load_documents()
 
 if not documents:
-    raise RuntimeError("No documents loaded from HF Dataset")
+    print("⚠️ No text documents found in dataset. PDFs must be converted to text.")
 
-embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
+embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL)
 
-vectordb = Chroma.from_documents(
-    documents=documents,
-    embedding=embeddings
-)
-
-
-# ----------------------------
-# Load LLM (NO device_map)
-# ----------------------------
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype="auto"
-)
+vectordb = None
+if documents:
+    vectordb = Chroma.from_documents(
+        documents,
+        embedding=embeddings,
+        persist_directory=CHROMA_DIR
+    )
 
 llm = pipeline(
     "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=256,
-    temperature=0.2
+    model="microsoft/Phi-3-mini-4k-instruct",
+    trust_remote_code=True,
+    max_new_tokens=256
 )
 
-
-# ----------------------------
-# 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:
+    if not vectordb:
         return {
-            "answer": "No relevant documents found.",
-            "status": status
+            "answer": "Knowledge base is empty. Please upload text documents to the dataset.",
+            "status": ["No text documents loaded"]
         }
 
+    docs = vectordb.similarity_search(question, k=3)
     context = "\n\n".join(d.page_content for d in docs)
 
-    prompt = f"""Use the context below to answer the question.
+    prompt = f"""Answer the question using only the context.
 
 Context:
 {context}
@@ -86,12 +61,12 @@ Question:
 
 Answer:"""
 
-    status.append("Generating answer…")
     result = llm(prompt)[0]["generated_text"]
 
-    answer = result.split("Answer:")[-1].strip()
-
     return {
-        "answer": answer,
-        "status": status
+        "answer": result.split("Answer:")[-1].strip(),
+        "status": [
+            f"Loaded {len(documents)} documents",
+            f"Retrieved {len(docs)} chunks"
+        ]
     }