import gradio as gr
import os
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, pipeline
from langchain.prompts import PromptTemplate
import time
# 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": "Compact 1.1B parameter model, fast but less powerful",
"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": "Compact Microsoft model with strong instruction following",
"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 specialized for code understanding",
"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):
"""Initialize the model once and cache it"""
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 from memory if any
if MODEL_CACHE["model"] is not None:
del MODEL_CACHE["model"]
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
# Handle T5 models separately
if 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",
low_cpu_mem_usage=True
)
else:
# Load tokenizer and model with appropriate configuration
MODEL_CACHE["tokenizer"] = AutoTokenizer.from_pretrained(model_name)
MODEL_CACHE["model"] = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=model_info["dtype"],
device_map="auto",
low_cpu_mem_usage=True,
trust_remote_code=True
)
return MODEL_CACHE["tokenizer"], MODEL_CACHE["model"], model_info.get("is_t5", False)
def create_llm_pipeline(model_key):
"""Create a new pipeline using the specified model"""
tokenizer, model, is_t5 = initialize_model_once(model_key)
# Create appropriate pipeline based on model type
if is_t5:
pipe = pipeline(
"text2text-generation",
model=model,
tokenizer=tokenizer,
max_new_tokens=256,
temperature=0.3,
top_p=0.9,
return_full_text=False,
)
else:
pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_new_tokens=256,
temperature=0.3,
top_p=0.9,
top_k=30,
repetition_penalty=1.2,
return_full_text=False,
)
# Wrap pipeline in HuggingFacePipeline for LangChain compatibility
return HuggingFacePipeline(pipeline=pipe)
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:
# Handle file from Gradio
file_path = file.name if hasattr(file, 'name') else str(file)
self.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
except Exception as e:
return f"Error membaca CSV: {str(e)}"
# Load document with reduced chunk size for better memory usage
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 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:
return f"Error creating vector database: {str(e)}"
# Create custom chain
try:
self.chain = create_conversational_chain(db, self.csv_file_path, self.model_key)
print(f"Chain created successfully using model: {self.model_key}")
except Exception as e:
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"""
if model_key == self.model_key:
return f"Model {model_key} sudah digunakan."
self.model_key = model_key
# If we have an active session with a file already loaded, recreate the chain
if self.csv_file_path:
try:
# Load existing database
db_path = f"{self.user_dir}/db_faiss"
embeddings = HuggingFaceEmbeddings(
model_name='sentence-transformers/all-MiniLM-L6-v2',
model_kwargs={'device': 'cpu'}
)
db = FAISS.load_local(db_path, embeddings)
# Create new chain with the selected model
self.chain = create_conversational_chain(db, self.csv_file_path, self.model_key)
return f"Model berhasil diubah ke {model_key}."
except Exception as e:
return f"Error mengubah model: {str(e)}"
else:
return f"Model diubah ke {model_key}. Silakan upload file CSV untuk memulai."
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)
# Get model choices
model_choices = list(MODEL_CONFIG.keys())
default_model = "DeepSeek Coder Instruct" # Default model
gr.HTML("Chat with CSV using AI Models
")
gr.HTML("Asisten analisis CSV untuk berbagai kebutuhan
")
with gr.Row():
with gr.Column(scale=1):
file_input = gr.File(
label="Upload CSV Anda",
file_types=[".csv"]
)
# Model selection accordion BEFORE process button
with gr.Accordion("Pilih Model AI", open=True):
model_dropdown = gr.Dropdown(
label="Model",
choices=model_choices,
value=default_model
)
model_info = gr.Markdown(
value=f"**{default_model}**: {MODEL_CONFIG[default_model]['description']}"
)
# Process button AFTER the accordion
process_button = gr.Button("Proses CSV")
with gr.Column(scale=2):
chatbot_interface = gr.Chatbot(
label="Riwayat Chat",
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
def handle_process_file(file, model_key, sess_id):
chatbot = ChatBot(sess_id, model_key)
result = chatbot.process_file(file)
return chatbot, [(None, result)]
process_button.click(
fn=handle_process_file,
inputs=[file_input, model_dropdown, session_id],
outputs=[chatbot_state, chatbot_interface]
)
# Change model handler
def handle_model_change(model_key, chatbot, sess_id):
if chatbot is None:
chatbot = ChatBot(sess_id, model_key)
return chatbot, [(None, f"Model diatur ke {model_key}. Silakan upload file CSV.")]
result = chatbot.change_model(model_key)
return chatbot, chatbot.chat_history + [(None, result)]
model_dropdown.change(
fn=handle_model_change,
inputs=[model_dropdown, chatbot_state, session_id],
outputs=[chatbot_state, chatbot_interface]
)
# 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)