Spaces:
Running
Running
| import os | |
| from typing import Dict | |
| from diacritization_evaluation import der, wer | |
| import torch | |
| from torch import nn | |
| from torch import optim | |
| from torch.cuda.amp import autocast | |
| from torch.utils.tensorboard.writer import SummaryWriter | |
| from tqdm.notebook import tqdm | |
| from tqdm import trange | |
| from diacritization_evaluation import util | |
| from .config_manager import ConfigManager | |
| from .dataset import load_iterators | |
| from .diacritizer import CBHGDiacritizer, Seq2SeqDiacritizer | |
| from .options import OptimizerType | |
| import gdown | |
| class Trainer: | |
| def run(self): | |
| raise NotImplementedError | |
| class GeneralTrainer(Trainer): | |
| def __init__(self, config_path: str, model_kind: str) -> None: | |
| self.config_path = config_path | |
| self.model_kind = model_kind | |
| self.config_manager = ConfigManager( | |
| config_path=config_path, model_kind=model_kind | |
| ) | |
| self.config = self.config_manager.config | |
| self.losses = [] | |
| self.lr = 0 | |
| self.pad_idx = 0 | |
| self.criterion = nn.CrossEntropyLoss(ignore_index=self.pad_idx) | |
| self.set_device() | |
| self.config_manager.create_remove_dirs() | |
| self.text_encoder = self.config_manager.text_encoder | |
| self.start_symbol_id = self.text_encoder.start_symbol_id | |
| self.summary_manager = SummaryWriter(log_dir=self.config_manager.log_dir) | |
| self.model = self.config_manager.get_model() | |
| self.optimizer = self.get_optimizer() | |
| self.model = self.model.to(self.device) | |
| self.load_model(model_path=self.config.get("train_resume_model_path")) | |
| self.load_diacritizer() | |
| self.initialize_model() | |
| def set_device(self): | |
| if self.config.get("device"): | |
| self.device = self.config["device"] | |
| else: | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| def load_diacritizer(self): | |
| if self.model_kind in ["cbhg", "baseline"]: | |
| self.diacritizer = CBHGDiacritizer(self.config_path, self.model_kind) | |
| elif self.model_kind in ["seq2seq", "tacotron_based"]: | |
| self.diacritizer = Seq2SeqDiacritizer(self.config_path, self.model_kind) | |
| def initialize_model(self): | |
| if self.global_step > 1: | |
| return | |
| if self.model_kind == "transformer": | |
| print("Initializing using xavier_uniform_") | |
| self.model.apply(initialize_weights) | |
| def load_model(self, model_path: str = None, load_optimizer: bool = True): | |
| with open( | |
| self.config_manager.base_dir / f"{self.model_kind}_network.txt", "w" | |
| ) as file: | |
| file.write(str(self.model)) | |
| if model_path is None: | |
| last_model_path = self.config_manager.get_last_model_path() | |
| if last_model_path is None: | |
| self.global_step = 1 | |
| return | |
| else: | |
| last_model_path = model_path | |
| print(f"loading from {last_model_path}") | |
| saved_model = torch.load(last_model_path, torch.device(self.config.get("device"))) | |
| self.model.load_state_dict(saved_model["model_state_dict"]) | |
| if load_optimizer: | |
| self.optimizer.load_state_dict(saved_model["optimizer_state_dict"]) | |
| self.global_step = saved_model["global_step"] + 1 | |
| class DiacritizationTester(GeneralTrainer): | |
| def __init__(self, config_path: str, model_kind: str, model_path: str) -> None: | |
| # if config_path == 'config/test.yml' or config_path == "Arabic_Diacritization/config/test.yml": | |
| # print("Exporting the pretrained models ... ") | |
| # url = 'https://drive.google.com/uc?id=12aYNY7cbsLNzhdPdC2K3u1sgrb1lpzwO' | |
| # gdown.cached_download(url,'model.zip', quiet=False, postprocess=gdown.extractall) | |
| self.config_path = config_path | |
| self.model_kind = model_kind | |
| self.config_manager = ConfigManager( | |
| config_path=config_path, model_kind=model_kind | |
| ) | |
| self.config = self.config_manager.config | |
| # print(self.config) | |
| self.pad_idx = 0 | |
| self.criterion = nn.CrossEntropyLoss(ignore_index=self.pad_idx) | |
| self.set_device() | |
| self.text_encoder = self.config_manager.text_encoder | |
| self.start_symbol_id = self.text_encoder.start_symbol_id | |
| self.model = self.config_manager.get_model() | |
| self.model = self.model.to(self.device) | |
| self.load_model(model_path=model_path, load_optimizer=False) | |
| self.load_diacritizer() | |
| self.diacritizer.set_model(self.model) | |
| self.initialize_model() | |
| def collate_fn(self, data): | |
| """ | |
| Padding the input and output sequences | |
| """ | |
| def merge(sequences): | |
| lengths = [len(seq) for seq in sequences] | |
| padded_seqs = torch.zeros(len(sequences), max(lengths)).long() | |
| for i, seq in enumerate(sequences): | |
| end = lengths[i] | |
| padded_seqs[i, :end] = seq[:end] | |
| return padded_seqs, lengths | |
| data.sort(key=lambda x: len(x[0]), reverse=True) | |
| # separate source and target sequences | |
| src_seqs, trg_seqs, original = zip(*data) | |
| # merge sequences (from tuple of 1D tensor to 2D tensor) | |
| src_seqs, src_lengths = merge(src_seqs) | |
| trg_seqs, trg_lengths = merge(trg_seqs) | |
| batch = { | |
| "original": original, | |
| "src": src_seqs, | |
| "target": trg_seqs, | |
| "lengths": torch.LongTensor(src_lengths), # src_lengths = trg_lengths | |
| } | |
| return batch | |
| def get_batch(self, sentence): | |
| data = self.text_encoder.clean(sentence) | |
| text, inputs, diacritics = util.extract_haraqat(data) | |
| inputs = torch.Tensor(self.text_encoder.input_to_sequence("".join(inputs))) | |
| diacritics = torch.Tensor(self.text_encoder.target_to_sequence(diacritics)) | |
| batch = self.collate_fn([(inputs, diacritics, text)]) | |
| return batch | |
| def infer(self, sentence): | |
| self.model.eval() | |
| batch = self.get_batch(sentence) | |
| predicted = self.diacritizer.diacritize_batch(batch) | |
| return predicted[0] | |