import gradio as gr
from transformers import pipeline
import numpy as np
import os
from huggingface_hub import login
import spaces

HF_TOKEN = os.environ.get("HF_TOKEN")
if HF_TOKEN:
    login(token=HF_TOKEN)

MODEL_ID = "badrex/JASRv1.1"
transcriber = pipeline("automatic-speech-recognition", model=MODEL_ID)

@spaces.GPU
def transcribe(audio):
    sr, y = audio
    # convert to mono if stereo
    if y.ndim > 1:
        y = y.mean(axis=1)
    y = y.astype(np.float32)
    y /= np.max(np.abs(y))
    return transcriber({"sampling_rate": sr, "raw": y})["text"]


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

    
demo = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(),
    outputs="text",
    title="<div>JASR v1.1 🎙️ <br>Speech Recognition for Dialectal Arabic</div>",
    description="""
        <div class="centered-content">
            <div>
                <p>
                Developed with ❤ by <a href="https://badrex.github.io/" style="color: #2563eb;">Badr al-Absi</a>  
                </p>
                <br>
                <p style="font-size: 15px; line-height: 1.8;">
                Marhaban 👋🏼
                <br>
                <br>
                 This is a demo for JASR, pronounced <i>Jāsir</i> [جاسِر], a Transformer-based automatic speech recognition (ASR) system for dialectal Arabic.
                 The current running instance is optimized for the regional dialects of <i>Jazirat al-Arab</i>, or the Arabian Peninsula. 
                 JASR is still under active development.
                <br>                   
                <p style="font-size: 15px; line-height: 1.8;">
                Simply <strong>upload an audio file</strong> 📤 or <strong>record yourself speaking</strong> 🎙️⏺️ to try out the model!
                </p>
            </div>
        </div>
        """,
    examples=examples if examples else None,
    example_labels=[
        "Kuwait Theatre",
        "Saudi Radio Poetry", 
        "News Report (MSA)", 
        "San3ani Arabic male", 
        "San3ani Arabic female", 
        "Khaleeji Theatre",
        "TEDx KSA",
        "Yousif Saif Football Commentary", 
        "Khaleeji Theatre 2",
        "TV Drama", 
        "KSA Theatre", 
        "TV Drama 2",
        "Radio Jeddah (KSA)", 
        "Omani Theatre", 
        "Khaleeji Drama", 
        "Radio News", 
        "TEDx KSA 2",
        "Radio Jeddah (KSA) 2",
    ],
    cache_examples=False,  
    examples_per_page=18,
    flagging_mode=None,
)

if __name__ == "__main__":
    demo.launch()