app.py

import gradio as gr
from transformers import pipeline
import numpy as np
import os

# Load the model
print("Loading model...")
model_id = "badrex/mms-300m-arabic-dialect-identifier"
try:
    classifier = pipeline("audio-classification", model=model_id)
    print("Model loaded successfully")
except Exception as e:
    print(f"Error loading model: {e}")
    
# Define dialect mapping
dialect_mapping = {
    "MSA": "Modern Standard Arabic",
    "Egyptian": "Egyptian Arabic",
    "Gulf": "Gulf Arabic",
    "Levantine": "Levantine Arabic",
    "Maghrebi": "Maghrebi Arabic"
}

def predict_dialect(audio):
    try:
        # The audio input from Gradio is a tuple of (sample_rate, audio_array)
        if audio is None:
            return {"Error": 1.0}
            
        sr, audio_array = audio
        
        # Process the audio input
        if len(audio_array.shape) > 1:
            audio_array = audio_array.mean(axis=1)  # Convert stereo to mono
        
        print(f"Processing audio: sample rate={sr}, shape={audio_array.shape}")
        
        # Classify the dialect
        predictions = classifier({"sampling_rate": sr, "raw": audio_array})
        
        # Format results for display
        results = {}
        for pred in predictions:
            dialect_name = dialect_mapping.get(pred['label'], pred['label'])
            results[dialect_name] = float(pred['score'])
        
        return results
    except Exception as e:
        print(f"Error in prediction: {e}")
        return {"Error": 1.0}

# Find example files
example_files = []
examples_dir = "examples"
if os.path.exists(examples_dir):
    for filename in os.listdir(examples_dir):
        if filename.endswith((".wav", ".mp3", ".ogg")):
            example_files.append(os.path.join(examples_dir, filename))
            
    print(f"Found {len(example_files)} example files")
else:
    print("Examples directory not found")

# Examples with labels
examples = []
if example_files:
    for file in example_files:
        basename = os.path.basename(file)
        dialect = basename.split("_")[0] if "_" in basename else basename.split(".")[0]
        label = dialect_mapping.get(dialect, dialect.capitalize())
        examples.append([file, f"{label} Sample"])

# Create the Gradio interface
demo = gr.Interface(
    fn=predict_dialect,
    inputs=gr.Audio(),
    outputs=gr.Label(num_top_classes=5, label="Predicted Dialect"),
    title="Arabic Dialect Identifier",
    description="""This demo identifies Arabic dialects from speech audio.
    Upload an audio file or record your voice speaking Arabic to see which dialect it matches.
    The model identifies: Modern Standard Arabic (MSA), Egyptian, Gulf, Levantine, and Maghrebi dialects.""",
    examples=examples if examples else None,
    examples_per_page=5,
    flagging_mode=None  # Updated from allow_flagging
)

# Launch the app
demo.launch()