import gradio as gr
from transformers import pipeline

# Load model - use the correct Hugging Face model ID
# Remove the '/spaces/' prefix and use just the username/model-name format
classifier = pipeline("audio-classification", model="Hnin/wav2vec2-base-finetuned-ks")

def predict(audio):
    if audio is None:
        return {"Error": "No audio provided"}
    
    try:
        preds = classifier(audio)
        return {p["label"]: p["score"] for p in preds}
    except Exception as e:
        return {"Error": f"Prediction failed: {str(e)}"}

# Gradio UI
gr.Interface(
    fn=predict,
    inputs=gr.Audio(sources=["microphone"], type="filepath"),  # Updated parameter name
    outputs=gr.Label(num_top_classes=3),
    title="🔊 Keyword Spotting",
    description="Upload an audio file or record from microphone for keyword spotting classification",
    examples=["mp3-output-ttsfree(dot)com (4).mp3"]  # Make sure these files exist
).launch()