class_docstr: Token Classification Model with BERT, applicable for tasks such as Named Entity Recognition
class_docstr_tok: ["Token","Classification","Model","with","BERT",",","applicable","for","tasks","such","as","Named","Entity","Recognition"]
code: def test_step(self, batch, batch_idx):
      (
          input_ids,
          input_type_ids,
          input_mask,
          subtokens_mask,
          loss_mask,
          labels,
      ) = batch
      logits = self(
          input_ids=input_ids,
          token_type_ids=input_type_ids,
          attention_mask=input_mask,
      )
      val_loss = self.loss(
          logits=logits, labels=labels, loss_mask=loss_mask
      )
  
      subtokens_mask = subtokens_mask > 0.5
  
      preds = torch.argmax(logits, axis=-1)[subtokens_mask]
      labels = labels[subtokens_mask]
      tp, fn, fp, _ = self.classification_report(preds, labels)
  
      loss = {"test_loss": val_loss, "tp": tp, "fn": fn, "fp": fp}
      self.test_step_outputs.append(loss)
      return loss
  
imports: import os
  from typing import List, Optional, Union
  import torch
  from omegaconf import DictConfig, OmegaConf
  from pytorch_lightning import Trainer
  from torch.utils.data import DataLoader
  from nemo.collections.common.losses import CrossEntropyLoss
  from nemo.collections.nlp.data.data_utils.data_preprocessing import get_labels_to_labels_id_mapping
  from nemo.collections.nlp.data.token_classification.token_classification_dataset import BertTokenClassificationDataset, BertTokenClassificationInferDataset
  from nemo.collections.nlp.data.token_classification.token_classification_utils import get_label_ids
  from nemo.collections.nlp.metrics.classification_report import ClassificationReport
  from nemo.collections.nlp.models.nlp_model import NLPModel
  from nemo.collections.nlp.modules.common import TokenClassifier
  from nemo.collections.nlp.parts.utils_funcs import get_classification_report, plot_confusion_matrix, tensor2list
  from nemo.core.classes.common import PretrainedModelInfo, typecheck
  from nemo.utils import logging
label_desc: Use image node_id 8 for calling the TokenClassificationModel obj's underlying member method code with example usage: obj.test_step(batch, batch_idx) and returns: loss