disiniStutterModel's picture
Update app.py
10fc44a verified
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()