|
from enum import Enum |
|
import os |
|
from pathlib import Path |
|
import shutil |
|
import subprocess |
|
from typing import Any, Dict |
|
|
|
import ruamel.yaml |
|
import torch |
|
|
|
from poetry_diacritizer.models.baseline import BaseLineModel |
|
from poetry_diacritizer.models.cbhg import CBHGModel |
|
from poetry_diacritizer.models.gpt import GPTModel |
|
from poetry_diacritizer.models.seq2seq import Decoder as Seq2SeqDecoder, Encoder as Seq2SeqEncoder, Seq2Seq |
|
from poetry_diacritizer.models.tacotron_based import ( |
|
Decoder as TacotronDecoder, |
|
Encoder as TacotronEncoder, |
|
Tacotron, |
|
) |
|
|
|
from poetry_diacritizer.options import AttentionType, LossType, OptimizerType |
|
from poetry_diacritizer.util.text_encoders import ( |
|
ArabicEncoderWithStartSymbol, |
|
BasicArabicEncoder, |
|
TextEncoder, |
|
) |
|
|
|
|
|
class ConfigManager: |
|
"""Co/home/almodhfer/Projects/daicritization/temp_results/CA_MSA/cbhg-new/model-10.ptnfig Manager""" |
|
|
|
def __init__(self, config_path: str, model_kind: str): |
|
available_models = ["baseline", "cbhg", "seq2seq", "tacotron_based", "gpt"] |
|
if model_kind not in available_models: |
|
raise TypeError(f"model_kind must be in {available_models}") |
|
self.config_path = Path(config_path) |
|
self.model_kind = model_kind |
|
self.yaml = ruamel.yaml.YAML() |
|
self.config: Dict[str, Any] = self._load_config() |
|
self.git_hash = self._get_git_hash() |
|
self.session_name = ".".join( |
|
[ |
|
self.config["data_type"], |
|
self.config["session_name"], |
|
f"{model_kind}", |
|
] |
|
) |
|
|
|
self.data_dir = Path( |
|
os.path.join(self.config["data_directory"], self.config["data_type"]) |
|
) |
|
self.base_dir = Path( |
|
os.path.join(self.config["log_directory"], self.session_name) |
|
) |
|
self.log_dir = Path(os.path.join(self.base_dir, "logs")) |
|
self.prediction_dir = Path(os.path.join(self.base_dir, "predictions")) |
|
self.plot_dir = Path(os.path.join(self.base_dir, "plots")) |
|
self.models_dir = Path(os.path.join(self.base_dir, "models")) |
|
if "sp_model_path" in self.config: |
|
self.sp_model_path = self.config["sp_model_path"] |
|
else: |
|
self.sp_model_path = None |
|
self.text_encoder: TextEncoder = self.get_text_encoder() |
|
self.config["len_input_symbols"] = len(self.text_encoder.input_symbols) |
|
self.config["len_target_symbols"] = len(self.text_encoder.target_symbols) |
|
if self.model_kind in ["seq2seq", "tacotron_based"]: |
|
self.config["attention_type"] = AttentionType[self.config["attention_type"]] |
|
self.config["optimizer"] = OptimizerType[self.config["optimizer_type"]] |
|
|
|
def _load_config(self): |
|
with open(self.config_path, "rb") as model_yaml: |
|
_config = self.yaml.load(model_yaml) |
|
return _config |
|
|
|
@staticmethod |
|
def _get_git_hash(): |
|
try: |
|
return ( |
|
subprocess.check_output(["git", "describe", "--always"]) |
|
.strip() |
|
.decode() |
|
) |
|
except Exception as e: |
|
print(f"WARNING: could not retrieve git hash. {e}") |
|
|
|
def _check_hash(self): |
|
try: |
|
git_hash = ( |
|
subprocess.check_output(["git", "describe", "--always"]) |
|
.strip() |
|
.decode() |
|
) |
|
if self.config["git_hash"] != git_hash: |
|
print( |
|
f"""WARNING: git hash mismatch. Current: {git_hash}. |
|
Config hash: {self.config['git_hash']}""" |
|
) |
|
except Exception as e: |
|
print(f"WARNING: could not check git hash. {e}") |
|
|
|
@staticmethod |
|
def _print_dict_values(values, key_name, level=0, tab_size=2): |
|
tab = level * tab_size * " " |
|
print(tab + "-", key_name, ":", values) |
|
|
|
def _print_dictionary(self, dictionary, recursion_level=0): |
|
for key in dictionary.keys(): |
|
if isinstance(key, dict): |
|
recursion_level += 1 |
|
self._print_dictionary(dictionary[key], recursion_level) |
|
else: |
|
self._print_dict_values( |
|
dictionary[key], key_name=key, level=recursion_level |
|
) |
|
|
|
def print_config(self): |
|
print("\nCONFIGURATION", self.session_name) |
|
self._print_dictionary(self.config) |
|
|
|
def update_config(self): |
|
self.config["git_hash"] = self._get_git_hash() |
|
|
|
def dump_config(self): |
|
self.update_config() |
|
_config = {} |
|
for key, val in self.config.items(): |
|
if isinstance(val, Enum): |
|
_config[key] = val.name |
|
else: |
|
_config[key] = val |
|
with open(self.base_dir / "config.yml", "w") as model_yaml: |
|
self.yaml.dump(_config, model_yaml) |
|
|
|
def create_remove_dirs( |
|
self, |
|
clear_dir: bool = False, |
|
clear_logs: bool = False, |
|
clear_weights: bool = False, |
|
clear_all: bool = False, |
|
): |
|
self.base_dir.mkdir(exist_ok=True, parents=True) |
|
self.plot_dir.mkdir(exist_ok=True) |
|
self.prediction_dir.mkdir(exist_ok=True) |
|
if clear_dir: |
|
delete = input(f"Delete {self.log_dir} AND {self.models_dir}? (y/[n])") |
|
if delete == "y": |
|
shutil.rmtree(self.log_dir, ignore_errors=True) |
|
shutil.rmtree(self.models_dir, ignore_errors=True) |
|
if clear_logs: |
|
delete = input(f"Delete {self.log_dir}? (y/[n])") |
|
if delete == "y": |
|
shutil.rmtree(self.log_dir, ignore_errors=True) |
|
if clear_weights: |
|
delete = input(f"Delete {self.models_dir}? (y/[n])") |
|
if delete == "y": |
|
shutil.rmtree(self.models_dir, ignore_errors=True) |
|
self.log_dir.mkdir(exist_ok=True) |
|
self.models_dir.mkdir(exist_ok=True) |
|
|
|
def get_last_model_path(self): |
|
""" |
|
Given a checkpoint, get the last save model name |
|
Args: |
|
checkpoint (str): the path where models are saved |
|
""" |
|
models = os.listdir(self.models_dir) |
|
models = [model for model in models if model[-3:] == ".pt"] |
|
if len(models) == 0: |
|
return None |
|
_max = max(int(m.split(".")[0].split("-")[0]) for m in models) |
|
model_name = f"{_max}-snapshot.pt" |
|
last_model_path = os.path.join(self.models_dir, model_name) |
|
|
|
return last_model_path |
|
|
|
def load_model(self, model_path: str = None): |
|
""" |
|
loading a model from path |
|
Args: |
|
checkpoint (str): the path to the model |
|
name (str): the name of the model, which is in the path |
|
model (Tacotron): the model to load its save state |
|
optimizer: the optimizer to load its saved state |
|
""" |
|
|
|
model = self.get_model() |
|
|
|
with open(self.base_dir / f"{self.model_kind}_network.txt", "w") as file: |
|
file.write(str(model)) |
|
|
|
if model_path is None: |
|
last_model_path = self.get_last_model_path() |
|
if last_model_path is None: |
|
return model, 1 |
|
else: |
|
last_model_path = model_path |
|
|
|
saved_model = torch.load(last_model_path) |
|
out = model.load_state_dict(saved_model["model_state_dict"]) |
|
print(out) |
|
global_step = saved_model["global_step"] + 1 |
|
return model, global_step |
|
|
|
def get_model(self, ignore_hash=False): |
|
if not ignore_hash: |
|
self._check_hash() |
|
if self.model_kind == "cbhg": |
|
return self.get_cbhg() |
|
|
|
elif self.model_kind == "seq2seq": |
|
return self.get_seq2seq() |
|
|
|
elif self.model_kind == "tacotron_based": |
|
return self.get_tacotron_based() |
|
|
|
elif self.model_kind == "baseline": |
|
return self.get_baseline() |
|
|
|
elif self.model_kind == "gpt": |
|
return self.get_gpt() |
|
|
|
def get_gpt(self): |
|
model = GPTModel( |
|
self.config["base_model_path"], |
|
freeze=self.config["freeze"], |
|
n_layer=self.config["n_layer"], |
|
use_lstm=self.config["use_lstm"], |
|
) |
|
return model |
|
|
|
def get_baseline(self): |
|
model = BaseLineModel( |
|
embedding_dim=self.config["embedding_dim"], |
|
inp_vocab_size=self.config["len_input_symbols"], |
|
targ_vocab_size=self.config["len_target_symbols"], |
|
layers_units=self.config["layers_units"], |
|
use_batch_norm=self.config["use_batch_norm"], |
|
) |
|
|
|
return model |
|
|
|
def get_cbhg(self): |
|
model = CBHGModel( |
|
embedding_dim=self.config["embedding_dim"], |
|
inp_vocab_size=self.config["len_input_symbols"], |
|
targ_vocab_size=self.config["len_target_symbols"], |
|
use_prenet=self.config["use_prenet"], |
|
prenet_sizes=self.config["prenet_sizes"], |
|
cbhg_gru_units=self.config["cbhg_gru_units"], |
|
cbhg_filters=self.config["cbhg_filters"], |
|
cbhg_projections=self.config["cbhg_projections"], |
|
post_cbhg_layers_units=self.config["post_cbhg_layers_units"], |
|
post_cbhg_use_batch_norm=self.config["post_cbhg_use_batch_norm"], |
|
) |
|
|
|
return model |
|
|
|
def get_seq2seq(self): |
|
encoder = Seq2SeqEncoder( |
|
embedding_dim=self.config["encoder_embedding_dim"], |
|
inp_vocab_size=self.config["len_input_symbols"], |
|
layers_units=self.config["encoder_units"], |
|
use_batch_norm=self.config["use_batch_norm"], |
|
) |
|
|
|
decoder = TacotronDecoder( |
|
self.config["len_target_symbols"], |
|
start_symbol_id=self.text_encoder.start_symbol_id, |
|
embedding_dim=self.config["decoder_embedding_dim"], |
|
encoder_dim=self.config["encoder_dim"], |
|
decoder_units=self.config["decoder_units"], |
|
decoder_layers=self.config["decoder_layers"], |
|
attention_type=self.config["attention_type"], |
|
attention_units=self.config["attention_units"], |
|
is_attention_accumulative=self.config["is_attention_accumulative"], |
|
use_prenet=self.config["use_decoder_prenet"], |
|
prenet_depth=self.config["decoder_prenet_depth"], |
|
teacher_forcing_probability=self.config["teacher_forcing_probability"], |
|
) |
|
|
|
model = Tacotron(encoder=encoder, decoder=decoder) |
|
|
|
return model |
|
|
|
def get_tacotron_based(self): |
|
encoder = TacotronEncoder( |
|
embedding_dim=self.config["encoder_embedding_dim"], |
|
inp_vocab_size=self.config["len_input_symbols"], |
|
prenet_sizes=self.config["prenet_sizes"], |
|
use_prenet=self.config["use_encoder_prenet"], |
|
cbhg_gru_units=self.config["cbhg_gru_units"], |
|
cbhg_filters=self.config["cbhg_filters"], |
|
cbhg_projections=self.config["cbhg_projections"], |
|
) |
|
|
|
decoder = TacotronDecoder( |
|
self.config["len_target_symbols"], |
|
start_symbol_id=self.text_encoder.start_symbol_id, |
|
embedding_dim=self.config["decoder_embedding_dim"], |
|
encoder_dim=self.config["encoder_dim"], |
|
decoder_units=self.config["decoder_units"], |
|
decoder_layers=self.config["decoder_layers"], |
|
attention_type=self.config["attention_type"], |
|
attention_units=self.config["attention_units"], |
|
is_attention_accumulative=self.config["is_attention_accumulative"], |
|
use_prenet=self.config["use_decoder_prenet"], |
|
prenet_depth=self.config["decoder_prenet_depth"], |
|
teacher_forcing_probability=self.config["teacher_forcing_probability"], |
|
) |
|
|
|
model = Tacotron(encoder=encoder, decoder=decoder) |
|
|
|
return model |
|
|
|
def get_text_encoder(self): |
|
"""Getting the class of TextEncoder from config""" |
|
if self.config["text_cleaner"] not in [ |
|
"basic_cleaners", |
|
"valid_arabic_cleaners", |
|
None, |
|
]: |
|
raise Exception(f"cleaner is not known {self.config['text_cleaner']}") |
|
|
|
if self.config["text_encoder"] == "BasicArabicEncoder": |
|
text_encoder = BasicArabicEncoder( |
|
cleaner_fn=self.config["text_cleaner"], sp_model_path=self.sp_model_path |
|
) |
|
elif self.config["text_encoder"] == "ArabicEncoderWithStartSymbol": |
|
text_encoder = ArabicEncoderWithStartSymbol( |
|
cleaner_fn=self.config["text_cleaner"], sp_model_path=self.sp_model_path |
|
) |
|
else: |
|
raise Exception( |
|
f"the text encoder is not found {self.config['text_encoder']}" |
|
) |
|
|
|
return text_encoder |
|
|
|
def get_loss_type(self): |
|
try: |
|
loss_type = LossType[self.config["loss_type"]] |
|
except: |
|
raise Exception(f"The loss type is not correct {self.config['loss_type']}") |
|
return loss_type |
|
|
|
|
|
if __name__ == "__main__": |
|
config_path = "config/tacotron-base-config.yml" |
|
model_kind = "tacotron" |
|
config = ConfigManager(config_path=config_path, model_kind=model_kind) |
|
|
