ort_common.py

import multiprocessing
import os
import tarfile as tf

import numpy as np
import torch
import torchaudio
import torchaudio.compliance.kaldi as kaldi
import yaml

try:
    from swig_decoders import (map_batch, ctc_beam_search_decoder_batch,
                               TrieVector, PathTrie)
except ModuleNotFoundError:
    class PathTrie:
        pass

    class TrieVector(list):
        pass

    def _log_add(*values):
        values = [value for value in values if value != -float("inf")]
        if not values:
            return -float("inf")
        max_value = max(values)
        return max_value + np.log(sum(np.exp(value - max_value)
                                      for value in values))

    def _map_sentence(sent, vocabulary, greedy=False, blank_id=0):
        mapped = []
        prev = None
        for token in sent:
            token = int(token)
            if greedy and token == prev:
                prev = token
                continue
            prev = token
            if token == blank_id or token < 0 or token >= len(vocabulary):
                continue
            piece = vocabulary[token]
            if piece.startswith("<") and piece.endswith(">"):
                continue
            mapped.append(piece)
        return "".join(mapped)

    def map_batch(batch_sents, vocabulary, num_processes, greedy=False,
                  blank_id=0):
        del num_processes
        return [_map_sentence(sent, vocabulary, greedy, blank_id)
                for sent in batch_sents]

    def _ctc_prefix_beam_search(log_probs_seq, log_probs_idx, beam_size,
                                blank_id):
        beam = {(): (0.0, -float("inf"))}
        for frame_probs, frame_ids in zip(log_probs_seq, log_probs_idx):
            next_beam = {}
            for prefix, (prob_blank, prob_non_blank) in beam.items():
                for prob, token in zip(frame_probs, frame_ids):
                    token = int(token)
                    prob = float(prob)
                    next_prob_blank, next_prob_non_blank = next_beam.get(
                        prefix, (-float("inf"), -float("inf")))
                    if token == blank_id:
                        next_beam[prefix] = (
                            _log_add(next_prob_blank, prob_blank + prob,
                                     prob_non_blank + prob),
                            next_prob_non_blank,
                        )
                        continue

                    last = prefix[-1] if prefix else None
                    if token == last:
                        next_beam[prefix] = (
                            next_prob_blank,
                            _log_add(next_prob_non_blank,
                                     prob_non_blank + prob),
                        )
                        new_prefix = prefix + (token, )
                        nb_blank, nb_non_blank = next_beam.get(
                            new_prefix, (-float("inf"), -float("inf")))
                        next_beam[new_prefix] = (
                            nb_blank,
                            _log_add(nb_non_blank, prob_blank + prob),
                        )
                    else:
                        new_prefix = prefix + (token, )
                        nb_blank, nb_non_blank = next_beam.get(
                            new_prefix, (-float("inf"), -float("inf")))
                        next_beam[new_prefix] = (
                            nb_blank,
                            _log_add(nb_non_blank, prob_blank + prob,
                                     prob_non_blank + prob),
                        )
            beam = dict(sorted(
                next_beam.items(),
                key=lambda item: _log_add(item[1][0], item[1][1]),
                reverse=True)[:beam_size])
        return [(_log_add(prob_blank, prob_non_blank), prefix)
                for prefix, (prob_blank, prob_non_blank) in sorted(
                    beam.items(),
                    key=lambda item: _log_add(item[1][0], item[1][1]),
                    reverse=True)]

    def ctc_beam_search_decoder_batch(batch_log_probs_seq,
                                      batch_log_probs_idx,
                                      batch_root_trie,
                                      batch_start,
                                      beam_size,
                                      num_processes,
                                      blank_id=0,
                                      space_id=-1,
                                      cutoff_prob=0.999,
                                      ext_scorer=None):
        del batch_root_trie, batch_start, num_processes, space_id
        del cutoff_prob, ext_scorer
        return [
            _ctc_prefix_beam_search(log_probs_seq, log_probs_idx, beam_size,
                                    blank_id)
            for log_probs_seq, log_probs_idx in zip(batch_log_probs_seq,
                                                    batch_log_probs_idx)
        ]


def load_config(config_path):
    with open(config_path, "r") as fin:
        return yaml.load(fin, Loader=yaml.FullLoader)


