import gradio as gr
import torch
from transformers import pipeline

# Cấu hình mô hình
model_id = "kotoba-tech/kotoba-whisper-bilingual-v1.0"
device = 0 if torch.cuda.is_available() else -1
pipe = pipeline("automatic-speech-recognition", model=model_id, device=device, chunk_length_s=30)

def transcribe(audio):
    result = pipe(audio, return_timestamps=True)["text"]
    return result

# Giao diện Gradio
with gr.Blocks() as demo:
    gr.Markdown("## 🎤 Kotoba Whisper Bilingual Demo")
    with gr.Tab("Tải lên tệp âm thanh"):
        audio_file = gr.Audio(type="filepath")
        output_text = gr.Textbox(label="Văn bản chuyển đổi")
        transcribe_button = gr.Button("Chuyển đổi")
        transcribe_button.click(fn=transcribe, inputs=audio_file, outputs=output_text)
    with gr.Tab("Ghi âm trực tiếp"):
        audio_mic = gr.Audio(sources="microphone", type="filepath")
        output_text_mic = gr.Textbox(label="Văn bản chuyển đổi")
        transcribe_button_mic = gr.Button("Chuyển đổi")
        transcribe_button_mic.click(fn=transcribe, inputs=audio_mic, outputs=output_text_mic)

demo.launch()