chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,244 @@
|
||||
#!/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 concurrent.futures import ThreadPoolExecutor
|
||||
import logging
|
||||
from omegaconf import MISSING
|
||||
import os
|
||||
import torch
|
||||
from typing import Optional
|
||||
import warnings
|
||||
|
||||
|
||||
from dataclasses import dataclass
|
||||
from fairseq.dataclass import FairseqDataclass
|
||||
from .kaldi_initializer import KaldiInitializerConfig, initalize_kaldi
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class KaldiDecoderConfig(FairseqDataclass):
|
||||
hlg_graph_path: Optional[str] = None
|
||||
output_dict: str = MISSING
|
||||
|
||||
kaldi_initializer_config: Optional[KaldiInitializerConfig] = None
|
||||
|
||||
acoustic_scale: float = 0.5
|
||||
max_active: int = 10000
|
||||
beam_delta: float = 0.5
|
||||
hash_ratio: float = 2.0
|
||||
|
||||
is_lattice: bool = False
|
||||
lattice_beam: float = 10.0
|
||||
prune_interval: int = 25
|
||||
determinize_lattice: bool = True
|
||||
prune_scale: float = 0.1
|
||||
max_mem: int = 0
|
||||
phone_determinize: bool = True
|
||||
word_determinize: bool = True
|
||||
minimize: bool = True
|
||||
|
||||
num_threads: int = 1
|
||||
|
||||
|
||||
class KaldiDecoder(object):
|
||||
def __init__(
|
||||
self,
|
||||
cfg: KaldiDecoderConfig,
|
||||
beam: int,
|
||||
nbest: int = 1,
|
||||
):
|
||||
try:
|
||||
from kaldi.asr import FasterRecognizer, LatticeFasterRecognizer
|
||||
from kaldi.base import set_verbose_level
|
||||
from kaldi.decoder import (
|
||||
FasterDecoder,
|
||||
FasterDecoderOptions,
|
||||
LatticeFasterDecoder,
|
||||
LatticeFasterDecoderOptions,
|
||||
)
|
||||
from kaldi.lat.functions import DeterminizeLatticePhonePrunedOptions
|
||||
from kaldi.fstext import read_fst_kaldi, SymbolTable
|
||||
except:
|
||||
warnings.warn(
|
||||
"pykaldi is required for this functionality. Please install from https://github.com/pykaldi/pykaldi"
|
||||
)
|
||||
|
||||
# set_verbose_level(2)
|
||||
|
||||
self.acoustic_scale = cfg.acoustic_scale
|
||||
self.nbest = nbest
|
||||
|
||||
if cfg.hlg_graph_path is None:
|
||||
assert (
|
||||
cfg.kaldi_initializer_config is not None
|
||||
), "Must provide hlg graph path or kaldi initializer config"
|
||||
cfg.hlg_graph_path = initalize_kaldi(cfg.kaldi_initializer_config)
|
||||
|
||||
assert os.path.exists(cfg.hlg_graph_path), cfg.hlg_graph_path
|
||||
|
||||
if cfg.is_lattice:
|
||||
self.dec_cls = LatticeFasterDecoder
|
||||
opt_cls = LatticeFasterDecoderOptions
|
||||
self.rec_cls = LatticeFasterRecognizer
|
||||
else:
|
||||
assert self.nbest == 1, "nbest > 1 requires lattice decoder"
|
||||
self.dec_cls = FasterDecoder
|
||||
opt_cls = FasterDecoderOptions
|
||||
self.rec_cls = FasterRecognizer
|
||||
|
||||
self.decoder_options = opt_cls()
|
||||
self.decoder_options.beam = beam
|
||||
self.decoder_options.max_active = cfg.max_active
|
||||
self.decoder_options.beam_delta = cfg.beam_delta
|
||||
self.decoder_options.hash_ratio = cfg.hash_ratio
|
||||
|
||||
if cfg.is_lattice:
|
||||
self.decoder_options.lattice_beam = cfg.lattice_beam
|
||||
self.decoder_options.prune_interval = cfg.prune_interval
|
||||
self.decoder_options.determinize_lattice = cfg.determinize_lattice
|
||||
self.decoder_options.prune_scale = cfg.prune_scale
|
||||
det_opts = DeterminizeLatticePhonePrunedOptions()
|
||||
det_opts.max_mem = cfg.max_mem
|
||||
det_opts.phone_determinize = cfg.phone_determinize
|
||||
det_opts.word_determinize = cfg.word_determinize
|
||||
det_opts.minimize = cfg.minimize
|
||||
self.decoder_options.det_opts = det_opts
|
||||
|
||||
self.output_symbols = {}
|
||||
with open(cfg.output_dict, "r") as f:
|
||||
for line in f:
|
||||
items = line.rstrip().split()
|
||||
assert len(items) == 2
|
||||
self.output_symbols[int(items[1])] = items[0]
|
||||
|
||||
logger.info(f"Loading FST from {cfg.hlg_graph_path}")
|
||||
self.fst = read_fst_kaldi(cfg.hlg_graph_path)
|
||||
self.symbol_table = SymbolTable.read_text(cfg.output_dict)
|
||||
|
||||
self.executor = ThreadPoolExecutor(max_workers=cfg.num_threads)
|
||||
|
||||
def generate(self, models, sample, **unused):
|
||||
"""Generate a batch of inferences."""
|
||||
# model.forward normally channels prev_output_tokens into the decoder
|
||||
# separately, but SequenceGenerator directly calls model.encoder
|
||||
encoder_input = {
|
||||
k: v for k, v in sample["net_input"].items() if k != "prev_output_tokens"
|
||||
}
|
||||
emissions, padding = self.get_emissions(models, encoder_input)
|
||||
return self.decode(emissions, padding)
|
||||
|
||||
def get_emissions(self, models, encoder_input):
|
||||
"""Run encoder and normalize emissions"""
|
||||
model = models[0]
|
||||
|
||||
all_encoder_out = [m(**encoder_input) for m in models]
|
||||
|
||||
if len(all_encoder_out) > 1:
|
||||
|
||||
if "encoder_out" in all_encoder_out[0]:
|
||||
encoder_out = {
|
||||
"encoder_out": sum(e["encoder_out"] for e in all_encoder_out)
|
||||
/ len(all_encoder_out),
|
||||
"encoder_padding_mask": all_encoder_out[0]["encoder_padding_mask"],
|
||||
}
|
||||
padding = encoder_out["encoder_padding_mask"]
|
||||
else:
|
||||
encoder_out = {
|
||||
"logits": sum(e["logits"] for e in all_encoder_out)
|
||||
/ len(all_encoder_out),
|
||||
"padding_mask": all_encoder_out[0]["padding_mask"],
|
||||
}
|
||||
padding = encoder_out["padding_mask"]
|
||||
else:
|
||||
encoder_out = all_encoder_out[0]
|
||||
padding = (
|
||||
encoder_out["padding_mask"]
|
||||
if "padding_mask" in encoder_out
|
||||
else encoder_out["encoder_padding_mask"]
|
||||
)
|
||||
|
||||
if hasattr(model, "get_logits"):
|
||||
emissions = model.get_logits(encoder_out, normalize=True)
|
||||
else:
|
||||
emissions = model.get_normalized_probs(encoder_out, log_probs=True)
|
||||
|
||||
return (
|
||||
emissions.cpu().float().transpose(0, 1),
|
||||
padding.cpu() if padding is not None and padding.any() else None,
|
||||
)
|
||||
|
||||
def decode_one(self, logits, padding):
|
||||
from kaldi.matrix import Matrix
|
||||
|
||||
decoder = self.dec_cls(self.fst, self.decoder_options)
|
||||
asr = self.rec_cls(
|
||||
decoder, self.symbol_table, acoustic_scale=self.acoustic_scale
|
||||
)
|
||||
|
||||
if padding is not None:
|
||||
logits = logits[~padding]
|
||||
|
||||
mat = Matrix(logits.numpy())
|
||||
|
||||
out = asr.decode(mat)
|
||||
|
||||
if self.nbest > 1:
|
||||
from kaldi.fstext import shortestpath
|
||||
from kaldi.fstext.utils import (
|
||||
convert_compact_lattice_to_lattice,
|
||||
convert_lattice_to_std,
|
||||
convert_nbest_to_list,
|
||||
get_linear_symbol_sequence,
|
||||
)
|
||||
|
||||
lat = out["lattice"]
|
||||
|
||||
sp = shortestpath(lat, nshortest=self.nbest)
|
||||
|
||||
sp = convert_compact_lattice_to_lattice(sp)
|
||||
sp = convert_lattice_to_std(sp)
|
||||
seq = convert_nbest_to_list(sp)
|
||||
|
||||
results = []
|
||||
for s in seq:
|
||||
_, o, w = get_linear_symbol_sequence(s)
|
||||
words = list(self.output_symbols[z] for z in o)
|
||||
results.append(
|
||||
{
|
||||
"tokens": words,
|
||||
"words": words,
|
||||
"score": w.value,
|
||||
"emissions": logits,
|
||||
}
|
||||
)
|
||||
return results
|
||||
else:
|
||||
words = out["text"].split()
|
||||
return [
|
||||
{
|
||||
"tokens": words,
|
||||
"words": words,
|
||||
"score": out["likelihood"],
|
||||
"emissions": logits,
|
||||
}
|
||||
]
|
||||
|
||||
def decode(self, emissions, padding):
|
||||
if padding is None:
|
||||
padding = [None] * len(emissions)
|
||||
|
||||
ret = list(
|
||||
map(
|
||||
lambda e, p: self.executor.submit(self.decode_one, e, p),
|
||||
emissions,
|
||||
padding,
|
||||
)
|
||||
)
|
||||
return ret
|
||||
Reference in New Issue
Block a user