def load_vocab(vocab_path):
    vocabulary = []
    char_dict = {}
    with open(vocab_path, "r") as fin:
        for line in fin:
            arr = line.strip().split()
            assert len(arr) == 2
            char_dict[int(arr[1])] = arr[0]
            vocabulary.append(arr[0])
    return vocabulary, char_dict


def compute_feats(audio_file: str, sr=16000) -> np.ndarray:
    try:
        import soundfile as sf
        waveform, sample_rate = sf.read(audio_file, dtype="int16",
                                        always_2d=True)
        waveform = torch.from_numpy(waveform.T).to(torch.float)
    except ModuleNotFoundError:
        waveform, sample_rate = torchaudio.load(audio_file, normalize=True)
        waveform = waveform.to(torch.float)
    if sample_rate != sr:
        waveform = torchaudio.transforms.Resample(
            orig_freq=sample_rate, new_freq=sr)(waveform)
    feats = kaldi.fbank(waveform,
                        num_mel_bins=80,
                        frame_length=25,
                        frame_shift=10,
                        energy_floor=0.0,
                        sample_frequency=sr)
    return feats.unsqueeze(0).numpy()


def pad_array_along_axis(array, pad_width, axis, mode="constant", **kwargs):
    if array.shape[axis] >= pad_width:
        return array
    full_pad_width = [(0, 0)] * array.ndim
    full_pad_width[axis] = (0, pad_width - array.shape[axis])
    return np.pad(array, pad_width=full_pad_width, mode=mode, **kwargs)


def ctc_decoding(beam_log_probs,
                 beam_log_probs_idx,
                 encoder_out_lens,
                 vocabulary,
                 mode="ctc_prefix_beam_search"):
    beam_size = beam_log_probs.shape[-1]
    batch_size = beam_log_probs.shape[0]
    num_processes = min(multiprocessing.cpu_count(), batch_size)
    hyps = []
    score_hyps = []

    if mode == "ctc_greedy_search":
        log_probs_idx = beam_log_probs_idx[:, :, 0]
        batch_sents = []
        for idx, seq in enumerate(log_probs_idx):
            batch_sents.append(seq[0:encoder_out_lens[idx]].tolist())
        hyps = map_batch(batch_sents, vocabulary, num_processes, True, 0)
    elif mode in ("ctc_prefix_beam_search", "attention_rescoring"):
        batch_log_probs_seq_list = beam_log_probs.tolist()
        batch_log_probs_idx_list = beam_log_probs_idx.tolist()
        batch_len_list = encoder_out_lens.tolist()
        batch_log_probs_seq = []
        batch_log_probs_ids = []
        batch_start = []
        batch_root = TrieVector()
        root_dict = {}
        for i in range(len(batch_len_list)):
            num_sent = batch_len_list[i]
            batch_log_probs_seq.append(batch_log_probs_seq_list[i][0:num_sent])
            batch_log_probs_ids.append(batch_log_probs_idx_list[i][0:num_sent])
            root_dict[i] = PathTrie()
            batch_root.append(root_dict[i])
            batch_start.append(True)
        score_hyps = ctc_beam_search_decoder_batch(batch_log_probs_seq,
                                                   batch_log_probs_ids,
                                                   batch_root,
                                                   batch_start,
                                                   beam_size,
                                                   num_processes,
                                                   0, -2, 0.99999)
        if mode == "ctc_prefix_beam_search":
            for cand_hyps in score_hyps:
                hyps.append(cand_hyps[0][1])
            hyps = map_batch(hyps, vocabulary, num_processes, False, 0)
    return hyps, score_hyps


