Update finetune.py

finetune.py

@@ -191,7 +191,34 @@ def finetune_model(lang, task, tempdir_path, log_every, max_epoch, scheduler, wa
     gr.Info("Finished collect stats, starting training.")
     log(tempdir_path, "Finished collect stats, starting training...")
     trainer.train()
-    gr.Info("Finished Fine-tuning! Archiving experiment files...")
+    gr.Info("Finished Fine-tuning!")
+
+    gr.Info("Start generating output for test set!")
+    log(tempdir_path, "Start generating output for test set!")
+
+    del trainer
+    model = Speech2Text(
+        "assets/owsm_ebf_v3.1_base/config.yaml",
+        "assets/owsm_ebf_v3.1_base/owsm_v3.1_base.trained.pth",
+        device="cuda" if torch.cuda.is_available() else "cpu",
+        token_type="bpe",
+        bpemodel="assets/owsm_ebf_v3.1_base/bpe.model",
+        beam_size=5,
+        ctc_weight=0.0,
+        lang_sym=f"<{lang}>",
+        task_sym=f"<{task}>",
+    )
+    model.s2t_model.eval()
+    d = torch.load(f"{tempdir_path}/exp/finetune/valid.acc.ave.pth")
+    model.s2t_model.load_state_dict(d)
+    
+    hyp = ""
+    with open(f"{tempdir_path}/hyp.txt", "w") as f_hyp:
+        for i in range(len(test_list)):
+            data = test_list[i]
+            out = model(librosa.load(data['audio_path'], sr=16000)[0])[0][3]
+            f_hyp.write(out + '\n')
+            hyp += out + '\n'
 
     log(tempdir_path, "Finished fine-tuning.")
     log(tempdir_path, "Start archiving experiment files...")
@@ -224,34 +251,7 @@ def finetune_model(lang, task, tempdir_path, log_every, max_epoch, scheduler, wa
 
     gr.Info("Finished generating result file in zip!")
     log(tempdir_path, "Finished generating result file in zip!")
-
-    gr.Info("Start generating output for test set!")
-    log(tempdir_path, "Start generating output for test set!")
-
-    del trainer
-    model = Speech2Text(
-        "assets/owsm_ebf_v3.1_base/config.yaml",
-        "assets/owsm_ebf_v3.1_base/owsm_v3.1_base.trained.pth",
-        device="cuda" if torch.cuda.is_available() else "cpu",
-        token_type="bpe",
-        bpemodel="assets/owsm_ebf_v3.1_base/bpe.model",
-        beam_size=5,
-        ctc_weight=0.0,
-        lang_sym=f"<{lang}>",
-        task_sym=f"<{task}>",
-    )
-    model.s2t_model.eval()
-    d = torch.load(f"{tempdir_path}/exp/finetune/valid.acc.ave.pth")
-    model.s2t_model.load_state_dict(d)
     
-    hyp = ""
-    with open(f"{tempdir_path}/hyp.txt", "w") as f_hyp:
-        for i in range(len(test_list)):
-            data = test_list[i]
-            out = model(librosa.load(data['audio_path'], sr=16000)[0])[0][3]
-            f_hyp.write(out + '\n')
-            hyp += out + '\n'
-
     return [f"{tempdir_path}/finetune.zip", f"{tempdir_path}/ref.txt", f"{tempdir_path}/base.txt", f"{tempdir_path}/hyp.txt"], ref, base, hyp