import gradio as gr
import torch
import librosa
import numpy as np
from transformers import AutoModelForAudioClassification, AutoFeatureExtractor

# ========= 1. LOAD MODEL (التعديل هنا) =========
# استخدمنا AutoModel عشان هو اللي بيعرف يقرأ ملفات .safetensors و .json
model_path = "." # النقطة تعني الفولدر الحالي اللي فيه الملفات
model = AutoModelForAudioClassification.from_pretrained(model_path)
feature_extractor = AutoFeatureExtractor.from_pretrained(model_path)
model.eval()

# ========= 2. PREPROCESS AUDIO (تعديل بسيط ليناسب Transformers) =========
def predict(audio_path):
    # تحميل الصوت
    audio, sr = librosa.load(audio_path, sr=feature_extractor.sampling_rate)
    
    # تحويل الصوت لـ Tensors مناسبة للموديل
    inputs = feature_extractor(audio, sampling_rate=sr, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        
    # الحصول على التوقع (Label)
    predicted_class_ids = torch.argmax(logits, dim=-1).item()
    prediction = model.config.id2label[predicted_class_ids]
    
    return prediction

# ========= 3. UI (زي ما هو مع تعديل بسيط في الـ fn) =========
interface = gr.Interface(
    fn=predict,
    inputs=gr.Audio(type="filepath", label="Upload Audio 🎤"),
    outputs=gr.Textbox(label="Prediction"),
    title="Audio AI Model 🎧"
)

interface.launch()