Update app.py

app.py

@@ -1,7 +1,8 @@
 import streamlit as st
-from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import pipeline, AutoModelForCTC, AutoProcessor, Wav2Vec2ForCTC
 import pyttsx3
 import random
+import torchaudio
 
 # Set up the app
 st.set_page_config(page_title="ChiliDron", layout="wide", page_icon="🎨")
@@ -13,13 +14,13 @@ def load_storytelling_model():
 
 @st.cache_resource
 def load_phoneme_model():
-    model = AutoModelForSeq2SeqLM.from_pretrained("mirfan899/kids_phoneme_sm_model")
-    tokenizer = AutoTokenizer.from_pretrained("mirfan899/kids_phoneme_sm_model")
-    return model, tokenizer
+    model = Wav2Vec2ForCTC.from_pretrained("mirfan899/kids_phoneme_sm_model")
+    processor = AutoProcessor.from_pretrained("mirfan899/kids_phoneme_sm_model")
+    return model, processor
 
 # Load models
 story_generator = load_storytelling_model()
-phoneme_model, phoneme_tokenizer = load_phoneme_model()
+phoneme_model, phoneme_processor = load_phoneme_model()
 
 # Initialize TTS engine
 def initialize_tts():
@@ -81,17 +82,16 @@ elif options == "Phoneme Practice":
     st.header("🔤 Phoneme Practice")
     st.write("Learn how to pronounce words correctly!")
 
-    word_to_practice = st.text_input("Enter a word to practice its phonemes:")
-
-    if st.button("Get Phonemes"):
-        if word_to_practice:
-            with st.spinner("Analyzing phonemes..."):
-                inputs = phoneme_tokenizer.encode(word_to_practice, return_tensors="pt")
-                outputs = phoneme_model.generate(inputs)
-                phonemes = phoneme_tokenizer.decode(outputs[0], skip_special_tokens=True)
-                st.write(f"Phonemes for '{word_to_practice}': {phonemes}")
-        else:
-            st.warning("Please enter a word to practice.")
+    uploaded_audio = st.file_uploader("Upload an audio file to analyze phonemes", type=["wav", "mp3"])
+
+    if uploaded_audio:
+        with st.spinner("Analyzing phonemes..."):
+            audio_input, _ = torchaudio.load(uploaded_audio)
+            input_values = phoneme_processor(audio_input, return_tensors="pt", sampling_rate=16000).input_values
+            logits = phoneme_model(input_values).logits
+            predicted_ids = torch.argmax(logits, dim=-1)
+            transcription = phoneme_processor.batch_decode(predicted_ids)
+            st.write(f"Transcription: {transcription[0]}")
 
 # Guided Meditation for Kids
 elif options == "Guided Meditation":