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: Copy pasted from distillbert (transformers/examples/distillation/)
docstr_tok: ["Copy","pasted","from","distillbert","(","transformers\/examples\/distillation\/",")"]
code: def calc_ce_loss(self, mask, s_logits, t_logits):
      """Copy pasted from distillbert (transformers/examples/distillation/)"""
      # mask has False at padding_idx
      sel_mask = mask[:, :, None].expand_as(s_logits)
      vocab_size = s_logits.size(-1)
      s_logits_slct = torch.masked_select(
          s_logits, sel_mask
      )  # (bs * seq_length * voc_size) modulo the 1s in mask
      t_logits_slct = torch.masked_select(
          t_logits, sel_mask
      )  # (bs * seq_length * voc_size) modulo the 1s in mask
      s_logits_slct = s_logits_slct.view(
          -1, vocab_size
      )  # (bs * seq_length, voc_size) modulo the 1s in mask
      t_logits_slct = t_logits_slct.view(
          -1, vocab_size
      )  # (bs * seq_length, voc_size) modulo the 1s in mask
      assert t_logits_slct.size() == s_logits_slct.size()
      loss_ce = (
          self.ce_loss_fct(
              nn.functional.log_softmax(
                  s_logits_slct / self.temperature, dim=-1
              ),
              nn.functional.softmax(
                  t_logits_slct / self.temperature, dim=-1
              ),
          )
          * (self.temperature) ** 2
      )
      return loss_ce
  
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 2 for calling the SummarizationDistiller obj's underlying member method code with example usage: obj.calc_ce_loss(mask, s_logits, t_logits) and returns: loss_ce