class_docstr: Transformer-based (duplex) decoder model for TN/ITN.
docstr: Lightning calls this inside the training loop with the data from the training dataloader
  passed in as `batch`.
code: def training_step(self, batch, batch_idx):
      """
      Lightning calls this inside the training loop with the data from the training dataloader
      passed in as `batch`.
      """
      # tarred dataset contains batches, and the first dimension of size 1 added by the DataLoader
      # (batch_size is set to 1) is redundant
      if batch["input_ids"].ndim == 3:
          batch = {k: v.squeeze(dim=0) for k, v in batch.items()}
  
      # Apply Transformer
      train_loss = self.forward(
          input_ids=batch["input_ids"],
          decoder_input_ids=batch["decoder_input_ids"],
          attention_mask=batch["attention_mask"],
          labels=batch["labels"],
      )
  
      lr = self._optimizer.param_groups[0]["lr"]
      self.log("train_loss", train_loss)
      self.log("lr", lr, prog_bar=True)
      return {"loss": train_loss, "lr": lr}
  
imports: import json
  import os
  from collections import defaultdict
  from math import ceil
  from typing import Dict, List, Optional, Union
  import torch
  from omegaconf import DictConfig
  from pytorch_lightning import Trainer
  from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, DataCollatorForSeq2Seq
  import nemo.collections.nlp.data.text_normalization.constants
  from nemo.collections.common.tokenizers.moses_tokenizers import MosesProcessor
  from nemo.collections.nlp.data.text_normalization import TextNormalizationTestDataset
  from nemo.collections.nlp.data.text_normalization.decoder_dataset import TarredTextNormalizationDecoderDataset, TextNormalizationDecoderDataset
  from nemo.collections.nlp.models.duplex_text_normalization.utils import get_formatted_string
  from nemo.collections.nlp.models.nlp_model import NLPModel
  from nemo.core.classes.common import PretrainedModelInfo, typecheck
  from nemo.core.neural_types import ChannelType, LabelsType, LossType, MaskType, NeuralType
  from nemo.utils import logging
label_desc: Use image node_id 6 for calling the DuplexDecoderModel obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: dict