Update app.py

app.py

@@ -1,78 +1,65 @@
-# app.py
-import os
-import torch
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from streamlit_webrtc import webrtc_streamer, AudioProcessorBase
-import tempfile
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import whisper
+from streamlit_webrtc import webrtc_streamer, AudioProcessorBase
+import torch
 
-# -----------------------------
-# 🛠 SETUP
-# -----------------------------
-st.set_page_config(page_title="🧠 Agentic AI Bot", layout="centered")
-os.makedirs("offload", exist_ok=True)
+# ----------------------------- SETUP -----------------------------
+st.set_page_config(page_title="🧠 Talkative AI Bot", layout="centered")
 
-# -----------------------------
-# 🧠 Load Whisper Model (for mic)
-# -----------------------------
+# ----------------------------- LOAD MODELS -----------------------------
+# Load Whisper model for speech-to-text
 @st.cache_resource
 def load_whisper():
-    return whisper.load_model("base")
+    try:
+        model = whisper.load_model("base")
+        return model
+    except Exception as e:
+        st.error(f"An error occurred while loading Whisper model: {e}")
+        return None
 
-whisper_model = load_whisper()
-
-# -----------------------------
-# 🤖 Load LLM (LLaMA-2)
-# -----------------------------
+# Load DistilGPT-2 model for generating responses
 @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",
-        torch_dtype=torch.float16,
-        offload_folder="offload",
-        use_auth_token=True
-    )
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-    return pipe
-
-pipe = load_llm()
+def load_language_model():
+    try:
+        tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
+        model = AutoModelForCausalLM.from_pretrained("distilgpt2")
+        return model, tokenizer
+    except Exception as e:
+        st.error(f"An error occurred while loading Language model: {e}")
+        return None, None
 
-# -----------------------------
-# 🎤 Microphone Input
-# -----------------------------
+# ----------------------------- FUNCTION TO HANDLE SPEECH -----------------------------
 class AudioProcessor(AudioProcessorBase):
     def __init__(self):
-        self.result = ""
+        self.whisper_model = load_whisper()
+
+    def transform(self, audio_frame):
+        # Convert audio frame to audio file and get text transcription
+        result = self.whisper_model.transcribe(audio_frame)
+        return result['text']
 
-    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
+# ----------------------------- FUNCTION TO GENERATE RESPONSE -----------------------------
+def generate_response(user_input):
+    model, tokenizer = load_language_model()
+    if model and tokenizer:
+        inputs = tokenizer(user_input, return_tensors="pt")
+        outputs = model.generate(inputs['input_ids'], max_length=100)
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
+    return "Sorry, I couldn't process that."
 
-st.subheader("🎙️ Speak to the Agent")
-audio_ctx = webrtc_streamer(key="mic", audio_processor_factory=AudioProcessor)
+# ----------------------------- STREAMLIT UI -----------------------------
+st.title("🧠 Talkative AI Bot")
+st.write("Talk to the bot using your microphone, and it will respond!")
 
-# -----------------------------
-# 🧠 Chat Interface
-# -----------------------------
-st.subheader("💬 Ask a Question")
-user_input = ""
-if audio_ctx and audio_ctx.audio_processor:
-    user_input = audio_ctx.audio_processor.result
+# Streamlit WebRTC for speech-to-text
+webrtc_streamer(key="example", audio_processor_factory=AudioProcessor)
 
-user_text = st.text_input("Or type your question here:", user_input)
+# Input text for chatbot
+user_input = st.text_input("Type something for the bot:")
 
-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)
+# Handle text input and generate response
+if user_input:
+    response = generate_response(user_input)
+    st.write(f"Bot: {response}")