class_docstr: k2 Mixin class that simplifies the construction of various models with k2-based losses.
  
  It does the following:
      -   Sets up the graph loss and decoder (methods _init_k2 and update_k2_modules).
      -   Registers external graphs, if needed.
      -   Augments forward(...) with optional graph decoding to get accurate predictions.
class_docstr_tok: ["k2","Mixin","class","that","simplifies","the","construction","of","various","models","with","k2-based","losses",".","It","does","the","following",":","-","Sets","up","the","graph","loss","and","decoder","(","methods","_init_k2","and","update_k2_modules",")",".","-","Registers","external","graphs",",","if","needed",".","-","Augments","forward","(","...",")","with","optional","graph","decoding","to","get","accurate","predictions","."]
docstr: Helper function to initialize or update k2 loss and transcribe_decoder.
  
  Args:
      input_cfg: DictConfig to take new parameters from. Schema is expected as in
          nemo.collections.asr.models.configs.k2_sequence_models_config.GraphModuleConfig
docstr_tok: ["Helper","function","to","initialize","or","update","k2","loss","and","transcribe_decoder",".","Args",":","input_cfg",":","DictConfig","to","take","new","parameters","from",".","Schema","is","expected","as","in","nemo.collections.asr.models.configs.k2_sequence_models_config.GraphModuleConfig"]
code: def update_k2_modules(self, input_cfg: DictConfig):
      """
      Helper function to initialize or update k2 loss and transcribe_decoder.
  
      Args:
          input_cfg: DictConfig to take new parameters from. Schema is expected as in
              nemo.collections.asr.models.configs.k2_sequence_models_config.GraphModuleConfig
      """
      del self.loss
      if hasattr(self, "transcribe_decoder"):
          del self.transcribe_decoder
  
      if hasattr(self, "joint"):
          # RNNT
          num_classes = self.joint.num_classes_with_blank - 1
      else:
          # CTC, MMI, ...
          num_classes = self.decoder.num_classes_with_blank - 1
          remove_consecutive = input_cfg.backend_cfg.get(
              "topo_with_self_loops", True
          ) and input_cfg.backend_cfg.get(
              "topo_type", "default"
          ) not in [
              "forced_blank",
              "identity",
          ]
          self._wer.remove_consecutive = remove_consecutive
  
      from nemo.collections.asr.losses.lattice_losses import LatticeLoss
  
      self.loss = LatticeLoss(
          num_classes=num_classes,
          reduction=self._cfg.get("ctc_reduction", "mean_batch"),
          backend="k2",
          criterion_type=input_cfg.get("criterion_type", "ml"),
          loss_type=input_cfg.get("loss_type", "ctc"),
          split_batch_size=input_cfg.get("split_batch_size", 0),
          graph_module_cfg=input_cfg.backend_cfg,
      )
  
      criterion_type = self.loss.criterion_type
      self.use_graph_lm = criterion_type == "map"
      transcribe_training = input_cfg.get("transcribe_training", False)
      if transcribe_training and criterion_type == "ml":
          logging.warning(
              f"""You do not need to use transcribe_training=`{transcribe_training}` 
                              with criterion_type=`{criterion_type}`. transcribe_training will be set to False."""
          )
          transcribe_training = False
      self.transcribe_training = transcribe_training
      if self.use_graph_lm:
          from nemo.collections.asr.modules.graph_decoder import (
              ViterbiDecoderWithGraph,
          )
  
          self.transcribe_decoder = ViterbiDecoderWithGraph(
              num_classes=num_classes,
              backend="k2",
              dec_type="token_lm",
              return_type="1best",
              return_ilabels=True,
              output_aligned=True,
              split_batch_size=input_cfg.get("split_batch_size", 0),
              graph_module_cfg=input_cfg.backend_cfg,
          )
  
file_path: NeMo/nemo/collections/asr/parts/k2/classes.py
imports: from abc import ABC
  from dataclasses import dataclass, field
  from typing import Any, Optional, Tuple
  import torch
  from omegaconf import DictConfig
  from nemo.utils import logging
label_desc: Use image node_id 2 for calling the ASRK2Mixin obj's underlying member method code with example usage: obj.update_k2_modules(input_cfg) without return types