def create_model(embeddings, d_model, d_ff, num_heads, num_layers, rpr_k, d_k, activation, checkpoint_name, device): if len(rpr_k) == 0 or rpr_k[0] < 1: rpr_k = [None] else: rpr_k = listify(rpr_k) logger.info("Creating tied encoder decoder model") hps = { "dsz": d_model, "hsz": d_model, "d_ff": d_ff, "dropout": 0.0, "num_heads": num_heads, "layers": num_layers, "encoder_type": "transformer", "decoder_type": "transformer", "src_lengths_key": "x_lengths", "d_k": d_k, "activation": activation, "rpr_k": rpr_k } model = TiedEmbeddingsSeq2SeqModel({'x': embeddings}, None, **hps) if checkpoint_name.endswith('npz'): load_transformer_seq2seq_npz(model, checkpoint_name) else: model.load_state_dict( torch.load(checkpoint_name, map_location=torch.device(device))) print(model) return model
def create(cls, src_embeddings, tgt_embedding, **kwargs): model = cls(src_embeddings, tgt_embedding, **kwargs) checkpoint_name = kwargs.get('checkpoint') if checkpoint_name is not None: if checkpoint_name.endswith('npz'): load_transformer_seq2seq_npz(model, checkpoint_name) else: model.load_state_dict(torch.load(checkpoint_name)) logger.info(model) return model
def reload_from_checkpoint(model_type, restart_from, restart_tick_type, model, steps_per_epoch): if os.path.isdir(restart_from): restart_from, _ = find_latest_checkpoint(restart_from) print(f'Latest checkpoint: {restart_from}') vec = restart_from.split("-") try: step_num = int(vec[-1].split(".")[0]) except: step_num = 0 start_epoch = 0 if restart_tick_type: tick_type = restart_tick_type else: tick_type = vec[-2] if restart_from.endswith('.npz'): # If its a seq2seq load either from a seq2seq or from a TLM encoder if model_type == 'encoder-decoder': try: load_transformer_seq2seq_npz(model, restart_from) except: print( 'Model file not recognized as seq2seq model, attempting to load as LM for encoder, reset step' ) load_seq2seq_enc_from_tlm_npz(model, restart_from) step_num = 0 tick_type = 'ignore' else: try: load_transformer_de_npz(model, restart_from) # If its a dual-encoder, assuming we have model.transformer and model.embeddings, we can load directly # from a Transformer Language Model except: print( 'Model file not recognized as a dual encoder model, attempting to load as LM for encoder, reset step' ) load_tlm_npz(model, restart_from) step_num = 0 tick_type = 'ignore' else: model.load_state_dict(torch.load(restart_from)) if tick_type == 'epoch': start_epoch = step_num step_num = start_epoch * steps_per_epoch elif tick_type == 'step': start_epoch = step_num // steps_per_epoch else: logger.warning( f"The previous tick was {step_num} but command-line specifies to ignore, setting to 0" ) step_num = 0 start_epoch = 0 return step_num, start_epoch