code: def plot_confusion_matrix(
      labels: List[int],
      preds: List[int],
      graph_fold: str,
      label_ids: Dict[str, int] = None,
      normalize: bool = False,
      prefix: str = "",
  ):
      """
      Plot confusion matrix.
      Args:
        labels: ground truth labels
        preds: model predictions
        graph_fold: path to a folder to store confusion matrix plot
        label_ids: str label to id map, for example: {'O': 0, 'LOC': 1}
        normalize: whether to normalize confusion matrix
        prefix: prefix for the plot name
      """
      if label_ids is None:
          _plot_confusion_matrix(labels, preds, graph_fold)
  
      else:
          # remove labels from label_ids that don't appear in the dev set
          used_labels = set(labels) | set(preds)
          label_ids = {
              k: label_ids[k]
              for k, v in label_ids.items()
              if v in used_labels
          }
  
          ids_to_labels = {label_ids[k]: k for k in label_ids}
          classes = [
              ids_to_labels[id] for id in sorted(label_ids.values())
          ]
  
          title = "Confusion_matrix"
          cm = confusion_matrix(labels, preds)
          if normalize:
              sums = cm.sum(axis=1)[:, np.newaxis]
              sums = np.where(sums == 0, 1, sums)
              cm = cm.astype("float") / sums
              title = "Normalized_" + title
  
          fig = plt.figure()
          ax = fig.add_subplot(111)
  
          cax = ax.matshow(cm)
  
          ax.set_xticks(np.arange(-1, len(classes)))
          ax.set_yticks(np.arange(-1, len(classes)))
          ax.set_xticklabels([""] + classes, rotation=90)
          ax.set_yticklabels([""] + classes)
          ax.set_ylabel("True")
          ax.set_xlabel("Predicted")
  
          os.makedirs(graph_fold, exist_ok=True)
          fig.colorbar(cax)
  
          title = (prefix + title).strip()
          fig_name = os.path.join(
              graph_fold, title + "_" + time.strftime("%Y%m%d-%H%M%S")
          )
          plt.savefig(fig_name)
          logging.info(f"Confusion matrix saved to {fig_name}")
  
imports: import os
  import time
  from typing import Callable, Dict, List, Optional, Union
  import numpy
  import torch
  import torch.nn.functional
  from matplotlib import pyplot
  from sklearn.metrics import classification_report, confusion_matrix
  from torch import Tensor
  from nemo.collections.nlp.modules.common.megatron.utils import erf_gelu
  from nemo.collections.nlp.modules.common.megatron.utils import openai_gelu
  from nemo.collections.nlp.modules.common.megatron.utils import squared_relu
  from nemo.utils import logging
label_desc: Use image node_id 4 for calling a global function with example usage: plot_confusion_matrix(labels, preds, graph_fold, label_ids, normalize, prefix) without return types