Cohere ASR chunk reassembly can crash on empty chunks and accept inconsistent metadata

[AliOsm]

AI-generated implementation disclosure: The proposed implementation was developed with AI assistance after three days of implementation, testing, benchmarking, and human review in our production Cohere Arabic batch-inference project. Any eventual PR will repeat this disclosure at the beginning. I am opening this issue first to follow the repository's agentic-contribution policy and will wait for maintainer approval before opening a PR.

Summary

CohereAsrProcessor._reassemble_chunk_texts() assumes that decoded texts and audio_chunk_index are complete and internally consistent. On current main, an all-empty chunk group from a supported processor path raises IndexError. Inconsistent metadata can also silently drop, overwrite, or synthesize outputs.

Would you accept a scoped hardening change that validates the reassembly metadata before allocating outputs and preserves the exact output of every valid input?

Reachable empty-chunk failure

Using the current Cohere processor with a valid 40-second waveform produces two rows for one sample:

audio_chunk_index = [(0, 0), (0, 1)]

If both generated rows decode to empty strings, which is possible for EOS-only output, the current helper raises:

CohereAsrProcessor._reassemble_chunk_texts(["", ""], audio_chunk_index)
# IndexError: list index out of range

This is a normal processor metadata shape, not malformed test-only input.

Other current behavior

Input inconsistency	Current result
Text/metadata length mismatch	Extra rows are silently ignored or represented as empty output
Duplicate unchunked row	Later text silently overwrites earlier text
Duplicate chunk index	Both texts are silently concatenated
Same sample represented directly and as chunks	One representation is silently discarded
Missing sample index	An empty sample is synthesized
Missing chunk index	Noncontiguous chunks are joined
Negative sample index	Python negative indexing can write another sample
Very large sparse sample index	Output allocation is based on the largest index

Proposed behavior

• Require decoded text and metadata counts to match.
• Return [] for empty inputs and [""] for an all-empty chunk group.
• Require every metadata entry to be a (sample_idx, chunk_idx) pair.
• Require non-negative integer-like indices and reject booleans.
• Reject duplicate direct rows, duplicate chunks, and mixed direct/chunked representations.
• Require sample and per-sample chunk indices to be contiguous from zero.
• Validate contiguity before output allocation, so sparse hostile indices cannot request a giant list.
• Preserve existing valid whitespace and separator behavior, including no-space Chinese/Japanese assembly.
• Keep the helper signature and all public APIs unchanged.

The exact upstream implementation is ready for review before PR creation:

• Branch: AliOsm:harden-cohere-asr-chunk-reassembly
• Commit: 05eb46aade
• Production implementation: strict reassembly in our batch-inference script

Compatibility evidence

The proposed helper matched current main exactly across 10,000 deterministic randomized valid metadata maps covering:

• 1-32 samples and 1-8 chunks per sample
• mixed direct and chunked samples
• shuffled processor rows
• empty and whitespace-only chunks
• separators " ", "", and "|"

All valid cases handled by the old implementation produced byte-identical outputs. All-empty chunk groups were excluded from parity comparison because current main crashes on them; the proposed output is an empty string for that sample.

One deliberate limitation remains: with the existing arguments, the helper cannot infer a completely absent trailing sample or trailing chunk if both its decoded text and metadata row are missing. The proposed validation detects count mismatches and leading/internal gaps without claiming to detect information that was never passed.

Performance

Pure-Python microbenchmark on an AMD Ryzen 5 5600X with Python 3.12.12, one pinned CPU core, and 30 alternating trials:

Workload	Current median	Strict median	Added cost
24 direct rows	1.662 us	4.632 us	2.970 us
24 samples / 48 chunks	16.983 us	28.189 us	11.206 us
500 samples / 1,000 chunks	340.569 us	561.698 us	221.129 us

The largest measured addition is 0.221 ms for 1,000 decoded rows. This postprocessing cost is negligible beside token decoding and ASR inference; no timing assertions are added to the test suite.

Validation completed

PYTHONPATH=src pytest -q tests/models/cohere_asr/test_processing_cohere_asr.py
  7 passed

PYTHONPATH=src pytest -q   tests/models/cohere_asr/test_modeling_cohere_asr.py   tests/models/cohere_asr/test_processing_cohere_asr.py
  132 passed, 131 skipped

make style
  4/4 checks passed

make check-repo
  24/24 checks passed

ty check tests/models/cohere_asr/test_processing_cohere_asr.py
  All checks passed

The focused tests cover valid ordering and whitespace parity, no-space separators, empty inputs, all-empty chunks, both length-mismatch directions, malformed pairs, invalid index types, negative indices, duplicates, mixed representations, missing indices, and trillion-scale sparse indices.

Duplicate-work audit

I searched open and closed issues/PRs for _reassemble_chunk_texts, audio_chunk_index, Cohere chunk reassembly, long-form decoding, and empty/duplicate/missing chunks. I also checked all currently open PR file lists; none modifies src/transformers/models/cohere_asr/processing_cohere_asr.py.

Adjacent work is unrelated:

• Add cohere asr #45023 introduced Cohere ASR and the current helper.
• 🚨 Modularize ProcessorMixin into smaller components #45493 refactored processor plumbing without changing reassembly behavior.
• [fix] fix requirements audio feature and proc #47113 adds feature-extractor dependency guards.
• fix: apply channel averaging correctly in audio feature extractors #45549 fixes feature-extractor channel averaging.
• Draft 🚨🚧 FeatureExtractor → AudioProcessor #44394 is a broad audio processor migration and does not touch this helper.
• Add repeated n-gram stopping criterion #47271 adds a general repeated n-gram stopping criterion.

May I open the PR?

cc @eustlb @vasqu