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
from typing import List
import tempfile

# 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 = {
    # bilingual zh-en with char+BPE
    "csukuangfj/k2fsa-zipformer-bilingual-zh-en-t": {
        "tokens": "data/lang_char_bpe/tokens.txt",
        "encoder_fp32": "exp/96/encoder-epoch-99-avg-1.onnx",
        "encoder_int8": "exp/96/encoder-epoch-99-avg-1.int8.onnx",
        "decoder_fp32": "exp/96/decoder-epoch-99-avg-1.onnx",
        "decoder_int8": "exp/96/decoder-epoch-99-avg-1.int8.onnx",
        "joiner_fp32": "exp/96/joiner-epoch-99-avg-1.onnx",
        "joiner_int8": "exp/96/joiner-epoch-99-avg-1.int8.onnx",
        "modeling_unit":"cjkchar+bpe",
        "bpe_vocab":   "data/lang_char_bpe/bpe.vocab",
    },
    # mixed Chinese+English (char+BPE)
    "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",
        "modeling_unit":"cjkchar+bpe",
        "bpe_vocab":   "data/lang_char_bpe/bpe.vocab",
    },
    # Korean-only (CJK chars)
    "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",
        "modeling_unit":"cjkchar",
        "bpe_vocab":   None,
    },
    # multi Chinese (Hans) (CJK chars)
    "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",
        "modeling_unit":"cjkchar",
        "bpe_vocab":   None,
    },
    # wenetspeech streaming (CJK chars)
    "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",
        "modeling_unit":"cjkchar",
        "bpe_vocab":   None,
    },
    # English-only (BPE)
    "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",
        "modeling_unit":"bpe",
        "bpe_vocab":   None,
    },
    "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",
        "modeling_unit":"bpe",
        "bpe_vocab":   None,
    },
    "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",
        "modeling_unit":"bpe",
        "bpe_vocab":   None,
    },
    # older bilingual zh-en (cjkchar+BPE) – no bpe.vocab shipped
    "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",
        "modeling_unit":"cjkchar+bpe",
        "bpe_vocab":   None,
    },
    # French-only (BPE)
    "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",
        "modeling_unit":"bpe",
        "bpe_vocab":   None,
    },
    # Chinese-only small (CJK chars)
    "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",
        "modeling_unit":"cjkchar",
        "bpe_vocab":   None,
    },
    # English-only 20M (BPE)
    "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",
        "modeling_unit":"bpe",
        "bpe_vocab":   None,
    },
}

# 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,
    hotwords: List[str] = None,
    hotwords_score: float = 0.0,
):
    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))

    # ——— Download BPE vocab if this model has one ———
    modeling_unit = entry.get("modeling_unit")
    bpe_rel_path  = entry.get("bpe_vocab")
    bpe_vocab_path = None
    if bpe_rel_path:
        try:
            bpe_vocab_path = hf_hub_download(
                repo_id=model_id,
                filename=bpe_rel_path,
                cache_dir=str(CACHE_DIR),
            )
            print(f"[DEBUG asr_worker] Downloaded bpe_vocab: {bpe_vocab_path}")
        except Exception as e:
            print(f"[WARNING asr_worker] Could not download bpe_vocab '{bpe_rel_path}': {e}")
            bpe_vocab_path = None

    # ——— Decide whether to use beam search with hotword biasing ———
    use_beam = (hotwords and hotwords_score > 0.0) and bpe_vocab_path
    if use_beam:
            # Write hotword list to a temp file (one entry per line)
            tf = tempfile.NamedTemporaryFile(
                mode="w", delete=False, suffix=".txt", dir=str(CACHE_DIR)
            )
            for w in hotwords:
                tf.write(f"{w}\n")
            tf.flush()
            tf.close()
            hotwords_file_path = tf.name
            print(f"[DEBUG asr_worker] Written {len(hotwords)} hotwords to {hotwords_file_path} with score {hotwords_score}")

            # Create beam-search recognizer with biasing :contentReference[oaicite:0]{index=0}
            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="modified_beam_search",
                hotwords_file=hotwords_file_path,
                hotwords_score=hotwords_score,
                modeling_unit=modeling_unit,
                bpe_vocab=bpe_vocab_path,
            )

    # ——— Fallback to original greedy-search (no hotword biasing) ———
    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