Fix Fun-ASR-Nano CTC batch fallback (#3187)

funasr/models/fun_asr_nano/model.py

@@ -796,6 +796,50 @@ class FunASRNano(nn.Module):
             )
         return results, {}
 
+    @staticmethod
+    def _slice_batch_value(value, index):
+        if value is None:
+            return None
+        if isinstance(value, torch.Tensor) and value.ndim > 0 and value.shape[0] > index:
+            return value[index : index + 1]
+        if isinstance(value, list) and len(value) > index:
+            return [value[index]]
+        if isinstance(value, tuple) and len(value) > index:
+            return (value[index],)
+        return value
+
+    @staticmethod
+    def _merge_inference_meta(target, source):
+        for name, value in source.items():
+            if isinstance(value, (int, float)):
+                target[name] = target.get(name, 0.0) + value
+            elif name not in target:
+                target[name] = value
+
+    def _inference_llm_ctc_sequential(
+        self, data_in, data_lengths, key, tokenizer, frontend, **kwargs
+    ):
+        """Run multi-segment input one segment at a time when CTC timestamps are active."""
+        if key is not None and len(key) > 0 and isinstance(key[0], (list, tuple)):
+            key = list(key[0])
+
+        results = []
+        meta_data = {}
+        for i, data_i in enumerate(data_in):
+            key_i = [key[i]] if key is not None and i < len(key) else None
+            data_lengths_i = self._slice_batch_value(data_lengths, i)
+            results_i, meta_i = self.inference_llm(
+                [data_i],
+                data_lengths=data_lengths_i,
+                key=key_i,
+                tokenizer=tokenizer,
+                frontend=frontend,
+                **kwargs,
+            )
+            results.extend(results_i)
+            self._merge_inference_meta(meta_data, meta_i)
+        return results, meta_data
+
     def inference_llm(
         self,
         data_in,
@@ -818,8 +862,12 @@ class FunASRNano(nn.Module):
         # Only batch when CTC timestamps are not needed; the batched path does not
         # produce ctc_timestamps, so fall back to the single-sample path when a CTC
         # decoder is loaded (preserves timestamp behavior).
-        if len(data_in) > 1 and self.ctc_decoder is None:
-            return self._inference_llm_batch(
+        if len(data_in) > 1:
+            if self.ctc_decoder is None:
+                return self._inference_llm_batch(
+                    data_in, data_lengths, key, tokenizer, frontend, **kwargs
+                )
+            return self._inference_llm_ctc_sequential(
                 data_in, data_lengths, key, tokenizer, frontend, **kwargs
             )
         inputs_embeds, contents, batch, source_ids, meta_data = self.inference_prepare(

tests/test_fun_asr_nano_ctc_batch_fallback.py

@@ -0,0 +1,103 @@
+from types import SimpleNamespace
+
+import torch
+
+from funasr.models.fun_asr_nano import model as nano_model
+
+
+class _FakeLLM:
+    def __init__(self):
+        self.config = SimpleNamespace(pad_token_id=None, eos_token_id=0)
+        self.calls = 0
+
+    def to(self, _dtype):
+        return self
+
+    def generate(self, **_kwargs):
+        self.calls += 1
+        return torch.tensor([[10 + self.calls]], dtype=torch.long)
+
+
+class _FakeTokenizer:
+    def batch_decode(self, generated_ids, **_kwargs):
+        return [f"text-{int(generated_ids[0, -1])}"]
+
+
+class _FakeCTCDecoder:
+    def __call__(self, encoder_out, encoder_out_lens):
+        logits = torch.tensor(
+            [[[0.0, 2.0, 0.0], [0.0, 0.0, 3.0]]],
+            dtype=torch.float32,
+        )
+        return logits, encoder_out_lens
+
+
+class _FakeCTC:
+    def log_softmax(self, decoder_out):
+        return decoder_out
+
+
+class _FakeCTCTokenizer:
+    def decode(self, token_ids):
+        return "".join(str(token_id) for token_id in token_ids)
+
+    def encode(self, text):
+        return [1] if text else []
+
+
+def test_ctc_decoder_multi_segment_input_uses_single_segment_fallback(monkeypatch):
+    instance = object.__new__(nano_model.FunASRNano)
+    instance.ctc_decoder = _FakeCTCDecoder()
+    instance.ctc = _FakeCTC()
+    instance.ctc_tokenizer = _FakeCTCTokenizer()
+    instance.blank_id = 0
+    instance.llm = _FakeLLM()
+
+    prepare_calls = []
+
+    def fake_inference_prepare(
+        data_in,
+        data_lengths=None,
+        key=None,
+        tokenizer=None,
+        frontend=None,
+        **_kwargs,
+    ):
+        if len(data_in) > 1:
+            raise NotImplementedError("batch decoding is not implemented")
+        prepare_calls.append((data_in[0], key[0]))
+        return (
+            torch.zeros(1, 2, 4),
+            {"assistant": [f"label-{data_in[0]}"]},
+            {"attention_mask": torch.ones(1, 2, dtype=torch.long)},
+            torch.empty(1, 0, dtype=torch.long),
+            {
+                "encoder_out": torch.zeros(1, 2, 3),
+                "encoder_out_lens": torch.tensor([2]),
+                "batch_data_time": 1.0,
+            },
+        )
+
+    monkeypatch.setattr(instance, "inference_prepare", fake_inference_prepare)
+    monkeypatch.setattr(
+        nano_model,
+        "forced_align",
+        lambda _logits, target_ids, _blank_id: [
+            {"token": int(target_ids[0]), "start_time": 0.0, "end_time": 1.0}
+        ],
+    )
+
+    results, meta = instance.inference_llm(
+        ["seg-a", "seg-b"],
+        key=["key-a", "key-b"],
+        tokenizer=_FakeTokenizer(),
+        frontend=None,
+        device="cpu",
+        llm_dtype="fp32",
+    )
+
+    assert prepare_calls == [("seg-a", "key-a"), ("seg-b", "key-b")]
+    assert [result["key"] for result in results] == ["key-a", "key-b"]
+    assert [result["text"] for result in results] == ["text-11", "text-12"]
+    assert all("timestamps" in result for result in results)
+    assert meta["batch_data_time"] == 2.0