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 grouped_batch_sampler import GroupedBatchSampler, create_lengths_groups
  from lm_seqs_dataset import LmSeqsDataset
  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 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)