shekkari21/codereviewer

1import os2import torch.nn as nn3import torch4import torch.nn.functional as F5from torch.nn import CrossEntropyLoss, BCEWithLogitsLoss6import numpy as np7from utils import MyTokenizer8from transformers import (9    RobertaConfig,10    RobertaModel,11    RobertaTokenizer,12    BartConfig,13    BartForConditionalGeneration,14    BartTokenizer,15    T5Config,16    T5ForConditionalGeneration,17    T5Tokenizer,18)19import logging20 21logger = logging.getLogger(__name__)22 23 24class ReviewerModel(T5ForConditionalGeneration):25 26    def __init__(self, config):27        super().__init__(config)28        self.cls_head = nn.Linear(self.config.d_model, 2, bias=True)29        self.init()30 31    def init(self):32        nn.init.xavier_uniform_(self.lm_head.weight)33        factor = self.config.initializer_factor34        self.cls_head.weight.data.normal_(mean=0.0, \35            std=factor * ((self.config.d_model) ** -0.5))36        self.cls_head.bias.data.zero_()37 38    def forward(39        self, *argv, **kwargs40    ):41        r"""42        Doc from Huggingface transformers:43        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):44            Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[-100, 0, ...,45            config.vocab_size - 1]`. All labels set to ``-100`` are ignored (masked), the loss is only computed for46            labels in ``[0, ..., config.vocab_size]``47        Returns:48        Examples::49            >>> from transformers import T5Tokenizer, T5ForConditionalGeneration50            >>> tokenizer = T5Tokenizer.from_pretrained('t5-small')51            >>> model = T5ForConditionalGeneration.from_pretrained('t5-small')52            >>> # training53            >>> input_ids = tokenizer('The <extra_id_0> walks in <extra_id_1> park', return_tensors='pt').input_ids54            >>> labels = tokenizer('<extra_id_0> cute dog <extra_id_1> the <extra_id_2>', return_tensors='pt').input_ids55            >>> outputs = model(input_ids=input_ids, labels=labels)56            >>> loss = outputs.loss57            >>> logits = outputs.logits58            >>> # inference59            >>> input_ids = tokenizer("summarize: studies have shown that owning a dog is good for you", return_tensors="pt").input_ids  # Batch size 160            >>> outputs = model.generate(input_ids)61            >>> print(tokenizer.decode(outputs[0], skip_special_tokens=True))62            >>> # studies have shown that owning a dog is good for you.63        """64        if "cls" in kwargs:65            assert (66                "input_ids" in kwargs and \67                "labels" in kwargs and \68                "attention_mask" in kwargs69            )70            return self.cls(71                input_ids=kwargs["input_ids"],72                labels=kwargs["labels"],73                attention_mask=kwargs["attention_mask"],74            )75        if "input_labels" in kwargs:76            assert (77                "input_ids" in kwargs and \78                "input_labels" in kwargs and \79                "decoder_input_ids" in kwargs and \80                "attention_mask" in kwargs and \81                "decoder_attention_mask" in kwargs82            ), "Please give these arg keys."83            input_ids = kwargs["input_ids"]84            input_labels = kwargs["input_labels"]85            decoder_input_ids = kwargs["decoder_input_ids"]86            attention_mask = kwargs["attention_mask"]87            decoder_attention_mask = kwargs["decoder_attention_mask"]88            if "encoder_loss" not in kwargs:89                encoder_loss = True90            else:91                encoder_loss = kwargs["encoder_loss"]92            return self.review_forward(input_ids, input_labels, decoder_input_ids, attention_mask, decoder_attention_mask, encoder_loss)93        return super().forward(*argv, **kwargs)94 95    def cls(96        self,97        input_ids,98        labels,99        attention_mask,100    ):101        encoder_outputs = self.encoder( \102            input_ids=input_ids,103            attention_mask=attention_mask,104            output_attentions=False,105            return_dict=False106        )107        hidden_states = encoder_outputs[0]108        first_hidden = hidden_states[:, 0, :]109        first_hidden = nn.Dropout(0.3)(first_hidden)110        logits = self.cls_head(first_hidden)111        loss_fct = CrossEntropyLoss()112        if labels != None:113            loss = loss_fct(logits, labels)114            return loss115        return logits116 117    def review_forward(118        self,119        input_ids,120        input_labels,121        decoder_input_ids,122        attention_mask,123        decoder_attention_mask,124        encoder_loss=True125    ):126        encoder_outputs = self.encoder( \127            input_ids=input_ids,128            attention_mask=attention_mask,129            output_attentions=False,130            return_dict=False131        )132        hidden_states = encoder_outputs[0]133        decoder_inputs = self._shift_right(decoder_input_ids)134        # Decode135        decoder_outputs = self.decoder(136            input_ids=decoder_inputs,137            attention_mask=decoder_attention_mask,138            encoder_hidden_states=hidden_states,139            encoder_attention_mask=attention_mask,140            output_attentions=False,141            return_dict=False142        )143        sequence_output = decoder_outputs[0]144        if self.config.tie_word_embeddings: # this is True default145            sequence_output = sequence_output * (self.model_dim ** -0.5)146        if encoder_loss:147            # print(self.encoder.get_input_embeddings().weight.shape)148            cls_logits = nn.functional.linear(hidden_states, self.encoder.get_input_embeddings().weight)149            # cls_logits = self.cls_head(hidden_states)150        lm_logits = self.lm_head(sequence_output)151        if decoder_input_ids is not None:152            lm_loss_fct = CrossEntropyLoss(ignore_index=0)      # Warning: PAD_ID should be 0153            loss = lm_loss_fct(lm_logits.view(-1, lm_logits.size(-1)), decoder_input_ids.view(-1))154            if encoder_loss and input_labels is not None:155                cls_loss_fct = CrossEntropyLoss(ignore_index=-100)156                loss += cls_loss_fct(cls_logits.view(-1, cls_logits.size(-1)), input_labels.view(-1))157            return loss158        return cls_logits, lm_logits159 160def get_model_size(model):161    model_parameters = filter(lambda p: p.requires_grad, model.parameters())162    model_size = sum([np.prod(p.size()) for p in model_parameters])163    return "{}M".format(round(model_size / 1e6))164 165 166def build_or_load_gen_model(args):167    config_class, model_class, tokenizer_class = T5Config, ReviewerModel, RobertaTokenizer168    169    config = config_class.from_pretrained(args.model_name_or_path)170    tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)171    model = model_class.from_pretrained(args.model_name_or_path, config=config)172 173    tokenizer.special_dict = {174        f"<e{i}>" : tokenizer.get_vocab()[f"<e{i}>"] for i in range(99, -1, -1)175    }176 177    tokenizer.mask_id = tokenizer.get_vocab()["<mask>"]178    tokenizer.bos_id = tokenizer.get_vocab()["<s>"]179    tokenizer.pad_id = tokenizer.get_vocab()["<pad>"]180    tokenizer.eos_id = tokenizer.get_vocab()["</s>"]181    tokenizer.msg_id = tokenizer.get_vocab()["<msg>"]182    tokenizer.keep_id = tokenizer.get_vocab()["<keep>"]183    tokenizer.add_id = tokenizer.get_vocab()["<add>"]184    tokenizer.del_id = tokenizer.get_vocab()["<del>"]185    tokenizer.start_id = tokenizer.get_vocab()["<start>"]186    tokenizer.end_id = tokenizer.get_vocab()["<end>"]187 188    logger.info(189        "Finish loading model [%s] from %s",190        get_model_size(model),191        args.model_name_or_path,192    )193 194    if args.load_model_path is not None:195        model_path = os.path.join(args.load_model_path, "pytorch_model.bin")196        logger.info("Reload model from {}".format(model_path))197        try:198            model.load_state_dict(torch.load(model_path, map_location="cpu"))199        except RuntimeError:200            saved = model.cls_head201            model.cls_head = None202            model.load_state_dict(torch.load(model_path, map_location="cpu"))203            model.cls_head = saved204        model.to(args.local_rank)205 206    return config, model, tokenizer207 208 209