def make_decoder_inputs(encoder_out,
                        encoder_out_lens,
                        beam_log_probs,
                        beam_log_probs_idx,
                        vocabulary,
                        sos,
                        eos,
                        decoder_len):
    _, score_hyps = ctc_decoding(beam_log_probs, beam_log_probs_idx,
                                 encoder_out_lens, vocabulary,
                                 "attention_rescoring")
    ignore_id = -1
    beam_size = beam_log_probs.shape[-1]
    batch_size = beam_log_probs.shape[0]
    ctc_score, all_hyps = [], []
    for hyps in score_hyps:
        cur_len = len(hyps)
        if len(hyps) < beam_size:
            hyps += (beam_size - cur_len) * [(-float("inf"), (0,))]
        cur_ctc_score = []
        for hyp in hyps:
            cur_ctc_score.append(hyp[0])
            all_hyps.append(list(hyp[1]))
        ctc_score.append(cur_ctc_score)
    ctc_score = np.array(ctc_score, dtype=np.float32)

    max_len = decoder_len - 2
    hyps_pad_sos_eos = np.ones((batch_size, beam_size, max_len + 2),
                               dtype=np.int64) * ignore_id
    r_hyps_pad_sos_eos = np.ones((batch_size, beam_size, max_len + 2),
                                 dtype=np.int64) * ignore_id
    hyps_lens_sos = np.ones((batch_size, beam_size), dtype=np.int32)
    k = 0
    for i in range(batch_size):
        for j in range(beam_size):
            cand = all_hyps[k][:max_len]
            length = len(cand) + 2
            hyps_pad_sos_eos[i][j][0:length] = [sos] + cand + [eos]
            r_hyps_pad_sos_eos[i][j][0:length] = [sos] + cand[::-1] + [eos]
            hyps_lens_sos[i][j] = len(cand) + 1
            k += 1

    if decoder_len > encoder_out.shape[1]:
        encoder_out = np.pad(encoder_out,
                             [(0, 0),
                              (0, decoder_len - encoder_out.shape[1]),
                              (0, 0)],
                             mode="constant",
                             constant_values=0)
    elif decoder_len < encoder_out.shape[1]:
        encoder_out = encoder_out[:, :decoder_len, :]

    return {
        "encoder_out": encoder_out,
        "encoder_out_lens": np.full(batch_size,
                                    fill_value=decoder_len,
                                    dtype=np.int32),
        "hyps_pad_sos_eos": hyps_pad_sos_eos.astype(np.int32),
        "hyps_lens_sos": hyps_lens_sos,
        "r_hyps_pad_sos_eos": r_hyps_pad_sos_eos.astype(np.int32),
        "ctc_score": ctc_score,
    }, all_hyps


def make_offline_inputs(feats, seq_len):
    feats = feats[:, :seq_len, :]
    speech_lengths = np.array([feats.shape[1]], dtype=np.int32)
    if feats.shape[1] < seq_len:
        feats = pad_array_along_axis(feats, pad_width=seq_len, axis=1)
    return {"speech": feats, "speech_lengths": speech_lengths}


def make_online_initial_state(configs,
                              batch_size=1,
                              decoding_chunk_size=16,
                              num_decoding_left_chunks=5):
    subsampling = 4
    context = 7
    stride = subsampling * decoding_chunk_size
    decoding_window = (decoding_chunk_size - 1) * subsampling + context
    required_cache_size = decoding_chunk_size * num_decoding_left_chunks

    output_size = configs["encoder_conf"]["output_size"]
    num_layers = configs["encoder_conf"]["num_blocks"]
    cnn_module_kernel = configs["encoder_conf"].get("cnn_module_kernel", 1) - 1
    head = configs["encoder_conf"]["attention_heads"]
    d_k = configs["encoder_conf"]["output_size"] // head

    state = {
        "att_cache": np.zeros((batch_size, num_layers, head,
                               required_cache_size, d_k * 2),
                              dtype=np.float32),
        "cnn_cache": np.zeros((batch_size, num_layers, output_size,
                               cnn_module_kernel),
                              dtype=np.float32),
        "cache_mask": np.zeros((batch_size, 1, required_cache_size),
                               dtype=np.float32),
        "offset": np.zeros((batch_size, 1), dtype=np.int32),
    }
    params = {
        "batch_size": batch_size,
        "context": context,
        "stride": stride,
        "decoding_window": decoding_window,
    }
    return state, params


