Update app.py

app.py

@@ -6,9 +6,8 @@ import pandas as pd
 from langchain.document_loaders.csv_loader import CSVLoader
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
-from langchain.llms import HuggingFacePipeline
+from langchain.llms import CTransformers
 from langchain.chains import ConversationalRetrievalChain
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 # Global model cache
 MODEL_CACHE = {
@@ -20,36 +19,19 @@ MODEL_CACHE = {
 os.makedirs("user_data", exist_ok=True)
 
 def initialize_model_once():
-    """Initialize model once using pipeline API"""
+    """Initialize model once using CTransformers API"""
     with MODEL_CACHE["init_lock"]:
         if MODEL_CACHE["model"] is None:
-            # Load model from Hugging Face Hub
-            model_id = "meta-llama/Llama-2-7b-chat-hf"
-            
-            # Tokenizer
-            tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ.get("HF_TOKEN"))
-            
-            # Model with low precision
-            model = AutoModelForCausalLM.from_pretrained(
-                model_id, 
-                token=os.environ.get("HF_TOKEN"),
-                device_map="auto",
-                load_in_8bit=True  # Quantize model to 8-bit precision
-            )
-            
-            # Create pipeline
-            pipe = pipeline(
-                "text-generation",
-                model=model,
-                tokenizer=tokenizer,
+            # Load Mistral-7B-Instruct-v0.2.Q4_K_M.gguf model
+            MODEL_CACHE["model"] = CTransformers(
+                model="TheBloke/Mistral-7B-Instruct-v0.2-GGUF",
+                model_file="mistral-7b-instruct-v0.2.Q4_K_M.gguf",
+                model_type="mistral",
                 max_new_tokens=512,
                 temperature=0.2,
                 top_p=0.9,
                 repetition_penalty=1.2
             )
-            
-            # Create LangChain wrapper
-            MODEL_CACHE["model"] = HuggingFacePipeline(pipeline=pipe)
     
     return MODEL_CACHE["model"]
 
@@ -91,7 +73,7 @@ class ChatBot:
                 db_path = f"{self.user_dir}/db_faiss"
                 embeddings = HuggingFaceEmbeddings(
                     model_name='sentence-transformers/all-MiniLM-L6-v2',
-                    model_kwargs={'device': 'auto'}
+                    model_kwargs={'device': 'cpu'}  # Explicitly set to CPU
                 )
                 
                 db = FAISS.from_documents(data, embeddings)
@@ -105,7 +87,8 @@ class ChatBot:
                 llm = initialize_model_once()
                 self.chain = ConversationalRetrievalChain.from_llm(
                     llm=llm, 
-                    retriever=db.as_retriever(search_kwargs={"k": 4})
+                    retriever=db.as_retriever(search_kwargs={"k": 4}),
+                    return_source_documents=True
                 )
                 print("Chain created successfully")
             except Exception as e:
@@ -115,7 +98,7 @@ class ChatBot:
             file_info = f"CSV berhasil dimuat dengan {df.shape[0]} baris dan {len(df.columns)} kolom. Kolom: {', '.join(df.columns.tolist())}"
             self.chat_history.append(("System", file_info))
             
-            return "File CSV berhasil diproses! Anda dapat mulai chat dengan model Llama 2."
+            return "File CSV berhasil diproses! Anda dapat mulai chat dengan Mistral 7B."
         except Exception as e:
             import traceback
             print(traceback.format_exc())
@@ -131,6 +114,15 @@ class ChatBot:
             
             # Update chat history
             answer = result["answer"]
+            
+            # Optional: Add source info to answer
+            sources = result.get("source_documents", [])
+            if sources:
+                source_text = "\n\nSumber:\n"
+                for i, doc in enumerate(sources[:2], 1):  # Limit to top 2 sources
+                    source_text += f"{i}. {doc.page_content[:100]}...\n"
+                answer += source_text
+            
             self.chat_history.append((message, answer))
             
             return answer
@@ -141,12 +133,12 @@ class ChatBot:
 
 # UI Code dan handler functions sama seperti sebelumnya
 def create_gradio_interface():
-    with gr.Blocks(title="Chat with CSV using Llama2 🦙") as interface:
+    with gr.Blocks(title="Chat with CSV using Mistral 7B") as interface:
         session_id = gr.State(lambda: str(uuid.uuid4()))
         chatbot_state = gr.State(lambda: None)
         
-        gr.HTML("<h1 style='text-align: center;'>Chat with CSV using Llama2 🦙</h1>")
-        gr.HTML("<h3 style='text-align: center;'>Asisten analisis CSV yang powerfull</h3>")
+        gr.HTML("<h1 style='text-align: center;'>Chat with CSV using Mistral 7B</h1>")
+        gr.HTML("<h3 style='text-align: center;'>Asisten analisis CSV yang powerful</h3>")
         
         with gr.Row():
             with gr.Column(scale=1):
@@ -158,11 +150,11 @@ def create_gradio_interface():
                 
                 with gr.Accordion("Informasi Model", open=False):
                     gr.Markdown("""
-                    **Model**: Llama-2-7b-chat-hf
+                    **Model**: Mistral-7B-Instruct-v0.2-GGUF
                     
                     **Fitur**:
-                    - Dioptimalkan untuk analisis data dan percakapan
-                    - Menggunakan API Hugging Face untuk efisiensi
+                    - GGUF model yang dioptimalkan untuk CPU
+                    - Efisien untuk analisis data dan percakapan
                     - Manajemen sesi per pengguna
                     """)