import streamlit as st
from transformers import pipeline
import os

# Define your Hugging Face token; secure this appropriately
HF_TOKEN = os.getenv('HF_TOKEN')

# Set up Hugging Face pipeline for audio classification using the specified model
model_name = "MIT/ast-finetuned-audioset-10-10-0.4593"

@st.cache(allow_output_mutation=True)
def load_model(token, model_name):
    return pipeline("audio-classification", model=model_name, use_auth_token=token)

audio_classifier = load_model(HF_TOKEN, model_name)

# Pre-uploaded audio files
audio_files = {
    "Labrador Barking": "labrador-barking.mp3",
    "Tolling Bell": "tolling-bell.mp3",
    "Airplane Landing": "airplane-landing.mp3",
    "Old Car Engine": "old-car-engine.mp3",
    "Hard Shoes": "hard_shoes.mp3",
    "Alien Spaceship": "alien-spaceship.mp3",
}

# Streamlit UI
st.title("Audio Classification with Pre-uploaded Files")

# Audio file selection
selected_audio_name = st.selectbox("Select an audio file", list(audio_files.keys()))
audio_file_path = audio_files[selected_audio_name]

# Perform classification
if st.button("Classify"):
    # Read audio file
    with open(audio_file_path, "rb") as audio_file:
        audio_bytes = audio_file.read()

    results = audio_classifier(audio_bytes)

    # Displaying results
    st.write("Classification Results:")
    for result in results:
        label = result['label']
        score = round(result['score'], 4)  # Adjust rounding as needed
        st.write(f"Label: {label}, Score: {score}")