Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,33 +3,36 @@ from transformers import pipeline
 from gtts import gTTS
 import os
 
-# Load small public instruction model
+# Load a public, PyTorch-compatible, conversational model
 @st.cache_resource(show_spinner="Loading AI Buddy...")
 def load_llm():
-    return pipeline("text2text-generation",
-                    model="google/flan-t5-small",
-                    tokenizer="google/flan-t5-small")
+    return pipeline("text-generation",
+                    model="declare-lab/flan-alpaca-base",
+                    tokenizer="declare-lab/flan-alpaca-base",
+                    max_new_tokens=100,
+                    do_sample=True,
+                    temperature=0.7)
 
 llm = load_llm()
 
-# Text-to-speech
+# Text-to-speech function
 def speak(text, filename="response.mp3"):
     tts = gTTS(text)
     tts.save(filename)
     audio_file = open(filename, "rb")
-    st.audio(audio_bytes := audio_file.read(), format="audio/mp3")
+    st.audio(audio_file.read(), format="audio/mp3")
     os.remove(filename)
 
 # Streamlit UI
 st.set_page_config(page_title="AI Learning Buddy", page_icon="🧸")
 st.title("🧸 AI Learning Buddy (Ages 4–7)")
+st.markdown("Ask anything fun or educational and hear your buddy talk!")
 
-st.markdown("Type a fun question and hear the AI Buddy reply!")
-
-user_input = st.text_input("Ask your question:")
+user_input = st.text_input("What would you like to ask?")
 
 if st.button("Ask the Buddy") and user_input:
     prompt = f"Explain to a 5-year-old: {user_input}"
     result = llm(prompt)[0]["generated_text"]
-    st.markdown(f"**AI Buddy says:** {result}")
-    speak(result)
+    answer = result.split(":")[-1].strip()
+    st.markdown(f"**AI Buddy says:** {answer}")
+    speak(answer)