Update app.py

app.py

@@ -15,28 +15,28 @@ import time
 
 
 auth_token = os.environ.get('TOKEN')
-model1 = WhisperForConditionalGeneration.from_pretrained("rohitp1/kkkh_whisper_small_distillation_att_loss_libri360_epochs_100_batch_4_concat_dataset",
-                                                         use_auth_token=auth_token)
 
-tokenizer1 = WhisperTokenizer.from_pretrained("rohitp1/kkkh_whisper_small_distillation_att_loss_libri360_epochs_100_batch_4_concat_dataset",
-                                                         use_auth_token=auth_token)
 
-feat_ext1 = WhisperFeatureExtractor.from_pretrained("rohitp1/kkkh_whisper_small_distillation_att_loss_libri360_epochs_100_batch_4_concat_dataset",
-                                                         use_auth_token=auth_token)
+M1 = "rohitp1/kkkh_whisper_small_distillation_att_loss_libri360_epochs_100_batch_4_concat_dataset"
+M2 = "rohitp1/dgx2_whisper_small_finetune_teacher_babble_noise_libri_360_hours_50_epochs_batch_8"
+M3 = "rohitp1/subhadeep_whisper_small_finetune_teacher_no_noise_libri_360_hours_100_epochs_batch_8"
 
+model1 = WhisperForConditionalGeneration.from_pretrained(M1, use_auth_token=auth_token)
+tokenizer1 = WhisperTokenizer.from_pretrained(M1, use_auth_token=auth_token)
+feat_ext1 = WhisperFeatureExtractor.from_pretrained(M1, use_auth_token=auth_token)
 
-model2 = WhisperForConditionalGeneration.from_pretrained("rohitp1/dgx2_whisper_small_finetune_teacher_babble_noise_libri_360_hours_50_epochs_batch_8",
-                                                         use_auth_token=auth_token)
+model2 = WhisperForConditionalGeneration.from_pretrained(M2, use_auth_token=auth_token)
+tokenizer2 = WhisperTokenizer.from_pretrained(M2, use_auth_token=auth_token)
+feat_ext2 = WhisperFeatureExtractor.from_pretrained(M2, use_auth_token=auth_token)
 
-tokenizer2 = WhisperTokenizer.from_pretrained("rohitp1/dgx2_whisper_small_finetune_teacher_babble_noise_libri_360_hours_50_epochs_batch_8",
-                                                         use_auth_token=auth_token)
-
-feat_ext2 = WhisperFeatureExtractor.from_pretrained("rohitp1/dgx2_whisper_small_finetune_teacher_babble_noise_libri_360_hours_50_epochs_batch_8",
-                                                         use_auth_token=auth_token)
+model3 = WhisperForConditionalGeneration.from_pretrained(M3, use_auth_token=auth_token)
+tokenizer3 = WhisperTokenizer.from_pretrained(M3, use_auth_token=auth_token)
+feat_ext3 = WhisperFeatureExtractor.from_pretrained(M3, use_auth_token=auth_token)
 
 
 p1 = pipeline('automatic-speech-recognition', model=model1, tokenizer=tokenizer1, feature_extractor=feat_ext1)
 p2 = pipeline('automatic-speech-recognition', model=model2, tokenizer=tokenizer2, feature_extractor=feat_ext2)
+p3 = pipeline('automatic-speech-recognition', model=model3, tokenizer=tokenizer3, feature_extractor=feat_ext3)
 
 def transcribe(mic_input, upl_input, model_type):
     if mic_input:
@@ -44,8 +44,10 @@ def transcribe(mic_input, upl_input, model_type):
     else:
         audio = upl_input
     time.sleep(3)
-    if model_type =='Finetuned':
+    if model_type == 'NoisyFinetuned':
         text = p2(audio)["text"]
+    elif model_type == 'CleanFinetuned':
+        text = p3(audio)["text"]
     else:
         text = p1(audio)["text"]
     # state = text + " "
@@ -77,7 +79,7 @@ def transcribe(mic_input, upl_input, model_type):
 # demo.launch()
 
 def clear_inputs_and_outputs():
-    return [None, None, "RobustDistillation", None]
+    return [None, None, "CleanFinetuned", None]
 
 # Main function
 if __name__ == "__main__":
@@ -100,7 +102,7 @@ if __name__ == "__main__":
                 )
 
                 with gr.Row():
-                    model_type = gr.inputs.Dropdown(["RobustDistillation", "Finetuned"], label='Model Type')
+                    model_type = gr.inputs.Dropdown(["RobustDistillation", "NoisyFinetuned", "CleanFinetuned"], label='Model Type')
 
                 with gr.Row():
                     clr_btn = gr.Button(value="Clear", variant="secondary")