import gradio as gr
import torch
import librosa
from transformers import pipeline

print("🔄 Loading model...")
pipe = pipeline(
    "automatic-speech-recognition",
    model="AigizK/wav2vec2-large-mms-1b-tatar",
    device=0 if torch.cuda.is_available() else -1
)
print("✅ Model loaded!")

def transcribe(audio_path):
    if not audio_path: 
        return "⚠️ Загрузите файл"
    try:
        # wav2vec2 строго требует 16kHz mono float32
        audio, _ = librosa.load(audio_path, sr=16000, mono=True, dtype="float32")
        result = pipe(audio, generate_kwargs={"language": "tatar", "max_new_tokens": 256})["text"]
        return result.replace("[UNK]", "й")
    except Exception as e:
        return f"❌ Ошибка: {e}"

# Gradio 6.0: убран allow_flagging, theme перенесён в launch()
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🗣️ Tatar Speech-to-Text Demo")
    gr.Markdown("Model: `wav2vec2-large-mms-1b-tatar` • Audio will be auto-resampled to 16kHz")
    
    with gr.Row():
        audio_input = gr.Audio(type="filepath", label="🎤 Upload audio (.wav/.mp3)")
        text_output = gr.Textbox(label="📝 Transcription", lines=4, interactive=False)
    
    btn = gr.Button("▶️ Transcribe", variant="primary")
    btn.click(fn=transcribe, inputs=audio_input, outputs=text_output)
    
    gr.Examples(
        examples=["./sample.mp3"],  # Optional: add your sample file path
        inputs=audio_input,
        outputs=text_output,
        fn=transcribe,
        cache_examples=False
    )

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