app.py

import gradio as gr
import gc
import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Force CPU only
import uuid
import threading
import pandas as pd
import torch
from langchain.document_loaders import CSVLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.llms import HuggingFacePipeline
from langchain.chains import LLMChain
from transformers import AutoTokenizer, AutoModelForCausalLM, T5Tokenizer, T5ForConditionalGeneration, BitsAndBytesConfig, pipeline
from langchain.prompts import PromptTemplate
from llama_cpp import Llama
import re
import datetime
import warnings
warnings.filterwarnings('ignore')

# Global model cache
MODEL_CACHE = {
    "model": None,
    "tokenizer": None,
    "init_lock": threading.Lock(),
    "model_name": None
}

# Create directories for user data
os.makedirs("user_data", exist_ok=True)

# Model configuration dictionary
MODEL_CONFIG = {
    "Llama 2 Chat": {
        "name": "TheBloke/Llama-2-7B-Chat-GGUF",
        "description": "Llama 2 7B Chat model with good general performance",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "TinyLlama Chat": {
        "name": "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF",
        "description": "Model ringan dengan 1.1B parameter, cepat dan ringan",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "Mistral Instruct": {
        "name": "TheBloke/Mistral-7B-Instruct-v0.2-GGUF",
        "description": "7B instruction-tuned model with excellent reasoning",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "Phi-4 Mini Instruct": {
        "name": "microsoft/Phi-4-mini-instruct",
        "description": "Model yang ringan dari Microsoft cocok untuk tugas instruksional",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "DeepSeek Coder Instruct": {
        "name": "deepseek-ai/deepseek-coder-1.3b-instruct",
        "description": "1.3B model untuk kode dan analisis data",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "DeepSeek Lite Chat": {
        "name": "deepseek-ai/DeepSeek-V2-Lite-Chat",
        "description": "Light but powerful chat model from DeepSeek",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "Qwen2.5 Coder Instruct": {
        "name": "Qwen/Qwen2.5-Coder-3B-Instruct-GGUF",
        "description": "3B model specialized for code and technical applications",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "DeepSeek Distill Qwen": {
        "name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
        "description": "1.5B distilled model with good balance of speed and quality",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32
    },
    "Flan T5 Small": {
        "name": "google/flan-t5-small",
        "description": "Lightweight T5 model optimized for instruction following",
        "dtype": torch.float16 if torch.cuda.is_available() else torch.float32,
        "is_t5": True
    }
}

def initialize_model_once(model_key):
    with MODEL_CACHE["init_lock"]:
        current_model = MODEL_CACHE["model_name"]
        if MODEL_CACHE["model"] is None or current_model != model_key:
            # Clear previous model
            if MODEL_CACHE["model"] is not None:
                del MODEL_CACHE["model"]
                if MODEL_CACHE["tokenizer"] is not None:
                    del MODEL_CACHE["tokenizer"]
                torch.cuda.empty_cache() if torch.cuda.is_available() else None
                
            model_info = MODEL_CONFIG[model_key]
            model_name = model_info["name"]
            MODEL_CACHE["model_name"] = model_key

            try:
                print(f"Loading model: {model_name}")
                
                # Periksa apakah ini model GGUF
                if "GGUF" in model_name:
                    # Download model file terlebih dahulu jika belum ada
                    from huggingface_hub import hf_hub_download
                    try:
                        # Coba temukan file GGUF di repo
                        repo_id = model_name
                        model_path = hf_hub_download(
                            repo_id=repo_id, 
                            filename="model.gguf"  # Nama file dapat berbeda
                        )
                    except Exception as e:
                        print(f"Couldn't find model.gguf, trying other filenames: {str(e)}")
                        # Coba cari file GGUF dengan nama lain
                        import requests
                        from huggingface_hub import list_repo_files
                        
                        files = list_repo_files(repo_id)
                        gguf_files = [f for f in files if f.endswith('.gguf')]
                        
                        if not gguf_files:
                            raise ValueError(f"No GGUF files found in {repo_id}")
                        
                        # Gunakan file GGUF pertama yang ditemukan
                        model_path = hf_hub_download(repo_id=repo_id, filename=gguf_files[0])
                    
                    # Load model GGUF dengan llama-cpp-python
                    MODEL_CACHE["model"] = Llama(
                        model_path=model_path,
                        n_ctx=2048,  # Konteks yang lebih kecil untuk penghematan memori
                        n_batch=512,
                        n_threads=2  # Sesuaikan dengan 2 vCPU
                    )
                    MODEL_CACHE["tokenizer"] = None  # GGUF tidak membutuhkan tokenizer terpisah
                    MODEL_CACHE["is_gguf"] = True
                
                # Handle T5 models
                elif model_info.get("is_t5", False):
                    MODEL_CACHE["tokenizer"] = T5Tokenizer.from_pretrained(model_name)
                    MODEL_CACHE["model"] = T5ForConditionalGeneration.from_pretrained(
                        model_name,
                        torch_dtype=model_info["dtype"],
                        device_map="auto" if torch.cuda.is_available() else None,
                        low_cpu_mem_usage=True
                    )
                    MODEL_CACHE["is_gguf"] = False
                    
                # Handle standard HF models
                else:
                    quantization_config = BitsAndBytesConfig(
                        load_in_4bit=True,
                        bnb_4bit_compute_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                        bnb_4bit_quant_type="nf4",
                        bnb_4bit_use_double_quant=True
                    )
                    MODEL_CACHE["tokenizer"] = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
                    MODEL_CACHE["model"] = AutoModelForCausalLM.from_pretrained(
                        model_name,
                        quantization_config=quantization_config,
                        torch_dtype=model_info["dtype"],
                        device_map="auto" if torch.cuda.is_available() else None,
                        low_cpu_mem_usage=True,
                        trust_remote_code=True
                    )
                    MODEL_CACHE["is_gguf"] = False
                    
                print(f"Model {model_name} loaded successfully")
            except Exception as e:
                import traceback
                print(f"Error loading model {model_name}: {str(e)}")
                print(traceback.format_exc())
                raise RuntimeError(f"Failed to load model {model_name}: {str(e)}")
                
    return MODEL_CACHE["tokenizer"], MODEL_CACHE["model"], MODEL_CACHE.get("is_gguf", False)

def create_llm_pipeline(model_key):
    """Create a new pipeline using the specified model"""
    try:
        print(f"Creating pipeline for model: {model_key}")
        tokenizer, model, is_gguf = initialize_model_once(model_key)

        if model is None:
            raise ValueError(f"Model is None for {model_key}")

        # For GGUF models from llama-cpp-python
        if is_gguf:
            # Buat adaptor untuk menggunakan model GGUF seperti HF pipeline
            from langchain.llms import LlamaCpp
            llm = LlamaCpp(
                model_path=model.model_path,
                temperature=0.3,
                max_tokens=128,
                top_p=0.9,
                n_ctx=2048,
                streaming=False
            )
            return llm
            
        # Create appropriate pipeline for HF models
        elif getattr(model_info, "is_t5", False):
            print("Creating T5 pipeline")
            pipe = pipeline(
                "text2text-generation",
                model=model,
                tokenizer=tokenizer,
                max_new_tokens=128,
                temperature=0.3,
                top_p=0.9,
                return_full_text=False,
            )
        else:
            print("Creating causal LM pipeline")
            pipe = pipeline(
                "text-generation",
                model=model,
                tokenizer=tokenizer,
                max_new_tokens=128,
                temperature=0.3,
                top_p=0.9,
                top_k=30,
                repetition_penalty=1.2,
                return_full_text=False,
            )
        
        print("Pipeline created successfully")
        return HuggingFacePipeline(pipeline=pipe)
    except Exception as e:
        import traceback
        print(f"Error creating pipeline: {str(e)}")
        print(traceback.format_exc())

def handle_model_loading_error(model_key, session_id):
    """Handle model loading errors by providing alternative model suggestions"""
    suggested_models = [
        "DeepSeek Coder Instruct",  # 1.3B model
        "Phi-4 Mini Instruct",      # Light model
        "TinyLlama Chat",           # 1.1B model 
        "Flan T5 Small"             # Lightweight T5
    ]
    
    # Remove the current model from suggestions if it's in the list
    if model_key in suggested_models:
        suggested_models.remove(model_key)
    
    suggestions = ", ".join(suggested_models[:3])  # Only show top 3 suggestions
    return None, f"Tidak dapat memuat model {model_key}. Silakan coba model lain seperti: {suggestions}"
        
def create_conversational_chain(db, file_path, model_key):
    llm = create_llm_pipeline(model_key)
    
    # Load the file into pandas to enable code execution for data analysis
    df = pd.read_csv(file_path)
    
    # Create improved prompt template that focuses on direct answers, not code
    template = """
    Berikut ini adalah informasi tentang file CSV:
    
    Kolom-kolom dalam file: {columns}
    
    Beberapa baris pertama:
    {sample_data}
    
    Konteks tambahan dari vector database:
    {context}
    
    Pertanyaan: {question}
    
    INSTRUKSI PENTING:
    1. Jangan tampilkan kode Python, berikan jawaban langsung dalam Bahasa Indonesia.
    2. Jika pertanyaan terkait statistik data (rata-rata, maksimum dll), lakukan perhitungan dan berikan hasilnya.
    3. Jawaban harus singkat, jelas dan akurat berdasarkan data yang ada.
    4. Gunakan format yang sesuai untuk angka (desimal 2 digit untuk nilai non-integer).
    5. Jangan menyebutkan proses perhitungan, fokus pada hasil akhir.
    
    Jawaban:
    """
    
    PROMPT = PromptTemplate(
        template=template,
        input_variables=["columns", "sample_data", "context", "question"]
    )
    
    # Create retriever
    retriever = db.as_retriever(search_kwargs={"k": 3})  # Reduced k for better performance
    
    # Process query with better error handling
    def process_query(query, chat_history):
        try:
            # Get information from dataframe for context
            columns_str = ", ".join(df.columns.tolist())
            sample_data = df.head(2).to_string()  # Reduced to 2 rows for performance
            
            # Get context from vector database
            docs = retriever.get_relevant_documents(query)
            context = "\n\n".join([doc.page_content for doc in docs])
            
            # Dynamically calculate answers for common statistical queries
            def preprocess_query():
                query_lower = query.lower()
                result = None
                
                # Handle statistical queries directly
                if "rata-rata" in query_lower or "mean" in query_lower or "average" in query_lower:
                    for col in df.columns:
                        if col.lower() in query_lower and pd.api.types.is_numeric_dtype(df[col]):
                            try:
                                result = f"Rata-rata {col} adalah {df[col].mean():.2f}"
                            except:
                                pass
                
                elif "maksimum" in query_lower or "max" in query_lower or "tertinggi" in query_lower:
                    for col in df.columns:
                        if col.lower() in query_lower and pd.api.types.is_numeric_dtype(df[col]):
                            try:
                                result = f"Nilai maksimum {col} adalah {df[col].max():.2f}"
                            except:
                                pass
                
                elif "minimum" in query_lower or "min" in query_lower or "terendah" in query_lower:
                    for col in df.columns:
                        if col.lower() in query_lower and pd.api.types.is_numeric_dtype(df[col]):
                            try:
                                result = f"Nilai minimum {col} adalah {df[col].min():.2f}"
                            except:
                                pass
                
                elif "total" in query_lower or "jumlah" in query_lower or "sum" in query_lower:
                    for col in df.columns:
                        if col.lower() in query_lower and pd.api.types.is_numeric_dtype(df[col]):
                            try:
                                result = f"Total {col} adalah {df[col].sum():.2f}"
                            except:
                                pass
                
                elif "baris" in query_lower or "jumlah data" in query_lower or "row" in query_lower:
                    result = f"Jumlah baris data adalah {len(df)}"
                
                elif "kolom" in query_lower or "field" in query_lower:
                    if "nama" in query_lower or "list" in query_lower or "sebutkan" in query_lower:
                        result = f"Kolom dalam data: {', '.join(df.columns.tolist())}"
                
                return result
            
            # Try direct calculation first
            direct_answer = preprocess_query()
            if direct_answer:
                return {"answer": direct_answer}
            
            # If no direct calculation, use the LLM
            chain = LLMChain(llm=llm, prompt=PROMPT)
            raw_result = chain.run(
                columns=columns_str,
                sample_data=sample_data,
                context=context,
                question=query
            )
            
            # Clean the result
            cleaned_result = raw_result.strip()
            
            # If result is empty after cleaning, use a fallback
            if not cleaned_result:
                return {"answer": "Tidak dapat memproses jawaban. Silakan coba pertanyaan lain."}
                
            return {"answer": cleaned_result}
        except Exception as e:
            import traceback
            print(f"Error in process_query: {str(e)}")
            print(traceback.format_exc())
            return {"answer": f"Terjadi kesalahan saat memproses pertanyaan: {str(e)}"}
    
    return process_query

class ChatBot:
    def __init__(self, session_id, model_key="DeepSeek Coder Instruct"):
        self.session_id = session_id
        self.chat_history = []
        self.chain = None
        self.user_dir = f"user_data/{session_id}"
        self.csv_file_path = None
        self.model_key = model_key
        os.makedirs(self.user_dir, exist_ok=True)

    def process_file(self, file, model_key=None):
        if model_key:
            self.model_key = model_key
        
        if file is None:
            return "Mohon upload file CSV terlebih dahulu."

        try:
            print(f"Processing file using model: {self.model_key}")
            # Handle file from Gradio
            file_path = file.name if hasattr(file, 'name') else str(file)
            self.csv_file_path = file_path
            print(f"CSV file path: {file_path}")

            # Copy to user directory
            user_file_path = f"{self.user_dir}/uploaded.csv"

            # Verify the CSV can be loaded
            try:
                df = pd.read_csv(file_path)
                print(f"CSV verified: {df.shape[0]} rows, {len(df.columns)} columns")

                # Save a copy in user directory
                df.to_csv(user_file_path, index=False)
                self.csv_file_path = user_file_path
                print(f"CSV saved to {user_file_path}")
            except Exception as e:
                print(f"Error reading CSV: {str(e)}")
                return f"Error membaca CSV: {str(e)}"

            # Load document with reduced chunk size for better memory usage
            try:
                loader = CSVLoader(file_path=user_file_path, encoding="utf-8", csv_args={
                    'delimiter': ','})
                data = loader.load()
                print(f"Documents loaded: {len(data)}")
            except Exception as e:
                print(f"Error loading documents: {str(e)}")
                return f"Error loading documents: {str(e)}"

            # Create vector database with optimized settings
            try:
                db_path = f"{self.user_dir}/db_faiss"
            
                # Use CPU-friendly embeddings with smaller dimensions
                embeddings = HuggingFaceEmbeddings(
                    model_name='sentence-transformers/all-MiniLM-L6-v2',
                    model_kwargs={'device': 'cpu'}
                )

                db = FAISS.from_documents(data, embeddings)
                db.save_local(db_path)
                print(f"Vector database created at {db_path}")
            except Exception as e:
                print(f"Error creating vector database: {str(e)}")
                return f"Error creating vector database: {str(e)}"

            # Create custom chain
            try:
                print(f"Creating conversation chain with model: {self.model_key}")
                self.chain = create_conversational_chain(db, self.csv_file_path, self.model_key)
                print("Chain created successfully")
            except Exception as e:
                import traceback
                print(f"Error creating chain: {str(e)}")
                print(traceback.format_exc())
                return f"Error creating chain: {str(e)}"

            # Add basic file info to chat history for context
            file_info = f"CSV berhasil dimuat dengan {df.shape[0]} baris dan {len(df.columns)} kolom menggunakan model {self.model_key}. Kolom: {', '.join(df.columns.tolist())}"
            self.chat_history.append(("System", file_info))

            return f"File CSV berhasil diproses dengan model {self.model_key}! Anda dapat mulai chat dengan model untuk analisis data."
        except Exception as e:
            import traceback
            print(traceback.format_exc())
            return f"Error pemrosesan file: {str(e)}"

    def change_model(self, model_key):
        """Change the model being used and recreate the chain if necessary"""
        try:
            if model_key == self.model_key:
                return f"Model {model_key} sudah digunakan."
            
            print(f"Changing model from {self.model_key} to {model_key}")
            self.model_key = model_key
        
            # If we have an active session with a file already loaded, recreate the chain
            if self.csv_file_path and os.path.exists(self.csv_file_path):
                try:
                    # Load existing database
                    db_path = f"{self.user_dir}/db_faiss"
                    if not os.path.exists(db_path):
                        return f"Error: Database tidak ditemukan. Silakan upload file CSV kembali."
                
                    print(f"Loading embeddings from {db_path}")
                    embeddings = HuggingFaceEmbeddings(
                        model_name='sentence-transformers/all-MiniLM-L6-v2',
                        model_kwargs={'device': 'cpu'}
                    )
                
                    # Tambahkan flag allow_dangerous_deserialization=True
                    db = FAISS.load_local(db_path, embeddings, allow_dangerous_deserialization=True)
                    print(f"FAISS database loaded successfully")
                
                    # Create new chain with the selected model
                    print(f"Creating new conversation chain with {model_key}")
                    self.chain = create_conversational_chain(db, self.csv_file_path, self.model_key)
                    print(f"Chain created successfully")
                
                    # Add notification to chat history
                    self.chat_history.append(("System", f"Model berhasil diubah ke {model_key}."))
                
                    return f"Model berhasil diubah ke {model_key}."
                except Exception as e:
                    import traceback
                    error_trace = traceback.format_exc()
                    print(f"Detailed error in change_model: {error_trace}")
                    return f"Error mengubah model: {str(e)}"
            else:
                # Just update the model key if no file is loaded yet
                print(f"No CSV file loaded yet, just updating model preference to {model_key}")
                return f"Model diubah ke {model_key}. Silakan upload file CSV untuk memulai."
        except Exception as e:
            import traceback
            error_trace = traceback.format_exc()
            print(f"Unexpected error in change_model: {error_trace}")
            return f"Error tidak terduga saat mengubah model: {str(e)}"

    def chat(self, message, history):
        if self.chain is None:
            return "Mohon upload file CSV terlebih dahulu."

        try:
            # Process the question with the chain
            result = self.chain(message, self.chat_history)

            # Get the answer with fallback
            answer = result.get("answer", "Maaf, tidak dapat menghasilkan jawaban. Silakan coba pertanyaan lain.")

            # Ensure we never return empty
            if not answer or answer.strip() == "":
                answer = "Maaf, tidak dapat menghasilkan jawaban yang sesuai. Silakan coba pertanyaan lain."
            
            # Update internal chat history
            self.chat_history.append((message, answer))

            # Return just the answer for Gradio
            return answer
        except Exception as e:
            import traceback
            print(traceback.format_exc())
            return f"Error: {str(e)}"

# UI Code
def create_gradio_interface():
    with gr.Blocks(title="Chat with CSV using AI Models") as interface:
        session_id = gr.State(lambda: str(uuid.uuid4()))
        chatbot_state = gr.State(lambda: None)
        model_selected = gr.State(lambda: False)  # Track if model is already in use
        
        # Get model choices
        model_choices = list(MODEL_CONFIG.keys())
        default_model = "DeepSeek Coder Instruct"  # Default model

        gr.HTML("<h1 style='text-align: center;'>Chat with CSV using AI Models</h1>")
        gr.HTML("<h3 style='text-align: center;'>Asisten analisis CSV untuk berbagai kebutuhan</h3>")

        with gr.Row():
            with gr.Column(scale=1):
                with gr.Group():
                    gr.Markdown("### Langkah 1: Pilih Model AI")
                    model_dropdown = gr.Dropdown(
                        label="Model",
                        choices=model_choices,
                        value=default_model,
                        interactive=True
                    )
                    model_info = gr.Markdown(
                        value=f"**{default_model}**: {MODEL_CONFIG[default_model]['description']}"
                    )
                
                with gr.Group():
                    gr.Markdown("### Langkah 2: Unggah dan Proses CSV")
                    file_input = gr.File(
                        label="Upload CSV Anda",
                        file_types=[".csv"]
                    )
                    process_button = gr.Button("Proses CSV")
                
                reset_button = gr.Button("Reset Sesi (Untuk Ganti Model)")

            with gr.Column(scale=2):
                chatbot_interface = gr.Chatbot(
                    label="Riwayat Chat",
                    # type="messages",
                    height=400
                )
                message_input = gr.Textbox(
                    label="Ketik pesan Anda",
                    placeholder="Tanyakan tentang data CSV Anda...",
                    lines=2
                )
                submit_button = gr.Button("Kirim")
                clear_button = gr.Button("Bersihkan Chat")

        # Update model info when selection changes
        def update_model_info(model_key):
            return f"**{model_key}**: {MODEL_CONFIG[model_key]['description']}"
            
        model_dropdown.change(
            fn=update_model_info,
            inputs=[model_dropdown],
            outputs=[model_info]
        )
        
        # Process file handler - disables model selection after file is processed
        def handle_process_file(file, model_key, sess_id):
            if file is None:
                return None, None, False, "Mohon upload file CSV terlebih dahulu."
        
            try:
                chatbot = ChatBot(sess_id, model_key)
                result = chatbot.process_file(file)
                return chatbot, True, [(None, result)]
            except Exception as e:
                import traceback
                print(f"Error processing file with {model_key}: {str(e)}")
                print(traceback.format_exc())

        process_button.click(
            fn=handle_process_file,
            inputs=[file_input, model_dropdown, session_id],
            outputs=[chatbot_state, model_selected, chatbot_interface]
        ).then(
            # Disable model dropdown after processing file
            fn=lambda selected: gr.update(interactive=not selected),
            inputs=[model_selected],
            outputs=[model_dropdown]
        )

        # Reset handler - enables model selection again
        def reset_session():
            return None, False, [], gr.update(interactive=True)
            
        reset_button.click(
            fn=reset_session,
            inputs=[],
            outputs=[chatbot_state, model_selected, chatbot_interface, model_dropdown]
        )

        # Chat handlers
        def user_message_submitted(message, history, chatbot, sess_id):
            history = history + [(message, None)]
            return history, "", chatbot, sess_id

        def bot_response(history, chatbot, sess_id):
            if chatbot is None:
                chatbot = ChatBot(sess_id)
                history[-1] = (history[-1][0], "Mohon upload file CSV terlebih dahulu.")
                return chatbot, history

            user_message = history[-1][0]
            response = chatbot.chat(user_message, history[:-1])
            history[-1] = (user_message, response)
            return chatbot, history

        submit_button.click(
            fn=user_message_submitted,
            inputs=[message_input, chatbot_interface, chatbot_state, session_id],
            outputs=[chatbot_interface, message_input, chatbot_state, session_id]
        ).then(
            fn=bot_response,
            inputs=[chatbot_interface, chatbot_state, session_id],
            outputs=[chatbot_state, chatbot_interface]
        )

        message_input.submit(
            fn=user_message_submitted,
            inputs=[message_input, chatbot_interface, chatbot_state, session_id],
            outputs=[chatbot_interface, message_input, chatbot_state, session_id]
        ).then(
            fn=bot_response,
            inputs=[chatbot_interface, chatbot_state, session_id],
            outputs=[chatbot_state, chatbot_interface]
        )

        # Clear chat handler
        def handle_clear_chat(chatbot):
            if chatbot is not None:
                chatbot.chat_history = []
            return chatbot, []

        clear_button.click(
            fn=handle_clear_chat,
            inputs=[chatbot_state],
            outputs=[chatbot_state, chatbot_interface]
        )

    return interface

# Launch the interface
demo = create_gradio_interface()
demo.launch(share=True)