Update app.py

app.py

@@ -1,54 +1,78 @@
+# app.py
+import os
+import torch
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
-import os
+from streamlit_webrtc import webrtc_streamer, AudioProcessorBase
+import tempfile
+import whisper
 
-# ✅ Must be first Streamlit command
+# -----------------------------
+# 🛠 SETUP
+# -----------------------------
 st.set_page_config(page_title="🧠 Agentic AI Bot", layout="centered")
+os.makedirs("offload", exist_ok=True)
 
+# -----------------------------
+# 🧠 Load Whisper Model (for mic)
+# -----------------------------
 @st.cache_resource
-def load_model():
-    model_id = "meta-llama/Llama-2-7b-hf"
-    
-    # Offload folder for Hugging Face Space
-    offload_dir = "/tmp/offload"
-    os.makedirs(offload_dir, exist_ok=True)
+def load_whisper():
+    return whisper.load_model("base")
 
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
+whisper_model = load_whisper()
 
+# -----------------------------
+# 🤖 Load LLM (LLaMA-2)
+# -----------------------------
+@st.cache_resource
+def load_llm():
+    model_id = "meta-llama/Llama-2-7b-hf"
+    tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=True)
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         device_map="auto",
-        offload_folder=offload_dir,
-        low_cpu_mem_usage=True
+        torch_dtype=torch.float16,
+        offload_folder="offload",
+        use_auth_token=True
     )
-
     pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
     return pipe
 
-# Load the model
-try:
-    with st.spinner("Loading model..."):
-        pipe = load_model()
-except Exception as e:
-    st.error(f"An error occurred while loading the model: {e}")
-    st.stop()
-
-# UI
-st.title("🧠 Agentic AI Assistant")
-st.markdown("Talk to your LLaMA 2-powered AI assistant.")
-
-user_input = st.text_area("Enter your question:", height=150)
-
-if st.button("Ask"):
-    if user_input.strip() == "":
-        st.warning("Please enter a message.")
-    else:
-        with st.spinner("Thinking..."):
-            try:
-                response = pipe(user_input, max_new_tokens=256, do_sample=True, temperature=0.7)[0]["generated_text"]
-                # Extract only new response portion
-                answer = response[len(user_input):].strip()
-                st.success(answer)
-            except Exception as e:
-                st.error(f"An error occurred while generating a response: {e}")
+pipe = load_llm()
+
+# -----------------------------
+# 🎤 Microphone Input
+# -----------------------------
+class AudioProcessor(AudioProcessorBase):
+    def __init__(self):
+        self.result = ""
+
+    def recv(self, frame):
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+            f.write(frame.to_ndarray().tobytes())
+            audio_path = f.name
+        try:
+            text = whisper_model.transcribe(audio_path)["text"]
+            self.result = text
+        except:
+            self.result = "[Could not transcribe audio]"
+        return frame
+
+st.subheader("🎙️ Speak to the Agent")
+audio_ctx = webrtc_streamer(key="mic", audio_processor_factory=AudioProcessor)
+
+# -----------------------------
+# 🧠 Chat Interface
+# -----------------------------
+st.subheader("💬 Ask a Question")
+user_input = ""
+if audio_ctx and audio_ctx.audio_processor:
+    user_input = audio_ctx.audio_processor.result
+
+user_text = st.text_input("Or type your question here:", user_input)
+
+if st.button("Ask") and user_text.strip():
+    with st.spinner("Thinking..."):
+        result = pipe(user_text, max_new_tokens=200, do_sample=True)[0]["generated_text"]
+        st.success(result)