Update app.py

app.py

@@ -3,11 +3,14 @@ import os
 import uuid
 import threading
 import pandas as pd
+import numpy as np
 from langchain.document_loaders.csv_loader import CSVLoader
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 from langchain.llms import CTransformers
-from langchain.chains import ConversationalRetrievalChain
+from langchain.agents import create_pandas_dataframe_agent
+from langchain.chains import LLMChain
+from langchain.prompts import PromptTemplate
 
 # Global model cache
 MODEL_CACHE = {
@@ -22,15 +25,16 @@ def initialize_model_once():
     """Initialize model once using CTransformers API"""
     with MODEL_CACHE["init_lock"]:
         if MODEL_CACHE["model"] is None:
-            # Load Mistral-7B-Instruct-v0.2.Q4_K_M.gguf model
+            # Load Phi-2 model (smaller than Mistral)
             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",
+                model="TheBloke/phi-2-GGUF",
+                model_file="phi-2.Q4_K_M.gguf",
+                model_type="phi2",
                 max_new_tokens=512,
-                temperature=0.2,
+                temperature=0.1,
                 top_p=0.9,
-                repetition_penalty=1.2
+                repetition_penalty=1.1,
+                context_length=2048
             )
     
     return MODEL_CACHE["model"]
@@ -38,8 +42,9 @@ def initialize_model_once():
 class ChatBot:
     def __init__(self, session_id):
         self.session_id = session_id
+        self.csv_info = None
+        self.df = None
         self.chat_history = []
-        self.chain = None
         self.user_dir = f"user_data/{session_id}"
         os.makedirs(self.user_dir, exist_ok=True)
         
@@ -50,95 +55,148 @@ class ChatBot:
         try:
             # Handle file from Gradio
             file_path = file.name if hasattr(file, 'name') else str(file)
+            file_name = os.path.basename(file_path)
             
-            # Verify and save CSV
+            # Load and save CSV directly with pandas
             try:
-                df = pd.read_csv(file_path)
+                self.df = pd.read_csv(file_path)
                 user_file_path = f"{self.user_dir}/uploaded.csv"
-                df.to_csv(user_file_path, index=False)
-                print(f"CSV verified: {df.shape[0]} rows, {len(df.columns)} columns")
-            except Exception as e:
-                return f"Error membaca CSV: {str(e)}"
-            
-            # Load document
-            try:
-                loader = CSVLoader(file_path=file_path, encoding="utf-8", csv_args={'delimiter': ','})
-                data = loader.load()
-                print(f"Documents loaded: {len(data)}")
-            except Exception as e:
-                return f"Error loading documents: {str(e)}"
-            
-            # Create vector database
-            try:
-                db_path = f"{self.user_dir}/db_faiss"
-                embeddings = HuggingFaceEmbeddings(
-                    model_name='sentence-transformers/all-MiniLM-L6-v2',
-                    model_kwargs={'device': 'cpu'}  # Explicitly set to CPU
-                )
+                self.df.to_csv(user_file_path, index=False)
                 
-                db = FAISS.from_documents(data, embeddings)
-                db.save_local(db_path)
-                print(f"Vector database created at {db_path}")
+                # Store CSV info
+                self.csv_info = {
+                    "filename": file_name,
+                    "rows": self.df.shape[0],
+                    "columns": self.df.shape[1],
+                    "column_names": self.df.columns.tolist(),
+                }
+                
+                print(f"CSV verified: {self.df.shape[0]} rows, {len(self.df.columns)} columns")
             except Exception as e:
-                return f"Error creating vector database: {str(e)}"
+                return f"Error membaca CSV: {str(e)}"
             
-            # Create LLM and chain
+            # Create query translator
             try:
                 llm = initialize_model_once()
