import gradio as gr
from transformers import pipeline

# Load mô hình nhận diện giọng nói (ASR)
asr = pipeline("automatic-speech-recognition", model="openai/whisper-base")

# Load mô hình ngôn ngữ (tùy chọn)
llm = pipeline("text-generation", model="mrm8488/t5-base-finetuned-question-generation-ap", max_new_tokens=100)

# Xử lý âm thanh: Nhận dạng giọng nói
def transcribe(audio_file):
    if audio_file is None:
        return "Không có file âm thanh."
    result = asr(audio_file)
    text = result["text"]
    return text

# Gửi văn bản vào LLM (nếu muốn)
def ask_llm(text):
    if not text:
        return "Chưa có văn bản đầu vào"
    response = llm(text)
    return response[0]["generated_text"]

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# Voice Assistant Demo")
    gr.Markdown("### Gửi file âm thanh để nhận diện giọng nói bằng Whisper")

    with gr.Row():
        audio_input = gr.Audio(label="Chọn file âm thanh", type="filepath")
        btn_transcribe = gr.Button("Nhận diện giọng nói")
    
    transcript_output = gr.Textbox(label="Kết quả nhận diện")

    btn_transcribe.click(fn=transcribe, inputs=audio_input, outputs=transcript_output)

    gr.Markdown("### Gửi văn bản vào mô hình LLM")
    llm_input = gr.Textbox(label="Văn bản đầu vào")
    btn_llm = gr.Button("Gửi vào LLM")
    llm_output = gr.Textbox(label="Kết quả LLM")

    btn_llm.click(fn=ask_llm, inputs=llm_input, outputs=llm_output)

demo.launch()