Rollback: stable local RAG

rag.py

@@ -1,90 +1,86 @@
 # rag.py
 
 import os
-from typing import List, Tuple
-
-from langchain_community.document_loaders import PyPDFLoader
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+from datasets import load_dataset
+from transformers import pipeline
+from langchain.schema import Document
 from langchain_community.vectorstores import Chroma
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain.schema import Document
-from transformers import pipeline
 
-from config import (
-    KB_DIR,
-    CHROMA_DIR,
-    EMBEDDING_MODEL,
-    LLM_MODEL,
-)
+from config import HF_DATASET_REPO, EMBEDDING_MODEL, LLM_MODEL
+
 
 # -----------------------------
-# Load documents
+# Load documents from HF dataset
 # -----------------------------
-def load_documents() -> List[Document]:
-    docs = []
+def load_documents():
+    documents = []
 
-    if not os.path.exists(KB_DIR):
-        print(f"⚠️ KB_DIR not found: {KB_DIR}")
-        return docs
+    try:
+        ds = load_dataset(HF_DATASET_REPO, split="train")
+    except Exception as e:
+        print(f"❌ Could not load dataset: {e}")
+        return []
 
-    for file in os.listdir(KB_DIR):
-        if file.lower().endswith(".pdf"):
-            loader = PyPDFLoader(os.path.join(KB_DIR, file))
-            docs.extend(loader.load())
+    # Expecting dataset rows like: { "text": "..." }
+    for row in ds:
+        text = row.get("text")
+        if text and isinstance(text, str):
+            documents.append(Document(page_content=text))
 
-    return docs
+    print(f"✅ Loaded {len(documents)} documents from dataset")
+    return documents
 
 
 # -----------------------------
-# Build vector DB (once)
+# Embeddings
 # -----------------------------
-documents = load_documents()
-
-splitter = RecursiveCharacterTextSplitter(
-    chunk_size=800,
-    chunk_overlap=100
-)
-
-chunks = splitter.split_documents(documents)
-
 embeddings = HuggingFaceEmbeddings(
     model_name=EMBEDDING_MODEL
 )
 
-vectordb = Chroma.from_documents(
-    documents=chunks,
-    embedding=embeddings,
-    persist_directory=CHROMA_DIR
-)
 
-retriever = vectordb.as_retriever(search_kwargs={"k": 3})
+# -----------------------------
+# Vector DB (safe creation)
+# -----------------------------
+documents = load_documents()
+
+if not documents:
+    print("⚠️ No documents loaded. Vector DB will be disabled.")
+    vectordb = None
+else:
+    vectordb = Chroma.from_documents(
+        documents=documents,
+        embedding=embeddings
+    )
+    print("✅ Vector DB initialized")
+
 
 # -----------------------------
-# LLM (CORRECT task)
+# LLM Pipeline (CPU safe)
 # -----------------------------
-llm = pipeline(
-    "text2text-generation",
+qa_pipeline = pipeline(
+    task="text-generation",
     model=LLM_MODEL,
-    device=-1
+    max_new_tokens=256
 )
 
+
 # -----------------------------
-# RAG call
+# RAG Query Function
 # -----------------------------
-def ask_rag_with_status(question: str) -> Tuple[str, list]:
-    status = []
-
-    if vectordb._collection.count() == 0:
-        return "Knowledge base is empty.", ["No documents indexed"]
+def ask_rag_with_status(question: str):
+    if vectordb is None:
+        return "Knowledge base is empty.", "NO_KB"
 
-    docs = retriever.get_relevant_documents(question)
+    docs = vectordb.similarity_search(question, k=3)
 
-    status.append(f"Retrieved {len(docs)} chunks")
+    if not docs:
+        return "No relevant documents found.", "NO_MATCH"
 
     context = "\n\n".join(d.page_content for d in docs)
 
-    prompt = f"""
-Answer the question using ONLY the context below.
+    prompt = f"""Use the context below to answer the question.
 
 Context:
 {context}
@@ -92,9 +88,9 @@ Context:
 Question:
 {question}
 
-Answer:
-"""
+Answer:"""
 
-    result = llm(prompt, max_new_tokens=256)[0]["generated_text"]
+    result = qa_pipeline(prompt)
+    answer = result[0]["generated_text"]
 
-    return result.strip(), status
+    return answer, "OK"