Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,17 +1,18 @@
-app_code = '''import streamlit as st
+import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 
 st.set_page_config(
-    page_title="AI Agent Chatbot V2",
+    page_title="AI Agent Chatbot",
     page_icon="🤖",
     layout="wide"
 )
 
 BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
-LORA_REPO = "Redfire-1234/AI-agent-v2"  # Updated repo name
+LORA_REPO = "Redfire-1234/AI-agent-v2"  # Change to AI-agent if using old model
 
+# Initialize session state
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
@@ -38,7 +39,7 @@ def generate_response(tokenizer, model, user_input):
             add_generation_prompt=True
         )
     except:
-        prompt = f"User: {user_input}\\nAssistant:"
+        prompt = f"User: {user_input}\nAssistant:"
     
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
@@ -57,61 +58,95 @@ def generate_response(tokenizer, model, user_input):
     reply = tokenizer.decode(new_tokens, skip_special_tokens=True)
     return reply.strip()
 
-# UI
+def clear_chat():
+    st.session_state.messages = []
+
+# Header
 col1, col2 = st.columns([6, 1])
 with col1:
-    st.title("🤖 AI Agent Chatbot V2")
-    st.caption("✨ Improved with Natural Conversations")
+    st.title("🤖 AI Agent Chatbot")
+    st.caption("Powered by Qwen 2.5 + LoRA Fine-tuning")
 with col2:
-    if st.button("🗑️ New Chat", use_container_width=True):
-        st.session_state.messages = []
+    if st.button("🗑️ New Chat", use_container_width=True, type="secondary"):
+        clear_chat()
         st.rerun()
 
-with st.spinner("Loading model..."):
+# Load model
+with st.spinner("Loading model... (first time takes 2-3 minutes)"):
     tokenizer, model = load_model()
 
+# Display chat history
 chat_container = st.container()
 with chat_container:
     if len(st.session_state.messages) == 0:
-        st.info("👋 **Welcome!** Ask me anything.\\n\\n⚠️ Type 'q' to end conversation")
+        st.info("👋 **Welcome to AI Agent Chatbot!**\n\n💬 Ask me anything and keep the conversation going.\n\n⚠️ **To end the conversation, simply type:** `q`")
     
     for message in st.session_state.messages:
-        with st.chat_message(message["role"], avatar="👤" if message["role"]=="user" else "🤖"):
-            st.write(message["content"])
+        if message["role"] == "user":
+            with st.chat_message("user", avatar="👤"):
+                st.write(message["content"])
+        else:
+            with st.chat_message("assistant", avatar="🤖"):
+                st.write(message["content"])
 
-user_input = st.chat_input("Type your message... (Type 'q' to end)")
+# Chat input
+user_input = st.chat_input("Type your message here... (Type 'q' to end conversation)")
 
 if user_input:
+    # Check if user wants to quit
     if user_input.strip().lower() == 'q':
         st.session_state.messages.append({
             "role": "assistant",
-            "content": "👋 Goodbye! Click 'New Chat' to start fresh."
+            "content": "👋 Goodbye! Click 'New Chat' to start a fresh conversation."
         })
         st.rerun()
     
-    st.session_state.messages.append({"role": "user", "content": user_input})
+    # Add user message
+    st.session_state.messages.append({
+        "role": "user",
+        "content": user_input
+    })
     
+    # Generate response
     with st.spinner("Thinking..."):
         reply = generate_response(tokenizer, model, user_input)
     
-    st.session_state.messages.append({"role": "assistant", "content": reply})
+    # Add bot response
+    st.session_state.messages.append({
+        "role": "assistant",
+        "content": reply
+    })
+    
     st.rerun()
 
+# Sidebar
 with st.sidebar:
     st.header("ℹ️ About")
-    st.write("**Version:** 2.0 - Natural Conversations")
-    st.write("**Base Model:** Qwen 2.5 1.5B")
-    st.write("**Fine-tuned:** Yes (LoRA)")
+    st.write(f"**Base Model:** Qwen 2.5 1.5B")
+    st.write(f"**LoRA Adapter:** AI-agent-v2")
     
     st.divider()
     
-    st.header("📊 Stats")
+    st.header("📊 Chat Stats")
     st.metric("Messages", len(st.session_state.messages))
+    st.metric("User Messages", len([m for m in st.session_state.messages if m["role"] == "user"]))
+    st.metric("Bot Messages", len([m for m in st.session_state.messages if m["role"] == "assistant"]))
     
     st.divider()
     
-    st.info("✨ **New in V2:**\\n- Natural greetings\\n- Better responses\\n- No more artifacts")
-'''
-
-with open("app.py", "w") as f:
-    f.write(app_code)
+    st.header("💡 Tips")
+    st.info("""
+    - Type your question and press Enter
+    - Type 'q' to end the conversation
+    - Click 'New Chat' to start fresh
+    - All messages are saved in this session
+    """)
+    
+    st.divider()
+    
+    if st.button("🗑️ Clear History", use_container_width=True, type="primary"):
+        clear_chat()
+        st.rerun()
+    
+    st.divider()
+    st.caption("Made with ❤️ using Streamlit")