refactor(aero_realtime): add speech span labels

Author: kcz358

src/lmms_engine/datasets/processor/aero_realtime_processor.py

@@ -20,12 +20,9 @@
 (assistant text segments are placed at specific temporal positions during
 video playback).
 
-The processor builds ``text_stream_ids`` with the delay mechanism:
-- ``<|rt_start|>`` at the first audio position
-- ``<|rt_pad|>`` for audio silence positions before the delay boundary
-- ``<|rt_speak|>`` at the delay boundary
-- After ``<|rt_speak|>``: ``<|rt_pad|>`` for normal QA, or actual text tokens
-  at the appropriate audio positions for realtime data
+The processor builds ``text_stream_ids`` on the audio timeline. ``<|rt_pad|>``
+is silence context only; labels supervise ``<|rt_speak|>``, speech span
+boundaries (``<|rt_start|>`` / ``<|rt_end|>``), and speech text tokens.
 """
 
 from typing import Dict, List, Optional
@@ -50,9 +47,9 @@ class AeroRealtimeDataProcessor(Qwen3_VLDataProcessor):
     """Data processor for AeroRealtime training.
 
     Builds ``input_ids``, ``text_stream_ids``, and ``labels`` for the
-    dual-stream training design.  Handles:
-    - Normal video QA: audio timeline filled with ``<|rt_pad|>`` after delay
-    - Realtime training: text tokens placed at temporal positions on audio tokens
+    realtime audio-stream training design.  Handles:
+    - Normal video QA: audio timeline filled with ``<|rt_pad|>`` context
+    - Realtime training: boundary and text labels on audio tokens
     - Image-only: standard scatter (no text_stream_ids)
     - Audio extraction from video for audio-vision fusion
     """
@@ -135,6 +132,10 @@ def rt_pad_id(self):
     def rt_speak_id(self):
         return self.tokenizer.convert_tokens_to_ids(self.processor.rt_speak_token)
 
+    @property
+    def rt_end_id(self):
+        return self.tokenizer.convert_tokens_to_ids(self.processor.rt_end_token)
+
     # ------------------------------------------------------------------
     # Main process entry point
     # ------------------------------------------------------------------
@@ -334,10 +335,9 @@ def _build_normal_qa_ids_and_labels(
         """Build input_ids, text_stream_ids, and labels from HF messages.
 
         For normal video QA the text_stream_ids only differ from input_ids
-        on audio pad positions:
-            - all ``<|audio_pad|>`` slots -> ``<|rt_pad|>``
-            - first ``<|audio_pad|>`` -> ``<|rt_start|>``
-            - delayed ``<|audio_pad|>`` -> ``<|rt_speak|>``
+        on audio pad positions, where all ``<|audio_pad|>`` slots become
+        ``<|rt_pad|>`` context. Normal QA keeps standard assistant labels;
+        realtime span labels are built by ``_build_realtime_ids_and_labels``.
 
         Video placeholders and envelope boundary tokens keep their original
         ids; vision features replace video placeholder embeddings in the model.
