code: def __init__(
      self,
      params: dict,
      dataset: LmSeqsDataset,
      token_probs: torch.tensor,
      student: nn.Module,
      teacher: nn.Module,
  ):
      logger.info("Initializing Distiller")
      self.params = params
      self.dump_path = params.dump_path
      self.multi_gpu = params.multi_gpu
      self.fp16 = params.fp16
  
      self.student = student
      self.teacher = teacher
  
      self.student_config = student.config
      self.vocab_size = student.config.vocab_size
  
      if params.n_gpu <= 1:
          sampler = RandomSampler(dataset)
      else:
          sampler = DistributedSampler(dataset)
  
      if params.group_by_size:
          groups = create_lengths_groups(
              lengths=dataset.lengths, k=params.max_model_input_size
          )
          sampler = GroupedBatchSampler(
              sampler=sampler,
              group_ids=groups,
              batch_size=params.batch_size,
          )
      else:
          sampler = BatchSampler(
              sampler=sampler,
              batch_size=params.batch_size,
              drop_last=False,
          )
  
      self.dataloader = DataLoader(
          dataset=dataset,
          batch_sampler=sampler,
          collate_fn=dataset.batch_sequences,
      )
  
      self.temperature = params.temperature
      assert self.temperature > 0.0
  
      self.alpha_ce = params.alpha_ce
      self.alpha_mlm = params.alpha_mlm
      self.alpha_clm = params.alpha_clm
      self.alpha_mse = params.alpha_mse
      self.alpha_cos = params.alpha_cos
  
      self.mlm = params.mlm
      if self.mlm:
          logger.info("Using MLM loss for LM step.")
          self.mlm_mask_prop = params.mlm_mask_prop
          assert 0.0 <= self.mlm_mask_prop <= 1.0
          assert (
              params.word_mask + params.word_keep + params.word_rand
              == 1.0
          )
          self.pred_probs = torch.FloatTensor(
              [params.word_mask, params.word_keep, params.word_rand]
          )
          self.pred_probs = (
              self.pred_probs.to(f"cuda:{params.local_rank}")
              if params.n_gpu > 0
              else self.pred_probs
          )
          self.token_probs = (
              token_probs.to(f"cuda:{params.local_rank}")
              if params.n_gpu > 0
              else token_probs
          )
          if self.fp16:
              self.pred_probs = self.pred_probs.half()
              self.token_probs = self.token_probs.half()
      else:
          logger.info("Using CLM loss for LM step.")
  
      self.epoch = 0
      self.n_iter = 0
      self.n_total_iter = 0
      self.n_sequences_epoch = 0
      self.total_loss_epoch = 0
      self.last_loss = 0
      self.last_loss_ce = 0
      self.last_loss_mlm = 0
      self.last_loss_clm = 0
      if self.alpha_mse > 0.0:
          self.last_loss_mse = 0
      if self.alpha_cos > 0.0:
          self.last_loss_cos = 0
      self.last_log = 0
  
      self.ce_loss_fct = nn.KLDivLoss(reduction="batchmean")
      self.lm_loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
      if self.alpha_mse > 0.0:
          self.mse_loss_fct = nn.MSELoss(reduction="sum")
      if self.alpha_cos > 0.0:
          self.cosine_loss_fct = nn.CosineEmbeddingLoss(
              reduction="mean"
          )
  
      logger.info("--- Initializing model optimizer")
      assert params.gradient_accumulation_steps >= 1
      self.num_steps_epoch = len(self.dataloader)
      num_train_optimization_steps = (
          int(
              self.num_steps_epoch
              / params.gradient_accumulation_steps
              * params.n_epoch
          )
          + 1
      )
  
      no_decay = ["bias", "LayerNorm.weight"]
      optimizer_grouped_parameters = [
          {
              "params": [
                  p
                  for n, p in student.named_parameters()
                  if not any(nd in n for nd in no_decay)
                  and p.requires_grad
              ],
              "weight_decay": params.weight_decay,
          },
          {
              "params": [
                  p
                  for n, p in student.named_parameters()
                  if any(nd in n for nd in no_decay) and p.requires_grad
              ],
              "weight_decay": 0.0,
          },
      ]
      logger.info(
          "------ Number of trainable parameters (student): %i"
          % sum(
              [
                  p.numel()
                  for p in self.student.parameters()
                  if p.requires_grad
              ]
          )
      )
      logger.info(
          "------ Number of parameters (student): %i"
          % sum([p.numel() for p in self.student.parameters()])
      )
      self.optimizer = AdamW(
          optimizer_grouped_parameters,
          lr=params.learning_rate,
          eps=params.adam_epsilon,
          betas=(0.9, 0.98),
      )
  
      warmup_steps = math.ceil(
          num_train_optimization_steps * params.warmup_prop
      )
      self.scheduler = get_linear_schedule_with_warmup(
          self.optimizer,
          num_warmup_steps=warmup_steps,
          num_training_steps=num_train_optimization_steps,
      )
  
      if self.fp16:
          try:
              from apex import amp
          except ImportError:
              raise ImportError(
                  "Please install apex from https://www.github.com/nvidia/apex to use fp16 training."
              )
          logger.info(
              f"Using fp16 training: {self.params.fp16_opt_level} level"
          )
          self.student, self.optimizer = amp.initialize(
              self.student,
              self.optimizer,
              opt_level=self.params.fp16_opt_level,
          )
          self.teacher = self.teacher.half()
  
      if self.multi_gpu:
          if self.fp16:
              from apex.parallel import DistributedDataParallel
  
              logger.info(
                  "Using apex.parallel.DistributedDataParallel for distributed training."
              )
              self.student = DistributedDataParallel(self.student)
          else:
              from torch.nn.parallel import DistributedDataParallel
  
              logger.info(
                  "Using nn.parallel.DistributedDataParallel for distributed training."
              )
              self.student = DistributedDataParallel(
                  self.student,
                  device_ids=[params.local_rank],
                  output_device=params.local_rank,
                  find_unused_parameters=True,
              )
  
      self.is_master = params.is_master
      if self.is_master:
          logger.info("--- Initializing Tensorboard")
          self.tensorboard = SummaryWriter(
              log_dir=os.path.join(self.dump_path, "log", "train")
          )
          self.tensorboard.add_text(
              tag="config/training",
              text_string=str(self.params),
              global_step=0,
          )
          self.tensorboard.add_text(
              tag="config/student",
              text_string=str(self.student_config),
              global_step=0,
          )
  
imports: import math
  import os
  import time
  import psutil
  import torch
  from torch import nn
  from torch.optim import AdamW
  from torch.utils.data import BatchSampler, DataLoader, RandomSampler
  from torch.utils.data.distributed import DistributedSampler
  from tqdm import tqdm
  from grouped_batch_sampler import GroupedBatchSampler, create_lengths_groups
  from lm_seqs_dataset import LmSeqsDataset
  from transformers import get_linear_schedule_with_warmup
  from utils import logger
label_desc: Use image node_id 1 to create a new Distiller object  with example: obj = Distiller(params, dataset, token_probs, student, teacher)