#%%
from transformers import pipeline
import gradio as gr
import os

#%%
whisper = pipeline(model='/mnt/projects/whisper/WhisperANSP/Models/whisper-large-v2-atco2-asr-atcosim-ANSP-3h1m', task='automatic-speech-recognition')
bert_atco_ner = pipeline(model='Jzuluaga/bert-base-ner-atc-en-atco2-1h')

#%%
def transcribe(audio_file, audio_mic):
    if audio_mic is not None:
        return whisper(audio_mic)['text']
    elif audio_file is not None:
        return whisper(audio_file)['text']
    else:
        return 'There was no audio to transcribe...'

#%%
def extractCallSignCommand(transcription):
    if type(transcription) is str:
        result = bert_atco_ner(transcription)
        callsigns = []
        commands = []
        values = []
        for item in result:
            if 'callsign' in item['entity']:
                callsigns.append(item['word'])
            if 'command' in item['entity']:
                commands.append(item['word'])
            if 'value' in item['entity']:
                values.append(item['word'])
                
        return 'Callsigns: ' + ', '.join(callsigns) + '\nCommands: ' + ', '.join(commands) + '\nValues: ' + ', '.join(values)
    else:
        return 'There was no transcription to extract a callsign or command from...'

#%%
def transcribeAndExtract(audio_file, audio_mic, transcribe_only):
    transcription = transcribe(audio_file, audio_mic)
    if not transcribe_only:
        callSignCommandValues = extractCallSignCommand(transcription)
    else:
        callSignCommandValues = ''
    return transcription, callSignCommandValues

#%%
iface = gr.Interface(
        fn=transcribeAndExtract,
        inputs=[gr.Audio(source='upload', type='filepath', interactive=True), gr.Audio(source='microphone', type='filepath'), gr.Checkbox(label='Transcribe only', default=False)],
        outputs=[gr.Text(label='Transcription'), gr.Text(label='Callsigns, commands and values')],
        title='Whisper Large v2 - ATCO2-ATCOSIM-ANSP',
        description='This demo will transcribe ATC audio files by using the Whisper Large v2 model fine-tuned on the ATCO2, ATCOSIM and ANSP datasets. Further it uses a Named Entity Recognition model to extract callsigns, commands and values from the transcription. This model is based on Google\'s BERT model and fine-tuned on the ATCO2 dataset.',
)

#%%
iface.launch(server_name='0.0.0.0', server_port=9000)