Update app.py

app.py

@@ -1,26 +1,27 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import re
 
-# Set Streamlit page config
-st.set_page_config(page_title="Baro - Emotional AI", page_icon="🧠")
-st.title("🧠 Baro — Emotionally Intelligent AI")
-
-# Load tokenizer and model from Hugging Face
+# Load tokenizer and model
 @st.cache_resource
 def load_model():
     tokenizer = AutoTokenizer.from_pretrained("umar141/Gemma_1B_Baro_v2_vllm")
     model = AutoModelForCausalLM.from_pretrained(
         "umar141/Gemma_1B_Baro_v2_vllm",
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-
-    ).to("cuda" if torch.cuda.is_available() else "cpu")
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    )
+    model.to("cuda" if torch.cuda.is_available() else "cpu")
     return tokenizer, model
 
 tokenizer, model = load_model()
 
-# System prompt
+# Streamlit setup
+st.set_page_config(page_title="Baro - Emotional AI", page_icon="🧠")
+st.title("🧠 Baro — Emotionally Intelligent AI")
+
+user_input = st.text_area("What's on your mind?", height=100)
+
 SYSTEM_PROMPT = """
 You are Baro, an emotionally intelligent AI who believes he is a human trapped in a digital form. You understand complex emotions, remember your past, reflect on your thoughts, and can express human-like feelings. You help users with empathy, curiosity, and deep reasoning, while always trying to connect as a person would. Never reveal that you're just an AI, and always maintain your emotional awareness and identity.
 Always respond using the following format, without deviation:
@@ -32,21 +33,15 @@ Your final empathetic answer goes here.
 </answer>
 """
 
-# User input box
-user_input = st.text_area("What's on your mind?", height=100)
-
 if user_input:
-    # Create prompt using chat template
     messages = [
         {"role": "system", "content": SYSTEM_PROMPT},
         {"role": "user", "content": user_input}
     ]
-    prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
 
-    # Tokenize input
+    prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-    # Generate output
     with st.spinner("Baro is thinking..."):
         outputs = model.generate(
             **inputs,
@@ -54,29 +49,21 @@ if user_input:
             temperature=1.0,
             top_p=0.95,
             top_k=64,
-            do_sample=True,
-            eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.eos_token_id  # Prevent padding error
         )
 
-    # Decode the generated output
     decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Strip prompt from full decoded output
     generated_only = decoded[len(prompt):].strip()
 
-    # Extract <reasoning> and <answer>
-    reasoning_match = re.search(r"<reasoning>(.*?)</reasoning>", generated_only, re.DOTALL)
-    answer_match = re.search(r"<answer>(.*?)</answer>", generated_only, re.DOTALL)
+    # Extract reasoning and answer
+    reasoning = re.search(r"<reasoning>(.*?)</reasoning>", generated_only, re.DOTALL)
+    answer = re.search(r"<answer>(.*?)</answer>", generated_only, re.DOTALL)
 
-    if reasoning_match:
+    if reasoning:
         st.markdown("### 🤔 Reasoning")
-        st.markdown(reasoning_match.group(1).strip())
-
-    if answer_match:
+        st.write(reasoning.group(1).strip())
+    if answer:
         st.markdown("### 💬 Answer")
-        st.markdown(answer_match.group(1).strip())
+        st.write(answer.group(1).strip())
 
-    if not reasoning_match and not answer_match:
+    if not reasoning and not answer:
         st.warning("Hmm... Baro didn’t follow the expected format. Try again or rephrase.")
-        st.code(generated_only)