import gradio as gr
import torch
from transformers import HubertForSequenceClassification, HubertProcessor

# Load the model and processor from Hugging Face
model = HubertForSequenceClassification.from_pretrained("HareemFatima/distilhubert-finetuned-stutterdetection")
processor = HubertProcessor.from_pretrained("HareemFatima/distilhubert-finetuned-stutterdetection")

# Define a function for stutter detection
def detect_stutter(audio):
    # Preprocess the audio
    inputs = processor(audio, sampling_rate=16000, return_tensors="pt", padding=True)
    
    # Get model predictions
    with torch.no_grad():
        logits = model(**inputs).logits
    predicted_class = logits.argmax(-1).item()
    
    # Map prediction to stutter type
    stutter_types = {0: "Non Stutter", 1: "Beginner Stutter", 2: "Middle Stutter", 3: "End Stutter"}
    return stutter_types.get(predicted_class, "Unknown Stutter")

# Create Gradio interface
iface = gr.Interface(fn=detect_stutter, inputs=gr.Audio(source="microphone", type="numpy"), outputs="text")

# Launch the interface
iface.launch()