import gradio as gr
import librosa
import numpy as np
from transformers import pipeline

asr_pipe = pipeline(
    "automatic-speech-recognition", model="Monsia/wav2ve-dyula-lr-3e-4-aug2-2"
)
mt_pipe = pipeline("text2text-generation", model="Monsia/nllb200-mt-dyu-fr-1e-05-Aug-4")


def load_and_fix_data(input_file):
    speech, sample_rate = librosa.load(input_file)
    if len(speech.shape) > 1:
        speech = speech[:, 0] + speech[:, 1]
    if sample_rate != 16000:
        speech = librosa.resample(speech, orig_sr=sample_rate, target_sr=16000)
    return speech


def koumankan_ast(path):
    speech_array = load_and_fix_data(path)
    text_dyu = asr_pipe(speech_array)["text"]
    text_fr = mt_pipe(text_dyu)[0]["generated_text"]
    return text_dyu, text_fr


iface = gr.Interface(
    fn=koumankan_ast,
    inputs=gr.Audio(
        sources=["microphone", "upload"], type="filepath", label="Record/ Drop audio"
    ),
    outputs=[gr.Textbox(label="Dioula"), gr.Textbox(label="Français")],
    title="Koumankan: AST Dioula vers Français",
    description="Le projet Koumankan vise à améliorer la disponibilité et la qualité des données vocales et textuelles pour les langues africaines, et à développer des modèles de reconnaissance vocale et de traduction automatique pour ces langues. Ici nous vous présentons Koumankan AST Dioula vers Français.",
)
iface.launch()