class_docstr: Pretraining machinery using SCRIPT from Nijkamp et al., 2021. Always running sparse prediction.
class_docstr_tok: ["Pretraining","machinery","using","SCRIPT","from","Nijkamp","et","al.",",","2021",".","Always","running","sparse","prediction","."]
code: def forward(
      self,
      input_ids,
      attention_mask: Optional[torch.Tensor] = None,
      labels: Optional[torch.Tensor] = None,
  ):
      loss = torch.tensor(
          0.0, dtype=torch.float, device=input_ids.device
      )
  
      outputs = self.encoder(input_ids, attention_mask)
      outputs = outputs.view(-1, outputs.shape[-1])
  
      if labels is not None:
          # ## Generation pass ##
          labels = labels.view(-1)
          mask_positions = labels.view(-1) != self.loss_fn.ignore_index
          num_masks_guaranteed = round(
              self.sparse_prediction * labels.shape[0]
          )
          indices = torch.argsort(mask_positions.int())[
              -num_masks_guaranteed:
          ]
  
          # sparse outputs for prediction
          outputs = outputs[indices]
          labels = labels[indices]
  
          logits = self.decoder(
              self.prediction_head(outputs)
          )  # sparse logits
          loss += self.loss_fn(logits, labels)
  
          # ## Discrimination pass ##
          resampled_token_ids = self._gumbel_sample(logits.detach())
          discriminator_input_ids = input_ids.clone().view(-1)
          discriminator_input_ids[indices] = resampled_token_ids
  
          critic_labels = (
              input_ids.view(-1) != discriminator_input_ids
          ).to(outputs.dtype)
  
          outputs = self.encoder(
              discriminator_input_ids.view_as(input_ids), attention_mask
          ).view(-1, outputs.shape[-1])
          disc_logits = self.decoder(
              self.prediction_head(outputs)
          )  # full logits
          binary_logits = self._get_binary_logits(disc_logits)
  
          # ELECTRA-type discriminator:
          loss += (
              self.ALPHA
              * torch.nn.functional.binary_cross_entropy_with_logits(
                  binary_logits, critic_labels
              )
          )
  
      else:
          logits = self.decoder(self.prediction_head(outputs))
          loss += outputs.new_zeros((1,))
  
      return {"loss": loss, "logits": logits}
  
imports: import torch
  from transformers import PretrainedConfig, PreTrainedModel
  from transformers import AutoConfig, AutoModel, AutoModelForMaskedLM, AutoModelForSequenceClassification, AutoModelForTokenClassification
  from typing import Optional
  from omegaconf import OmegaConf
  from .fff import FFF
  from .components import _get_norm_fn, _get_nonlin_fn, EmbeddingComponent, PoolingComponent, PredictionHeadComponent, GLU, get_extended_attention_mask, _init_module
  from .attention import get_attention_mechanism
label_desc: Use image node_id 3 for calling the ScriptableLMForSCRIPTTraining obj's underlying member method code with example usage: obj.forward(input_ids, attention_mask, labels) and returns: dict