class_docstr: Supports T5, Bart, Pegasus and other models that inherit from Bart.
class_docstr_tok: ["Supports","T5",",","Bart",",","Pegasus","and","other","models","that","inherit","from","Bart","."]
docstr: Compute the loss for a batch
docstr_tok: ["Compute","the","loss","for","a","batch"]
code: def _step(self, batch: dict) -> tuple:
      """Compute the loss for a batch"""
      pad_token_id = self.tokenizer.pad_token_id
      input_ids, src_mask, labels = (
          batch["input_ids"],
          batch["attention_mask"],
          batch["labels"],
      )
      if isinstance(self.model, T5ForConditionalGeneration):
          decoder_input_ids = self.model._shift_right(labels)
      else:
          decoder_input_ids = shift_tokens_right(labels, pad_token_id)
  
      # noinspection PyCallingNonCallable
      student_outputs = self(
          input_ids,
          attention_mask=src_mask,
          decoder_input_ids=decoder_input_ids,
          output_hidden_states=self.do_calc_hidden_loss,
          output_attentions=False,
          use_cache=False,
      )
      lm_logits = student_outputs["logits"]
  
      # Same cross entropy vs. label smoothing logic as finetune.py
      assert lm_logits.shape[-1] == self.model.config.vocab_size
      if self.hparams.label_smoothing == 0:
          # Same behavior as modeling_bart.py, besides ignoring pad_token_id
          loss_fct = nn.CrossEntropyLoss(ignore_index=pad_token_id)
          student_lm_loss = loss_fct(
              lm_logits.view(-1, lm_logits.shape[-1]), labels.view(-1)
          )
      else:
          lprobs = nn.functional.log_softmax(lm_logits, dim=-1)
          student_lm_loss, _ = label_smoothed_nll_loss(
              lprobs,
              labels,
              self.hparams.label_smoothing,
              ignore_index=pad_token_id,
          )
  
      def zero_tensor():
          return torch.tensor(0.0).type_as(student_lm_loss)
  
      teacher_enc_outputs = student_outputs[
          "encoder_last_hidden_state"
      ]  # use this unless self.different_base_models
      hid_loss_enc, hid_loss_dec = zero_tensor(), zero_tensor()
      if self.different_encoder:  # compute encoder hidden state loss
          all_teacher_encoder_outputs = self.teacher.get_encoder()(
              input_ids,
              attention_mask=src_mask,
              output_hidden_states=self.do_calc_hidden_loss,
          )
          if self.different_base_models:
              teacher_enc_outputs = all_teacher_encoder_outputs[
                  "last_hidden_state"
              ]
          elif self.do_calc_hidden_loss:
              hid_loss_enc = self.calc_hidden_loss(
                  src_mask,
                  student_outputs["encoder_hidden_states"],
                  all_teacher_encoder_outputs["hidden_states"],
                  self.e_matches,
                  normalize_hidden=self.hparams.normalize_hidden,
              )
  
      teacher_outputs = self.teacher(
          input_ids,
          attention_mask=src_mask,
          encoder_outputs=(teacher_enc_outputs,),
          decoder_input_ids=decoder_input_ids,
          output_hidden_states=self.do_calc_hidden_loss,
          use_cache=False,  # since we are not passing labels, never let this default to True
      )
      dec_mask = decoder_input_ids.ne(pad_token_id)
      loss_ce = self.calc_ce_loss(
          dec_mask, lm_logits, teacher_outputs["logits"]
      )
      if (
          self.do_calc_hidden_loss
      ):  # Intermediate supervision of decoder hidden states
          hid_loss_dec = self.calc_hidden_loss(
              dec_mask,
              student_outputs["decoder_hidden_states"],
              teacher_outputs["decoder_hidden_states"],
              self.d_matches,
              normalize_hidden=self.hparams.normalize_hidden,
          )
  
      blended_loss = (
          self.alpha_ce * loss_ce
          + self.alpha_mlm * student_lm_loss
          + self.hparams.alpha_hid * (hid_loss_enc + hid_loss_dec)
      )
      return (
          blended_loss,
          loss_ce,
          student_lm_loss,
          hid_loss_enc,
          hid_loss_dec,
      )
  
imports: import argparse
  import gc
  import os
  import sys
  from pathlib import Path
  from typing import List
  import pytorch_lightning
  import torch
  from finetune import SummarizationModule, TranslationModule
  from finetune import main
  from make_student import create_student_by_copying_alternating_layers, get_layers_to_supervise
  from torch import nn
  from transformers import AutoModelForSeq2SeqLM, MBartTokenizer, T5ForConditionalGeneration
  from transformers.models.bart.modeling_bart import shift_tokens_right
  from utils import calculate_bleu, check_output_dir, freeze_params, label_smoothed_nll_loss, use_task_specific_params
  from lightning_base import generic_train
label_desc: Use image node_id 4 for calling the SummarizationDistiller obj's underlying member method code with example usage: obj._step(batch) and returns: blended_loss,  loss_ce,  student_lm_loss,  hid_loss_enc,  hid_loss_dec, torch