model.py (2018B)
1 import logging 2 import numpy as np 3 import torch 4 import torch.nn as nn 5 from transformers import AutoConfig, AutoModelForPreTraining 6 import ipdb 7 8 logger = logging.getLogger(__name__) 9 10 class GenerativeModel(nn.Module): 11 def __init__(self, config, tokenizer): 12 super().__init__() 13 self.tokenizer = tokenizer 14 logger.info(f'Loading pre-trained model {config.model_name}') 15 self.max_length = config.max_length 16 self.model_config = AutoConfig.from_pretrained(config.model_name, cache_dir=config.cache_dir) 17 self.model = AutoModelForPreTraining.from_pretrained(config.model_name, cache_dir=config.cache_dir, config=self.model_config) 18 self.model.config.no_repeat_ngram_size = 0 19 self.model.resize_token_embeddings(len(self.tokenizer)) 20 21 def forward(self, batch): 22 outputs = self.model(input_ids=batch.enc_idxs[:, :self.max_length], 23 attention_mask=batch.enc_attn[:, :self.max_length], 24 decoder_input_ids=batch.dec_idxs, 25 decoder_attention_mask=batch.dec_attn, 26 labels=batch.lbl_idxs, 27 return_dict=True) 28 29 loss = outputs['loss'] 30 31 return loss 32 33 def predict(self, batch, num_beams=4, max_length=50): 34 self.eval() 35 with torch.no_grad(): 36 outputs = self.model.generate(input_ids=batch.enc_idxs[:, :self.max_length], 37 attention_mask=batch.enc_attn[:, :self.max_length], 38 num_beams=num_beams, 39 max_length=max_length) 40 41 final_output = [] 42 for bid in range(len(batch.enc_idxs)): 43 output_sentence = self.tokenizer.decode(outputs[bid], skip_special_tokens=True, clean_up_tokenization_spaces=True) 44 final_output.append(output_sentence) 45 self.train() 46 47 return final_output 48