Not-Grim-Refer/CodeBERT_CodeReviewer

1import gradio as gr  2import requests3from torch import nn4from torch.nn import CrossEntropyLoss5from transformers import AutoTokenizer, T5ForConditionalGeneration, AutoModelForSeq2SeqLM, T5Config6import torch7 8MAX_SOURCE_LENGTH = 5129 10 11class ReviewerModel(T5ForConditionalGeneration):12 13    def __init__(self, config):14        super().__init__(config)15        self.cls_head = nn.Linear(config.d_model, 2, bias=True) 16        # Fixed typo: config not self.config17        self.init()18 19    def init(self):20        nn.init.xavier_uniform_(self.lm_head.weight)21        factor = self.config.initializer_factor22        self.cls_head.weight.data.normal_(mean=0.0, std=factor * (self.config.d_model ** -0.5))23        # Fixed exponentiation operator 24        self.cls_head.bias.data.zero_()25 26    def forward(27            self, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, labels=None):28        29        # Simplified method signature to only include necessary arguments30        31        if labels is not None:32            # Added validation check for seq2seq case33            34            outputs = super().forward(35                input_ids=input_ids, 36                attention_mask=attention_mask,37                decoder_input_ids=decoder_input_ids,38                decoder_attention_mask=decoder_attention_mask,39                labels=labels40            )41            # Call super forward method with correct arguments42            43            return outputs44        45        # Removed unnecessary conditional logic46        # Return super() forward directly for generation case47        48