def make_online_encoder_input(feats, cur, params, state):
    batch_size = params["batch_size"]
    decoding_window = params["decoding_window"]
    end = min(cur + decoding_window, feats.shape[1])
    chunk_xs = feats[:, cur:end, :]
    if chunk_xs.shape[1] < decoding_window:
        chunk_xs = pad_array_along_axis(chunk_xs,
                                        pad_width=decoding_window,
                                        axis=1)
    chunk_xs = chunk_xs.astype(np.float32)
    chunk_lens = np.full(batch_size,
                         fill_value=chunk_xs.shape[1],
                         dtype=np.int32)
    return {
        "chunk_xs": chunk_xs,
        "chunk_lens": chunk_lens,
        "offset": state["offset"],
        "att_cache": state["att_cache"],
        "cnn_cache": state["cnn_cache"],
        "cache_mask": state["cache_mask"],
    }


def update_online_state(state, outputs):
    state["offset"] = outputs[4]
    state["att_cache"] = outputs[5]
    state["cnn_cache"] = outputs[6]
    state["cache_mask"] = outputs[7]


def save_calibration_inputs(calib_data_path, inputs, sample_id):
    for input_name, data in inputs.items():
        data_path = os.path.join(calib_data_path, input_name)
        os.makedirs(data_path, exist_ok=True)
        np.save(os.path.join(data_path, f"{sample_id}.npy"), data)


def pack_calibration_dataset(calib_data_path):
    for input_name in sorted(os.listdir(calib_data_path)):
        data_path = os.path.join(calib_data_path, input_name)
        if not os.path.isdir(data_path):
            continue
        tar_path = os.path.join(calib_data_path, input_name + ".tar.gz")
        with tf.open(tar_path, "w:gz") as tf_file:
            tf_file.add(data_path, arcname=input_name)


