class_docstr: Base class for NLP augmentation
docstr: Augmenting
  
  Args:
      batch: current batch
  
  Returns:
      augmented batch
code: def forward(self, batch: Dict) -> Dict:
      """Augmenting
  
      Args:
          batch: current batch
  
      Returns:
          augmented batch
      """
  
      if self.cfg.augmentation.token_mask_probability > 0:
          input_ids = batch["input_ids"].clone()
          # special_mask = ~batch["special_tokens_mask"].clone().bool()
          mask = (
              torch.bernoulli(
                  torch.full(
                      input_ids.shape,
                      float(
                          self.cfg.augmentation.token_mask_probability
                      ),
                  )
              )
              .to(input_ids.device)
              .bool()
              # & special_mask
          ).bool()
          input_ids[mask] = self.cfg._tokenizer_mask_token_id
          batch["input_ids"] = input_ids.clone()
          batch["attention_mask"][mask] = 0
          if batch["labels"].shape[1] == batch["input_ids"].shape[1]:
              batch["labels"][mask] = -100
  
      return batch
  
file_path: h2o-llmstudio/llm_studio/src/augmentations/nlp_aug.py
imports: import logging
  from abc import abstractmethod
  from typing import Any, Dict
  import torch
  from torch import nn
label_desc: Use image node_id 2 for calling the BaseNLPAug obj's underlying member method code with example usage: obj.forward(batch) and returns: batch