Evaluating code quality before comment generation.

app.py

@@ -1,12 +1,151 @@
 import gradio as gr
 import requests
-from transformers import AutoTokenizer, T5ForConditionalGeneration, AutoModelForSeq2SeqLM
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers import AutoTokenizer, T5ForConditionalGeneration, AutoModelForSeq2SeqLM, T5Config
 import torch
 
-
 MAX_SOURCE_LENGTH = 512
 
 
+class ReviewerModel(T5ForConditionalGeneration):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.cls_head = nn.Linear(self.config.d_model, 2, bias=True)
+        self.init()
+
+    def init(self):
+        nn.init.xavier_uniform_(self.lm_head.weight)
+        factor = self.config.initializer_factor
+        self.cls_head.weight.data.normal_(mean=0.0, \
+                                          std=factor * ((self.config.d_model) ** -0.5))
+        self.cls_head.bias.data.zero_()
+
+    def forward(
+            self, *argv, **kwargs
+    ):
+        r"""
+        Doc from Huggingface transformers:
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
+            Labels for computing the sequence classification/regression loss. Indices should be in :obj:`[-100, 0, ...,
+            config.vocab_size - 1]`. All labels set to ``-100`` are ignored (masked), the loss is only computed for
+            labels in ``[0, ..., config.vocab_size]``
+        Returns:
+        Examples::
+            >>> from transformers import T5Tokenizer, T5ForConditionalGeneration
+            >>> tokenizer = T5Tokenizer.from_pretrained('t5-small')
+            >>> model = T5ForConditionalGeneration.from_pretrained('t5-small')
+            >>> # training
+            >>> input_ids = tokenizer('The <extra_id_0> walks in <extra_id_1> park', return_tensors='pt').input_ids
+            >>> labels = tokenizer('<extra_id_0> cute dog <extra_id_1> the <extra_id_2>', return_tensors='pt').input_ids
+            >>> outputs = model(input_ids=input_ids, labels=labels)
+            >>> loss = outputs.loss
+            >>> logits = outputs.logits
+            >>> # inference
+            >>> input_ids = tokenizer("summarize: studies have shown that owning a dog is good for you", return_tensors="pt").input_ids  # Batch size 1
+            >>> outputs = model.generate(input_ids)
+            >>> print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+            >>> # studies have shown that owning a dog is good for you.
+        """
+        if "cls" in kwargs:
+            assert (
+                    "input_ids" in kwargs and \
+                    "labels" in kwargs and \
+                    "attention_mask" in kwargs
+            )
+            return self.cls(
+                input_ids=kwargs["input_ids"],
+                labels=kwargs["labels"],
+                attention_mask=kwargs["attention_mask"],
+            )
+        if "input_labels" in kwargs:
+            assert (
+                    "input_ids" in kwargs and \
+                    "input_labels" in kwargs and \
+                    "decoder_input_ids" in kwargs and \
+                    "attention_mask" in kwargs and \
+                    "decoder_attention_mask" in kwargs
+            ), "Please give these arg keys."
+            input_ids = kwargs["input_ids"]
+            input_labels = kwargs["input_labels"]
+            decoder_input_ids = kwargs["decoder_input_ids"]
+            attention_mask = kwargs["attention_mask"]
+            decoder_attention_mask = kwargs["decoder_attention_mask"]
+            if "encoder_loss" not in kwargs:
+                encoder_loss = True
+            else:
+                encoder_loss = kwargs["encoder_loss"]
+            return self.review_forward(input_ids, input_labels, decoder_input_ids, attention_mask,
+                                       decoder_attention_mask, encoder_loss)
+        return super().forward(*argv, **kwargs)
+
+    def cls(
+            self,
+            input_ids,
+            labels,
+            attention_mask,
+    ):
+        encoder_outputs = self.encoder( \
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_attentions=False,
+            return_dict=False
+        )
+        hidden_states = encoder_outputs[0]
+        first_hidden = hidden_states[:, 0, :]
+        first_hidden = nn.Dropout(0.3)(first_hidden)
+        logits = self.cls_head(first_hidden)
+        loss_fct = CrossEntropyLoss()
+        if labels != None:
+            loss = loss_fct(logits, labels)
+            return loss
+        return logits
+
+    def review_forward(
+            self,
+            input_ids,
+            input_labels,
+            decoder_input_ids,
+            attention_mask,
+            decoder_attention_mask,
+            encoder_loss=True
+    ):
+        encoder_outputs = self.encoder( \
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_attentions=False,
+            return_dict=False
+        )
+        hidden_states = encoder_outputs[0]
+        decoder_inputs = self._shift_right(decoder_input_ids)
+        # Decode
+        decoder_outputs = self.decoder(
+            input_ids=decoder_inputs,
+            attention_mask=decoder_attention_mask,
+            encoder_hidden_states=hidden_states,
+            encoder_attention_mask=attention_mask,
+            output_attentions=False,
+            return_dict=False
+        )
+        sequence_output = decoder_outputs[0]
+        if self.config.tie_word_embeddings:  # this is True default
+            sequence_output = sequence_output * (self.model_dim ** -0.5)
+        if encoder_loss:
+            # print(self.encoder.get_input_embeddings().weight.shape)
+            cls_logits = nn.functional.linear(hidden_states, self.encoder.get_input_embeddings().weight)
+            # cls_logits = self.cls_head(hidden_states)
+        lm_logits = self.lm_head(sequence_output)
+        if decoder_input_ids is not None:
+            lm_loss_fct = CrossEntropyLoss(ignore_index=0)  # Warning: PAD_ID should be 0
+            loss = lm_loss_fct(lm_logits.view(-1, lm_logits.size(-1)), decoder_input_ids.view(-1))
+            if encoder_loss and input_labels is not None:
+                cls_loss_fct = CrossEntropyLoss(ignore_index=-100)
+                loss += cls_loss_fct(cls_logits.view(-1, cls_logits.size(-1)), input_labels.view(-1))
+            return loss
+        return cls_logits, lm_logits
+
+
 def prepare_models():
     tokenizer = AutoTokenizer.from_pretrained("microsoft/codereviewer")
 
@@ -24,7 +163,8 @@ def prepare_models():
     tokenizer.start_id = tokenizer.get_vocab()["<start>"]
     tokenizer.end_id = tokenizer.get_vocab()["<end>"]
 
-    model = AutoModelForSeq2SeqLM.from_pretrained("microsoft/codereviewer")
+    config = T5Config.from_pretrained("microsoft/codereviewer")
+    model = ReviewerModel.from_pretrained("microsoft/codereviewer", config=config)
 
     model.eval()
     return tokenizer, model
@@ -104,6 +244,19 @@ def review_commit(user="p4vv37", repository="ueflow", commit="610a8c7b02b946bc9e
         for diff in fd.diffs:
             inputs = torch.tensor([encode_diff(tokenizer, diff, msg, source)], dtype=torch.long).to("cpu")
             inputs_mask = inputs.ne(tokenizer.pad_id)
+            logits = model(
+                input_ids=inputs,
+                cls=True,
+                attention_mask=inputs_mask,
+                labels=None,
+                use_cache=True,
+                num_beams=5,
+                early_stopping=True,
+                max_length=100
+            )
+            needs_review = torch.argmax(logits, dim=-1).cpu().numpy()[0]
+            if not needs_review:
+                continue
             preds = model.generate(inputs,
                                    attention_mask=inputs_mask,
                                    use_cache=True,