import streamlit as st
from transformers import pipeline

# Load the model pipeline
model = pipeline("audio-classification", model="HareemFatima/distilhubert-finetuned-stutterdetection")

# Streamlit app
def main():
    st.title("Stutter Classification App")
    audio_input = st.audio("Capture Audio", format="audio/wav", start_recording=True, channels=1)
    if st.button("Stop Recording"):
        # Assuming the recording is saved as "recording.wav"
        recording_path = "recording.wav"
        # Call the model pipeline to classify the audio
        prediction = model(recording_path)
        # Get the predicted label
        predicted_label = prediction[0]["label"]
        # Map the label to the corresponding stutter type
        if predicted_label == 0:
            stutter_type = "nonstutter"
        elif predicted_label == 1:
            stutter_type = "prolongation"
        elif predicted_label == 2:
            stutter_type = "repetition"
        elif predicted_label == 3:
            stutter_type = "blocks"
        else:
            stutter_type = "Unknown"
        st.write("Predicted Stutter Type:", stutter_type)

if __name__ == "__main__":
    main()