Added model version dropdown

app.py

@@ -13,14 +13,26 @@ login(token=os.environ['hf_token'])
 #     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.'
+# 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')
 
 #%%
-def transcribe(audio_file, audio_mic):
+def transcribe(audio_file, audio_mic, model_version):
+    if model_version == 'local':
+        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.'
+    elif model_version == 'large-v2':
+        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.'
+    elif model_version == 'large-v3':
+        whisper = pipeline(model='jlvdoorn/whisper-large-v3-atco2-asr-atcosim')
+        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:
@@ -48,8 +60,8 @@ def extractCallSignCommand(transcription):
         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)
+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:
@@ -59,10 +71,16 @@ def transcribeAndExtract(audio_file, audio_mic, transcribe_only):
 #%%
 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)],
+        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=['local', 'large-v2', 'large-v3'], value='large-v3', label='Whisper model version'),
+            ],
         outputs=[gr.Text(label='Transcription'), gr.Text(label='Callsigns, commands and values')],
-        title=ttl,
-        description=dis,
+        title='Whisper',
+        description='Transcribe and extract',
 )
 
 #%%