app/asr_worker.py

import os
from pathlib import Path
import numpy as np
import sherpa_onnx
import scipy.signal
from opencc import OpenCC
from huggingface_hub import hf_hub_download

# Ensure Hugging Face cache is in a user-writable directory
CACHE_DIR = Path(__file__).parent / "hf_cache"
os.makedirs(CACHE_DIR, exist_ok=True)

converter = OpenCC('s2t')

# Streaming Zipformer model registry: paths relative to repo root
STREAMING_ZIPFORMER_MODELS = {
    "pfluo/k2fsa-zipformer-chinese-english-mixed": {
        "tokens": "data/lang_char_bpe/tokens.txt",
        "encoder_fp32": "exp/encoder-epoch-99-avg-1.onnx",
        "encoder_int8": "exp/encoder-epoch-99-avg-1.int8.onnx",
        "decoder_fp32": "exp/decoder-epoch-99-avg-1.onnx",
        "decoder_int8": None,
        "joiner_fp32": "exp/joiner-epoch-99-avg-1.onnx",
        "joiner_int8": "exp/joiner-epoch-99-avg-1.int8.onnx",
    },
    "k2-fsa/sherpa-onnx-streaming-zipformer-korean-2024-06-16": {
        "tokens": "tokens.txt",
        "encoder_fp32": "encoder-epoch-99-avg-1.onnx",
        "encoder_int8": "encoder-epoch-99-avg-1.int8.onnx",
        "decoder_fp32": "decoder-epoch-99-avg-1.onnx",
        "decoder_int8": "decoder-epoch-99-avg-1.int8.onnx",
        "joiner_fp32": "joiner-epoch-99-avg-1.onnx",
        "joiner_int8": "joiner-epoch-99-avg-1.int8.onnx",
    },
    "k2-fsa/sherpa-onnx-streaming-zipformer-multi-zh-hans-2023-12-12": {
        "tokens": "tokens.txt",
        "encoder_fp32": "encoder-epoch-20-avg-1-chunk-16-left-128.onnx",
        "encoder_int8": "encoder-epoch-20-avg-1-chunk-16-left-128.int8.onnx",
        "decoder_fp32": "decoder-epoch-20-avg-1-chunk-16-left-128.onnx",
        "decoder_int8": "decoder-epoch-20-avg-1-chunk-16-left-128.int8.onnx",
        "joiner_fp32": "joiner-epoch-20-avg-1-chunk-16-left-128.onnx",
        "joiner_int8": "joiner-epoch-20-avg-1-chunk-16-left-128.int8.onnx",
    },
    "pkufool/icefall-asr-zipformer-streaming-wenetspeech-20230615": {
        "tokens": "data/lang_char/tokens.txt",
        "encoder_fp32": "exp/encoder-epoch-12-avg-4-chunk-16-left-128.onnx",
        "encoder_int8": "exp/encoder-epoch-12-avg-4-chunk-16-left-128.int8.onnx",
        "decoder_fp32": "exp/decoder-epoch-12-avg-4-chunk-16-left-128.onnx",
        "decoder_int8": "exp/decoder-epoch-12-avg-4-chunk-16-left-128.int8.onnx",
        "joiner_fp32": "exp/joiner-epoch-12-avg-4-chunk-16-left-128.onnx",
        "joiner_int8": "exp/joiner-epoch-12-avg-4-chunk-16-left-128.int8.onnx",
    },
    "csukuangfj/sherpa-onnx-streaming-zipformer-en-2023-06-26": {
        "tokens": "tokens.txt",
        "encoder_fp32": "encoder-epoch-99-avg-1-chunk-16-left-128.onnx",
        "encoder_int8": "encoder-epoch-99-avg-1-chunk-16-left-128.int8.onnx",
        "decoder_fp32": "decoder-epoch-99-avg-1-chunk-16-left-128.onnx",
        "decoder_int8": None,
        "joiner_fp32": "joiner-epoch-99-avg-1-chunk-16-left-128.onnx",
        "joiner_int8": "joiner-epoch-99-avg-1-chunk-16-left-128.int8.onnx",
    },
    "csukuangfj/sherpa-onnx-streaming-zipformer-en-2023-06-21": {
        "tokens": "tokens.txt",
        "encoder_fp32": "encoder-epoch-99-avg-1.onnx",
        "encoder_int8": "encoder-epoch-99-avg-1.int8.onnx",
        "decoder_fp32": "decoder-epoch-99-avg-1.onnx",
        "decoder_int8": "decoder-epoch-99-avg-1.int8.onnx",
        "joiner_fp32": "joiner-epoch-99-avg-1.onnx",
        "joiner_int8": "joiner-epoch-99-avg-1.int8.onnx",
    },
    "csukuangfj/sherpa-onnx-streaming-zipformer-en-2023-02-21": {
        "tokens": "tokens.txt",
        "encoder_fp32": "encoder-epoch-99-avg-1.onnx",
        "encoder_int8": "encoder-epoch-99-avg-1.int8.onnx",
        "decoder_fp32": "decoder-epoch-99-avg-1.onnx",
        "decoder_int8": "decoder-epoch-99-avg-1.int8.onnx",
        "joiner_fp32": "joiner-epoch-99-avg-1.onnx",
        "joiner_int8": "joiner-epoch-99-avg-1.int8.onnx",
    },
    "csukuangfj/sherpa-onnx-streaming-zipformer-bilingual-zh-en-2023-02-20": {
        "tokens": "tokens.txt",
        "encoder_fp32": "encoder-epoch-99-avg-1.onnx",
        "encoder_int8": "encoder-epoch-99-avg-1.int8.onnx",
        "decoder_fp32": "decoder-epoch-99-avg-1.onnx",
        "decoder_int8": "decoder-epoch-99-avg-1.int8.onnx",
        "joiner_fp32": "joiner-epoch-99-avg-1.onnx",
        "joiner_int8": "joiner-epoch-99-avg-1.int8.onnx",
    },
    "shaojieli/sherpa-onnx-streaming-zipformer-fr-2023-04-14": {
        "tokens": "tokens.txt",
        "encoder_fp32": "encoder-epoch-29-avg-9-with-averaged-model.onnx",
        "encoder_int8": "encoder-epoch-29-avg-9-with-averaged-model.int8.onnx",
        "decoder_fp32": "decoder-epoch-29-avg-9-with-averaged-model.onnx",
        "decoder_int8": "decoder-epoch-29-avg-9-with-averaged-model.int8.onnx",
        "joiner_fp32": "joiner-epoch-29-avg-9-with-averaged-model.onnx",
        "joiner_int8": "joiner-epoch-29-avg-9-with-averaged-model.int8.onnx",
    },
    "csukuangfj/sherpa-onnx-streaming-zipformer-small-bilingual-zh-en-2023-02-16": {
        "tokens": "tokens.txt",
        "encoder_fp32": "encoder-epoch-99-avg-1.onnx",
        "encoder_int8": "encoder-epoch-99-avg-1.int8.onnx",
        "decoder_fp32": "decoder-epoch-99-avg-1.onnx",
        "decoder_int8": "decoder-epoch-99-avg-1.int8.onnx",
        "joiner_fp32": "joiner-epoch-99-avg-1.onnx",
        "joiner_int8": "joiner-epoch-99-avg-1.int8.onnx",
    },
    "csukuangfj/sherpa-onnx-streaming-zipformer-zh-14M-2023-02-23": {
        "tokens": "tokens.txt",
        "encoder_fp32": "encoder-epoch-99-avg-1.onnx",
        "encoder_int8": "encoder-epoch-99-avg-1.int8.onnx",
        "decoder_fp32": "decoder-epoch-99-avg-1.onnx",
        "decoder_int8": "decoder-epoch-99-avg-1.int8.onnx",
        "joiner_fp32": "joiner-epoch-99-avg-1.onnx",
        "joiner_int8": "joiner-epoch-99-avg-1.int8.onnx",
    },
    "csukuangfj/sherpa-onnx-streaming-zipformer-en-20M-2023-02-17": {
        "tokens": "tokens.txt",
        "encoder_fp32": "encoder-epoch-99-avg-1.onnx",
        "encoder_int8": "encoder-epoch-99-avg-1.int8.onnx",
        "decoder_fp32": "decoder-epoch-99-avg-1.onnx",
        "decoder_int8": "decoder-epoch-99-avg-1.int8.onnx",
        "joiner_fp32": "joiner-epoch-99-avg-1.onnx",
        "joiner_int8": "joiner-epoch-99-avg-1.int8.onnx",
    },
}

