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"]
docstr: Setup loss
     Setup or update loss.
  
  Args:
      class_balancing: whether to use class weights during training
docstr_tok: ["Setup","loss","Setup","or","update","loss",".","Args",":","class_balancing",":","whether","to","use","class","weights","during","training"]
code: def setup_loss(self, class_balancing: str = None):
      """Setup loss
         Setup or update loss.
  
      Args:
          class_balancing: whether to use class weights during training
      """
      if class_balancing not in ["weighted_loss", None]:
          raise ValueError(
              f"Class balancing {class_balancing} is not supported. Choose from: [null, weighted_loss]"
          )
      if class_balancing == "weighted_loss" and self.class_weights:
          # you may need to increase the number of epochs for convergence when using weighted_loss
          loss = CrossEntropyLoss(
              logits_ndim=3, weight=self.class_weights
          )
          logging.debug(f"Using {class_balancing} class balancing.")
      else:
          loss = CrossEntropyLoss(logits_ndim=3)
          logging.debug(f"Using CrossEntropyLoss class balancing.")
      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 3 for calling the TokenClassificationModel obj's underlying member method code with example usage: obj.setup_loss(class_balancing) and returns: loss