import gradio as gr
import torch
from transformers import pipeline

# -----------------------------
# 1. Load Whisper model on CPU
# -----------------------------
MODEL_NAME = "Phonepadith/whisper-3-large-turbo-lao-finetuned-v1"  # Change if needed

device = "cpu"
torch_dtype = torch.float32

asr_pipeline = pipeline(
    "automatic-speech-recognition",
    model=MODEL_NAME,
    tokenizer=MODEL_NAME,
    feature_extractor=MODEL_NAME,
    device=-1,  # Force CPU
    torch_dtype=torch_dtype,
)

# -----------------------------
# 2. Transcription function
# -----------------------------
def transcribe(audio_file):
    if audio_file is None:
        return "❌ Please upload an audio file."
    try:
        result = asr_pipeline(audio_file)
        text = result["text"]
        return f"🧩 Lao Transcription:\n\n{text}"
    except Exception as e:
        return f"❌ Error: {str(e)}"

# -----------------------------
# 3. Gradio Interface
# -----------------------------
title = "🎧 Lao Whisper Test (CPU)"
description = """
ທົດສອບຮູບແບບຈຳລອງ Whisper ທີ່ປັບປຸງໃຫ້ຮັບຮູ້ພາສາລາວ 🇱🇦  
Upload an audio file (`.wav`, `.mp3`, `.m4a`, `.flac`) and get the Lao transcription.
"""

iface = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(type="filepath", label="🎵 Upload Lao Speech Audio"),
    outputs=gr.Textbox(label="🧾 Transcription Result", lines=8),
    title=title,
    description=description,
    allow_flagging="never",
)

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