Update app.py

app.py

@@ -1,207 +1,206 @@
-import streamlit as st
+import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import time
 import os
-from threading import Thread
-import matplotlib.pyplot as plt
-import pandas as pd
 import json
+import pandas as pd
+import matplotlib.pyplot as plt
 import re
+from threading import Thread
+import numpy as np
 from io import StringIO
 
-# Set page configuration
-st.set_page_config(
-    page_title="Advanced Mistral-7B Chatbot",
-    page_icon="🤖",
-    layout="wide",
-    initial_sidebar_state="expanded"
-)
+# Global variables to store model, tokenizer and pipe
+MODEL = None
+TOKENIZER = None
+PIPE = None
+MODEL_LOADING = False
+MODEL_LOADED = False
+
+# Store chat history for different chat sessions
+CHATS = {"Main Chat": []}
+CURRENT_CHAT = "Main Chat"
 
-# Custom CSS for better appearance
-st.markdown("""
-<style>
-    .chat-message {
-        padding: 1.5rem; 
-        border-radius: 0.8rem; 
-        margin-bottom: 1rem; 
-        display: flex;
-        flex-direction: column;
-    }
-    .chat-message.user {
-        background-color: #e0f7fa;
-        border-left: 5px solid #039be5;
-    }
-    .chat-message.assistant {
-        background-color: #f1f8e9;
-        border-left: 5px solid #7cb342;
-    }
-    .chat-message .content {
-        margin-top: 0.5rem;
-    }
-    .sidebar-content {
-        padding: 1rem;
-    }
-    .sidebar-header {
-        font-weight: bold;
-        margin-bottom: 0.5rem;
-    }
-</style>
-""", unsafe_allow_html=True)
+# System prompt and generation config
+SYSTEM_PROMPT = "You are a helpful AI assistant based on the Mistral-7B-Instruct model. You specialize in creating structured JSON data for automation workflows like n8n. When asked, format JSON properly with correct indentation and structure."
+GENERATE_CONFIG = {
+    "max_new_tokens": 512,
+    "temperature": 0.7,
+    "top_p": 0.95,
+    "top_k": 50,
+    "repetition_penalty": 1.1,
+    "do_sample": True
+}
 
-# Initialize session state variables
-if 'messages' not in st.session_state:
-    st.session_state.messages = []
-if 'model' not in st.session_state:
-    st.session_state.model = None
-if 'tokenizer' not in st.session_state:
-    st.session_state.tokenizer = None
-if 'pipe' not in st.session_state:
-    st.session_state.pipe = None
-if 'loading_model' not in st.session_state:
-    st.session_state.loading_model = False
-if 'model_loaded' not in st.session_state:
-    st.session_state.model_loaded = False
-if 'system_prompt' not in st.session_state:
-    st.session_state.system_prompt = "You are a helpful AI assistant based on the Mistral-7B-Instruct model. Answer the user's questions to the best of your ability."
-if 'generate_config' not in st.session_state:
-    st.session_state.generate_config = {
-        "max_new_tokens": 512,
-        "temperature": 0.7,
-        "top_p": 0.95,
-        "top_k": 50,
-        "repetition_penalty": 1.1,
-        "do_sample": True
-    }
-if 'chats' not in st.session_state:
-    st.session_state.chats = {"Main Chat": []}
-if 'current_chat' not in st.session_state:
-    st.session_state.current_chat = "Main Chat"
-if 'file_data' not in st.session_state:
-    st.session_state.file_data = None
-if 'analyzed_data' not in st.session_state:
-    st.session_state.analyzed_data = None
+# File data storage
+FILE_DATA = None
+ANALYZED_DATA = None
 
 # Function to load the model in background
 def load_model_in_background():
+    global MODEL, TOKENIZER, PIPE, MODEL_LOADING, MODEL_LOADED
     try:
-        st.session_state.loading_model = True
+        MODEL_LOADING = True
+        
         # Model identifier
         model_id = "mistralai/Mistral-7B-Instruct-v0.3"
         
         # Load tokenizer
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        st.session_state.tokenizer = tokenizer
+        TOKENIZER = AutoTokenizer.from_pretrained(model_id)
         
         # Configure model loading (with lower precision for efficiency)
