import torch
import scipy
import gradio as gr
from transformers import set_seed
from datasets import load_dataset, Audio

import  goai_stt, goai_tts, goai_traduction

#language_list = ['mos', 'fra', 'eng']
device = 0 if torch.cuda.is_available() else "cpu"

demo = gr.Blocks()

goai_stt = gr.Interface(
    fn = goai_stt.goai_stt,
    inputs=[
        gr.Audio(sources=["microphone", "upload"], type="filepath"),
        device
    ],
    outputs="text",
    title="Speech-to-text"
)

goai_tts = gr.Interface(
    fn=goai_tts.goai_tts,
    inputs=[
        gr.Text(label="Input text"),
        device
    ],
    outputs=[
        gr.Audio(label="Generated Audio", type="numpy")
    ],
    title="Text-to-speech"
)

goai_traduction = gr.Interface(
    fn=goai_traduction.goai_traduction,
    inputs=[
        gr.Textbox(label="Text", placeholder="Yaa sõama"),
        gr.Dropdown(label="Source Language", choices=["eng_Latn", "fra_Latn", "mos_Latn"]),
        gr.Dropdown(label="Target Language", choices=["eng_Latn", "fra_Latn", "mos_Latn"])
    ],
    outputs=["text"],
    examples=[["Yʋʋm a wãn la b kẽesd biig lekolle?", "mos_Latn", "fra_Latn"]],
    title="Translation Demo",
)


with demo:
    gr.TabbedInterface(
        [goai_traduction, goai_tts, goai_stt],
        ["Translation", "Text-to-speech", "Speech-to-text"],
    )

demo.launch()