| import argparse | |
| import json | |
| import sys | |
| from ScanDL2.scandl_module.original_scandl import sp_gaussian_diffusion as gd | |
| from ScanDL2.scandl_module.original_scandl.sp_gaussian_diffusion import ( | |
| SpacedDiffusion, | |
| space_timesteps, | |
| ) | |
| from ScanDL2.scandl_module.original_scandl.sp_transformer_model import TransformerNetModel | |
| sys.path.append("./") | |
| sys.path.append("../") | |
| def load_defaults_config(config_path: str): | |
| """ | |
| Load defaults for training args. | |
| """ | |
| with open(config_path, "r") as f: | |
| return json.load(f) | |
| def create_model_and_diffusion( | |
| hidden_t_dim, | |
| hidden_dim, | |
| vocab_size, | |
| config_name, | |
| use_plm_init, | |
| dropout, | |
| num_transformer_layers, | |
| num_transformer_heads, | |
| mask_padding, | |
| diffusion_steps, | |
| noise_schedule, | |
| learn_sigma, | |
| timestep_respacing, | |
| predict_xstart, | |
| rescale_timesteps, | |
| sigma_small, | |
| rescale_learned_sigmas, | |
| use_kl, | |
| one_noise_step, | |
| nll_in_loss, | |
| notes, | |
| **kwargs, | |
| ): | |
| model = TransformerNetModel( | |
| input_dims=hidden_dim, | |
| output_dims=(hidden_dim if not learn_sigma else hidden_dim * 2), | |
| hidden_t_dim=hidden_t_dim, | |
| num_transformer_layers=num_transformer_layers, | |
| num_transformer_heads=num_transformer_heads, | |
| one_noise_step=one_noise_step, | |
| mask_padding=mask_padding, | |
| dropout=dropout, | |
| config_name=config_name, | |
| vocab_size=vocab_size, | |
| init_pretrained=use_plm_init, | |
| ) | |
| betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps) | |
| if not timestep_respacing: | |
| timestep_respacing = [diffusion_steps] | |
| diffusion = SpacedDiffusion( | |
| use_timesteps=space_timesteps(diffusion_steps, timestep_respacing), | |
| betas=betas, | |
| rescale_timesteps=rescale_timesteps, | |
| predict_xstart=predict_xstart, | |
| learn_sigmas=learn_sigma, | |
| sigma_small=sigma_small, | |
| use_kl=use_kl, | |
| one_noise_step=one_noise_step, | |
| nll_in_loss=nll_in_loss, | |
| mask_padding=mask_padding, | |
| rescale_learned_sigmas=rescale_learned_sigmas, | |
| ) | |
| return model, diffusion | |
| def add_dict_to_argparser(parser, default_dict): | |
| for k, v in default_dict.items(): | |
| v_type = type(v) | |
| if v is None: | |
| v_type = str | |
| elif isinstance(v, bool): | |
| v_type = str2bool | |
| parser.add_argument(f"--{k}", default=v, type=v_type) | |
| def args_to_dict(args, keys): | |
| return {k: getattr(args, k) for k in keys} | |
| def str2bool(v): | |
| """ | |
| https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse | |
| """ | |
| if isinstance(v, bool): | |
| return v | |
| if v.lower() in ("yes", "true", "t", "y", "1"): | |
| return True | |
| elif v.lower() in ("no", "false", "f", "n", "0"): | |
| return False | |
| else: | |
| raise argparse.ArgumentTypeError("boolean value expected") | |