Spaces:
Running
Running
| #%% | |
| 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') | |
| # 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))] | |
| bert_atco_ner = pipeline(model='Jzuluaga/bert-base-ner-atc-en-atco2-1h') | |
| whisper_v2 = pipeline(model='jlvdoorn/whisper-large-v2-atco2-asr-atcosim') | |
| whisper_v3 = pipeline(model='jlvdoorn/whisper-large-v3-atco2-asr-atcosim') | |
| #%% | |
| def transcribe(audio, model_version): | |
| if model_version == 'large-v2': | |
| whisper = whisper_v2 | |
| ttl = 'Whisper Large v2 - ATCO2-ATCOSIM' | |
| dis = 'This demo will transcribe ATC audio files by using the Whisper Large v2 model fine-tuned on the ATCO2 and ATCOSIM datasets. \n \n Further it uses a Named Entity Recognition model to extract callsigns, commands and values from the transcription. \n This model is based on Google\'s BERT model and fine-tuned on the ATCO2 dataset.' | |
| elif model_version == 'large-v3': | |
| whisper = whisper_v3 | |
| ttl = 'Whisper Large v3 - ATCO2-ATCOSIM' | |
| dis = 'This demo will transcribe ATC audio files by using the Whisper Large v3 model fine-tuned on the ATCO2 and ATCOSIM datasets. \n \n Further it uses a Named Entity Recognition model to extract callsigns, commands and values from the transcription. \n This model is based on Google\'s BERT model and fine-tuned on the ATCO2 dataset.' | |
| if audio is not None: | |
| return whisper(audio)['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, transcribe_only, model_version): | |
| transcription = transcribe(audio, model_version) | |
| 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), | |
| gr.Dropdown(choices=['large-v2', 'large-v3'], value='large-v3', label='Whisper model version'), | |
| ], | |
| outputs=[gr.Text(label='Transcription'), gr.Text(label='Callsigns, commands and values')], | |
| title='Whisper ATC - Large v3', | |
| description='Transcribe and extract', | |
| # examples = examples, | |
| ) | |
| #%% | |
| file_iface = gr.Interface( | |
| fn = transcribeAndExtract, | |
| inputs = [gr.Audio(source='upload', type='filepath', 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', | |
| ) | |
| 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() |