class WenetONNXRunner:

    def __init__(self,
                 config_path,
                 vocab_path,
                 onnx_dir="onnx_model",
                 offline_seq_len=1024,
                 decoder_len=32,
                 decoding_chunk_size=16,
                 num_decoding_left_chunks=5,
                 batch_size=1,
                 providers=None):
        self.config_path = config_path
        self.vocab_path = vocab_path
        self.onnx_dir = onnx_dir
        self.offline_seq_len = offline_seq_len
        self.decoder_len = decoder_len
        self.decoding_chunk_size = decoding_chunk_size
        self.num_decoding_left_chunks = num_decoding_left_chunks
        self.batch_size = batch_size
        self.providers = providers or ["CPUExecutionProvider"]

        self.configs = load_config(config_path)
        self.vocabulary, self.char_dict = load_vocab(vocab_path)
        self.eos = self.sos = len(self.char_dict) - 1

        self._offline_encoder = None
        self._online_encoder = None
        self._decoder = None

    @property
    def offline_encoder(self):
        if self._offline_encoder is None:
            self._offline_encoder = self._new_session("encoder_offline.onnx")
        return self._offline_encoder

    @property
    def online_encoder(self):
        if self._online_encoder is None:
            self._online_encoder = self._new_session("encoder_online.onnx")
        return self._online_encoder

    @property
    def decoder(self):
        if self._decoder is None:
            self._decoder = self._new_session("decoder.onnx")
        return self._decoder

    def _new_session(self, filename):
        import onnxruntime as ort
        return ort.InferenceSession(os.path.join(self.onnx_dir, filename),
                                    providers=self.providers)

    def compute_feats(self, audio_file):
        return compute_feats(audio_file)

    def run_offline_encoder(self,
                            feats,
                            calib_data_path=None,
                            sample_id=0):
        encoder_input = make_offline_inputs(feats, self.offline_seq_len)
        if calib_data_path:
            save_calibration_inputs(calib_data_path, encoder_input, sample_id)
        encoder_out, encoder_out_lens, ctc_log_probs, beam_log_probs, beam_log_probs_idx = (
            self.offline_encoder.run(None, encoder_input))
        return {
            "encoder_out": encoder_out,
            "encoder_out_lens": encoder_out_lens,
            "ctc_log_probs": ctc_log_probs,
            "beam_log_probs": beam_log_probs,
            "beam_log_probs_idx": beam_log_probs_idx,
        }

    def run_online_encoder(self,
                           feats,
                           calib_data_path=None,
                           sample_prefix=0):
        state, online_params = make_online_initial_state(
            self.configs, self.batch_size, self.decoding_chunk_size,
            self.num_decoding_left_chunks)
        encoder_out = []
        beam_log_probs = []
        beam_log_probs_idx = []
        num_frames = feats.shape[1]

        for cur in range(0, num_frames - online_params["context"] + 1,
                         online_params["stride"]):
            encoder_input = make_online_encoder_input(feats, cur,
                                                      online_params, state)
            if calib_data_path:
                save_calibration_inputs(calib_data_path, encoder_input,
                                        f"{sample_prefix}_{cur}")
            outputs = self.online_encoder.run(None, encoder_input)
            chunk_log_probs, chunk_log_probs_idx, chunk_out, chunk_out_lens = outputs[:4]
            update_online_state(state, outputs)
            del chunk_out_lens
            encoder_out.append(chunk_out)
            beam_log_probs.append(chunk_log_probs)
            beam_log_probs_idx.append(chunk_log_probs_idx.astype(np.int32))

        return {
            "encoder_out": np.concatenate(encoder_out, axis=1),
            "encoder_out_lens": np.full(self.batch_size,
                                        fill_value=sum(
                                            out.shape[1]
                                            for out in encoder_out),
                                        dtype=np.int32),
            "beam_log_probs": np.concatenate(beam_log_probs, axis=1),
            "beam_log_probs_idx": np.concatenate(beam_log_probs_idx, axis=1),
            "num_chunks": len(encoder_out),
        }

    def ctc_decode(self, encoder_outputs, mode):
        return ctc_decoding(encoder_outputs["beam_log_probs"],
                            encoder_outputs["beam_log_probs_idx"],
                            encoder_outputs["encoder_out_lens"],
                            self.vocabulary, mode)

    def run_decoder(self,
                    encoder_outputs,
                    calib_data_path=None,
                    sample_id=0):
        decoder_input, all_hyps = make_decoder_inputs(
            encoder_outputs["encoder_out"],
            encoder_outputs["encoder_out_lens"],
            encoder_outputs["beam_log_probs"],
            encoder_outputs["beam_log_probs_idx"],
            self.vocabulary,
            self.sos,
            self.eos,
            self.decoder_len,
        )
        if calib_data_path:
            save_calibration_inputs(calib_data_path, decoder_input, sample_id)

        best_index = self.decoder.run(None, decoder_input)[0].astype(np.int32)
        beam_size = encoder_outputs["beam_log_probs"].shape[-1]
        num_processes = min(multiprocessing.cpu_count(), best_index.shape[0])
        best_sents = []
        k = 0
        for idx in best_index:
            best_sents.append(all_hyps[k:k + beam_size][idx])
            k += beam_size
        hyps = map_batch(best_sents, self.vocabulary, num_processes)
        return "".join(hyps)

    def transcribe(self,
                   audio_file,
                   online=False,
                   mode="ctc_prefix_beam_search",
                   calib_data_path=None):
        feats = self.compute_feats(audio_file)
        if online:
            encoder_outputs = self.run_online_encoder(
                feats, calib_data_path=calib_data_path, sample_prefix=0)
        else:
            encoder_outputs = self.run_offline_encoder(
                feats, calib_data_path=calib_data_path, sample_id=0)

        if mode == "attention_rescoring":
            result = self.run_decoder(encoder_outputs, calib_data_path, 0)
        else:
            hyps, _ = self.ctc_decode(encoder_outputs, mode)
            result = "".join(hyps) if hyps else ""

        if calib_data_path:
            pack_calibration_dataset(calib_data_path)
        return result

    def save_calibration_for_audio(self,
                                   audio_file,
                                   parts,
                                   calib_data_path,
                                   sample_id):
        counts = {"offline": 0, "online": 0, "decoder": 0}
        feats = self.compute_feats(audio_file)
        offline_outputs = None
        if "offline" in parts or "decoder" in parts:
            offline_outputs = self.run_offline_encoder(
                feats,
                calib_data_path=calib_data_path if "offline" in parts else None,
                sample_id=sample_id)
            if "offline" in parts:
                counts["offline"] += 1
        if "online" in parts:
            online_outputs = self.run_online_encoder(
                feats,
                calib_data_path=calib_data_path,
                sample_prefix=sample_id)
            counts["online"] += online_outputs["num_chunks"]
        if "decoder" in parts:
            self.run_decoder(offline_outputs,
                             calib_data_path=calib_data_path,
                             sample_id=sample_id)
            counts["decoder"] += 1
        return counts