bug fix

finetune.py

@@ -3,6 +3,7 @@ import sys
 from pathlib import Path
 import shutil
 import os
+import zipfile
 
 from espnet2.tasks.s2t import S2TTask
 from espnet2.text.sentencepiece_tokenizer import SentencepiecesTokenizer
@@ -18,24 +19,9 @@ import gradio as gr
 import librosa
 
 
-class Logger:
-    def __init__(self, filename):
-        self.terminal = sys.stdout
-        self.log = open(filename, "w")
-
-    def write(self, message):
-        self.terminal.write(message)
-        self.log.write(message)
-
-    def flush(self):
-        self.terminal.flush()
-        self.log.flush()
-
-    def isatty(self):
-        return False
-
-
-sys.stdout = Logger("output.log")
+def log(temp_dir, text):
+    with open(f"{temp_dir}/output.log", "a") as f:
+        f.write(text + "\n")
 
 
 def count_parameters(model):
@@ -128,11 +114,11 @@ class CustomFinetuneModel(ESPnetS2TModel):
 def finetune_model(lang, task, tempdir_path, log_every, max_epoch, scheduler, warmup_steps, optimizer, learning_rate, weight_decay):
     """Main function for finetuning the model."""
 
-    print("Start generating baseline...")
+    log(tempdir_path, "Start generating baseline...")
     gr.Info("Start generating baseline...")
-    baseline_model(lang, task, tempdir_path)
+    ref, base = baseline_model(lang, task, tempdir_path)
     
-    print("Start loading dataset...")
+    log(tempdir_path, "Start generating hypothesis...")
     gr.Info("Start Fine-tuning process...")
     if len(tempdir_path) == 0:
         raise gr.Error("Please upload a zip file first.")
@@ -153,11 +139,11 @@ def finetune_model(lang, task, tempdir_path, log_every, max_epoch, scheduler, wa
 
     # load dataset and define data_info
     train_dataset, test_dataset, test_list = get_dataset(tempdir_path, data_info)
-    print("Loaded dataset.")
+    log(tempdir_path, "Loading dataset...")
     gr.Info("Loaded dataset.")
 
     # load and update configuration
-    print("Setting up the training configuration...")
+    log(tempdir_path, "Setting up the training configuration...")
     pretrain_config = ez.config.from_yaml(
         "s2t",
         "assets/owsm_ebf_v3.1_base/config.yaml",
@@ -197,24 +183,38 @@ def finetune_model(lang, task, tempdir_path, log_every, max_epoch, scheduler, wa
         ngpu=1
     )
     gr.Info("start collect stats")
-    print("Start collect stats process...")
+    log(tempdir_path, "Start collect stats process...")
     trainer.collect_stats()
+
     gr.Info("Finished collect stats, starting training.")
-    print("Finished collect stats process. Start training.")
+    log(tempdir_path, "Finished collect stats, starting training...")
     trainer.train()
     gr.Info("Finished Fine-tuning! Archiving experiment files...")
-    print("Finished fine-tuning.")
-    print("Start archiving experiment files...")
-    print("Create zip file for the following files into `finetune.zip`:")
-    for f in glob.glob(f"{tempdir_path}/exp/finetune/*"):
-        print(f.replace(tempdir_path, ""))
 
-    shutil.make_archive(f"{tempdir_path}/finetune", 'zip', f"{tempdir_path}/exp")
+    log(tempdir_path, "Finished fine-tuning.")
+    log(tempdir_path, "Start archiving experiment files...")
+    log(tempdir_path, "Create zip file for the following files into `finetune.zip`:")
+    log(tempdir_path, "exp/s2t_stats_raw_bpe50000")
+    log(tempdir_path, "exp/finetune/tensorboard")
+    log(tempdir_path, "exp/finetune/images")
+    log(tempdir_path, "exp/finetune/train.log")
+    log(tempdir_path, "exp/finetune/config.yaml")
+    log(tempdir_path, "exp/finetune/valid.acc.ave.pth")
+
+    finetune_zip = zipfile.ZipFile(f"{tempdir_path}/finetune.zip", "w", zipfile.ZIP_DEFLATED)
+    finetune_zip.write(f"{tempdir_path}/exp/s2t_stats_raw_bpe50000")
+    finetune_zip.write(f"{tempdir_path}/exp/finetune/tensorboard")
+    finetune_zip.write(f"{tempdir_path}/exp/finetune/images")
+    finetune_zip.write(f"{tempdir_path}/exp/finetune/train.log")
+    finetune_zip.write(f"{tempdir_path}/exp/finetune/config.yaml")
+    finetune_zip.write(f"{tempdir_path}/exp/finetune/valid.acc.ave.pth")
+    finetune_zip.close()
+
     gr.Info("Finished generating result file in zip!")
-    print("Finished archiving experiment files.")
+    log(tempdir_path, "Finished generating result file in zip!")
 
-    print("Start generating test result...")
     gr.Info("Start generating output for test set!")
+    log(tempdir_path, "Start generating output for test set!")
 
     del trainer
     model = Speech2Text(
@@ -224,7 +224,7 @@ def finetune_model(lang, task, tempdir_path, log_every, max_epoch, scheduler, wa
         token_type="bpe",
         bpemodel="assets/owsm_ebf_v3.1_base/bpe.model",
         beam_size=5,
-        ctc_weight=0.3,
+        ctc_weight=0.0,
         lang_sym=f"<{lang}>",
         task_sym=f"<{task}>",
     )
@@ -240,12 +240,13 @@ def finetune_model(lang, task, tempdir_path, log_every, max_epoch, scheduler, wa
             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"], hyp
+    return [f"{tempdir_path}/finetune.zip", f"{tempdir_path}/ref.txt", f"{tempdir_path}/base.txt", f"{tempdir_path}/hyp.txt"], ref, base, hyp
 
 
 def baseline_model(lang, task, tempdir_path):
-    print("Start loading dataset...")
+    log(tempdir_path, "Start loading dataset...")
     if len(tempdir_path) == 0:
+        log(tempdir_path, "Please upload a zip file first.")
         raise gr.Error("Please upload a zip file first.")
 
     # define tokenizer
@@ -264,11 +265,11 @@ def baseline_model(lang, task, tempdir_path):
 
     # load dataset and define data_info
     train_dataset, test_dataset, test_list = get_dataset(tempdir_path, data_info)
-    print("Loaded dataset.")
+    log(tempdir_path, "Loaded dataset.")
     gr.Info("Loaded dataset.")
 
-    print("Loading pretrained model...")
     gr.Info("Loading pretrained model...")
+    log(tempdir_path, "Loading pretrained model...")
 
     model = Speech2Text(
         "assets/owsm_ebf_v3.1_base/config.yaml",