Fix batched inputs longer than target_length (clips mixed) and keep the last full segment
For inputs longer than target_length (1008 frames in the released configs, ~10.1 s), AudioTransformerMAE_Encoder.forward splits the spectrogram into segments and encodes all of them as one batch. This PR fixes two problems in that path.
1. Clips in a batch get mixed. torch.flatten(splits, 0, 1) orders the first axis as (segment, batch), but torch.reshape(x, (x.shape[0] // n_splits, -1, C)) reads it back as (batch, segment). With batch size > 1, each output row contains segments of other clips. Batch size 1 is not affected, which is why single-clip checks look fine. The fix reshapes as (segment, batch, time, C), swaps the first two axes and joins segments along time. The pip dasheng package encodes each segment separately and concatenates along time, so it already behaves this way.
2. The last segment is dropped when the length is an exact multiple of target_length. The else branch stacks splits[:-1] even though the last split is full, so e.g. 2 x 1008 input frames give 252 output frames instead of 504. This happens at any batch size. The fix stacks all splits.
This is the cause of RicherMans/Dasheng#16 (VL-top10, 18.64 s clips, batch 4, not reproducible with the HF checkpoint but reproducible with the pip package). Details: https://github.com/RicherMans/Dasheng/issues/16#issuecomment-6381097475 . Downstream effect on the XARES-LLM baseline: https://github.com/xiaomi-research/xares-llm/issues/7 (task1 0.628 -> 0.719, task2 0.246 -> 0.474 once the audio is chunked before the model).
Checks. Real checkpoints, fp32 on GPU, same patch for dasheng-base / dasheng-0.6B / dasheng-1.2B:
| check | base | 0.6B | 1.2B |
|---|---|---|---|
| batch of 1, 25 s: patched vs current, max abs diff | 0 | 0 | 0 |
| batch of 2, 25 s: each row vs the clip alone, max abs diff (current code in brackets) | 0 (11.6) | 0 (9.2) | 4e-5 (17.2) |
| batch of 2, 25 s: vs pip-style per-segment loop | 7e-6 | 2e-5 | 0 |
| exactly 2 x 1008 frames: output frames (current code in brackets) | 504 (252) | 504 (252) | 504 (252) |
| 5 s input: patched vs current | 0 | 0 | 0 |
Output scale is 5-15. Inputs of at most 1008 frames are unchanged. Inputs with batch size 1 are also unchanged, except when the length is an exact multiple of 1008 frames (problem 2).
(The same change is proposed for mispeech/dasheng-base, mispeech/dasheng-0.6B and mispeech/dasheng-1.2B, which share this code.)