funasr_onnx ContextualParaformer / SeacoParaformer return results only for the first batch of a file list #3837

import json
import wave
from pathlib import Path

import numpy as np
import torch

work = Path("hotword_repro")
work.mkdir(exist_ok=True)
tokens = ["<blank>", "<s>", "</s>", "你", "好", "魔", "搭", "<unk>"]


class Embedder(torch.nn.Module):  # hotword [N, 10] -> hw_embed [10, N, 512]
    def forward(self, hotword):
        weight = torch.ones(len(tokens), 512)
        return torch.nn.functional.embedding(hotword.long(), weight).transpose(0, 1)


class Backbone(torch.nn.Module):  # bias_embed [B, N, 512]; every frame predicts "你"
    def forward(self, speech, speech_lengths, bias_embed):
        bias = torch.tensor([0.0, 0.0, 0.0, 5.0, 0.0, 0.0, 0.0, 0.0])
        logits = speech[..., :1] * 0.0 + bias_embed.sum((1, 2), keepdim=True) * 0.0 + bias
        return logits, speech_lengths // 4


hotword = torch.zeros(3, 10, dtype=torch.int32)
speech, lengths = torch.zeros(2, 12, 4), torch.tensor([12, 8], dtype=torch.int32)
torch.onnx.export(
    Embedder(), (hotword,), str(work / "model_eb.onnx"), dynamo=False,
    input_names=["hotword"], output_names=["hw_embed"],
    dynamic_axes={"hotword": {0: "N"}, "hw_embed": {1: "N"}},
)
torch.onnx.export(
    Backbone(), (speech, lengths, torch.zeros(2, 3, 512)), str(work / "model.onnx"), dynamo=False,
    input_names=["speech", "speech_lengths", "bias_embed"], output_names=["logits", "token_num"],
    dynamic_axes={"speech": {0: "B", 1: "T"}, "speech_lengths": {0: "B"},
                  "bias_embed": {0: "B", 1: "N"}, "logits": {0: "B", 1: "T"}, "token_num": {0: "B"}},
)
(work / "tokens.json").write_text(json.dumps(tokens))
(work / "config.yaml").write_text(json.dumps({
    "frontend_conf": {"n_mels": 4, "lfr_m": 1, "lfr_n": 1, "dither": 0.0},
    "model_conf": {"predictor_bias": 1},
}))
(work / "am.mvn").write_text(
    "<AddShift> 4 4\n<LearnRateCoef> 1 [ 0 0 0 0 ]\n<Rescale> 4 4\n<LearnRateCoef> 1 [ 1 1 1 1 ]\n"
)
wavs = []
for index, samples in enumerate([4000, 8000, 12000]):  # 0.25 s, 0.5 s, 0.75 s
    tone = (np.sin(np.arange(samples) * 2 * np.pi * 440 / 16000) * 3000).astype("<i2")
    with wave.open(str(work / f"{index}.wav"), "wb") as output:
        output.setnchannels(1)
        output.setsampwidth(2)
        output.setframerate(16000)
        output.writeframes(tone.tobytes())
    wavs.append(str(work / f"{index}.wav"))

from funasr_onnx import ContextualParaformer

model = ContextualParaformer(str(work), batch_size=1)
print([model(wav, "魔搭 你好")[0]["preds"][0] for wav in wavs])  # one call per file
print([row["preds"][0] for row in model(wavs, "魔搭 你好")])  # one call with all three files