Update app.py

app.py

@@ -1,105 +1,54 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from gtts import gTTS
-import tempfile
-import os
 import torch
-import speech_recognition as sr
-import soundfile as sf
-import io
-import logging
+import os
 
-# Set up logging
-logging.basicConfig(filename="app_errors.log", level=logging.ERROR)
+# ✅ Must be first Streamlit command
+st.set_page_config(page_title="🧠 Agentic AI Bot", layout="centered")
 
-# Load model and tokenizer with error handling
 @st.cache_resource
 def load_model():
-    try:
-        # Load the tokenizer using AutoTokenizer for Llama-2
-        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", use_fast=True)
-        model = AutoModelForCausalLM.from_pretrained(
-            "meta-llama/Llama-2-7b-hf",
-            device_map="auto",
-            torch_dtype=torch.float16
-        )
-        pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-        return pipe
-    except Exception as e:
-        # Log the error
-        logging.error(f"An error occurred while loading the model: {e}")
-        st.error(f"An error occurred while loading the model: {e}")
-        return None
-
-pipe = load_model()
-
-# Check if the pipe is None (which means an error occurred)
-if pipe is None:
-    st.stop()  # Stop further execution if the model couldn't be loaded
-
-# Title and description
-import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-
-# ✅ Set page config FIRST
-st.set_page_config(page_title="🧠 Agentic AI Bot", layout="centered")
-
-# Now the rest of your app
-st.title("Welcome to the Agentic AI Bot")
-
-# Text-to-speech function
-def speak(text):
-    tts = gTTS(text)
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as fp:
-        tts.save(fp.name)
-        return fp.name
-
-# Speech-to-text function
-def transcribe(audio_bytes):
-    recognizer = sr.Recognizer()
-    with sr.AudioFile(io.BytesIO(audio_bytes)) as source:
-        audio_data = recognizer.record(source)
-        try:
-            return recognizer.recognize_google(audio_data)
-        except sr.UnknownValueError:
-            return "Sorry, I could not understand the audio."
-        except sr.RequestError:
-            return "Speech recognition service is unavailable."
-
-# Input mode selection
-mode = st.radio("Choose input mode:", ["Text", "Voice"])
-
-user_input = ""
-if mode == "Text":
-    user_input = st.text_input("Enter your query:")
-else:
-    audio = st.file_uploader("Upload your voice (WAV format only)", type=["wav"])
-    if audio is not None:
-        audio_bytes = audio.read()
-        user_input = transcribe(audio_bytes)
-        st.write(f"You said: {user_input}")
-
-# Run the assistant
-if user_input:
-    with st.spinner("Thinking..."):
-        try:
-            result = pipe(user_input, max_new_tokens=200, temperature=0.7, do_sample=True)
-            response = result[0]['generated_text']
-
-            # Trim prompt from response if repeated
-            if response.lower().startswith(user_input.lower()):
-                response = response[len(user_input):].strip()
-
-            st.subheader("🤖 Assistant's Response:")
-            st.write(response)
-
-            # Speak response
-            audio_path = speak(response)
-            st.audio(audio_path, format="audio/mp3")
-
-            # Cleanup
-            os.remove(audio_path)
-        except Exception as e:
-            logging.error(f"Error while generating response: {str(e)}")
-            st.error("Sorry, there was an issue generating a response.")
+    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)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        device_map="auto",
+        offload_folder=offload_dir,
+        low_cpu_mem_usage=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}")