ax_pretrained_infer.py

#!/usr/bin/env python3
# Copyright      2023  Xiaomi Corp.        (authors: Fangjun Kuang)

"""
This script loads ONNX models exported by ./export-onnx.py
and uses them to decode waves.

We use the pre-trained model from
https://huggingface.co/Zengwei/icefall-asr-librispeech-pruned-transducer-stateless7-streaming-2022-12-29
as an example to show how to use this file.

1. Download the pre-trained model

cd egs/librispeech/ASR

repo_url=https://huggingface.co/Zengwei/icefall-asr-librispeech-pruned-transducer-stateless7-streaming-2022-12-29
GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
repo=$(basename $repo_url)

pushd $repo
git lfs pull --include "data/lang_bpe_500/bpe.model"
git lfs pull --include "exp/pretrained.pt"
cd exp
ln -s pretrained.pt epoch-99.pt
popd

2. Export the model to ONNX

./pruned_transducer_stateless7_streaming/export-onnx.py \
  --bpe-model $repo/data/lang_bpe_500/bpe.model \
  --use-averaged-model 0 \
  --epoch 99 \
  --avg 1 \
  --decode-chunk-len 32 \
  --exp-dir $repo/exp/

It will generate the following 3 files in $repo/exp

  - encoder-epoch-99-avg-1.onnx
  - decoder-epoch-99-avg-1.onnx
  - joiner-epoch-99-avg-1.onnx

3. Run this file with the exported ONNX models

./pruned_transducer_stateless7_streaming/onnx_pretrained.py \
  --encoder-model-filename $repo/exp/encoder-epoch-99-avg-1.onnx \
  --decoder-model-filename $repo/exp/decoder-epoch-99-avg-1.onnx \
  --joiner-model-filename $repo/exp/joiner-epoch-99-avg-1.onnx \
  --tokens $repo/data/lang_bpe_500/tokens.txt \
  $repo/test_wavs/1089-134686-0001.wav

Note: Even though this script only supports decoding a single file,
the exported ONNX models do support batch processing.
"""

import argparse
import logging
from typing import Dict, List, Optional, Tuple

import k2
import numpy as np
# import onnxruntime as ort
import torch
# import torchaudio
from kaldifeat import FbankOptions, OnlineFbank, OnlineFeature
from axengine import InferenceSession



def get_parser():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument(
        "--encoder-model-filename",
        type=str,
        required=True,
        help="Path to the encoder onnx model. ",
    )

    parser.add_argument(
        "--decoder-model-filename",
        type=str,
        required=True,
        help="Path to the decoder onnx model. ",
    )

    parser.add_argument(
        "--joiner-model-filename",
        type=str,
        required=True,
        help="Path to the joiner onnx model. ",
    )

    parser.add_argument(
        "--tokens",
        type=str,
        help="""Path to tokens.txt.""",
    )

    parser.add_argument(
        "--sound-dir",
        type=str,
        help="The input sound file to transcribe. "
        "Supported formats are those supported by torchaudio.load(). "
        "For example, wav and flac are supported. "
        "The sample rate has to be 16kHz.",
    )

    return parser


