diff --git a/funasr/models/glm_asr/inference_vllm.py b/funasr/models/glm_asr/inference_vllm.py index e2aaf2c27..a8d4564fd 100644 --- a/funasr/models/glm_asr/inference_vllm.py +++ b/funasr/models/glm_asr/inference_vllm.py @@ -110,6 +110,63 @@ def prepare_glmasr_vllm_dir(model_dir: str) -> str: return output_dir +# Warn only once per process so batch loops do not spam the log. +_warned_dup_keys = False + + +def _dedup_keys(keys): + """Make result keys unique while preserving order and first-occurrence names. + + Each result key is derived from the audio file basename + (``os.path.splitext(os.path.basename(path))[0]``), so two inputs that live + in different directories but share a basename -- e.g. ``spk1/segment.wav`` + and ``spk2/segment.wav`` -- both map to ``"segment"``. A downstream + ``{r["key"]: r["text"]}`` mapping (the canonical FunASR result shape) would + then silently drop all but the last colliding entry, returning fewer + transcripts than inputs with no error. Appending a deterministic ``_N`` + suffix to later collisions keeps every transcript addressable. + + Args: + keys: Result keys in input order. + + Returns: + A new list of unique keys, same length and order as ``keys``. The first + occurrence of each key is preserved unchanged; the n-th repeat becomes + ``"_"`` (e.g. ``"seg"`` -> ``"seg"``, ``"seg_1"``, ``"seg_2"``). + """ + global _warned_dup_keys + + seen = set() + out = [] + collided = False + for key in keys: + if key not in seen: + seen.add(key) + out.append(key) + continue + # Find the first free "_" so a suffixed key cannot itself clash + # with an existing one (e.g. inputs "seg", "seg_1", "seg"). + collided = True + n = 1 + candidate = f"{key}_{n}" + while candidate in seen: + n += 1 + candidate = f"{key}_{n}" + seen.add(candidate) + out.append(candidate) + + if collided and not _warned_dup_keys: + logger.warning( + "Duplicate result keys from audio basenames were made unique with " + "'_N' suffixes (e.g. two files named 'segment.wav' in different " + "directories map to the same key); pass distinct filenames if you " + "rely on the basename as the result key." + ) + _warned_dup_keys = True + + return out + + class GLMASRVLLMEngine: """GLM-ASR with vLLM backend. @@ -263,13 +320,18 @@ class GLMASRVLLMEngine: t2 = time.perf_counter() logger.info(f"vLLM generation: {t2-t1:.3f}s") + raw_keys = [ + os.path.splitext(os.path.basename(x))[0] if isinstance(x, str) else f"sample_{i}" + for i, x in enumerate(inputs) + ] + keys = _dedup_keys(raw_keys) + results = [] for i, output in enumerate(outputs): token_ids = list(output.outputs[0].token_ids) text = self.tokenizer.decode(token_ids, skip_special_tokens=True) text = re.sub(r'\s+', ' ', text).strip() - key = os.path.splitext(os.path.basename(inputs[i]))[0] if isinstance(inputs[i], str) else f"sample_{i}" - results.append({"key": key, "text": text}) + results.append({"key": keys[i], "text": text}) return results diff --git a/tests/test_glm_asr_vllm_dedup_keys.py b/tests/test_glm_asr_vllm_dedup_keys.py new file mode 100644 index 000000000..0cd403a3e --- /dev/null +++ b/tests/test_glm_asr_vllm_dedup_keys.py @@ -0,0 +1,124 @@ +"""Unit tests for GLM-ASR vLLM result-key deduplication. + +These tests exercise ``GLMASRVLLMEngine.generate`` without a GPU or a real +vLLM installation: the vLLM entry points are stubbed in ``sys.modules`` and the +audio/encoder/engine collaborators are mocked, so only the result-key wiring is +under test. The pure ``_dedup_keys`` helper is also tested directly. +""" + +import sys +import types +import unittest +from unittest import mock + + +def _install_vllm_stub(): + """Install a minimal ``vllm`` stub so ``inference_vllm`` imports without GPU.""" + + class _SamplingParams: + def __init__(self, **kwargs): + self.kwargs = kwargs + + class _EmbedsPrompt: + def __init__(self, **kwargs): + self.kwargs = kwargs + + vllm_mod = types.ModuleType("vllm") + vllm_mod.SamplingParams = _SamplingParams + vllm_mod.LLM = object + inputs_mod = types.ModuleType("vllm.inputs") + inputs_mod.EmbedsPrompt = _EmbedsPrompt + data_mod = types.ModuleType("vllm.inputs.data") + data_mod.EmbedsPrompt = _EmbedsPrompt + + sys.modules["vllm"] = vllm_mod + sys.modules["vllm.inputs"] = inputs_mod + sys.modules["vllm.inputs.data"] = data_mod + + +class DedupKeysHelperTest(unittest.TestCase): + def setUp(self): + _install_vllm_stub() + from funasr.models.glm_asr import inference_vllm + + self.inference_vllm = inference_vllm + # Reset the warn-once flag so warning-related assertions are deterministic. + inference_vllm._warned_dup_keys = False + + def test_collision_free_input_unchanged(self): + keys = ["a", "b", "c"] + self.assertEqual(self.inference_vllm._dedup_keys(keys), ["a", "b", "c"]) + + def test_repeated_keys_get_deterministic_suffix(self): + keys = ["seg", "seg", "other", "seg"] + self.assertEqual( + self.inference_vllm._dedup_keys(keys), + ["seg", "seg_1", "other", "seg_2"], + ) + + def test_suffix_does_not_clash_with_existing_key(self): + # A naive "seg" -> "seg_1" scheme would re-collide with the literal + # "seg_1" input; the result must stay globally unique. + keys = ["seg", "seg_1", "seg"] + result = self.inference_vllm._dedup_keys(keys) + self.assertEqual(result, ["seg", "seg_1", "seg_2"]) + self.assertEqual(len(set(result)), len(result)) + + def test_empty_input(self): + self.assertEqual(self.inference_vllm._dedup_keys([]), []) + + def test_sentinel_keys_preserved(self): + keys = ["sample_0", "sample_1", "sample_0"] + self.assertEqual( + self.inference_vllm._dedup_keys(keys), + ["sample_0", "sample_1", "sample_0_1"], + ) + + def test_does_not_mutate_input(self): + keys = ["seg", "seg"] + self.inference_vllm._dedup_keys(keys) + self.assertEqual(keys, ["seg", "seg"]) + + +class DedupResultKeysTest(unittest.TestCase): + def setUp(self): + _install_vllm_stub() + from funasr.models.glm_asr.inference_vllm import GLMASRVLLMEngine + + # Build an engine without running __init__ (no model load / GPU needed). + engine = GLMASRVLLMEngine.__new__(GLMASRVLLMEngine) + engine.device = "cpu" + engine._encode_audio = mock.Mock(return_value="audio_embeds") + engine._build_prompt_embeds = mock.Mock( + return_value=mock.Mock(float=lambda: "embeds") + ) + engine.tokenizer = mock.Mock() + engine.tokenizer.decode = mock.Mock(return_value="hello world") + self.engine = engine + + def _set_outputs(self, n): + token_out = types.SimpleNamespace(token_ids=[1, 2, 3]) + outs = [types.SimpleNamespace(outputs=[token_out]) for _ in range(n)] + self.engine.vllm_engine = mock.Mock() + self.engine.vllm_engine.generate = mock.Mock(return_value=outs) + + def test_duplicate_basenames_get_unique_keys(self): + self._set_outputs(2) + results = self.engine.generate(["spk1/segment.wav", "spk2/segment.wav"]) + self.assertEqual([r["key"] for r in results], ["segment", "segment_1"]) + # No transcript is dropped when results are folded into a {key: text} dict. + self.assertEqual(len({r["key"] for r in results}), len(results)) + + def test_distinct_basenames_unchanged(self): + self._set_outputs(2) + results = self.engine.generate(["a.wav", "b.wav"]) + self.assertEqual([r["key"] for r in results], ["a", "b"]) + + def test_single_input_key(self): + self._set_outputs(1) + results = self.engine.generate("only.wav") + self.assertEqual(results, [{"key": "only", "text": "hello world"}]) + + +if __name__ == "__main__": + unittest.main()