@@ -363,20 +363,15 @@ def _build_normal_qa_ids_and_labels(
         text_stream_id = list(input_id)  # start as a copy of input_ids
 
         if has_video and has_audio:
-            # video + audio: only audio pads carry realtime stream tokens
+            # video + audio: only audio pads carry realtime stream context
             self.processor._fill_text_stream_video_audio(
                 stream=text_stream_id,
                 video_grid_thw=video_grid_thw,
                 video_metadata=video_metadata,
                 temporal_patch_size=getattr(self.processor.video_processor, "temporal_patch_size", 2),
-                vision_start_id=self.tokenizer.convert_tokens_to_ids(self.processor.vision_start_token),
-                vision_end_id=self.tokenizer.convert_tokens_to_ids(self.processor.vision_end_token),
                 audio_start_id=self.tokenizer.convert_tokens_to_ids(self.processor.audio_start_token),
                 audio_end_id=self.tokenizer.convert_tokens_to_ids(self.processor.audio_end_token),
-                audio_pad_id=self.audio_token_id,
-                rt_start_id=self.rt_start_id,
                 rt_pad_id=self.rt_pad_id,
-                rt_speak_id=self.rt_speak_id,
             )
         elif has_audio:
             # audio-only: single envelope per audio sample
@@ -454,14 +449,9 @@ def _build_realtime_ids_and_labels(
             video_grid_thw=video_grid_thw,
             video_metadata=video_metadata,
             temporal_patch_size=getattr(self.processor.video_processor, "temporal_patch_size", 2),
-            vision_start_id=self.tokenizer.convert_tokens_to_ids(self.processor.vision_start_token),
-            vision_end_id=self.tokenizer.convert_tokens_to_ids(self.processor.vision_end_token),
             audio_start_id=self.tokenizer.convert_tokens_to_ids(self.processor.audio_start_token),
             audio_end_id=self.tokenizer.convert_tokens_to_ids(self.processor.audio_end_token),
-            audio_pad_id=self.audio_token_id,
-            rt_start_id=self.rt_start_id,
             rt_pad_id=self.rt_pad_id,
-            rt_speak_id=self.rt_speak_id,
         )
 
         audio_positions = [idx for idx, token_id in enumerate(input_id) if token_id == self.audio_token_id]
@@ -474,29 +464,36 @@ def _build_realtime_ids_and_labels(
             raise ValueError(f"Audio position/time mismatch: {len(audio_positions)} != {len(audio_times)}")
 
         delay = getattr(self.processor, "delay_seconds", 2.0)
-        for pos, t_sec in zip(audio_positions, audio_times):
-            if t_sec >= delay:
-                target[pos] = self.rt_pad_id
+        speak_audio_idx = self._first_index_at_or_after(audio_times, delay)
+        if speak_audio_idx < len(audio_positions):
+            speak_pos = audio_positions[speak_audio_idx]
+            text_stream_id[speak_pos] = self.rt_speak_id
 
         assistant_segments = sorted(
             [seg for seg in (realtime_segments or []) if seg.get("role") == "assistant" and seg.get("text")],
             key=lambda item: float(item["time"]),
         )
-        event_times = sorted(float(seg["time"]) for seg in (realtime_segments or []))
+        occupied_audio_indices = {speak_audio_idx} if speak_audio_idx < len(audio_positions) else set()
         for segment in assistant_segments:
             start_time = float(segment["time"])
-            end_time = self._next_time_after(event_times, start_time)
             start_audio_idx = self._first_index_at_or_after(audio_times, start_time)
-            end_audio_idx = (
-                self._first_index_at_or_after(audio_times, end_time) if end_time is not None else len(audio_positions)
+            if speak_audio_idx < len(audio_positions):
+                start_audio_idx = max(start_audio_idx, speak_audio_idx + 1)
+            available_indices = self._next_available_indices(
+                start=start_audio_idx,
+                count=len(audio_positions),
+                limit=len(audio_positions),
+                occupied=occupied_audio_indices,
             )
-            if start_audio_idx < end_audio_idx and text_stream_id[audio_positions[start_audio_idx]] == self.rt_speak_id:
-                start_audio_idx += 1
-            token_ids = self._encode_realtime_text(segment["text"])
-            for offset, token_id in enumerate(token_ids[: max(0, end_audio_idx - start_audio_idx)]):
-                pos = audio_positions[start_audio_idx + offset]
+            if len(available_indices) < 2:
+                continue
+            text_token_budget = len(available_indices) - 2
+            token_ids = [self.rt_start_id] + self._encode_realtime_text(segment["text"])[:text_token_budget] + [self.rt_end_id]
+            for audio_idx, token_id in zip(available_indices, token_ids):
+                pos = audio_positions[audio_idx]
                 text_stream_id[pos] = token_id
                 target[pos] = token_id
+                occupied_audio_indices.add(audio_idx)
 
         input_tensor = torch.tensor(input_id, dtype=torch.long)
         text_stream_tensor = torch.tensor(text_stream_id, dtype=torch.long)
@@ -579,11 +576,14 @@ def _first_index_at_or_after(values: List[float], target: float) -> int:
         return len(values)
 
     @staticmethod
-    def _next_time_after(values: List[float], target: float) -> Optional[float]:
-        for value in values:
-            if value > target:
-                return value
-        return None
+    def _next_available_indices(start: int, count: int, limit: int, occupied: set) -> List[int]:
+        indices = []
+        idx = start
+        while idx < limit and len(indices) < count:
+            if idx not in occupied:
+                indices.append(idx)
+            idx += 1
+        return indices
 
     def get_qwen_template_labels(
         self,
@@ -682,12 +682,10 @@ def _expand_encode_id_video_tokens(
 
         - Without audio: per-frame Qwen3VL legacy expansion (delegated to
           parent).
-        - With audio: per-chunk envelope expansion matching the model
-          processor's path 5b layout::
+        - With audio: per-chunk separated vision/audio envelopes::
 
-            <t.t seconds><|vision_start|><|audio_start|>
-              <|video_pad|>×spatial <|audio_pad|>×N_t
-            <|audio_end|><|vision_end|>
+            <t.t seconds><|vision_start|><|video_pad|>×spatial<|vision_end|>
+            <|audio_start|><|audio_pad|>×N_t<|audio_end|>
         """
         if audio_per_chunk_per_video is None:
             return super()._expand_encode_id_video_tokens(
@@ -740,11 +738,11 @@ def _expand_encode_id_video_tokens(
                 n_audio_t = audio_per_chunk[t]
                 expanded_encode_id.extend(timestamp_token_ids)
                 expanded_encode_id.append(vision_start_id)
-                expanded_encode_id.append(audio_start_id)
                 expanded_encode_id.extend([self.video_token_id] * spatial)
+                expanded_encode_id.append(vision_end_id)
+                expanded_encode_id.append(audio_start_id)
                 expanded_encode_id.extend([self.audio_token_id] * n_audio_t)
                 expanded_encode_id.append(audio_end_id)
-                expanded_encode_id.append(vision_end_id)
 
             prev = pos + 2  # skip past original <|vision_end|>
 

src/lmms_engine/models/aero_realtime/configuration_aero_realtime.py

@@ -79,6 +79,8 @@ class AeroRealtimeConfig(PretrainedConfig):
         rt_speak_token_index (`int`, *optional*, defaults to `151674`):
             Token index for ``<|rt_speak|>`` — delay boundary marker after which
             the model may begin producing text.
+        rt_end_token_index (`int`, *optional*, defaults to `151675`):
+            Token index for ``<|rt_end|>`` — closes one realtime speech span.
         delay_seconds (`float`, *optional*, defaults to `2.0`):
             Delay in seconds before the model is allowed to speak. Converted to
             a number of vision tokens based on the video's temporal resolution.
@@ -122,6 +124,7 @@ def __init__(
         rt_start_token_index=151672,
         rt_pad_token_index=151673,
         rt_speak_token_index=151674,
+        rt_end_token_index=151675,
         delay_seconds=2.0,
         tie_word_embeddings=False,
         **kwargs,
@@ -138,6 +141,7 @@ def __init__(
         self.rt_start_token_index = rt_start_token_index
         self.rt_pad_token_index = rt_pad_token_index
         self.rt_speak_token_index = rt_speak_token_index
+        self.rt_end_token_index = rt_end_token_index
         self.delay_seconds = delay_seconds
 
         # Aliases expected by qwen3_vl_get_rope_index (shared RoPE helper)

src/lmms_engine/models/aero_realtime/modeling_aero_realtime.py

@@ -516,8 +516,8 @@ def forward(
         """Forward pass for AeroRealtime.
 
         Audio and video are kept as **separate** token streams in the input
-        sequence (per-chunk envelope ``[VS][AS][video_pad×S][audio_pad×N]
-        [AE][VE]``) so time alignment is expressed entirely through token
+        sequence (per-chunk envelope ``[VS][video_pad×S][VE][AS]
+        [audio_pad×N][AE]``) so time alignment is expressed entirely through token
         order and RoPE.  Vision features replace vision placeholders; audio
         features are added to the realtime text stream on audio placeholders.
 
@@ -536,9 +536,9 @@ def forward(
             ``audio_token_index`` positions.
 
         **Video + Audio**: video placeholders receive pure vision features.
-        ``text_stream_ids`` carries realtime markers (``<|rt_start|>``,
-        ``<|rt_pad|>``, ``<|rt_speak|>``) only at audio positions, where audio
-        features are added to the realtime text embeddings.
+        ``text_stream_ids`` carries realtime markers (``<|rt_speak|>``,
+        ``<|rt_start|>``, ``<|rt_end|>``, and speech text) only at audio
+        positions, where audio features are added to the realtime text embeddings.
 
         Pipeline:
             1. Embed ``text_stream_ids`` (if provided) or ``input_ids``.
@@ -553,8 +553,8 @@ def forward(
                 position masks for image/video/audio features.
             text_stream_ids: Parallel text-stream token ids.
                 Shape ``[batch_size, seq_len]``.  At audio positions contains
-                ``<|rt_start|>``, ``<|rt_pad|>``, ``<|rt_speak|>``, or actual
-                text tokens; mirrors ``input_ids`` elsewhere.
+                ``<|rt_pad|>``, ``<|rt_speak|>``, speech boundary tokens, or
+                actual text tokens; mirrors ``input_ids`` elsewhere.
                 If not provided, falls back to ``input_ids``.
             pixel_values: Image pixel values (flat across batch).
             image_grid_thw: Grid info per image. ``[num_images, 3]``.