class OnnxModel:
    def __init__(
        self,
        encoder_model_filename: str,
        decoder_model_filename: str,
        joiner_model_filename: str,
    ):
        # session_opts = ort.SessionOptions()
        # session_opts.inter_op_num_threads = 1
        # session_opts.intra_op_num_threads = 1

        # self.session_opts = session_opts

        self.init_encoder(encoder_model_filename)
        self.init_decoder(decoder_model_filename)
        self.init_joiner(joiner_model_filename)

    def init_encoder(self, encoder_model_filename: str):
        self.encoder = InferenceSession(
            encoder_model_filename,
            #sess_options=self.session_opts,      
        )
        print("==========Encoder input==========")
        for i in self.encoder.get_inputs():
            print(i)
        print("==========Encoder output==========")
        for i in self.encoder.get_outputs():
            print(i)

        self.init_encoder_states()

    def init_encoder_states(self, batch_size: int = 1):
        # encoder_meta = self.encoder.get_modelmeta().custom_metadata_map
        # print(encoder_meta)

        # model_type = encoder_meta["model_type"]
        # assert model_type == "zipformer", model_type

        # decode_chunk_len = int(encoder_meta["decode_chunk_len"])
        # T = int(encoder_meta["T"])

        # num_encoder_layers = encoder_meta["num_encoder_layers"]
        # encoder_dims = encoder_meta["encoder_dims"]
        # attention_dims = encoder_meta["attention_dims"]
        # cnn_module_kernels = encoder_meta["cnn_module_kernels"]
        # left_context_len = encoder_meta["left_context_len"]

        model_type = "zipformer"
        assert model_type == "zipformer", model_type

        decode_chunk_len = 96
        T = 103

        num_encoder_layers = "2,2,2,2,2"
        encoder_dims = "256,256,256,256,256"
        attention_dims = "192,192,192,192,192"
        cnn_module_kernels = "31,31,31,31,31"
        left_context_len = "192,96,48,24,96"

        def to_int_list(s):
            return list(map(int, s.split(",")))

        num_encoder_layers = to_int_list(num_encoder_layers)
        encoder_dims = to_int_list(encoder_dims)
        attention_dims = to_int_list(attention_dims)
        cnn_module_kernels = to_int_list(cnn_module_kernels)
        left_context_len = to_int_list(left_context_len)

        # logging.info(f"decode_chunk_len: {decode_chunk_len}")
        # logging.info(f"T: {T}")
        # logging.info(f"num_encoder_layers: {num_encoder_layers}")
        # logging.info(f"encoder_dims: {encoder_dims}")
        # logging.info(f"attention_dims: {attention_dims}")
        # logging.info(f"cnn_module_kernels: {cnn_module_kernels}")
        # logging.info(f"left_context_len: {left_context_len}")

        num_encoders = len(num_encoder_layers)

        cached_len = []
        cached_avg = []
        cached_key = []
        cached_val = []
        cached_val2 = []
        cached_conv1 = []
        cached_conv2 = []

        N = batch_size

        for i in range(num_encoders):
            cached_len.append(torch.zeros(num_encoder_layers[i], N, dtype=torch.int64))
            cached_avg.append(torch.zeros(num_encoder_layers[i], N, encoder_dims[i]))
            cached_key.append(
                torch.zeros(
                    num_encoder_layers[i], left_context_len[i], N, attention_dims[i]
                )
            )
            cached_val.append(
                torch.zeros(
                    num_encoder_layers[i],
                    left_context_len[i],
                    N,
                    attention_dims[i] // 2,
                )
            )
            cached_val2.append(
                torch.zeros(
                    num_encoder_layers[i],
                    left_context_len[i],
                    N,
                    attention_dims[i] // 2,
                )
            )
            cached_conv1.append(
                torch.zeros(
                    num_encoder_layers[i], N, encoder_dims[i], cnn_module_kernels[i] - 1
                )
            )
            cached_conv2.append(
                torch.zeros(
                    num_encoder_layers[i], N, encoder_dims[i], cnn_module_kernels[i] - 1
                )
            )

        self.cached_len = cached_len
        self.cached_avg = cached_avg
        self.cached_key = cached_key
        self.cached_val = cached_val
        self.cached_val2 = cached_val2
        self.cached_conv1 = cached_conv1
        self.cached_conv2 = cached_conv2

        self.num_encoders = num_encoders

        self.segment = T
        self.offset = decode_chunk_len

    def init_decoder(self, decoder_model_filename: str):
        self.decoder = InferenceSession(
            decoder_model_filename,
            # sess_options=self.session_opts,
        )
        print("==========Decoder input==========")
        for i in self.decoder.get_inputs():
            print(i)
        print("==========Decoder output==========")
        for i in self.decoder.get_outputs():
            print(i)

        # decoder_meta = self.decoder.get_modelmeta().custom_metadata_map
        # self.context_size = int(decoder_meta["context_size"])
        # self.vocab_size = int(decoder_meta["vocab_size"])
        self.context_size = 2
        self.vocab_size = 6254

        logging.info(f"context_size: {self.context_size}")
        logging.info(f"vocab_size: {self.vocab_size}")

    def init_joiner(self, joiner_model_filename: str):
        self.joiner = InferenceSession(
            joiner_model_filename,
            # sess_options=self.session_opts,
        )

        print("==========Joiner input==========")
        for i in self.joiner.get_inputs():
            print(i)
        print("==========Joiner output==========")
        for i in self.joiner.get_outputs():
            print(i)

        # joiner_meta = self.joiner.get_modelmeta().custom_metadata_map
        # self.joiner_dim = int(joiner_meta["joiner_dim"])
        self.joiner_dim = 512

        logging.info(f"joiner_dim: {self.joiner_dim}")

    def _build_encoder_input_output(
        self,
        x: torch.Tensor,
    ) -> Tuple[Dict[str, np.ndarray], List[str]]:
        encoder_input = {"x": x.numpy()}
        encoder_output = ["encoder_out"]


        def build_states_input(states: List[torch.Tensor], name: str):
            for i, s in enumerate(states):
                if isinstance(s, torch.Tensor):
                    # Fix dtype for cached_len (int32)
                    if name == "cached_len":
                        encoder_input[f"{name}_{i}"] = s.to(torch.int32).numpy()
                    else:
                        encoder_input[f"{name}_{i}"] = s.numpy()
                else:
                    encoder_input[f"{name}_{i}"] = s

                encoder_output.append(f"new_{name}_{i}")

        build_states_input(self.cached_len, "cached_len")
        build_states_input(self.cached_avg, "cached_avg")
        build_states_input(self.cached_key, "cached_key")
        build_states_input(self.cached_val, "cached_val")
        build_states_input(self.cached_val2, "cached_val2")
        build_states_input(self.cached_conv1, "cached_conv1")
        build_states_input(self.cached_conv2, "cached_conv2")

        return encoder_input, encoder_output

    def _update_states(self, states: List[np.ndarray]):
        num_encoders = self.num_encoders

        self.cached_len = states[num_encoders * 0 : num_encoders * 1]
        self.cached_avg = states[num_encoders * 1 : num_encoders * 2]
        self.cached_key = states[num_encoders * 2 : num_encoders * 3]
        self.cached_val = states[num_encoders * 3 : num_encoders * 4]
        self.cached_val2 = states[num_encoders * 4 : num_encoders * 5]
        self.cached_conv1 = states[num_encoders * 5 : num_encoders * 6]
        self.cached_conv2 = states[num_encoders * 6 : num_encoders * 7]

    def run_encoder(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
          x:
            A 3-D tensor of shape (N, T, C)
        Returns:
          Return a 3-D tensor of shape (N, T', joiner_dim) where
          T' is usually equal to ((T-7)//2+1)//2
        """
        encoder_input, encoder_output_names = self._build_encoder_input_output(x)
        # 
        out = self.encoder.run(encoder_output_names, encoder_input)

        self._update_states(out[1:])

        return torch.from_numpy(out[0])

    def run_decoder(self, decoder_input: torch.Tensor) -> torch.Tensor:
            """
            Args:
                decoder_input:
                    A 2-D tensor of shape (N, context_size)
            Returns:
                Return a 2-D tensor of shape (N, joiner_dim)
            """
            # Fix dtype for decoder input (int32)
            decoder_input_np = decoder_input.to(torch.int32).numpy()
            out = self.decoder.run(
                    [self.decoder.get_outputs()[0].name],
                    {self.decoder.get_inputs()[0].name: decoder_input_np},
            )[0]

            return torch.from_numpy(out)

    def run_joiner(
        self, encoder_out: torch.Tensor, decoder_out: torch.Tensor
    ) -> torch.Tensor:
        """
        Args:
          encoder_out:
            A 2-D tensor of shape (N, joiner_dim)
          decoder_out:
            A 2-D tensor of shape (N, joiner_dim)
        Returns:
          Return a 2-D tensor of shape (N, vocab_size)
        """
        
        out = self.joiner.run(
            [self.joiner.get_outputs()[0].name],
            {
                self.joiner.get_inputs()[0].name: encoder_out.numpy(),
                self.joiner.get_inputs()[1].name: decoder_out.numpy(),
            },
        )[0]
        
        return torch.from_numpy(out)

def read_sound_files(
    filenames: List[str], expected_sample_rate: float
) -> List[torch.Tensor]:
    """Read a list of sound files into a list 1-D float32 torch tensors.
    Args:
      filenames:
        A list of sound filenames.
      expected_sample_rate:
        The expected sample rate of the sound files.
    Returns:
      Return a list of 1-D float32 torch tensors.
    """
    try:
        import soundfile as sf
    except ImportError:
        raise ImportError("Please install soundfile: pip install soundfile")
    try:
        import librosa
    except ImportError:
        librosa = None
    ans = []
    for f in filenames:
        data, sample_rate = sf.read(f)
        if data.ndim > 1:
            data = data[:, 0]
        if sample_rate != expected_sample_rate:
            if librosa is None:
                raise ImportError("Please install librosa: pip install librosa for resampling audio.")
            # librosa.resample expects float32
            data = data.astype(np.float32)
            data = librosa.resample(data, orig_sr=sample_rate, target_sr=expected_sample_rate)
            sample_rate = expected_sample_rate
        wave = torch.from_numpy(data.astype(np.float32)).contiguous()
        ans.append(wave)
    return ans


def create_streaming_feature_extractor() -> OnlineFeature:
    """Create a CPU streaming feature extractor.

    At present, we assume it returns a fbank feature extractor with
    fixed options. In the future, we will support passing in the options
    from outside.

    Returns:
      Return a CPU streaming feature extractor.
    """
    opts = FbankOptions()
    opts.device = "cpu"
    opts.frame_opts.dither = 0
    opts.frame_opts.snip_edges = False
    opts.frame_opts.samp_freq = 16000
    opts.mel_opts.num_bins = 80
    opts.mel_opts.high_freq = -400
    return OnlineFbank(opts)


def greedy_search(
    model: OnnxModel,
    encoder_out: torch.Tensor,
    context_size: int,
    decoder_out: Optional[torch.Tensor] = None,
    hyp: Optional[List[int]] = None,
) -> List[int]:
    """Greedy search in batch mode. It hardcodes --max-sym-per-frame=1.
    Args:
      model:
        The transducer model.
      encoder_out:
        A 3-D tensor of shape (1, T, joiner_dim)
      context_size:
        The context size of the decoder model.
      decoder_out:
        Optional. Decoder output of the previous chunk.
      hyp:
        Decoding results for previous chunks.
    Returns:
      Return the decoded results so far.
    """

    blank_id = 0

    if decoder_out is None:
        assert hyp is None, hyp
        hyp = [blank_id] * context_size
        decoder_input = torch.tensor([hyp], dtype=torch.int64)
        decoder_out = model.run_decoder(decoder_input)
    else:
        assert hyp is not None, hyp

    encoder_out = encoder_out.squeeze(0)
    T = encoder_out.size(0)
    for t in range(T):
        
        cur_encoder_out = encoder_out[t : t + 1]
        joiner_out = model.run_joiner(cur_encoder_out, decoder_out).squeeze(0)
        y = joiner_out.argmax(dim=0).item()
        if y != blank_id:
            hyp.append(y)
            decoder_input = hyp[-context_size:]
            decoder_input = torch.tensor([decoder_input], dtype=torch.int64)
            decoder_out = model.run_decoder(decoder_input)
    
    return hyp, decoder_out


@torch.no_grad()
def main():
    import time
    parser = get_parser()
    args = parser.parse_args()
    logging.info(vars(args))

    t0 = time.time()
    model = OnnxModel(
        encoder_model_filename=args.encoder_model_filename,
        decoder_model_filename=args.decoder_model_filename,
        joiner_model_filename=args.joiner_model_filename,
    )
    t1 = time.time()
    print(f" Model loading time: {t1 - t0:.3f} s")

    sample_rate = 16000

    # logging.info("Constructing Fbank computer")
    # online_fbank = create_streaming_feature_extractor()


    import os
    from glob import glob

    # 判断sound_dir是文件还是目录
    input_path = args.sound_dir
    if os.path.isdir(input_path):
        # 支持常见音频格式
        audio_exts = ["*.wav", "*.flac", "*.mp3", "*.ogg", "*.m4a"]
        audio_files = []
        for ext in audio_exts:
            audio_files.extend(glob(os.path.join(input_path, ext)))
        audio_files.sort()
    else:
        audio_files = [input_path]

    logging.info(f"Found {len(audio_files)} audio files.")

    rtf_list = []
    for idx, sound_file in enumerate(audio_files):
        model.init_encoder_states()  # 重置encoder缓存，保证每个音频独立识别
        t2 = time.time()
        waves = read_sound_files(
            filenames=[sound_file],
            expected_sample_rate=sample_rate,
        )[0]
        t3 = time.time()
        print(f"[{idx+1}/{len(audio_files)}] {os.path.basename(sound_file)} Audio loading time: {t3 - t2:.3f} s")

        tail_padding = torch.zeros(int(0.3 * sample_rate), dtype=torch.float32)
        wave_samples = torch.cat([waves, tail_padding])

        online_fbank = create_streaming_feature_extractor()

        num_processed_frames = 0
        segment = model.segment
        offset = model.offset

        context_size = model.context_size
        hyp = None
        decoder_out = None

        chunk = int(1 * sample_rate)  # 1 second
        start = 0
        t4 = time.time()
        while start < wave_samples.numel():
            end = min(start + chunk, wave_samples.numel())
            samples = wave_samples[start:end]
            start += chunk

            online_fbank.accept_waveform(
                sampling_rate=sample_rate,
                waveform=samples,
            )

            while online_fbank.num_frames_ready - num_processed_frames >= segment:
                frames = []
                for i in range(segment):
                    frames.append(online_fbank.get_frame(num_processed_frames + i))
                num_processed_frames += offset
                frames = torch.cat(frames, dim=0)
                frames = frames.unsqueeze(0)
                encoder_out = model.run_encoder(frames)
                
                hyp, decoder_out = greedy_search(
                    model,
                    encoder_out,
                    context_size,
                    decoder_out,
                    hyp,
                )
        t5 = time.time()
        print(f" Audio processing time: {t5 - t4:.3f} s")

        symbol_table = k2.SymbolTable.from_file(args.tokens)

        text = ""
        for i in hyp[context_size:]:
            text += symbol_table[i]
        text = text.replace("▁", " ").strip()

        logging.info(sound_file)
        logging.info(text)

        audio_duration = float(waves.shape[0]) / sample_rate
        total_time = (t3 - t2) + (t5 - t4)
        rtf = total_time / audio_duration if audio_duration > 0 else float('inf')
        print(f" Audio duration: {audio_duration:.3f} s")
        print(f" Total (load+process) time: {total_time:.3f} s")
        print(f" RTF (total_time/audio_duration): {rtf:.3f}")

        rtf_list.append(rtf)
        # logging.info("Decoding Done")

    if rtf_list:
        avg_rtf = sum(rtf_list) / len(rtf_list)
        print(f"\nAverage RTF over {len(rtf_list)} files: {avg_rtf:.3f}")
    else:
        print("No audio files processed.")


if __name__ == "__main__":
    formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"

    logging.basicConfig(format=formatter, level=logging.INFO)
    main()