chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,698 @@
|
||||
#!/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,
|
||||
)
|
||||
|
||||
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()
|
||||
Reference in New Issue
Block a user