docstr: Prepare the batch: from the token_ids and the lengths, compute the attention mask and the masked label for MLM.
  
  Input:
  ------
      batch: `Tuple`
          token_ids: `torch.tensor(bs, seq_length)` - The token ids for each of the sequence. It is padded.
          lengths: `torch.tensor(bs)` - The lengths of each of the sequences in the batch.
  
  Output:
  -------
      token_ids: `torch.tensor(bs, seq_length)` - The token ids after the modifications for MLM.
      attn_mask: `torch.tensor(bs, seq_length)` - The attention mask for the self-attention.
      mlm_labels: `torch.tensor(bs, seq_length)` - The masked language modeling labels. There is a -100 where there is nothing to predict.
docstr_tok: ["Prepare","the","batch",":","from","the","token_ids","and","the","lengths",",","compute","the","attention","mask","and","the","masked","label","for","MLM",".","Input",":","--","--","--","batch",":","`","Tuple","`","token_ids",":","`","torch.tensor","(","bs",",","seq_length",")","`","-","The","token","ids","for","each","of","the","sequence",".","It","is","padded",".","lengths",":","`","torch.tensor","(","bs",")","`","-","The","lengths","of","each","of","the","sequences","in","the","batch",".","Output",":","--","--","--","-","token_ids",":","`","torch.tensor","(","bs",",","seq_length",")","`","-","The","token","ids","after","the","modifications","for","MLM",".","attn_mask",":","`","torch.tensor","(","bs",",","seq_length",")","`","-","The","attention","mask","for","the","self-attention",".","mlm_labels",":","`","torch.tensor","(","bs",",","seq_length",")","`","-","The","masked","language","modeling","labels",".","There","is","a","-100","where","there","is","nothing","to","predict","."]
code: def prepare_batch_mlm(self, batch):
      """
      Prepare the batch: from the token_ids and the lengths, compute the attention mask and the masked label for MLM.
  
      Input:
      ------
          batch: `Tuple`
              token_ids: `torch.tensor(bs, seq_length)` - The token ids for each of the sequence. It is padded.
              lengths: `torch.tensor(bs)` - The lengths of each of the sequences in the batch.
  
      Output:
      -------
          token_ids: `torch.tensor(bs, seq_length)` - The token ids after the modifications for MLM.
          attn_mask: `torch.tensor(bs, seq_length)` - The attention mask for the self-attention.
          mlm_labels: `torch.tensor(bs, seq_length)` - The masked language modeling labels. There is a -100 where there is nothing to predict.
      """
      token_ids, lengths = batch
      token_ids, lengths = self.round_batch(
          x=token_ids, lengths=lengths
      )
      assert token_ids.size(0) == lengths.size(0)
  
      attn_mask = (
          torch.arange(
              token_ids.size(1), dtype=torch.long, device=lengths.device
          )
          < lengths[:, None]
      )
  
      bs, max_seq_len = token_ids.size()
      mlm_labels = token_ids.new(token_ids.size()).copy_(token_ids)
  
      x_prob = self.token_probs[token_ids.flatten()]
      n_tgt = math.ceil(self.mlm_mask_prop * lengths.sum().item())
      tgt_ids = torch.multinomial(
          x_prob / x_prob.sum(), n_tgt, replacement=False
      )
      pred_mask = torch.zeros(
          bs * max_seq_len, dtype=torch.bool, device=token_ids.device
      )  # previously `dtype=torch.uint8`, cf pytorch 1.2.0 compatibility
      pred_mask[tgt_ids] = 1
      pred_mask = pred_mask.view(bs, max_seq_len)
  
      pred_mask[
          token_ids == self.params.special_tok_ids["pad_token"]
      ] = 0
  
      # mask a number of words == 0 [8] (faster with fp16)
      if self.fp16:
          n1 = pred_mask.sum().item()
          if n1 > 8:
              pred_mask = pred_mask.view(-1)
              n2 = max(n1 % 8, 8 * (n1 // 8))
              if n2 != n1:
                  pred_mask[
                      torch.nonzero(pred_mask).view(-1)[: n1 - n2]
                  ] = 0
              pred_mask = pred_mask.view(bs, max_seq_len)
              assert (
                  pred_mask.sum().item() % 8 == 0
              ), pred_mask.sum().item()
  
      _token_ids_real = token_ids[pred_mask]
      _token_ids_rand = _token_ids_real.clone().random_(self.vocab_size)
      _token_ids_mask = _token_ids_real.clone().fill_(
          self.params.special_tok_ids["mask_token"]
      )
      probs = torch.multinomial(
          self.pred_probs, len(_token_ids_real), replacement=True
      )
      _token_ids = (
          _token_ids_mask * (probs == 0).long()
          + _token_ids_real * (probs == 1).long()
          + _token_ids_rand * (probs == 2).long()
      )
      token_ids = token_ids.masked_scatter(pred_mask, _token_ids)
  
      mlm_labels[
          ~pred_mask
      ] = (
          -100
      )  # previously `mlm_labels[1-pred_mask] = -1`, cf pytorch 1.2.0 compatibility
  
      # sanity checks
      assert 0 <= token_ids.min() <= token_ids.max() < self.vocab_size
  
      return token_ids, attn_mask, mlm_labels
  
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 2 for calling the Distiller obj's underlying member method code with example usage: obj.prepare_batch_mlm(batch) and returns: token_ids,  attn_mask,  mlm_labels