import gradio as gr
from transformers import pipeline

asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-tiny")

llm_pipe = pipeline(
    "text-generation",
    model="microsoft/Phi-3-mini-4k-instruct",
    device_map="auto",
    torch_dtype="auto",
)

tts_pipe = pipeline("text-to-speech", model="suno/bark")


def text_to_speech(text):
    speech = tts_pipe(text)
    audio = speech["audio"].squeeze()
    return speech["sampling_rate"], audio


def process_audio(audio_path):
    if audio_path is None:
        return "Please record something first.", None

    user_text = asr_pipe(audio_path)["text"].strip()

    messages = [
        {"role": "system", "content": "You are a helpful voice assistant. Answer in 1-3 short sentences."},
        {"role": "user", "content": user_text},
    ]
    output = llm_pipe(
        messages,
        max_new_tokens=100,
        do_sample=False,
    )
    reply = output[0]["generated_text"][-1]["content"].strip()

    sampling_rate, audio = text_to_speech(reply)
    return reply, (sampling_rate, audio)


demo = gr.Interface(
    fn=process_audio,
    inputs=gr.Audio(label="Audio Input", sources=["microphone"], type="filepath"),
    outputs=[
        gr.Textbox(label="Assistant reply"),
        gr.Audio(label="Spoken reply", type="numpy"),
    ],
    title="A Simple AI Voice Assistant",
    description="An AI Assistant That Will Answer Your Questions!",
)

demo.launch(inbrowser=True, debug=True)