-                self.chain = ConversationalRetrievalChain.from_llm(
-                    llm=llm, 
-                    retriever=db.as_retriever(search_kwargs={"k": 4}),
-                    return_source_documents=True
+                
+                query_template = """
+                Kamu adalah asisten yang mengubah pertanyaan natural language menjadi kode Python dengan pandas.
+                
+                Informasi tentang DataFrame:
+                - Nama kolom: {column_names}
+                - Jumlah baris: {num_rows}
+                - Sample data:
+                {sample_data}
+                
+                Pertanyaan pengguna: {question}
+                
+                Ubah pertanyaan tersebut menjadi kode pandas yang bisa dijalankan. Kode harus ringkas, efisien, dan menggunakan variabel 'df'.
+                Berikan HANYA kode python saja, tanpa backtick, tanpa penjelasan.
+                
+                Kode:
+                """
+                
+                self.query_chain = LLMChain(
+                    llm=llm,
+                    prompt=PromptTemplate(
+                        input_variables=["column_names", "num_rows", "sample_data", "question"],
+                        template=query_template
+                    )
                 )
-                print("Chain created successfully")
+                
+                print("Query translator created successfully")
             except Exception as e:
-                return f"Error creating chain: {str(e)}"
+                return f"Error creating query translator: {str(e)}"
             
             # Add file info to chat history
-            file_info = f"CSV berhasil dimuat dengan {df.shape[0]} baris dan {len(df.columns)} kolom. Kolom: {', '.join(df.columns.tolist())}"
+            file_info = f"CSV berhasil dimuat: {file_name} dengan {self.df.shape[0]} baris dan {len(self.df.columns)} kolom. Kolom: {', '.join(self.df.columns.tolist())}"
             self.chat_history.append(("System", file_info))
             
-            return "File CSV berhasil diproses! Anda dapat mulai chat dengan Mistral 7B."
+            return f"File CSV '{file_name}' berhasil diproses! Anda dapat mulai mengajukan pertanyaan tentang data."
         except Exception as e:
             import traceback
             print(traceback.format_exc())
             return f"Error pemrosesan file: {str(e)}"
 
     def chat(self, message, history):
-        if self.chain is None:
+        if self.df is None or self.query_chain is None:
             return "Mohon upload file CSV terlebih dahulu."
         
         try:
-            # Process with the chain
-            result = self.chain({"question": message, "chat_history": self.chat_history})
+            # Handle metadata questions directly
+            message_lower = message.lower()
+            if "nama file" in message_lower:
+                return f"Nama file CSV adalah: {self.csv_info['filename']}"
+            elif "nama kolom" in message_lower:
+                return f"Kolom dalam CSV: {', '.join(self.csv_info['column_names'])}"
+            elif "jumlah baris" in message_lower or "berapa baris" in message_lower:
+                return f"Jumlah baris dalam CSV: {self.csv_info['rows']}"
             
-            # Update chat history
-            answer = result["answer"]
+            # Get sample data for context
+            sample_str = self.df.head(3).to_string()
             
-            # 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))
+            # Translate question to pandas code
+            code_response = self.query_chain.run(
+                column_names=str(self.csv_info["column_names"]),
+                num_rows=self.csv_info["rows"],
+                sample_data=sample_str,
+                question=message
+            )
             
-            return answer
+            # Clean and execute the code
+            try:
+                code = code_response.strip()
+                # Add safety prefix to prevent malicious code
+                if not code.startswith("df"):
+                    code = "result = " + code
+                else:
+                    code = "result = " + code
+                    
+                # Create local context with the dataframe
+                locals_dict = {"df": self.df, "pd": pd, "np": np}
+                
+                # Execute the code
+                print(f"Executing code: {code}")
+                exec(code, {"pd": pd, "np": np}, locals_dict)
+                result = locals_dict.get("result", "No result returned")
+                
+                # Format the result
+                if isinstance(result, pd.DataFrame):
+                    if len(result) > 5:
+                        result_str = result.head(5).to_string() + f"\n\n[{len(result)} baris ditemukan]"
+                    else:
+                        result_str = result.to_string()
+                elif isinstance(result, (pd.Series, np.ndarray)):
+                    result_str = str(result)
+                else:
+                    result_str = str(result)
+                
+                # Build the response
+                response = f"Hasil analisis untuk pertanyaan: '{message}'\n\n"
+                response += f"Kode yang digunakan:\n```python\n{code}\n```\n\n"
+                response += f"Output:\n{result_str}"
+                
+                self.chat_history.append((message, response))
+                return response
+                
+            except Exception as e:
+                error_msg = f"Error mengeksekusi kode: {str(e)}\nKode yang dihasilkan:\n```python\n{code}\n```"
+                print(error_msg)
+                return error_msg
+                
         except Exception as e:
             import traceback
             print(traceback.format_exc())
             return f"Error: {str(e)}"
 
-# UI Code dan handler functions sama seperti sebelumnya
+# UI Code 
 def create_gradio_interface():
-    with gr.Blocks(title="Chat with CSV using Mistral 7B") as interface:
+    with gr.Blocks(title="CSV Data Analyzer") 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 Mistral 7B</h1>")
-        gr.HTML("<h3 style='text-align: center;'>Asisten analisis CSV yang powerful</h3>")
+        gr.HTML("<h1 style='text-align: center;'>CSV Data Analyzer</h1>")
+        gr.HTML("<h3 style='text-align: center;'>Ajukan pertanyaan tentang data CSV Anda</h3>")
         
         with gr.Row():
             with gr.Column(scale=1):
@@ -148,14 +206,12 @@ def create_gradio_interface():
                 )
                 process_button = gr.Button("Proses CSV")
                 
-                with gr.Accordion("Informasi Model", open=False):
+                with gr.Accordion("Contoh Pertanyaan", open=False):
                     gr.Markdown("""
-                    **Model**: Mistral-7B-Instruct-v0.2-GGUF
-                    
-                    **Fitur**:
-                    - GGUF model yang dioptimalkan untuk CPU
-                    - Efisien untuk analisis data dan percakapan
-                    - Manajemen sesi per pengguna
+                    - "Berapa jumlah data yang memiliki nilai Glucose di atas 150?"
+                    - "Bagaimana distribusi kolom Age?"
+                    - "Hitung nilai rata-rata dan standar deviasi untuk setiap kolom numerik"
+                    - "Buat tabel frekuensi untuk kolom Outcome"
                     """)
             
             with gr.Column(scale=2):
@@ -164,8 +220,8 @@ def create_gradio_interface():
                     height=400
                 )
                 message_input = gr.Textbox(
-                    label="Ketik pesan Anda",
-                    placeholder="Tanyakan tentang data CSV Anda...",
+                    label="Ketik pertanyaan Anda",
+                    placeholder="Contoh: Berapa jumlah data yang memiliki nilai Glucose di atas 150?",
                     lines=2
                 )
                 submit_button = gr.Button("Kirim")