text: #!/usr/bin/env python3
  
  # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  from dataclasses import dataclass
  import hydra
  from hydra.core.config_store import ConfigStore
  import logging
  from omegaconf import MISSING, OmegaConf
  import os
  import os.path as osp
  from pathlib import Path
  import subprocess
  from typing import Optional
  
  from fairseq.data.dictionary import Dictionary
  from fairseq.dataclass import FairseqDataclass
  
  script_dir = Path(__file__).resolve().parent
  config_path = script_dir / "config"
  
  
  logger = logging.getLogger(__name__)
  
  
  @dataclass
  class KaldiInitializerConfig(FairseqDataclass):
      data_dir: str = MISSING
      fst_dir: Optional[str] = None
      in_labels: str = MISSING
      out_labels: Optional[str] = None
      wav2letter_lexicon: Optional[str] = None
      lm_arpa: str = MISSING
      kaldi_root: str = MISSING
      blank_symbol: str = "<s>"
      silence_symbol: Optional[str] = None
  
  
  def create_units(fst_dir: Path, in_labels: str, vocab: Dictionary) -> Path:
      in_units_file = fst_dir / f"kaldi_dict.{in_labels}.txt"
      if not in_units_file.exists():
  
          logger.info(f"Creating {in_units_file}")
  
          with open(in_units_file, "w") as f:
              print("<eps> 0", file=f)
              i = 1
              for symb in vocab.symbols[vocab.nspecial :]:
                  if not symb.startswith("madeupword"):
                      print(f"{symb} {i}", file=f)
                      i += 1
      return in_units_file
  
  
  def create_lexicon(
      cfg: KaldiInitializerConfig,
      fst_dir: Path,
      unique_label: str,
      in_units_file: Path,
      out_words_file: Path,
  ) -> (Path, Path):
  
      disambig_in_units_file = fst_dir / f"kaldi_dict.{cfg.in_labels}_disambig.txt"
      lexicon_file = fst_dir / f"kaldi_lexicon.{unique_label}.txt"
      disambig_lexicon_file = fst_dir / f"kaldi_lexicon.{unique_label}_disambig.txt"
      if (
          not lexicon_file.exists()
          or not disambig_lexicon_file.exists()
          or not disambig_in_units_file.exists()
      ):
          logger.info(f"Creating {lexicon_file} (in units file: {in_units_file})")
  
          assert cfg.wav2letter_lexicon is not None or cfg.in_labels == cfg.out_labels
  
          if cfg.wav2letter_lexicon is not None:
              lm_words = set()
              with open(out_words_file, "r") as lm_dict_f:
                  for line in lm_dict_f:
                      lm_words.add(line.split()[0])
  
              num_skipped = 0
              total = 0
              with open(cfg.wav2letter_lexicon, "r") as w2l_lex_f, open(
                  lexicon_file, "w"
              ) as out_f:
                  for line in w2l_lex_f:
                      items = line.rstrip().split("\t")
                      assert len(items) == 2, items
                      if items[0] in lm_words:
                          print(items[0], items[1], file=out_f)
                      else:
                          num_skipped += 1
                          logger.debug(
                              f"Skipping word {items[0]} as it was not found in LM"
                          )
                      total += 1
              if num_skipped > 0:
                  logger.warning(
                      f"Skipped {num_skipped} out of {total} words as they were not found in LM"
                  )
          else:
              with open(in_units_file, "r") as in_f, open(lexicon_file, "w") as out_f:
                  for line in in_f:
                      symb = line.split()[0]
                      if symb != "<eps>" and symb != "<ctc_blank>" and symb != "<SIL>":
                          print(symb, symb, file=out_f)
  
          lex_disambig_path = (
              Path(cfg.kaldi_root) / "egs/wsj/s5/utils/add_lex_disambig.pl"
          )
          res = subprocess.run(
              [lex_disambig_path, lexicon_file, disambig_lexicon_file],
              check=True,
              capture_output=True,
          )
          ndisambig = int(res.stdout)
          disamib_path = Path(cfg.kaldi_root) / "egs/wsj/s5/utils/add_disambig.pl"
          res = subprocess.run(
              [disamib_path, "--include-zero", in_units_file, str(ndisambig)],
              check=True,
              capture_output=True,
          )
          with open(disambig_in_units_file, "wb") as f:
              f.write(res.stdout)
  
      return disambig_lexicon_file, disambig_in_units_file
  
  
  def create_G(
      kaldi_root: Path, fst_dir: Path, lm_arpa: Path, arpa_base: str
  ) -> (Path, Path):
  
      out_words_file = fst_dir / f"kaldi_dict.{arpa_base}.txt"
      grammar_graph = fst_dir / f"G_{arpa_base}.fst"
      if not grammar_graph.exists() or not out_words_file.exists():
          logger.info(f"Creating {grammar_graph}")
          arpa2fst = kaldi_root / "src/lmbin/arpa2fst"
          subprocess.run(
              [
                  arpa2fst,
                  "--disambig-symbol=#0",
                  f"--write-symbol-table={out_words_file}",
                  lm_arpa,
                  grammar_graph,
              ],
              check=True,
          )
      return grammar_graph, out_words_file
  
  
  def create_L(
      kaldi_root: Path,
      fst_dir: Path,
      unique_label: str,
      lexicon_file: Path,
      in_units_file: Path,
      out_words_file: Path,
  ) -> Path:
      lexicon_graph = fst_dir / f"L.{unique_label}.fst"
  
      if not lexicon_graph.exists():
          logger.info(f"Creating {lexicon_graph} (in units: {in_units_file})")
          make_lex = kaldi_root / "egs/wsj/s5/utils/make_lexicon_fst.pl"
          fstcompile = kaldi_root / "tools/openfst-1.6.7/bin/fstcompile"
          fstaddselfloops = kaldi_root / "src/fstbin/fstaddselfloops"
          fstarcsort = kaldi_root / "tools/openfst-1.6.7/bin/fstarcsort"
  
          def write_disambig_symbol(file):
              with open(file, "r") as f:
                  for line in f:
                      items = line.rstrip().split()
                      if items[0] == "#0":
                          out_path = str(file) + "_disamig"
                          with open(out_path, "w") as out_f:
                              print(items[1], file=out_f)
                              return out_path
  
              return None
  
          in_disambig_sym = write_disambig_symbol(in_units_file)
          assert in_disambig_sym is not None
          out_disambig_sym = write_disambig_symbol(out_words_file)
          assert out_disambig_sym is not None
  
          try:
              with open(lexicon_graph, "wb") as out_f:
                  res = subprocess.run(
                      [make_lex, lexicon_file], capture_output=True, check=True
                  )
                  assert len(res.stderr) == 0, res.stderr.decode("utf-8")
                  res = subprocess.run(
                      [
                          fstcompile,
                          f"--isymbols={in_units_file}",
                          f"--osymbols={out_words_file}",
                          "--keep_isymbols=false",
                          "--keep_osymbols=false",
                      ],
                      input=res.stdout,
                      capture_output=True,
                  )
                  assert len(res.stderr) == 0, res.stderr.decode("utf-8")
                  res = subprocess.run(
                      [fstaddselfloops, in_disambig_sym, out_disambig_sym],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [fstarcsort, "--sort_type=olabel"],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  out_f.write(res.stdout)
          except subprocess.CalledProcessError as e:
              logger.error(f"cmd: {e.cmd}, err: {e.stderr.decode('utf-8')}")
              os.remove(lexicon_graph)
              raise
          except AssertionError:
              os.remove(lexicon_graph)
              raise
  
      return lexicon_graph
  
  
  def create_LG(
      kaldi_root: Path,
      fst_dir: Path,
      unique_label: str,
      lexicon_graph: Path,
      grammar_graph: Path,
  ) -> Path:
      lg_graph = fst_dir / f"LG.{unique_label}.fst"
  
      if not lg_graph.exists():
          logger.info(f"Creating {lg_graph}")
  
          fsttablecompose = kaldi_root / "src/fstbin/fsttablecompose"
          fstdeterminizestar = kaldi_root / "src/fstbin/fstdeterminizestar"
          fstminimizeencoded = kaldi_root / "src/fstbin/fstminimizeencoded"
          fstpushspecial = kaldi_root / "src/fstbin/fstpushspecial"
          fstarcsort = kaldi_root / "tools/openfst-1.6.7/bin/fstarcsort"
  
          try:
              with open(lg_graph, "wb") as out_f:
                  res = subprocess.run(
                      [fsttablecompose, lexicon_graph, grammar_graph],
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [
                          fstdeterminizestar,
                          "--use-log=true",
                      ],
                      input=res.stdout,
                      capture_output=True,
                  )
                  res = subprocess.run(
                      [fstminimizeencoded],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [fstpushspecial],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [fstarcsort, "--sort_type=ilabel"],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  out_f.write(res.stdout)
          except subprocess.CalledProcessError as e:
              logger.error(f"cmd: {e.cmd}, err: {e.stderr.decode('utf-8')}")
              os.remove(lg_graph)
              raise
  
      return lg_graph
  
  
  def create_H(
      kaldi_root: Path,
      fst_dir: Path,
      disambig_out_units_file: Path,
      in_labels: str,
      vocab: Dictionary,
      blk_sym: str,
      silence_symbol: Optional[str],
  ) -> (Path, Path, Path):
      h_graph = (
          fst_dir / f"H.{in_labels}{'_' + silence_symbol if silence_symbol else ''}.fst"
      )
      h_out_units_file = fst_dir / f"kaldi_dict.h_out.{in_labels}.txt"
      disambig_in_units_file_int = Path(str(h_graph) + "isym_disambig.int")
      disambig_out_units_file_int = Path(str(disambig_out_units_file) + ".int")
      if (
          not h_graph.exists()
          or not h_out_units_file.exists()
          or not disambig_in_units_file_int.exists()
      ):
          logger.info(f"Creating {h_graph}")
          eps_sym = "<eps>"
  
          num_disambig = 0
          osymbols = []
  
          with open(disambig_out_units_file, "r") as f, open(
              disambig_out_units_file_int, "w"
          ) as out_f:
              for line in f:
                  symb, id = line.rstrip().split()
                  if line.startswith("#"):
                      num_disambig += 1
                      print(id, file=out_f)
                  else:
                      if len(osymbols) == 0:
                          assert symb == eps_sym, symb
                      osymbols.append((symb, id))
  
          i_idx = 0
          isymbols = [(eps_sym, 0)]
  
          imap = {}
  
          for i, s in enumerate(vocab.symbols):
              i_idx += 1
              isymbols.append((s, i_idx))
              imap[s] = i_idx
  
          fst_str = []
  
          node_idx = 0
          root_node = node_idx
  
          special_symbols = [blk_sym]
          if silence_symbol is not None:
              special_symbols.append(silence_symbol)
  
          for ss in special_symbols:
              fst_str.append("{} {} {} {}".format(root_node, root_node, ss, eps_sym))
  
          for symbol, _ in osymbols:
              if symbol == eps_sym or symbol.startswith("#"):
                  continue
  
              node_idx += 1
              # 1. from root to emitting state
              fst_str.append("{} {} {} {}".format(root_node, node_idx, symbol, symbol))
              # 2. from emitting state back to root
              fst_str.append("{} {} {} {}".format(node_idx, root_node, eps_sym, eps_sym))
              # 3. from emitting state to optional blank state
              pre_node = node_idx
              node_idx += 1
              for ss in special_symbols:
                  fst_str.append("{} {} {} {}".format(pre_node, node_idx, ss, eps_sym))
              # 4. from blank state back to root
              fst_str.append("{} {} {} {}".format(node_idx, root_node, eps_sym, eps_sym))
  
          fst_str.append("{}".format(root_node))
  
          fst_str = "\n".join(fst_str)
          h_str = str(h_graph)
          isym_file = h_str + ".isym"
  
          with open(isym_file, "w") as f:
              for sym, id in isymbols:
                  f.write("{} {}\n".format(sym, id))
  
          with open(h_out_units_file, "w") as f:
              for sym, id in osymbols:
                  f.write("{} {}\n".format(sym, id))
  
          with open(disambig_in_units_file_int, "w") as f:
              disam_sym_id = len(isymbols)
              for _ in range(num_disambig):
                  f.write("{}\n".format(disam_sym_id))
                  disam_sym_id += 1
  
          fstcompile = kaldi_root / "tools/openfst-1.6.7/bin/fstcompile"
          fstaddselfloops = kaldi_root / "src/fstbin/fstaddselfloops"
          fstarcsort = kaldi_root / "tools/openfst-1.6.7/bin/fstarcsort"
  
          try:
              with open(h_graph, "wb") as out_f:
                  res = subprocess.run(
                      [
                          fstcompile,
                          f"--isymbols={isym_file}",
                          f"--osymbols={h_out_units_file}",
                          "--keep_isymbols=false",
                          "--keep_osymbols=false",
                      ],
                      input=str.encode(fst_str),
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [
                          fstaddselfloops,
                          disambig_in_units_file_int,
                          disambig_out_units_file_int,
                      ],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [fstarcsort, "--sort_type=olabel"],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  out_f.write(res.stdout)
          except subprocess.CalledProcessError as e:
              logger.error(f"cmd: {e.cmd}, err: {e.stderr.decode('utf-8')}")
              os.remove(h_graph)
              raise
      return h_graph, h_out_units_file, disambig_in_units_file_int
  
  
  def create_HLGa(
      kaldi_root: Path,
      fst_dir: Path,
      unique_label: str,
      h_graph: Path,
      lg_graph: Path,
      disambig_in_words_file_int: Path,
  ) -> Path:
      hlga_graph = fst_dir / f"HLGa.{unique_label}.fst"
  
      if not hlga_graph.exists():
          logger.info(f"Creating {hlga_graph}")
  
          fsttablecompose = kaldi_root / "src/fstbin/fsttablecompose"
          fstdeterminizestar = kaldi_root / "src/fstbin/fstdeterminizestar"
          fstrmsymbols = kaldi_root / "src/fstbin/fstrmsymbols"
          fstrmepslocal = kaldi_root / "src/fstbin/fstrmepslocal"
          fstminimizeencoded = kaldi_root / "src/fstbin/fstminimizeencoded"
  
          try:
              with open(hlga_graph, "wb") as out_f:
                  res = subprocess.run(
                      [
                          fsttablecompose,
                          h_graph,
                          lg_graph,
                      ],
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [fstdeterminizestar, "--use-log=true"],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [fstrmsymbols, disambig_in_words_file_int],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [fstrmepslocal],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [fstminimizeencoded],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  out_f.write(res.stdout)
          except subprocess.CalledProcessError as e:
              logger.error(f"cmd: {e.cmd}, err: {e.stderr.decode('utf-8')}")
              os.remove(hlga_graph)
              raise
  
      return hlga_graph
  
  
  def create_HLa(
      kaldi_root: Path,
      fst_dir: Path,
      unique_label: str,
      h_graph: Path,
      l_graph: Path,
      disambig_in_words_file_int: Path,
  ) -> Path:
      hla_graph = fst_dir / f"HLa.{unique_label}.fst"
  
      if not hla_graph.exists():
          logger.info(f"Creating {hla_graph}")
  
          fsttablecompose = kaldi_root / "src/fstbin/fsttablecompose"
          fstdeterminizestar = kaldi_root / "src/fstbin/fstdeterminizestar"
          fstrmsymbols = kaldi_root / "src/fstbin/fstrmsymbols"
          fstrmepslocal = kaldi_root / "src/fstbin/fstrmepslocal"
          fstminimizeencoded = kaldi_root / "src/fstbin/fstminimizeencoded"
  
          try:
              with open(hla_graph, "wb") as out_f:
                  res = subprocess.run(
                      [
                          fsttablecompose,
                          h_graph,
                          l_graph,
                      ],
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [fstdeterminizestar, "--use-log=true"],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [fstrmsymbols, disambig_in_words_file_int],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [fstrmepslocal],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  res = subprocess.run(
                      [fstminimizeencoded],
                      input=res.stdout,
                      capture_output=True,
                      check=True,
                  )
                  out_f.write(res.stdout)
          except subprocess.CalledProcessError as e:
              logger.error(f"cmd: {e.cmd}, err: {e.stderr.decode('utf-8')}")
              os.remove(hla_graph)
              raise
  
      return hla_graph
  
  
  def create_HLG(
      kaldi_root: Path,
      fst_dir: Path,
      unique_label: str,
      hlga_graph: Path,
      prefix: str = "HLG",
  ) -> Path:
      hlg_graph = fst_dir / f"{prefix}.{unique_label}.fst"
      if not hlg_graph.exists():
          logger.info(f"Creating {hlg_graph}")
  
          add_self_loop = script_dir / "add-self-loop-simple"
          kaldi_src = kaldi_root / "src"
          kaldi_lib = kaldi_src / "lib"
  
          try:
              if not add_self_loop.exists():
                  fst_include = kaldi_root / "tools/openfst-1.6.7/include"
                  add_self_loop_src = script_dir / "add-self-loop-simple.cc"
  
                  subprocess.run(
                      [
                          "c++",
                          f"-I{kaldi_src}",
                          f"-I{fst_include}",
                          f"-L{kaldi_lib}",
                          add_self_loop_src,
                          "-lkaldi-base",
                          "-lkaldi-fstext",
                          "-o",
                          add_self_loop,
                      ],
                      check=True, #check=True
                  )
  
              my_env = os.environ.copy()
              my_env["LD_LIBRARY_PATH"] = f"{kaldi_lib}:{my_env['LD_LIBRARY_PATH']}"
  
              subprocess.run(
                  [
                      add_self_loop,
                      hlga_graph,
                      hlg_graph,
                  ],
                  check=True,
                  capture_output=True,
                  env=my_env,
              )
          except subprocess.CalledProcessError as e:
              logger.error(f"cmd: {e.cmd}, err: {e.stderr.decode('utf-8')}")
              raise
  
      return hlg_graph
  
  
  def initalize_kaldi(cfg: KaldiInitializerConfig) -> Path:
      if cfg.fst_dir is None:
          cfg.fst_dir = osp.join(cfg.data_dir, "kaldi")
      if cfg.out_labels is None:
          cfg.out_labels = cfg.in_labels
  
      kaldi_root = Path(cfg.kaldi_root)
      data_dir = Path(cfg.data_dir)
      fst_dir = Path(cfg.fst_dir)
      fst_dir.mkdir(parents=True, exist_ok=True)
  
      arpa_base = osp.splitext(osp.basename(cfg.lm_arpa))[0]
      unique_label = f"{cfg.in_labels}.{arpa_base}"
  
      with open(data_dir / f"dict.{cfg.in_labels}.txt", "r") as f:
          vocab = Dictionary.load(f)
  
      in_units_file = create_units(fst_dir, cfg.in_labels, vocab)
  
      grammar_graph, out_words_file = create_G(
          kaldi_root, fst_dir, Path(cfg.lm_arpa), arpa_base
      )
  
      disambig_lexicon_file, disambig_L_in_units_file = create_lexicon(
          cfg, fst_dir, unique_label, in_units_file, out_words_file
      )
  
      h_graph, h_out_units_file, disambig_in_units_file_int = create_H(
          kaldi_root,
          fst_dir,
          disambig_L_in_units_file,
          cfg.in_labels,
          vocab,
          cfg.blank_symbol,
          cfg.silence_symbol,
      )
      lexicon_graph = create_L(
          kaldi_root,
          fst_dir,
          unique_label,
          disambig_lexicon_file,
          disambig_L_in_units_file,
          out_words_file,
      )
      lg_graph = create_LG(
          kaldi_root, fst_dir, unique_label, lexicon_graph, grammar_graph
      )
      hlga_graph = create_HLGa(
          kaldi_root, fst_dir, unique_label, h_graph, lg_graph, disambig_in_units_file_int
      )
      hlg_graph = create_HLG(kaldi_root, fst_dir, unique_label, hlga_graph)
  
      # for debugging
      # hla_graph = create_HLa(kaldi_root, fst_dir, unique_label, h_graph, lexicon_graph, disambig_in_units_file_int)
      # hl_graph = create_HLG(kaldi_root, fst_dir, unique_label, hla_graph, prefix="HL_looped")
      # create_HLG(kaldi_root, fst_dir, "phnc", h_graph, prefix="H_looped")
  
      return hlg_graph
  
  
  @hydra.main(config_path=config_path, config_name="kaldi_initializer")
  def cli_main(cfg: KaldiInitializerConfig) -> None:
      container = OmegaConf.to_container(cfg, resolve=True, enum_to_str=True)
      cfg = OmegaConf.create(container)
      OmegaConf.set_struct(cfg, True)
      initalize_kaldi(cfg)
  
  
  if __name__ == "__main__":
  
      logging.root.setLevel(logging.INFO)
      logging.basicConfig(level=logging.INFO)
  
      try:
          from hydra._internal.utils import (
              get_args,
          )  # pylint: disable=import-outside-toplevel
  
          cfg_name = get_args().config_name or "kaldi_initializer"
      except ImportError:
          logger.warning("Failed to get config name from hydra args")
          cfg_name = "kaldi_initializer"
  
      cs = ConfigStore.instance()
      cs.store(name=cfg_name, node=KaldiInitializerConfig)
  
      cli_main()
  