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 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 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