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: MSE(student_hid, teacher_hid[matches]). Called "Intermediate supervision" in paper. Inspired by TinyBERT.
docstr_tok: ["MSE","(","student_hid",",","teacher_hid","[","matches","]",")",".","Called","``","Intermediate","supervision","''","in","paper",".","Inspired","by","TinyBERT","."]
code: def calc_hidden_loss(
      attention_mask,
      hidden_states,
      hidden_states_T,
      matches,
      normalize_hidden,
  ):
      """MSE(student_hid, teacher_hid[matches]). Called "Intermediate supervision" in paper. Inspired by TinyBERT."""
      msg = "expected list or tuple for hidden_states, got tensor of shape: "
      assert not isinstance(
          hidden_states, torch.Tensor
      ), f"{msg}{hidden_states.shape}"
      assert not isinstance(
          hidden_states_T, torch.Tensor
      ), f"{msg}{hidden_states_T.shape}"
      mask = attention_mask.to(hidden_states[0])
      valid_count = mask.sum() * hidden_states[0].size(-1)
      student_states = torch.stack(
          [hidden_states[i] for i in range(len(matches))]
      )
      teacher_states = torch.stack(
          [hidden_states_T[j] for j in matches]
      )
      assert (
          student_states.shape == teacher_states.shape
      ), f"{student_states.shape} != {teacher_states.shape}"
      if normalize_hidden:
          student_states = nn.functional.layer_norm(
              student_states, student_states.shape[1:]
          )
          teacher_states = nn.functional.layer_norm(
              teacher_states, teacher_states.shape[1:]
          )
      mse = nn.functional.mse_loss(
          student_states, teacher_states, reduction="none"
      )
      masked_mse = (
          mse * mask.unsqueeze(0).unsqueeze(-1)
      ).sum() / valid_count
      return masked_mse
  
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 5 for calling the SummarizationDistiller obj's underlying member method code with example usage: obj.calc_hidden_loss(attention_mask, hidden_states, hidden_states_T, matches, normalize_hidden) and returns: masked_mse