class_docstr: W-Transducer loss: RNN-T loss modification for training RNN-T model for the case
  when some text at the beginning/end of the utterance is missing.
  The resulting model behaves like the RNN-T model (no modification for decoding is required).
  For details see "Powerful and Extensible WFST Framework for RNN-Transducer Losses" paper
      https://ieeexplore.ieee.org/document/10096679
class_docstr_tok: ["W-Transducer","loss",":","RNN-T","loss","modification","for","training","RNN-T","model","for","the","case","when","some","text","at","the","beginning\/end","of","the","utterance","is","missing",".","The","resulting","model","behaves","like","the","RNN-T","model","(","no","modification","for","decoding","is","required",")",".","For","details","see","``","Powerful","and","Extensible","WFST","Framework","for","RNN-Transducer","Losses","''","paper","https",":","\/\/ieeexplore.ieee.org\/document\/10096679"]
docstr: Get unit schema (target text) graph for W-Transducer loss (Compose-Transducer).
  Forward arcs represent text labels.
  
  Example graph: text [1, 2], blank=0. Eps ids: 3, 4.
  
  graph::
  
          3:3:0                  0:0:1                  0:0:2
        +-------+              +-------+              +-------+
        v       |              v       |              v       |
      +-----------+  1:1:0   +-----------+  2:2:1   +-----------+  -1:-1:-1  #===#
      |     0     | -------> |     1     | -------> |     2     | ---------> H 3 H
      +-----------+          +-----------+          +-----------+            #===#
        ^ 0:0:0 |                                     ^ 4:4:2 |
        +-------+                                     +-------+
  
  Args:
      units_tensor: 1d tensor with text units
      vocab_size: number of total labels (vocab size including blank)
  
  Returns:
      unit schema graph (k2.Fsa).
      Labels: <unit>:<unit>:<unit_position> (k2.Fsa: labels, aux_labels, unit_positions)
docstr_tok: ["Get","unit","schema","(","target","text",")","graph","for","W-Transducer","loss","(","Compose-Transducer",")",".","Forward","arcs","represent","text","labels",".","Example","graph",":","text","[","1",",","2","]",",","blank=0",".","Eps","ids",":","3",",","4.","graph",":",":","3:3:0","0:0:1","0:0:2","+","--","--","--","-+","+","--","--","--","-+","+","--","--","--","-+","v","|","v","|","v","|","+","--","--","--","--","--","-+","1:1:0","+","--","--","--","--","--","-+","2:2:1","+","--","--","--","--","--","-+","-1",":","-1",":","-1","#","===","#","|","0","|","--","--","--","-",">","|","1","|","--","--","--","-",">","|","2","|","--","--","--","--","-",">","H","3","H","+","--","--","--","--","--","-+","+","--","--","--","--","--","-+","+","--","--","--","--","--","-+","#","===","#","^","0:0:0","|","^","4:4:2","|","+","--","--","--","-+","+","--","--","--","-+","Args",":","units_tensor",":","1d","tensor","with","text","units","vocab_size",":","number","of","total","labels","(","vocab","size","including","blank",")","Returns",":","unit","schema","graph","(","k2.Fsa",")",".","Labels",":","<","unit",">",":","<","unit",">",":","<","unit_position",">","(","k2.Fsa",":","labels",",","aux_labels",",","unit_positions",")"]
code: def get_unit_schema(
      self, units_tensor: torch.Tensor, vocab_size: int
  ) -> "k2.Fsa":
      """
      Get unit schema (target text) graph for W-Transducer loss (Compose-Transducer).
      Forward arcs represent text labels.
  
      Example graph: text [1, 2], blank=0. Eps ids: 3, 4.
  
      graph::
  
              3:3:0                  0:0:1                  0:0:2
            +-------+              +-------+              +-------+
            v       |              v       |              v       |
          +-----------+  1:1:0   +-----------+  2:2:1   +-----------+  -1:-1:-1  #===#
          |     0     | -------> |     1     | -------> |     2     | ---------> H 3 H
          +-----------+          +-----------+          +-----------+            #===#
            ^ 0:0:0 |                                     ^ 4:4:2 |
            +-------+                                     +-------+
  
      Args:
          units_tensor: 1d tensor with text units
          vocab_size: number of total labels (vocab size including blank)
  
      Returns:
          unit schema graph (k2.Fsa).
          Labels: <unit>:<unit>:<unit_position> (k2.Fsa: labels, aux_labels, unit_positions)
      """
  
      blank_id = self.blank
      start_eps_id = vocab_size
      end_eps_id = vocab_size + 1
      device = units_tensor.device
      text_len = units_tensor.shape[0]
  
      # arcs: scr, dest, label, score
      arcs = torch.zeros(
          ((text_len + 1) * 2 + 2, 4), dtype=torch.int32, device=device
      )
      text_indices = torch.arange(
          0, text_len + 1, dtype=torch.int32, device=device
      )
      # eps
      arcs[0, 2] = start_eps_id
      # blank labels
      arcs[1:-1:2, 0] = text_indices  # from state
      arcs[1:-1:2, 1] = text_indices  # to state
      arcs[1:-1:2, 2] = blank_id
  
      # text labels
      arcs[2:-1:2, 0] = text_indices  # from state
      arcs[2:-1:2, 1] = text_indices + 1  # to state
      arcs[2:-2:2, 2] = units_tensor  # labels: text
  
      arcs[-1] = arcs[-2]
      arcs[-2, 1] = text_len
      arcs[-2, 2] = end_eps_id
      arcs[
          -1, 2
      ] = (
          -1
      )  # last transition to final state, ilabel=-1 (special for k2)
      olabels = arcs[:, 2].detach().clone()  # same as ilabels
  
      fsa_text = k2.Fsa(arcs, olabels)
      fsa_text.unit_positions = torch.zeros_like(olabels)
      fsa_text.unit_positions[1:-1] = (
          text_indices.expand(2, -1).transpose(0, 1).flatten()
      )
      fsa_text.unit_positions[-1] = -1
      return fsa_text
  
file_path: NeMo/nemo/collections/asr/parts/k2/w_transducer.py
imports: from contextlib import nullcontext
  from typing import Union
  import torch
  import torch.nn.functional
  from nemo.collections.asr.parts.k2.graph_transducer import GraphRnntLoss, force_float32_context
  from nemo.core.utils.k2_guard import k2
  from nemo.utils.enum import PrettyStrEnum
label_desc: Use image node_id 2 for calling the GraphWTransducerLoss obj's underlying member method code with example usage: obj.get_unit_schema(units_tensor, vocab_size) and returns: fsa_text