Update app.py

app.py

@@ -3,9 +3,11 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 from datetime import datetime
 
-# Initialize session state for chat history
+# Initialize session state variables
 if 'messages' not in st.session_state:
     st.session_state.messages = []
+if "user_input_widget" not in st.session_state:
+    st.session_state.user_input_widget = ""
 
 @st.cache_resource
 def load_model():
@@ -46,7 +48,11 @@ def generate_response(prompt, model, tokenizer, history):
     return response
 
 def main():
-    st.set_page_config(page_title="AMD-OLMo Chatbot", layout="wide")
+    st.set_page_config(
+        page_title="AMD-OLMo Chatbot",
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
     
     # Custom CSS
     st.markdown("""
@@ -58,17 +64,31 @@ def main():
             background-color: #f0f2f6;
             padding: 20px;
             border-radius: 10px;
+            margin: 10px 0;
         }
         .chat-message {
-            padding: 10px;
+            padding: 15px;
             border-radius: 10px;
-            margin: 5px 0;
+            margin: 10px 0;
         }
         .user-message {
             background-color: #e6f3ff;
+            border-left: 5px solid #2e6da4;
         }
         .assistant-message {
             background-color: #f0f2f6;
+            border-left: 5px solid #5cb85c;
+        }
+        .stTextArea textarea {
+            font-size: 16px;
+        }
+        .timestamp {
+            font-size: 12px;
+            color: #666;
+            margin-top: 5px;
+        }
+        .st-emotion-cache-1v0mbdj.e115fcil1 {
+            margin-top: 20px;
         }
         </style>
     """, unsafe_allow_html=True)
@@ -79,35 +99,46 @@ def main():
     with tab1:
         st.title("AMD-OLMo-1B-SFT Model Information")
         
-        st.markdown("""
-        ## Model Overview
-        AMD-OLMo-1B-SFT is a state-of-the-art language model developed by AMD[1][2]. Key features include:
-        
-        ### Architecture
-        - **Base Model**: 1.2B parameters
-        - **Layers**: 16
-        - **Attention Heads**: 16
-        - **Hidden Size**: 2048
-        - **Context Length**: 2048
-        - **Vocabulary Size**: 50,280
-        
-        ### Training Details
-        - Pre-trained on 1.3 trillion tokens from Dolma v1.7
-        - Supervised fine-tuned (SFT) in two phases:
-          1. Tulu V2 dataset
-          2. OpenHermes-2.5, WebInstructSub, and Code-Feedback datasets
-        
-        ### Capabilities
-        - General text generation
-        - Question answering
-        - Code understanding
-        - Reasoning tasks
-        - Instruction following
-        
-        ### Hardware Requirements
-        - Optimized for AMD Instinct™ MI250 GPUs
-        - Training performed on 16 nodes with 4 GPUs each
-        """)
+        with st.container():
+            st.markdown("""
+            <div class="model-info">
+            <h2>Model Overview</h2>
+            AMD-OLMo-1B-SFT is a state-of-the-art language model developed by AMD. This model represents a significant advancement in AMD's AI capabilities.
+            
+            <h3>Architecture Specifications</h3>
+            
+            | Component | Specification |
+            |-----------|---------------|
+            | Parameters | 1.2B |
+            | Layers | 16 |
+            | Attention Heads | 16 |
+            | Hidden Size | 2048 |
+            | Context Length | 2048 |
+            | Vocabulary Size | 50,280 |
+            
+            <h3>Training Details</h3>
+            
+            - Pre-trained on 1.3 trillion tokens from Dolma v1.7
+            - Two-phase supervised fine-tuning (SFT):
+                1. Tulu V2 dataset
+                2. OpenHermes-2.5, WebInstructSub, and Code-Feedback datasets
+            
+            <h3>Key Capabilities</h3>
+            
+            - Natural language understanding and generation
+            - Context-aware responses
+            - Code understanding and generation
+            - Complex reasoning tasks
+            - Instruction following
+            - Multi-turn conversations
+            
+            <h3>Hardware Optimization</h3>
+            
+            - Optimized for AMD Instinct™ MI250 GPUs
+            - Distributed training across 16 nodes with 4 GPUs each
+            - Efficient inference on consumer hardware
+            </div>
+            """, unsafe_allow_html=True)
 
     with tab2:
         st.title("Chat with AMD-OLMo")
@@ -127,37 +158,54 @@ def main():
         with chat_container:
             for message in st.session_state.messages:
                 div_class = "user-message" if message["role"] == "user" else "assistant-message"
+                timestamp = message.get("timestamp", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
                 st.markdown(f"""
                     <div class="chat-message {div_class}">
                         <b>{message["role"].title()}:</b> {message["content"]}
+                        <div class="timestamp">{timestamp}</div>
                     </div>
                 """, unsafe_allow_html=True)
 
-        # User input
+        # User input section
         with st.container():
-            user_input = st.text_area("Your message:", key="user_input", height=100)
+            user_input = st.text_area(
+                "Your message:",
+                key="user_input_widget",
+                height=100,
+                placeholder="Type your message here..."
+            )
+            
             col1, col2, col3 = st.columns([1, 1, 4])
             
             with col1:
-                if st.button("Send"):
+                if st.button("Send", use_container_width=True):
                     if user_input.strip():
-                        # Add user message to history
-                        st.session_state.messages.append({"role": "user", "content": user_input})
+                        # Add user message to history with timestamp
+                        st.session_state.messages.append({
+                            "role": "user",
+                            "content": user_input,
+                            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+                        })
                         
                         # Generate response
-                        with st.spinner("Thinking..."):
+                        with st.spinner("Generating response..."):
                             response = generate_response(user_input, model, tokenizer, st.session_state.messages)
                         
-                        # Add assistant response to history
-                        st.session_state.messages.append({"role": "assistant", "content": response})
+                        # Add assistant response to history with timestamp
+                        st.session_state.messages.append({
+                            "role": "assistant",
+                            "content": response,
+                            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+                        })
                         
                         # Clear input
-                        st.session_state.user_input = ""
+                        st.session_state.user_input_widget = ""
                         st.experimental_rerun()
             
             with col2:
-                if st.button("Clear History"):
+                if st.button("Clear History", use_container_width=True):
                     st.session_state.messages = []
+                    st.session_state.user_input_widget = ""
                     st.experimental_rerun()
 
 if __name__ == "__main__":