UI fixes 5

app.py

@@ -1,188 +1,145 @@
 import streamlit as st
 import torch
 from model_handler import load_model_and_tokenizer, generate_response
-import time
 
 # Page configuration
 st.set_page_config(
     page_title="Phi-2 Fine-tuned Assistant",
     page_icon="🤖",
-    layout="centered",
+    layout="wide",
     initial_sidebar_state="collapsed"
 )
 
-# Custom CSS for chat interface - fixed scrolling issues
+# Clean, minimal CSS
 st.markdown("""
 <style>
-/* Reset Streamlit default spacing */
-.main .block-container {
-    padding-top: 1rem;
-    padding-bottom: 1rem;
-    max-width: 800px;
-}
-
-/* Fix header spacing */
-header {
-    visibility: hidden;
-}
-
-/* Hide main menu and footer */
-#MainMenu {visibility: hidden;}
-footer {visibility: hidden;}
-
-/* Chat bubbles */
-.user-bubble {
-    background-color: #e6f7ff;
-    border-radius: 15px;
-    padding: 10px 15px;
-    margin: 5px 0;
-    max-width: 80%;
-    margin-left: auto;
-    margin-right: 10px;
-    color: #000000;
-}
-.assistant-bubble {
-    background-color: #f0f0f0;
-    border-radius: 15px;
-    padding: 10px 15px;
-    margin: 5px 0;
-    max-width: 80%;
-    margin-left: 10px;
-    color: #000000;
-}
-
-/* Chat container - fixed height with single scrollbar */
-.chat-container {
-    display: flex;
-    flex-direction: column;
-    height: 60vh;
-    overflow-y: auto;
-    padding: 10px;
-    margin-bottom: 10px;
-    border-radius: 10px;
-    background-color: rgba(255, 255, 255, 0.05);
-}
-
-/* Input styling */
-.stTextInput>div>div>input {
-    border-radius: 20px;
-}
-.stButton>button {
-    border-radius: 20px;
-    width: 100%;
-}
-
-/* Ensure dark mode compatibility */
-.stApp {
-    background-color: #121212;
-}
-
-/* Title styling */
-h2 {
-    margin-top: 0 !important;
-    margin-bottom: 1rem !important;
-    text-align: center;
-}
-
-/* Loading indicator */
-.loading-container {
-    display: flex;
-    justify-content: center;
-    align-items: center;
-    height: 60vh;
-}
+    /* Remove all default padding and margins */
+    .block-container {
+        padding: 0 !important;
+        max-width: 100% !important;
+    }
+    
+    /* Hide header and footer */
+    header, footer, #MainMenu {
+        visibility: hidden;
+    }
+    
+    /* Main container */
+    .main-container {
+        max-width: 800px;
+        margin: 0 auto;
+        padding: 20px;
+    }
+    
+    /* Title */
+    .title {
+        text-align: center;
+        margin-bottom: 20px;
+    }
+    
+    /* Chat messages */
+    .user-message {
+        background-color: #e6f7ff;
+        border-radius: 15px;
+        padding: 10px 15px;
+        margin: 5px 0;
+        max-width: 80%;
+        margin-left: auto;
+        margin-right: 10px;
+        color: #000;
+    }
+    
+    .assistant-message {
+        background-color: #f0f0f0;
+        border-radius: 15px;
+        padding: 10px 15px;
+        margin: 5px 0;
+        max-width: 80%;
+        margin-left: 10px;
+        color: #000;
+    }
+    
+    /* Input area */
+    .input-area {
+        display: flex;
+        margin-top: 20px;
+    }
+    
+    /* Loading message */
+    .loading {
+        text-align: center;
+        padding: 20px;
+    }
 </style>
 """, unsafe_allow_html=True)
 
-# Initialize session state for chat history
+# Initialize session state
 if "messages" not in st.session_state:
     st.session_state.messages = []
-
-if "model_loaded" not in st.session_state:
-    st.session_state.model_loaded = False
-
-# App title
-st.markdown("<h2>Phi-2 Fine-tuned Assistant</h2>", unsafe_allow_html=True)
-
-# Load model (only once)
-if not st.session_state.model_loaded:
-    # Show loading indicator in the chat area
-    st.markdown('<div class="loading-container"><div>Loading the fine-tuned model... This may take a minute.</div></div>', unsafe_allow_html=True)
-    
-    # Load model in the background
-    model, tokenizer = load_model_and_tokenizer()
-    st.session_state.model = model
-    st.session_state.tokenizer = tokenizer
-    st.session_state.model_loaded = True
-    st.rerun()  # Rerun to refresh UI after loading
-
-# Chat interface (only show when model is loaded)
-if st.session_state.model_loaded:
-    # Display chat messages
-    st.markdown('<div class="chat-container" id="chat-container">', unsafe_allow_html=True)
     
-    # Display existing messages
-    for message in st.session_state.messages:
-        role = message["role"]
-        content = message["content"]
-        
-        if role == "user":
-            st.markdown(f'<div class="user-bubble">{content}</div>', unsafe_allow_html=True)
-        else:
-            st.markdown(f'<div class="assistant-bubble">{content}</div>', unsafe_allow_html=True)
-    
-    st.markdown('</div>', unsafe_allow_html=True)
+if "model" not in st.session_state:
+    st.session_state.model = None
     
-    # Auto-scroll to bottom using JavaScript
-    if st.session_state.messages:
-        st.markdown("""
-        <script>
-        function scrollToBottom() {
-            var chatContainer = document.getElementById('chat-container');
-            if (chatContainer) {
-                chatContainer.scrollTop = chatContainer.scrollHeight;
-            }
-        }
-        scrollToBottom();
-        </script>
-        """, unsafe_allow_html=True)
-    
-    # Chat input
-    col1, col2 = st.columns([5, 1])
+if "tokenizer" not in st.session_state:
+    st.session_state.tokenizer = None
+
+# Main app container
+st.markdown('<div class="main-container">', unsafe_allow_html=True)
+st.markdown('<h1 class="title">Phi-2 Fine-tuned Assistant</h1>', unsafe_allow_html=True)
+
+# Load model if not already loaded
+if st.session_state.model is None:
+    st.markdown('<div class="loading">Loading model... This may take a minute.</div>', unsafe_allow_html=True)
+    try:
+        model, tokenizer = load_model_and_tokenizer()
+        st.session_state.model = model
+        st.session_state.tokenizer = tokenizer
+        st.rerun()
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        st.stop()
+
+# Display chat messages
+for message in st.session_state.messages:
+    if message["role"] == "user":
+        st.markdown(f'<div class="user-message">{message["content"]}</div>', unsafe_allow_html=True)
+    else:
+        st.markdown(f'<div class="assistant-message">{message["content"]}</div>', unsafe_allow_html=True)
+
+# Input form with submit button
+with st.form(key="message_form", clear_on_submit=True):
+    user_input = st.text_input("Message:", key="user_input")
+    col1, col2 = st.columns([4, 1])
     with col1:
-        user_input = st.text_input("", placeholder="Type your message here...")
+        submit_button = st.form_submit_button("Send")
     with col2:
-        clear_button = st.button("Clear")
+        clear_button = st.form_submit_button("Clear Chat")
+
+# Handle form submission
+if submit_button and user_input:
+    # Add user message
+    st.session_state.messages.append({"role": "user", "content": user_input})
     
-    # Process input when user submits a message
-    if user_input:
-        # Add user message to chat history
-        st.session_state.messages.append({"role": "user", "content": user_input})
-        
-        # Generate response
-        try:
-            # Show typing indicator
-            temp_messages = st.session_state.messages.copy()
-            temp_messages.append({"role": "assistant", "content": "Assistant is typing..."})
-            
-            # Generate full response first (no streaming in this version to avoid UI issues)
-            full_response = ""
-            for token in generate_response(
-                st.session_state.model, 
-                st.session_state.tokenizer, 
-                user_input
-            ):
-                full_response += token
-            
-            # Add assistant's response to chat history
-            st.session_state.messages.append({"role": "assistant", "content": full_response})
-        except Exception as e:
-            st.session_state.messages.append({"role": "assistant", "content": f"I'm sorry, I encountered an error: {str(e)}"})
+    # Generate response
+    try:
+        full_response = ""
+        for token in generate_response(
+            st.session_state.model,
+            st.session_state.tokenizer,
+            user_input
+        ):
+            full_response += token
         
-        # Rerun to update the UI
-        st.rerun()
+        # Add assistant response
+        st.session_state.messages.append({"role": "assistant", "content": full_response})
+    except Exception as e:
+        st.session_state.messages.append({"role": "assistant", "content": f"Error: {str(e)}"})
+    
+    st.rerun()
 
-    # Clear chat when button is pressed
-    if clear_button:
-        st.session_state.messages = []
-        st.rerun()
+# Handle clear button
+if clear_button:
+    st.session_state.messages = []
+    st.rerun()
+
+st.markdown('</div>', unsafe_allow_html=True)