# Audio resampling utility
def resample_audio(audio: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
    return scipy.signal.resample_poly(audio, target_sr, orig_sr)

# Create an online recognizer for a given model and precision
# model_id: full HF repo ID
# precision: "int8" or "fp32"
def create_recognizer(model_id: str, precision: str):
    if model_id not in STREAMING_ZIPFORMER_MODELS:
        raise ValueError(f"Model '{model_id}' is not registered.")
    entry = STREAMING_ZIPFORMER_MODELS[model_id]

    tokens_file = entry['tokens']
    encoder_file = entry['encoder_int8'] if precision == 'int8' else entry['encoder_fp32']
    decoder_file = entry['decoder_fp32']
    joiner_file = entry['joiner_int8'] if precision == 'int8' else entry['joiner_fp32']

    tokens_path = hf_hub_download(repo_id=model_id, filename=tokens_file, cache_dir=str(CACHE_DIR))
    encoder_path = hf_hub_download(repo_id=model_id, filename=encoder_file, cache_dir=str(CACHE_DIR))
    decoder_path = hf_hub_download(repo_id=model_id, filename=decoder_file, cache_dir=str(CACHE_DIR))
    joiner_path = hf_hub_download(repo_id=model_id, filename=joiner_file, cache_dir=str(CACHE_DIR))

    return sherpa_onnx.OnlineRecognizer.from_transducer(
        tokens=tokens_path,
        encoder=encoder_path,
        decoder=decoder_path,
        joiner=joiner_path,
        provider="cpu",
        num_threads=1,
        sample_rate=16000,
        feature_dim=80,
        decoding_method="greedy_search"
    )

def stream_audio(raw_pcm_bytes, stream, recognizer, orig_sr):
    audio = np.frombuffer(raw_pcm_bytes, dtype=np.float32)
    if audio.size == 0:
        return "", 0.0

    resampled = resample_audio(audio, orig_sr, 16000)
    rms = float(np.sqrt(np.mean(resampled ** 2)))

    stream.accept_waveform(16000, resampled)
    if recognizer.is_ready(stream):
        recognizer.decode_streams([stream])
    result = recognizer.get_result(stream)
    return converter.convert(result), rms

def finalize_stream(stream, recognizer):
    tail = np.zeros(int(0.66 * 16000), dtype=np.float32)
    stream.accept_waveform(16000, tail)
    stream.input_finished()
    while recognizer.is_ready(stream):
        recognizer.decode_streams([stream])
    result = recognizer.get_result(stream)
    return converter.convert(result)