Update src/simple_rag.py

src/simple_rag.py

@@ -1,3 +1,5 @@
+# Modified RAG Pipeline for General Document Q&A (Khmer & English)
+
 import os
 import logging
 import torch
@@ -7,7 +9,6 @@ from langchain.schema import Document
 from langchain.vectorstores.chroma import Chroma
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.document_loaders import PyPDFDirectoryLoader
-from transformers import BitsAndBytesConfig
 
 logging.basicConfig(level=logging.INFO)
 
@@ -34,8 +35,8 @@ if use_gpu:
 else:
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
-        load_in_8bit=True,
-        device_map={"": "cpu"},  # Force CPU
+        device_map="cpu",  # Force CPU
+        torch_dtype=torch.float32,
     )
     
 
@@ -53,7 +54,8 @@ embedding_model = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-bas
 PROMPT_TEMPLATE = """
 You are a helpful assistant.
 Answer the question based ONLY on the context below.
-
+If the user asks in Khmer, respond in Khmer.
+If the user asks in English, respond in English.
 Use clear, concise sentences, no more than 50 words. Do not mention the existence of context.
 
 Context: