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'], atcosim[0]['audio'], atco2[1]['audio'], atcosim[1]['audio'], atco2[2]['audio'], atcosim[2]['audio']] | |
| examples_labels = ['Example ' + str(i) for i in len(examples)] | |
| ## Try to load a local model if available | |
| # try: | |
| # whisper = pipeline(model='/mnt/projects/whisper/WhisperANSP/Models/whisper-large-v2-atco2-asr-atcosim-ANSP-3h1m', task='automatic-speech-recognition') | |
| # ttl = 'Whisper Large v2 - ATCO2-ATCOSIM-ANSP' | |
| # dis = 'This demo will transcribe ATC audio files by using the Whisper Large v2 model fine-tuned on the ATCO2, ATCOSIM and ANSP 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.' | |
| # except: | |
| # whisper = pipeline(model='jlvdoorn/whisper-large-v2-atco2-asr-atcosim') | |
| # 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.' | |
| 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_file, audio_mic, 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_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, model_version): | |
| transcription = transcribe(audio_file, audio_mic, 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, | |
| ) | |
| #%% | |
| #iface.launch(server_name='0.0.0.0', server_port=9000) | |
| iface.launch() |