-        model = AutoModelForCausalLM.from_pretrained(
+        MODEL = AutoModelForCausalLM.from_pretrained(
             model_id,
             torch_dtype=torch.float16,
             device_map="auto",
             low_cpu_mem_usage=True
         )
-        st.session_state.model = model
         
         # Create text generation pipeline
-        st.session_state.pipe = pipeline(
+        PIPE = pipeline(
             "text-generation",
-            model=model,
-            tokenizer=tokenizer,
+            model=MODEL,
+            tokenizer=TOKENIZER,
             return_full_text=False
         )
         
-        st.session_state.loading_model = False
-        st.session_state.model_loaded = True
+        MODEL_LOADING = False
+        MODEL_LOADED = True
+        return "Model loaded successfully!"
     except Exception as e:
-        st.session_state.loading_model = False
-        st.error(f"Error loading model: {str(e)}")
-
-# Function to format and display chat messages
-def display_messages():
-    for message in st.session_state.chats[st.session_state.current_chat]:
-        with st.container():
-            col1, col2 = st.columns([1, 12])
-            
-            if message["role"] == "user":
-                with st.chat_message("user"):
-                    st.markdown(message["content"])
-            elif message["role"] == "assistant":
-                with st.chat_message("assistant"):
-                    st.markdown(message["content"])
-            elif message["role"] == "system":
-                with st.chat_message("system"):
-                    st.markdown(f"🔧 **System:** {message['content']}")
+        MODEL_LOADING = False
+        return f"Error loading model: {str(e)}"
 
 # Function to generate response using the model
-def generate_response(prompt):
-    if not st.session_state.model_loaded:
+def generate_response(prompt, chat_history, progress=gr.Progress()):
+    global MODEL, TOKENIZER, PIPE, CHATS, CURRENT_CHAT, SYSTEM_PROMPT, GENERATE_CONFIG, FILE_DATA, ANALYZED_DATA
+    
+    if not MODEL_LOADED:
         return "Model is still loading. Please wait a moment before sending messages."
     
     try:
-        messages = st.session_state.chats[st.session_state.current_chat]
+        # Use the current chat's history
+        messages = CHATS[CURRENT_CHAT]
         
         # Format conversation history in Mistral's chat format
         conversation = []
         
         # Add system prompt if it exists
-        if st.session_state.system_prompt:
-            conversation.append({"role": "system", "content": st.session_state.system_prompt})
+        if SYSTEM_PROMPT:
+            conversation.append({"role": "system", "content": SYSTEM_PROMPT})
         
         # Add previous messages
         for msg in messages:
             if msg["role"] != "system":  # Skip system messages in the history
                 conversation.append({"role": msg["role"], "content": msg["content"]})
         
+        # Handle file-related queries by including context
+        if ANALYZED_DATA is not None and any(keyword in prompt.lower() 
+                                             for keyword in ["file", "data", "analyze", "show", "tell me about", "json"]):
+            file_context = ""
+            if ANALYZED_DATA["type"] in ["csv", "excel"]:
+                # For structured data, provide summary info
+                summary = ANALYZED_DATA["summary"]
+                file_context = f"""I've analyzed the uploaded {ANALYZED_DATA["type"]} file with {summary["rows"]} rows and {summary["columns"]} columns.
+                The columns are: {', '.join(summary["column_names"])}.
+                Here's a sample of the data (first 5 rows): {json.dumps(summary["sample"])}
+                """
+            elif ANALYZED_DATA["type"] == "text":
+                # For text data, provide the content if not too large
+                summary = ANALYZED_DATA["summary"]
+                file_context = f"""I've analyzed the uploaded text file with {summary["word_count"]} words and {summary["line_count"]} lines.
+                Here's a preview of the content: {summary["preview"]}
+                """
+            elif ANALYZED_DATA["type"] == "json":
+                # For JSON data
+                summary = ANALYZED_DATA["summary"]
+                file_context = f"""I've analyzed the uploaded JSON file which contains a {summary["type"]}.
+                {"Keys: " + ', '.join(summary["keys"]) if summary["keys"] else ""}
+                {"Items: " + str(summary["length"]) if summary["length"] else ""}
+                Here's a preview: {summary["preview"]}
+                """
+                
+            # Enhance the user's query with file context
+            enhanced_prompt = f"{prompt}\n\nContext about the file: {file_context}"
+        else:
+            enhanced_prompt = prompt
+            
         # Add current prompt
-        conversation.append({"role": "user", "content": prompt})
+        conversation.append({"role": "user", "content": enhanced_prompt})
         
         # Convert to Mistral's chat format
-        formatted_prompt = st.session_state.tokenizer.apply_chat_template(
+        formatted_prompt = TOKENIZER.apply_chat_template(
             conversation, 
             tokenize=False,
             add_generation_prompt=True
         )
         
-        # Generate response
-        generation_config = st.session_state.generate_config
+        # Generate response with progress reporting
+        progress(0, desc="Generating response...")
         
-        response = st.session_state.pipe(
+        # Generate response
+        response = PIPE(
             formatted_prompt,
-            max_new_tokens=generation_config["max_new_tokens"],
-            temperature=generation_config["temperature"],
-            top_p=generation_config["top_p"],
-            top_k=generation_config["top_k"],
-            repetition_penalty=generation_config["repetition_penalty"],
-            do_sample=generation_config["do_sample"]
+            max_new_tokens=GENERATE_CONFIG["max_new_tokens"],
+            temperature=GENERATE_CONFIG["temperature"],
+            top_p=GENERATE_CONFIG["top_p"],
+            top_k=GENERATE_CONFIG["top_k"],
+            repetition_penalty=GENERATE_CONFIG["repetition_penalty"],
+            do_sample=GENERATE_CONFIG["do_sample"]
         )
         
-        # Extract and return generated text
+        progress(1, desc="Response generated!")
+        
+        # Extract generated text
         generated_text = response[0]["generated_text"]
-        return generated_text
+        
+        # Add user message to chat history
+        CHATS[CURRENT_CHAT].append({
+            "role": "user",
+            "content": prompt
+        })
+        
+        # Add assistant response to chat history
+        CHATS[CURRENT_CHAT].append({
+            "role": "assistant",
+            "content": generated_text
+        })
+        
+        # Update the chat history for the Gradio component
+        chat_history.append((prompt, generated_text))
+        
+        return chat_history
     
     except Exception as e:
-        return f"Error generating response: {str(e)}"
+        error_message = f"Error generating response: {str(e)}"
+        chat_history.append((prompt, error_message))
+        return chat_history
 
 # Function to create a new chat
 def create_new_chat(chat_name):
-    if chat_name and chat_name not in st.session_state.chats:
-        st.session_state.chats[chat_name] = []
-        st.session_state.current_chat = chat_name
-        return True
-    return False
+    global CHATS, CURRENT_CHAT
+    
+    if chat_name and chat_name not in CHATS:
+        CHATS[chat_name] = []
+        CURRENT_CHAT = chat_name
+        return f"Created new chat: {chat_name}"
+    return "Please enter a unique chat name"
 
 # Function to handle file upload and analysis
 def analyze_uploaded_file(file):
-    if file is None:
-        return None
+    global FILE_DATA, ANALYZED_DATA, CHATS, CURRENT_CHAT
     
-    file_extension = file.name.split('.')[-1].lower()
+    if file is None:
+        return "No file uploaded."
     
     try:
+        file_extension = file.name.split('.')[-1].lower()
+        
         if file_extension == 'csv':
-            data = pd.read_csv(file)
-            st.session_state.file_data = data
-            return {
+            data = pd.read_csv(file.name)
+            FILE_DATA = data
+            ANALYZED_DATA = {
                 "type": "csv",
                 "data": data,
                 "summary": {
@@ -214,9 +213,10 @@ def analyze_uploaded_file(file):
             }
         
         elif file_extension in ['txt', 'md']:
-            content = file.getvalue().decode('utf-8')
-            st.session_state.file_data = content
-            return {
+            with open(file.name, 'r', encoding='utf-8') as f:
+                content = f.read()
+            FILE_DATA = content
+            ANALYZED_DATA = {
                 "type": "text",
                 "data": content,
                 "summary": {
@@ -228,10 +228,11 @@ def analyze_uploaded_file(file):
             }
         
         elif file_extension == 'json':
-            content = file.getvalue().decode('utf-8')
-            data = json.loads(content)
-            st.session_state.file_data = data
-            return {
+            with open(file.name, 'r', encoding='utf-8') as f:
+                content = f.read()
+                data = json.loads(content)
+            FILE_DATA = data
+            ANALYZED_DATA = {
                 "type": "json",
                 "data": data,
                 "summary": {
@@ -243,9 +244,9 @@ def analyze_uploaded_file(file):
             }
         
         elif file_extension in ['xls', 'xlsx']:
-            data = pd.read_excel(file)
-            st.session_state.file_data = data
-            return {
+            data = pd.read_excel(file.name)
+            FILE_DATA = data
+            ANALYZED_DATA = {
                 "type": "excel",
                 "data": data,
                 "summary": {
@@ -258,349 +259,297 @@ def analyze_uploaded_file(file):
             }
         
         else:
-            return {
-                "type": "unsupported",
-                "error": f"File type .{file_extension} is not supported."
-            }
+            return f"File type .{file_extension} is not supported."
+        
+        # Add file summary to current chat as system message
+        file_summary = f"File analyzed: {file.name}\n"
+        if ANALYZED_DATA['type'] == 'csv' or ANALYZED_DATA['type'] == 'excel':
+            file_summary += f"- {ANALYZED_DATA['summary']['rows']} rows, {ANALYZED_DATA['summary']['columns']} columns\n"
+            file_summary += f"- Columns: {', '.join(ANALYZED_DATA['summary']['column_names'])}"
+        elif ANALYZED_DATA['type'] == 'text':
+            file_summary += f"- {ANALYZED_DATA['summary']['word_count']} words, {ANALYZED_DATA['summary']['line_count']} lines"
+        elif ANALYZED_DATA['type'] == 'json':
+            file_summary += f"- Type: {ANALYZED_DATA['summary']['type']}"
+            if ANALYZED_DATA['summary']['keys']:
+                file_summary += f"\n- Keys: {', '.join(ANALYZED_DATA['summary']['keys'])}"
+        
+        # Add system message to current chat
+        CHATS[CURRENT_CHAT].append({
+            "role": "system",
+            "content": file_summary
+        })
+        
+        return f"Successfully analyzed {ANALYZED_DATA['type']} file: {file.name}"
     
     except Exception as e:
-        return {
-            "type": "error",
-            "error": str(e)
-        }
-
-# Sidebar - Model Control and Settings
-with st.sidebar:
-    st.title("🤖 Mistral-7B Chatbot")
-    
-    # Model loading section
-    st.header("Model Control")
-    if not st.session_state.model_loaded and not st.session_state.loading_model:
-        if st.button("Load Mistral-7B Model"):
-            loading_thread = Thread(target=load_model_in_background)
-            loading_thread.start()
-    
-    if st.session_state.loading_model:
-        st.info("Loading model... This may take a few minutes.")
-        progress_bar = st.progress(0)
-        for i in range(100):
-            if not st.session_state.loading_model:
-                progress_bar.progress(100)
-                st.success("Model loaded successfully!")
-                break
-            progress_bar.progress(i + 1)
-            time.sleep(0.1)
-    
-    if st.session_state.model_loaded:
-        st.success("Model loaded and ready!")
-    
-    # Chat management
-    st.header("Chat Management")
-    
-    # Select chat
-    selected_chat = st.selectbox(
-        "Select Chat", 
-        options=list(st.session_state.chats.keys()),
-        index=list(st.session_state.chats.keys()).index(st.session_state.current_chat)
-    )
-    
-    if selected_chat != st.session_state.current_chat:
-        st.session_state.current_chat = selected_chat
-    
-    # Create new chat
-    new_chat_name = st.text_input("New Chat Name")
-    if st.button("Create New Chat"):
-        if create_new_chat(new_chat_name):
-            st.success(f"Created new chat: {new_chat_name}")
-        else:
-            st.error("Please enter a unique chat name")
-    
-    # Clear current chat
-    if st.button("Clear Current Chat"):
-        st.session_state.chats[st.session_state.current_chat] = []
-        st.success(f"Cleared chat: {st.session_state.current_chat}")
-    
-    # Advanced settings
-    st.header("Model Settings")
-    
-    with st.expander("System Prompt"):
-        system_prompt = st.text_area(
-            "Set the AI's behavior",
-            value=st.session_state.system_prompt,
-            height=100
-        )
-        if st.button("Update System Prompt"):
-            st.session_state.system_prompt = system_prompt
-            st.success("System prompt updated!")
-    
-    with st.expander("Generation Parameters"):
-        st.slider(
-            "Temperature",
-            min_value=0.1,
-            max_value=2.0,
-            value=st.session_state.generate_config["temperature"],
-            step=0.1,
-            key="temp_slider",
-            help="Higher = more creative, Lower = more focused"
-        )
-        st.session_state.generate_config["temperature"] = st.session_state["temp_slider"]
-        
-        st.slider(
-            "Max Tokens",
-            min_value=64,
-            max_value=2048,
-            value=st.session_state.generate_config["max_new_tokens"],
-            step=64,
-            key="max_tokens_slider",
-            help="Maximum number of tokens in the response"
-        )
-        st.session_state.generate_config["max_new_tokens"] = st.session_state["max_tokens_slider"]
-        
-        st.slider(
-            "Top P",
-            min_value=0.1,
-            max_value=1.0,
-            value=st.session_state.generate_config["top_p"],
-            step=0.05,
-            key="top_p_slider",
-            help="Nucleus sampling parameter"
-        )
-        st.session_state.generate_config["top_p"] = st.session_state["top_p_slider"]
-        
-        st.slider(
-            "Repetition Penalty",
-            min_value=1.0,
-            max_value=2.0,
-            value=st.session_state.generate_config["repetition_penalty"],
-            step=0.1,
-            key="rep_pen_slider",
-            help="Penalizes repetition (higher = less repetition)"
-        )
-        st.session_state.generate_config["repetition_penalty"] = st.session_state["rep_pen_slider"]
-    
-    # File upload and analysis
-    st.header("File Analysis")
-    uploaded_file = st.file_uploader("Upload a file to analyze", 
-                                      type=["csv", "txt", "json", "xlsx", "xls", "md"])
-    
-    if uploaded_file is not None:
-        if st.button("Analyze File"):
-            with st.spinner("Analyzing file..."):
-                analysis_result = analyze_uploaded_file(uploaded_file)
-                st.session_state.analyzed_data = analysis_result
-                
-                if analysis_result:
-                    if "error" in analysis_result:
-                        st.error(f"Error analyzing file: {analysis_result['error']}")
-                    else:
-                        st.success(f"Successfully analyzed {analysis_result['type']} file")
-                        
-                        # Add file summary to chat as system message
-                        file_summary = f"File analyzed: {uploaded_file.name}\n"
-                        if analysis_result['type'] == 'csv' or analysis_result['type'] == 'excel':
-                            file_summary += f"- {analysis_result['summary']['rows']} rows, {analysis_result['summary']['columns']} columns\n"
-                            file_summary += f"- Columns: {', '.join(analysis_result['summary']['column_names'])}"
-                        elif analysis_result['type'] == 'text':
-                            file_summary += f"- {analysis_result['summary']['word_count']} words, {analysis_result['summary']['line_count']} lines"
-                        
-                        # Add file summary to current chat
-                        st.session_state.chats[st.session_state.current_chat].append({
-                            "role": "system",
-                            "content": file_summary
-                        })
+        return f"Error analyzing file: {str(e)}"
 
-# Main chat interface
-st.header(f"💬 Chat: {st.session_state.current_chat}")
+# Function to update system prompt
+def update_system_prompt(new_prompt):
+    global SYSTEM_PROMPT
+    SYSTEM_PROMPT = new_prompt
+    return f"System prompt updated!"
 
-# Display chat messages
-display_messages()
+# Function to update generation parameters
+def update_generation_params(temp, max_tokens, top_p, rep_penalty):
+    global GENERATE_CONFIG
+    GENERATE_CONFIG["temperature"] = temp
+    GENERATE_CONFIG["max_new_tokens"] = max_tokens
+    GENERATE_CONFIG["top_p"] = top_p
+    GENERATE_CONFIG["repetition_penalty"] = rep_penalty
+    return f"Generation parameters updated!"
 
-# User input
-if user_prompt := st.chat_input("Type your message here...", disabled=not st.session_state.model_loaded and not st.session_state.loading_model):
-    # Add user message to chat history
-    st.session_state.chats[st.session_state.current_chat].append({
-        "role": "user",
-        "content": user_prompt
-    })
+# Function to display file data information
+def display_file_info():
+    global ANALYZED_DATA
     
-    # Display user message
-    with st.chat_message("user"):
-        st.markdown(user_prompt)
+    if ANALYZED_DATA is None:
+        return "No file has been analyzed yet."
     
-    # Generate and display assistant response
-    with st.chat_message("assistant"):
-        message_placeholder = st.empty()
-        message_placeholder.markdown("Thinking...")
+    file_info = f"## File Analysis Results\n\n"
+    file_info += f"**File Type:** {ANALYZED_DATA['type']}\n\n"
+    
+    if ANALYZED_DATA['type'] in ['csv', 'excel']:
+        summary = ANALYZED_DATA['summary']
+        file_info += f"**Rows:** {summary['rows']}\n"
+        file_info += f"**Columns:** {summary['columns']}\n"
+        file_info += f"**Column Names:** {', '.join(summary['column_names'])}\n\n"
         
-        # Handle file-related queries by including context
-        if st.session_state.analyzed_data is not None and any(keyword in user_prompt.lower() 
-                                                        for keyword in ["file", "data", "analyze", "show", "tell me about"]):
-            file_context = ""
-            if st.session_state.analyzed_data["type"] in ["csv", "excel"]:
-                # For structured data, provide summary info
-                summary = st.session_state.analyzed_data["summary"]
-                file_context = f"""I've analyzed the uploaded {st.session_state.analyzed_data["type"]} file with {summary["rows"]} rows and {summary["columns"]} columns.
-                The columns are: {', '.join(summary["column_names"])}.
-                Here's a sample of the data (first 5 rows): {json.dumps(summary["sample"])}
-                """
-            elif st.session_state.analyzed_data["type"] == "text":
-                # For text data, provide the content if not too large
-                summary = st.session_state.analyzed_data["summary"]
-                file_context = f"""I've analyzed the uploaded text file with {summary["word_count"]} words and {summary["line_count"]} lines.
-                Here's a preview of the content: {summary["preview"]}
-                """
-            elif st.session_state.analyzed_data["type"] == "json":
-                # For JSON data
-                summary = st.session_state.analyzed_data["summary"]
-                file_context = f"""I've analyzed the uploaded JSON file which contains a {summary["type"]}.
-                {"Keys: " + ', '.join(summary["keys"]) if summary["keys"] else ""}
-                {"Items: " + str(summary["length"]) if summary["length"] else ""}
-                Here's a preview: {summary["preview"]}
-                """
-                
-            # Enhance the user's query with file context
-            enhanced_prompt = f"{user_prompt}\n\nContext about the file: {file_context}"
-        else:
-            enhanced_prompt = user_prompt
-            
-        # Generate response
-        response = generate_response(enhanced_prompt)
+        # Sample data preview
+        file_info += "**Sample Data (First 5 rows):**\n"
+        sample_df = pd.DataFrame(summary['sample'])
+        file_info += sample_df.to_markdown()
         
-        # Update the placeholder with the complete response
-        message_placeholder.markdown(response)
+    elif ANALYZED_DATA['type'] == 'text':
+        summary = ANALYZED_DATA['summary']
+        file_info += f"**Length:** {summary['length']} characters\n"
+        file_info += f"**Word Count:** {summary['word_count']}\n"
+        file_info += f"**Line Count:** {summary['line_count']}\n\n"
+        file_info += "**Preview:**\n```\n" + summary['preview'] + "\n```"
         
-        # Add assistant response to chat history
-        st.session_state.chats[st.session_state.current_chat].append({
-            "role": "assistant",
-            "content": response
-        })
+    elif ANALYZED_DATA['type'] == 'json':
+        summary = ANALYZED_DATA['summary']
+        file_info += f"**Type:** {summary['type']}\n"
+        if summary['keys']:
+            file_info += f"**Keys:** {', '.join(summary['keys'])}\n"
+        if summary['length'] is not None:
+            file_info += f"**Length:** {summary['length']} items\n"
+        file_info += "\n**Preview:**\n```json\n" + summary['preview'] + "\n```"
+    
+    return file_info
+
+# Function to select current chat
+def select_chat(chat_name):
+    global CURRENT_CHAT
+    CURRENT_CHAT = chat_name
+    return f"Switched to chat: {chat_name}"
+
+# Function to clear current chat
+def clear_current_chat():
+    global CHATS, CURRENT_CHAT
+    CHATS[CURRENT_CHAT] = []
+    return f"Cleared chat: {CURRENT_CHAT}"
+
+# Function to load model and return status
+def load_model_button():
+    if MODEL_LOADED:
+        return "Model is already loaded and ready!"
+    elif MODEL_LOADING:
+        return "Model is currently loading... Please wait."
+    else:
+        thread = Thread(target=load_model_in_background)
+        thread.start()
+        return "Started loading the model. This may take a few minutes..."
+
+# Function to get available chats
+def get_available_chats():
+    global CHATS
+    return list(CHATS.keys())
 
-# Display data analysis tools when file is uploaded
-if st.session_state.analyzed_data is not None and st.session_state.file_data is not None:
-    st.header("📊 Data Analysis")
+# Main Gradio app
+def create_gradio_interface():
+    # CSS for better styling
+    css = """
+    .gradio-container {max-width: 100% !important; padding: 0}
+    .chat-message-user {background-color: #e0f7fa; padding: 12px; border-radius: 8px; margin-bottom: 8px}
+    .chat-message-bot {background-color: #f1f8e9; padding: 12px; border-radius: 8px; margin-bottom: 8px}
+    .file-info {border: 1px solid #ddd; padding: 15px; border-radius: 5px; margin-top: 10px}
+    """
     
-    if st.session_state.analyzed_data["type"] in ["csv", "excel"]:
-        data = st.session_state.file_data
+    # Setup tabs for different functionalities
+    with gr.Blocks(css=css) as app:
+        gr.Markdown("# 🤖 Advanced Mistral-7B-Instruct Chatbot")
         
-        # Data exploration tabs
-        tab1, tab2, tab3 = st.tabs(["Data Preview", "Statistics", "Visualization"])
-        
-        with tab1:
-            st.dataframe(data.head(10))
-            
-            # Simple search/filter
-            search_col = st.selectbox("Search column", [""] + list(data.columns))
-            if search_col:
-                search_term = st.text_input("Search term")
-                if search_term:
-                    filtered_data = data[data[search_col].astype(str).str.contains(search_term, case=False)]
-                    st.dataframe(filtered_data)
-        
-        with tab2:
-            st.write("### Numerical Statistics")
-            st.dataframe(data.describe())
-            
-            st.write("### Column Information")
-            col_info = pd.DataFrame({
-                "Data Type": data.dtypes,
-                "Non-Null Count": data.count(),
-                "Null Count": data.isna().sum(),
-                "Unique Values": [data[col].nunique() for col in data.columns]
-            })
-            st.dataframe(col_info)
-        
-        with tab3:
-            st.write("### Data Visualization")
-            
-            # Simple chart creation
-            chart_type = st.selectbox("Chart Type", ["Bar Chart", "Line Chart", "Scatter Plot", "Histogram"])
-            
-            if chart_type in ["Bar Chart", "Line Chart"]:
-                x_axis = st.selectbox("X-Axis", list(data.columns), key="x_axis_1")
-                y_axis = st.selectbox("Y-Axis", list(data.columns), key="y_axis_1")
-                
-                if x_axis and y_axis:
-                    fig, ax = plt.subplots(figsize=(10, 6))
+        with gr.Tab("Chat"):
+            with gr.Row():
+                with gr.Column(scale=3):
+                    chatbot = gr.Chatbot(
+                        [],
+                        elem_id="chatbot",
+                        height=500,
+                        bubble_full_width=False,
+                        avatar_images=(None, None),
+                    )
                     
-                    if chart_type == "Bar Chart":
-                        data.groupby(x_axis)[y_axis].mean().plot(kind='bar', ax=ax)
-                    else:  # Line Chart
-                        data.groupby(x_axis)[y_axis].mean().plot(kind='line', ax=ax)
+                    with gr.Row():
+                        msg = gr.Textbox(
+                            placeholder="Type your message here...",
+                            container=False,
+                            scale=8,
+                            autofocus=True
+                        )
+                        send_btn = gr.Button("Send", scale=1, variant="primary")
                     
-                    plt.xlabel(x_axis)
-                    plt.ylabel(y_axis)
-                    plt.title(f"{chart_type} of {y_axis} by {x_axis}")
-                    plt.xticks(rotation=45)
-                    plt.tight_layout()
-                    st.pyplot(fig)
-            
-            elif chart_type == "Scatter Plot":
-                x_axis = st.selectbox("X-Axis", list(data.columns), key="x_axis_2")
-                y_axis = st.selectbox("Y-Axis", list(data.columns), key="y_axis_2")
+                    with gr.Row():
+                        chat_selector = gr.Dropdown(
+                            choices=get_available_chats(),
+                            value=CURRENT_CHAT,
+                            label="Select Chat",
+                            interactive=True
+                        )
+                        
+                        with gr.Column(scale=1):
+                            new_chat_name = gr.Textbox(
+                                placeholder="New chat name",
+                                label="Create New Chat",
+                                container=True
+                            )
+                            with gr.Row():
+                                create_chat_btn = gr.Button("Create", variant="secondary")
+                                clear_chat_btn = gr.Button("Clear Current Chat", variant="secondary")
                 
-                if x_axis and y_axis:
-                    fig, ax = plt.subplots(figsize=(10, 6))
-                    ax.scatter(data[x_axis], data[y_axis])
-                    plt.xlabel(x_axis)
-                    plt.ylabel(y_axis)
-                    plt.title(f"Scatter Plot of {y_axis} vs {x_axis}")
-                    plt.tight_layout()
-                    st.pyplot(fig)
-            
-            elif chart_type == "Histogram":
-                column = st.selectbox("Column", list(data.columns))
-                bins = st.slider("Bins", min_value=5, max_value=100, value=20)
+                with gr.Column(scale=1):
+                    # Model Loading and Settings
+                    load_model_btn = gr.Button("Load Mistral-7B Model", variant="primary")
+                    model_status = gr.Textbox(label="Model Status", value="Not loaded", interactive=False)
+                    
+                    # System Prompt
+                    system_prompt_input = gr.Textbox(
+                        label="System Prompt", 
+                        value=SYSTEM_PROMPT,
+                        lines=4,
+                        placeholder="Enter system prompt to guide the AI's behavior..."
+                    )
+                    update_prompt_btn = gr.Button("Update System Prompt", variant="secondary")
+                    
+                    # Model Parameters
+                    gr.Markdown("### Generation Parameters")
+                    temperature = gr.Slider(
+                        minimum=0.1, 
+                        maximum=2.0, 
+                        value=GENERATE_CONFIG["temperature"], 
+                        step=0.1,
+                        label="Temperature"
+                    )
+                    max_tokens = gr.Slider(
+                        minimum=64, 
+                        maximum=2048, 
+                        value=GENERATE_CONFIG["max_new_tokens"], 
+                        step=64,
+                        label="Max Tokens"
+                    )
+                    top_p = gr.Slider(
+                        minimum=0.1, 
+                        maximum=1.0, 
+                        value=GENERATE_CONFIG["top_p"], 
+                        step=0.05,
+                        label="Top P"
+                    )
+                    rep_penalty = gr.Slider(
+                        minimum=1.0, 
+                        maximum=2.0, 
+                        value=GENERATE_CONFIG["repetition_penalty"], 
+                        step=0.1,
+                        label="Repetition Penalty"
+                    )
+                    update_params_btn = gr.Button("Update Parameters", variant="secondary")
+        
+        with gr.Tab("File Analysis"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    file_upload = gr.File(label="Upload a file to analyze")
+                    analyze_btn = gr.Button("Analyze File", variant="primary")
                 
-                if column:
-                    fig, ax = plt.subplots(figsize=(10, 6))
-                    ax.hist(data[column].dropna(), bins=bins)
-                    plt.xlabel(column)
-                    plt.ylabel("Frequency")
-                    plt.title(f"Histogram of {column}")
-                    plt.tight_layout()
-                    st.pyplot(fig)
-    
-    elif st.session_state.analyzed_data["type"] == "text":
-        text_data = st.session_state.file_data
+                with gr.Column(scale=2):
+                    file_analysis_output = gr.Markdown(label="File Analysis Results")
         
-        # Text analysis tools
-        tab1, tab2 = st.tabs(["Text Preview", "Text Analysis"])
+        # Set up event handlers
+        send_btn.click(
+            generate_response, 
+            inputs=[msg, chatbot], 
+            outputs=[chatbot], 
+            api_name="chat"
+        )
         
-        with tab1:
-            st.text_area("Content", text_data, height=300)
+        msg.submit(
+            generate_response, 
+            inputs=[msg, chatbot], 
+            outputs=[chatbot], 
+            api_name=False
+        )
         
-        with tab2:
-            # Basic text analysis
-            word_count = len(text_data.split())
-            char_count = len(text_data)
-            line_count = len(text_data.splitlines())
-            
-            col1, col2, col3 = st.columns(3)
-            col1.metric("Word Count", word_count)
-            col2.metric("Character Count", char_count)
-            col3.metric("Line Count", line_count)
-            
-            # Word frequency
-            if st.checkbox("Show Word Frequency"):
-                words = re.findall(r'\b\w+\b', text_data.lower())
-                word_freq = pd.Series(words).value_counts().head(20)
-                
-                st.write("### Most Common Words")
-                fig, ax = plt.subplots(figsize=(10, 6))
-                word_freq.plot(kind='bar', ax=ax)
-                plt.xlabel("Word")
-                plt.ylabel("Frequency")
-                plt.title("Top 20 Most Common Words")
-                plt.xticks(rotation=45)
-                plt.tight_layout()
-                st.pyplot(fig)
-    
-    elif st.session_state.analyzed_data["type"] == "json":
-        json_data = st.session_state.file_data
+        load_model_btn.click(
+            load_model_button, 
+            outputs=model_status, 
+            api_name="load_model"
+        )
+        
+        update_prompt_btn.click(
+            update_system_prompt, 
+            inputs=system_prompt_input, 
+            outputs=model_status, 
+            api_name="update_prompt"
+        )
+        
+        update_params_btn.click(
+            update_generation_params,
+            inputs=[temperature, max_tokens, top_p, rep_penalty],
+            outputs=model_status,
+            api_name="update_params"
+        )
+        
+        analyze_btn.click(
+            analyze_uploaded_file,
+            inputs=file_upload,
+            outputs=model_status,
+            api_name="analyze_file"
+        ).then(
+            display_file_info,
+            outputs=file_analysis_output,
+            api_name=False
+        )
         
-        # JSON explorer
-        st.json(json_data)
+        chat_selector.change(
+            select_chat,
+            inputs=chat_selector,
+            outputs=model_status,
+            api_name="select_chat"
+        )
+        
+        create_chat_btn.click(
+            create_new_chat,
+            inputs=new_chat_name,
+            outputs=model_status,
+            api_name="create_chat"
+        ).then(
+            get_available_chats,
+            outputs=chat_selector,
+            api_name=False
+        )
+        
+        clear_chat_btn.click(
+            clear_current_chat,
+            outputs=model_status,
+            api_name="clear_chat"
+        )
+        
+        # Refresh example messages every few seconds
+        chat_selector.change(lambda: [], outputs=chatbot)
+
+    return app
+
+# Launch the app
+demo = create_gradio_interface()
 
-# Footer
-st.markdown("---")
-st.markdown("Advanced Mistral-7B Chatbot | Built with Streamlit & Hugging Face")
+if __name__ == "__main__":
+    demo.launch()