Spaces:
Sleeping
Sleeping
| import glob | |
| 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 | |
| from espnet2.text.token_id_converter import TokenIDConverter | |
| from espnet2.s2t.espnet_model import ESPnetS2TModel | |
| from espnet2.bin.s2t_inference import Speech2Text | |
| import espnetez as ez | |
| import torch | |
| import numpy as np | |
| import logging | |
| import gradio as gr | |
| import librosa | |
| def log(temp_dir, text): | |
| with open(f"{temp_dir}/output.log", "a") as f: | |
| f.write(text + "\n") | |
| def count_parameters(model): | |
| return sum(p.numel() for p in model.parameters() if p.requires_grad) | |
| def get_dataset(data_path, data_info, test_count=10): | |
| # load data | |
| data = {} | |
| keys = [] | |
| with open(f"{data_path}/text", "r", encoding="utf-8") as f: | |
| for line in f.readlines(): | |
| audio_id, text = line.split(maxsplit=1) | |
| data[audio_id.strip()] = {"text": text.strip()} | |
| keys.append(audio_id.strip()) | |
| # load text_ctc data | |
| with open(f"{data_path}/text_ctc", "r", encoding="utf-8") as f: | |
| for line in f.readlines(): | |
| audio_id, text = line.split(maxsplit=1) | |
| data[audio_id.strip()]["text_ctc"] = text.strip() | |
| # load audio path | |
| for audio_path in glob.glob(f"{data_path}/audio/*"): | |
| audio_id = Path(audio_path).stem | |
| data[audio_id]["audio_path"] = audio_path | |
| # Convert to list | |
| data = [{ | |
| 'id': audio_id, | |
| 'text': data[audio_id]['text'], | |
| 'text_ctc': data[audio_id]['text_ctc'], | |
| 'audio_path': data[audio_id]['audio_path'], | |
| } for audio_id in keys] | |
| return ez.dataset.ESPnetEZDataset(data[test_count:], data_info), ez.dataset.ESPnetEZDataset(data[:test_count], data_info), data[:test_count] | |
| class CustomFinetuneModel(ESPnetS2TModel): | |
| def __init__(self, model, tempdir_path, log_every=500): | |
| super().__init__( | |
| vocab_size=model.vocab_size, | |
| token_list=model.token_list, | |
| frontend=model.frontend, | |
| specaug=model.specaug, | |
| normalize=model.normalize, | |
| preencoder=model.preencoder, | |
| encoder=model.encoder, | |
| postencoder=model.postencoder, | |
| decoder=model.decoder, | |
| ctc=model.ctc, | |
| ctc_weight=model.ctc_weight, | |
| interctc_weight=model.interctc_weight, | |
| ignore_id=model.ignore_id, | |
| lsm_weight=0.0, | |
| length_normalized_loss=False, | |
| report_cer=False, | |
| report_wer=False, | |
| sym_space="<space>", | |
| sym_blank="<blank>", | |
| sym_sos = "<sos>", | |
| sym_eos = "<eos>", | |
| sym_sop = "<sop>", # start of prev | |
| sym_na = "<na>", # not available | |
| extract_feats_in_collect_stats=model.extract_feats_in_collect_stats, | |
| ) | |
| self.iter_count = 0 | |
| self.log_every = log_every | |
| self.log_stats = { | |
| 'loss': 0.0, | |
| 'acc': 0.0 | |
| } | |
| self.tempdir_path = tempdir_path | |
| def forward(self, *args, **kwargs): | |
| out = super().forward(*args, **kwargs) | |
| self.log_stats['loss'] += out[1]['loss'].item() | |
| self.log_stats['acc'] += out[1]['acc'].item() | |
| self.iter_count += 1 | |
| if self.iter_count % self.log_every == 0: | |
| loss = self.log_stats['loss'] / self.log_every | |
| acc = self.log_stats['acc'] / self.log_every | |
| log(self.tempdir_path, f"[{self.iter_count}] - loss: {loss:.3f} - acc: {acc:.3f}") | |
| self.log_stats['loss'] = 0.0 | |
| self.log_stats['acc'] = 0.0 | |
| return out | |
| 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.""" | |
| log(tempdir_path, "Start generating baseline...") | |
| gr.Info("Start generating baseline...") | |
| ref, base = baseline_model(lang, task, tempdir_path) | |
| 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.") | |
| # define tokenizer | |
| tokenizer = SentencepiecesTokenizer("assets/owsm_ebf_v3.1_base/bpe.model") | |
| converter = TokenIDConverter("assets/owsm_ebf_v3.1_base/tokens.txt") | |
| def tokenize(text): | |
| return np.array(converter.tokens2ids(tokenizer.text2tokens(text))) | |
| data_info = { | |
| "speech": lambda d: librosa.load(d["audio_path"], sr=16000)[0], | |
| "text": lambda d: tokenize(f"<{lang}><{task}><notimestamps> {d['text']}"), | |
| "text_ctc": lambda d: tokenize(d["text_ctc"]), | |
| "text_prev": lambda d: tokenize("<na>"), | |
| } | |
| # load dataset and define data_info | |
| train_dataset, test_dataset, test_list = get_dataset(tempdir_path, data_info) | |
| log(tempdir_path, "Loading dataset...") | |
| gr.Info("Loaded dataset.") | |
| # load and update configuration | |
| log(tempdir_path, "Setting up the training configuration...") | |
| pretrain_config = ez.config.from_yaml( | |
| "s2t", | |
| "assets/owsm_ebf_v3.1_base/config.yaml", | |
| ) | |
| finetune_config = ez.config.update_finetune_config( | |
| "s2t", pretrain_config, "assets/owsm_ebf_v3.1_base/owsm_finetune_base.yaml" | |
| ) | |
| finetune_config['max_epoch'] = max_epoch | |
| finetune_config['optim'] = optimizer | |
| finetune_config['optim_conf']['lr'] = learning_rate | |
| finetune_config['optim_conf']['weight_decay'] = weight_decay | |
| finetune_config['scheduler'] = scheduler | |
| finetune_config['scheduler_conf']['warmup_steps'] = warmup_steps | |
| finetune_config['multiple_iterator'] = False | |
| finetune_config['num_iters_per_epoch'] = None | |
| finetune_config['multiprocessing_distributed'] = False | |
| def build_model_fn(args): | |
| model, _ = S2TTask.build_model_from_file( | |
| "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", | |
| ) | |
| model.train() | |
| log(tempdir_path, f'Trainable parameters: {count_parameters(model)}') | |
| model = CustomFinetuneModel(model, tempdir_path, log_every=log_every) | |
| return model | |
| trainer = ez.Trainer( | |
| task='s2t', | |
| train_config=finetune_config, | |
| train_dataset=train_dataset, | |
| valid_dataset=test_dataset, | |
| build_model_fn=build_model_fn, # provide the pre-trained model | |
| data_info=data_info, | |
| output_dir=f"{tempdir_path}/exp/finetune", | |
| stats_dir=f"{tempdir_path}/exp/stats", | |
| ngpu=1 | |
| ) | |
| gr.Info("start collect stats") | |
| log(tempdir_path, "Start collect stats process...") | |
| trainer.collect_stats() | |
| gr.Info("Finished collect stats, starting training.") | |
| log(tempdir_path, "Finished collect stats, starting training...") | |
| trainer.train() | |
| 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...") | |
| 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/config.yaml") | |
| log(tempdir_path, "exp/finetune/valid.acc.ave.pth") | |
| log(tempdir_path, "output.log") | |
| log(tempdir_path, "acc.png") | |
| log(tempdir_path, "loss.png") | |
| log(tempdir_path, "base.txt") | |
| log(tempdir_path, "hyp.txt") | |
| log(tempdir_path, "ref.txt") | |
| finetune_zip = zipfile.ZipFile(f"{tempdir_path}/finetune.zip", "w", zipfile.ZIP_DEFLATED) | |
| finetune_zip.write(f"{tempdir_path}/exp/stats") | |
| 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/config.yaml") | |
| finetune_zip.write(f"{tempdir_path}/exp/finetune/valid.acc.ave.pth") | |
| finetune_zip.write(f"{tempdir_path}/output.log") | |
| finetune_zip.write(f"{tempdir_path}/acc.png") | |
| finetune_zip.write(f"{tempdir_path}/loss.png") | |
| finetune_zip.write(f"{tempdir_path}/base.txt") | |
| finetune_zip.write(f"{tempdir_path}/hyp.txt") | |
| finetune_zip.write(f"{tempdir_path}/ref.txt") | |
| finetune_zip.close() | |
| gr.Info("Finished generating result file in zip!") | |
| log(tempdir_path, "Finished generating result file in zip!") | |
| 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): | |
| 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 | |
| tokenizer = SentencepiecesTokenizer("assets/owsm_ebf_v3.1_base/bpe.model") | |
| converter = TokenIDConverter("assets/owsm_ebf_v3.1_base/tokens.txt") | |
| def tokenize(text): | |
| return np.array(converter.tokens2ids(tokenizer.text2tokens(text))) | |
| data_info = { | |
| "speech": lambda d: librosa.load(d["audio_path"], sr=16000)[0], | |
| "text": lambda d: tokenize(f"<{lang}><{task}><notimestamps> {d['text']}"), | |
| "text_ctc": lambda d: tokenize(d["text_ctc"]), | |
| "text_prev": lambda d: tokenize("<na>"), | |
| } | |
| # load dataset and define data_info | |
| train_dataset, test_dataset, test_list = get_dataset(tempdir_path, data_info) | |
| log(tempdir_path, "Loaded dataset.") | |
| gr.Info("Loaded dataset.") | |
| gr.Info("Loading pretrained model...") | |
| log(tempdir_path, "Loading pretrained model...") | |
| 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.3, | |
| lang_sym=f"<{lang}>", | |
| task_sym=f"<{task}>", | |
| ) | |
| model.s2t_model.eval() | |
| base = "" | |
| ref = "" | |
| with open(f"{tempdir_path}/base.txt", "w") as f_base, open(f"{tempdir_path}/ref.txt", "w") as f_ref: | |
| for i in range(len(test_list)): | |
| data = test_list[i] | |
| f_ref.write(data['text'] + '\n') | |
| out = model(librosa.load(data['audio_path'], sr=16000)[0])[0][3] | |
| f_base.write(out + '\n') | |
| ref += data['text'] + '\n' | |
| base += out + '\n' | |
| return ref, base | |