Update app.py

app.py

@@ -1,63 +1,25 @@
-import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
-import torch
-import openai
-
-st.set_page_config(page_title="Baro - Emotional AI", layout="centered")
-
-st.title("🧠 Baro - Emotion-Aware AI (Gemma 1B)")
-st.markdown("Interact with your emotionally intelligent AI assistant fine-tuned on Gemma 1B.")
-
-# Inference Mode Selector
-mode = st.radio("Choose Inference Mode:", ["Transformers (local)", "vLLM API (remote)"])
-
-# Shared Input
-prompt = st.text_area("🗣️ Your Message:", height=200)
-
-if st.button("🔮 Generate Response"):
-    if not prompt.strip():
-        st.warning("Please enter a message.")
-    else:
-        with st.spinner("Baro is thinking..."):
-
-            # Transformers Mode
-            if mode == "Transformers (local)":
-                try:
-                    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,
-                        device_map="auto"
-                    )
-                    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
-                    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-
-                    _ = model.generate(
-                        input_ids=input_ids,
-                        max_new_tokens=200,
-                        do_sample=True,
-                        top_p=0.9,
-                        temperature=0.7,
-                        streamer=streamer,
-                    )
-                except Exception as e:
-                    st.error(f"Error in Transformers mode: {e}")
-
-            # vLLM API Mode
-            else:
-                api_url = st.text_input("vLLM Server Base URL", value="http://localhost:8000/v1")
-                if api_url:
-                    openai.api_key = "EMPTY"
-                    openai.base_url = api_url
-
-                    try:
-                        response = openai.ChatCompletion.create(
-                            model="umar141/Gemma_1B_Baro_v2_vllm",
-                            messages=[
-                                {"role": "system", "content": "You are Baro, an emotionally intelligent assistant."},
-                                {"role": "user", "content": prompt}
-                            ]
-                        )
-                        st.success(response.choices[0].message["content"])
-                    except Exception as e:
-                        st.error(f"Error in vLLM API mode: {e}")
+from transformers import GemmaForCausalLM, AutoTokenizer
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained("umar141/Gemma_1B_Baro_v2_vllm")
+
+# Load model
+model = GemmaForCausalLM.from_pretrained(
+    "umar141/Gemma_1B_Baro_v2_vllm",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto"
+)
+
+# Tokenize prompt
+input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+
+# Generate
+outputs = model.generate(
+    input_ids=input_ids,
+    max_new_tokens=200,
+    do_sample=True,
+    top_p=0.9,
+    temperature=0.7,
+)
+
+response = tokenizer.decode(outputs[0], skip_special_tokens=True)