Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,117 @@
-import altair as alt
-import numpy as np
-import pandas as pd
-import streamlit as st
-
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+app_code = '''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_icon="🤖",
+    layout="wide"
+)
+
+BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
+LORA_REPO = "Redfire-1234/AI-agent-v2"  # Updated repo name
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(LORA_REPO)
+    base_model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        low_cpu_mem_usage=True
+    )
+    model = PeftModel.from_pretrained(base_model, LORA_REPO)
+    model.eval()
+    return tokenizer, model
+
+def generate_response(tokenizer, model, user_input):
+    messages = [{"role": "user", "content": user_input}]
+    
+    try:
+        prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+    except:
+        prompt = f"User: {user_input}\\nAssistant:"
+    
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=200,
+            temperature=0.7,
+            do_sample=True,
+            top_p=0.9,
+            repetition_penalty=1.1,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    new_tokens = outputs[0][inputs['input_ids'].shape[1]:]
+    reply = tokenizer.decode(new_tokens, skip_special_tokens=True)
+    return reply.strip()
+
+# UI
+col1, col2 = st.columns([6, 1])
+with col1:
+    st.title("🤖 AI Agent Chatbot V2")
+    st.caption("✨ Improved with Natural Conversations")
+with col2:
+    if st.button("🗑️ New Chat", use_container_width=True):
+        st.session_state.messages = []
+        st.rerun()
+
+with st.spinner("Loading model..."):
+    tokenizer, model = load_model()
+
+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")
+    
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"], avatar="👤" if message["role"]=="user" else "🤖"):
+            st.write(message["content"])
+
+user_input = st.chat_input("Type your message... (Type 'q' to end)")
+
+if user_input:
+    if user_input.strip().lower() == 'q':
+        st.session_state.messages.append({
+            "role": "assistant",
+            "content": "👋 Goodbye! Click 'New Chat' to start fresh."
+        })
+        st.rerun()
+    
+    st.session_state.messages.append({"role": "user", "content": user_input})
+    
+    with st.spinner("Thinking..."):
+        reply = generate_response(tokenizer, model, user_input)
+    
+    st.session_state.messages.append({"role": "assistant", "content": reply})
+    st.rerun()
+
+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.divider()
+    
+    st.header("📊 Stats")
+    st.metric("Messages", len(st.session_state.messages))
+    
+    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)