Fix: remove device_map; CPU-safe Phi-3 load

rag.py

@@ -1,31 +1,35 @@
-from langchain_community.vectorstores import Chroma
-from langchain_community.embeddings import HuggingFaceEmbeddings
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_chroma import Chroma
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import PyPDFLoader
+import os
 
-from ingest import load_and_split_docs
-
-print("⏳ Loading documents...")
-
-documents = load_and_split_docs()
+MODEL_ID = "microsoft/Phi-3-mini-4k-instruct"
 
+print("⏳ Loading embeddings...")
 embeddings = HuggingFaceEmbeddings(
     model_name="sentence-transformers/all-MiniLM-L6-v2"
 )
 
-if documents:
-    vectorstore = Chroma.from_documents(
-        documents=documents,
-        embedding=embeddings
-    )
-    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
-else:
-    retriever = None
+print("⏳ Loading documents...")
+docs = []
+if os.path.exists("kb_docs"):
+    for f in os.listdir("kb_docs"):
+        if f.endswith(".pdf"):
+            loader = PyPDFLoader(os.path.join("kb_docs", f))
+            docs.extend(loader.load())
+
+splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
+splits = splitter.split_documents(docs)
+
+vectorstore = Chroma.from_documents(
+    splits,
+    embedding=embeddings,
+    persist_directory="./chroma_db"
+)
 
 print("⏳ Loading LLM...")
-
-MODEL_ID = "microsoft/Phi-3-mini-4k-instruct"
-
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID,
     trust_remote_code=True
@@ -33,37 +37,16 @@ tokenizer = AutoTokenizer.from_pretrained(
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    trust_remote_code=True,
-    torch_dtype=torch.float32,
-    device_map="cpu"
-)
-
-print("✅ RAG initialized.")
-
-def generate(prompt: str) -> str:
-    inputs = tokenizer(prompt, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=300,
-            temperature=0.2,
-            do_sample=True
-        )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    trust_remote_code=True
+)  # 👈 NO device_map, NO low_cpu_mem_usage
 
 def ask_rag_with_status(question: str):
-    if not retriever:
-        return {
-            "status": ["⚠️ No documents uploaded yet"],
-            "answer": "Please upload PDF files to the kb_docs folder and restart the Space."
-        }
-
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
     docs = retriever.get_relevant_documents(question)
+
     context = "\n\n".join(d.page_content for d in docs)
 
-    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,12 +54,18 @@ Context:
 Question:
 {question}
 
-Answer:
-"""
+Answer:"""
 
-    answer = generate(prompt)
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=256,
+        do_sample=True,
+        temperature=0.7
+    )
 
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return {
-        "status": ["🔍 Retrieved documents", "🧠 Generated answer"],
-        "answer": answer
+        "answer": answer,
+        "status": ["✅ Answer generated"]
     }