#%%
from huggingface_hub import login
from transformers import pipeline
from datasets import load_dataset
import gradio as gr
import os

#login(token=os.environ['hf_token'])

atco2 = load_dataset('jlvdoorn/atco2-asr', split='validation')
atcosim = load_dataset('jlvdoorn/atcosim', split='validation')

num_examples = 3
examples_atco2   = [ [{'sampling_rate': atco2[i]['audio']['sampling_rate'], 'raw': atco2[i]['audio']['array']}, False, 'large-v3'] for i in range(num_examples)]
#examples_atcosim = [ [{'sampling_rate': atcosim[i]['audio']['sampling_rate'], 'raw': atcosim[i]['audio']['array']}, False, 'large-v3'] for i in range(num_examples)]
examples = examples_atco2 #+ examples_atcosim
# examples = [atco2[0]['audio']['array'], atcosim[0]['audio']['array'], atco2[1]['audio']['array'], atcosim[1]['audio']['array'], atco2[2]['audio']['array'], atcosim[2]['audio']['array']]
# examples_labels = ['Example ' + str(i+1) for i in range(len(examples))]

whisper = pipeline(model='jlvdoorn/whisper-large-v3-atco2-asr-atcosim')
#%%
def transcribe(audio, model_version):
    if audio is not None:
        return whisper(audio)['text']
    else:
        return 'There was no audio to transcribe...'

#%%
file_iface = gr.Interface(
    fn = transcribeAndExtract,
    inputs = [gr.Audio(source='upload', interactive=True),
              gr.Checkbox(label='Transcribe only', default=False),
              gr.Dropdown(choices=['large-v2', 'large-v3'], value='large-v3', label='Whisper model version')
            ],

    outputs = [gr.Textbox(label='Transcription'), gr.Textbox(label='Callsigns, commands and values')],
    title = 'Whisper ATC - Large v3',
    description = 'Transcribe and extract',
    examples = examples,
)

mic_iface = gr.Interface(
    fn = transcribeAndExtract,
    inputs = [gr.Audio(source='microphone', type='filepath'),
              gr.Checkbox(label='Transcribe only', default=False),
              gr.Dropdown(choices=['large-v2', 'large-v3'], value='large-v3', label='Whisper model version')
            ],

    outputs = [gr.Textbox(label='Transcription'), gr.Textbox(label='Callsigns, commands and values')],
    title = 'Whisper ATC - Large v3',
    description = 'Transcribe and extract',
)
#%%
demo = gr.TabbedInterface([file_iface, mic_iface], ["File", "Microphone"])
demo.launch()