feat(aero_realtime): add realtime stream labels

Author: kcz358

src/lmms_engine/datasets/iterable/aero_realtime_iterable_dataset.py

@@ -68,12 +68,28 @@ def load_from_json(self, data, data_folder=None) -> Dict[str, torch.Tensor]:
         if isinstance(messages, str):
             messages = json.loads(messages)
 
+        is_realtime = bool(data.get("realtime", False))
+
         # First pass: collect media references and realtime segments
         for message in messages:
+            message_time = message.get("time")
+            if is_realtime and message_time is not None and message["role"] in ["user", "assistant"]:
+                text = self._extract_text_content(message.get("content", []))
+                if text:
+                    realtime_segments.append(
+                        {
+                            "time": float(message_time),
+                            "role": message["role"],
+                            "text": text,
+                        }
+                    )
+                continue
+
             for content in message["content"]:
-                if content["type"] == "image_url":
+                content_type = content.get("type")
+                if content_type == "image_url":
                     images_list.append(content["image_url"]["url"])
-                elif content["type"] == "video_url":
+                elif content_type == "video_url":
                     video_url = content["video_url"]["url"]
                     if data_folder is not None:
                         video_path = os.path.join(data_folder, video_url)
@@ -90,10 +106,11 @@ def load_from_json(self, data, data_folder=None) -> Dict[str, torch.Tensor]:
                     kwargs["video_metadata"] = video_metadata
                     kwargs["do_sample_frames"] = False
 
-                elif content["type"] == "realtime_text":
+                elif content_type == "realtime_text":
                     realtime_segments.append(
                         {
-                            "start_sec": content["start_sec"],
+                            "time": content["start_sec"],
+                            "role": "assistant",
                             "text": content["text"],
                         }
                     )
@@ -133,6 +150,16 @@ def load_from_json(self, data, data_folder=None) -> Dict[str, torch.Tensor]:
         )
         return inputs
 
+    @staticmethod
+    def _extract_text_content(content) -> str:
+        if isinstance(content, str):
+            return content
+        texts = []
+        for item in content:
+            if item and item.get("type") == "text" and item.get("text"):
+                texts.append(item["text"])
+        return "\n".join(texts)
+
     def _load_video_with_metadata(
         self,
         video_path: str,

src/lmms_engine/datasets/processor/aero_realtime_processor.py

@@ -249,17 +249,20 @@ def process(
                 metadata = _video_metadata[v_idx]
                 fps = metadata.fps if metadata.fps is not None else 24.0
                 grid_t = int(video_grid_thw[v_idx][0])
-                second_per_grid = temporal_patch_size / fps
+                curr_timestamp = self.processor._calculate_timestamps(
+                    metadata.frames_indices,
+                    fps,
+                    temporal_patch_size,
+                )
                 # Audio sample paired with this video by positional index
                 a_idx = v_idx if v_idx < len(num_audio_tokens_list) else 0
                 n_audio = num_audio_tokens_list[a_idx]
                 audio_duration = self.processor._get_audio_duration_seconds(audio_inputs["audio_attention_mask"][a_idx])
                 audio_rate = (n_audio / audio_duration) if audio_duration > 0 else 0.0
                 audio_per_chunk_per_video.append(
-                    self.processor._split_audio_across_chunks(
+                    self.processor._split_audio_across_chunk_times(
                         n_audio=n_audio,
-                        grid_t=grid_t,
-                        second_per_grid=second_per_grid,
+                        chunk_start_times=curr_timestamp[:grid_t],
                         audio_rate=audio_rate,
                     )
                 )
@@ -280,9 +283,18 @@ def process(
                 add_system_prompt=add_system_prompt,
             )
         else:
-            # TODO:
-            # Build realtime qa ids and labels
-            raise RuntimeError("Not implemented yet")
+            inputs = self._build_realtime_ids_and_labels(
+                hf_messages=hf_messages,
+                num_image_tokens=num_image_tokens,
+                num_video_tokens=num_video_tokens,
+                video_grid_thw=video_grid_thw,
+                video_metadata=_video_metadata,
+                audio_per_chunk_per_video=audio_per_chunk_per_video,
+                audio_attention_mask=audio_inputs.get("audio_attention_mask") if has_audio else None,
+                realtime_segments=realtime_segments,
+                system_message=system_message,
+                add_system_prompt=add_system_prompt,
+            )
 
         # ==============================================================
         # 6. Pack vision/audio tensors into output
@@ -398,6 +410,184 @@ def _build_normal_qa_ids_and_labels(
 
         return result
 
+    def _build_realtime_ids_and_labels(
+        self,
+        hf_messages,
+        num_image_tokens: Optional[List[int]],
+        num_video_tokens: Optional[List[int]],
+        video_grid_thw=None,
+        video_metadata=None,
+        audio_per_chunk_per_video: Optional[List[List[int]]] = None,
+        audio_attention_mask: Optional[torch.Tensor] = None,
+        realtime_segments: Optional[List[Dict]] = None,
+        system_message: str = "You are a helpful assistant",
+        add_system_prompt: bool = True,
+    ) -> dict:
+        if video_grid_thw is None or audio_per_chunk_per_video is None or audio_attention_mask is None:
+            raise ValueError("Realtime training requires both video and audio inputs.")
+
+        base_messages, timed_user_segments = self._build_realtime_base_messages(
+            hf_messages=hf_messages,
+            realtime_segments=realtime_segments or [],
+            video_grid_thw=video_grid_thw,
+            video_metadata=video_metadata,
+            audio_per_chunk_per_video=audio_per_chunk_per_video,
+            system_message=system_message,
+            add_system_prompt=add_system_prompt,
+        )
+
+        results = self.get_qwen_template_labels(
+            base_messages,
+            num_image_tokens,
+            num_video_tokens,
+            video_metadata,
+            video_grid_thw,
+            audio_per_chunk_per_video=audio_per_chunk_per_video,
+            timed_user_segments=timed_user_segments,
+            system_message=system_message,
+            add_system_prompt=False,
+        )
+        input_id = results["input_ids"].tolist()
+        text_stream_id = list(input_id)
+        target = [-100] * len(input_id)
+
+        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),
+            video_pad_id=self.video_token_id,
+            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]
+        audio_times = self._get_audio_position_times(
+            video_grid_thw=video_grid_thw,
+            video_metadata=video_metadata,
+            audio_per_chunk_per_video=audio_per_chunk_per_video,
+        )
+        if len(audio_positions) != len(audio_times):
+            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
+
+        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 []))
+        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 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]
+                text_stream_id[pos] = token_id
+                target[pos] = token_id
+
+        input_tensor = torch.tensor(input_id, dtype=torch.long)
+        text_stream_tensor = torch.tensor(text_stream_id, dtype=torch.long)
+        target_tensor = torch.tensor(target, dtype=torch.long)
+
+        return dict(
+            input_ids=input_tensor,
+            labels=target_tensor,
+            text_stream_ids=text_stream_tensor,
+        )
+
+    def _build_realtime_base_messages(
+        self,
+        hf_messages,
+        realtime_segments: List[Dict],
+        video_grid_thw,
+        video_metadata,
+        audio_per_chunk_per_video: List[List[int]],
+        system_message: str,
+        add_system_prompt: bool,
+    ):
+        messages = []
+        first_content = []
+        timed_user_segments = sorted(
+            [seg for seg in realtime_segments if seg.get("role") == "user" and seg.get("text")],
+            key=lambda item: float(item["time"]),
+        )
+
+        if add_system_prompt and (not hf_messages or hf_messages[0]["role"] != "system"):
+            messages.append({"role": "system", "content": [{"type": "text", "text": system_message}]})
+
+        for message in hf_messages:
+            if message["role"] == "system":
+                messages.append(message)
+                continue
+            if message.get("time") is not None:
+                continue
+            for content in message["content"]:
+                if content.get("type") in ["image", "video", "audio"]:
+                    first_content.append(content)
+
+        content = []
+        content.extend(first_content)
+
+        messages.append({"role": "user", "content": content})
+        return messages, timed_user_segments
+
+    def _get_chunk_start_times(self, video_grid_thw, video_metadata, audio_per_chunk_per_video: List[List[int]]):
+        times = []
+        for v_idx in range(len(video_grid_thw)):
+            metadata = video_metadata[v_idx]
+            fps = metadata.fps if metadata.fps is not None else 24.0
+            curr_timestamp = self.processor._calculate_timestamps(
+                metadata.frames_indices,
+                fps,
+                self.processor.video_processor.temporal_patch_size,
+            )
+            for t in range(len(audio_per_chunk_per_video[v_idx])):
+                times.append(curr_timestamp[t] if t < len(curr_timestamp) else curr_timestamp[-1])
+        return times
+
+    def _get_audio_position_times(self, video_grid_thw, video_metadata, audio_per_chunk_per_video: List[List[int]]):
+        times = []
+        chunk_times = self._get_chunk_start_times(video_grid_thw, video_metadata, audio_per_chunk_per_video)
+        chunk_idx = 0
+        for audio_per_chunk in audio_per_chunk_per_video:
+            for n_audio in audio_per_chunk:
+                times.extend([chunk_times[chunk_idx]] * n_audio)
+                chunk_idx += 1
+        return times
+
+    def _encode_realtime_text(self, text: str) -> List[int]:
+        return self.tokenizer.encode(text, add_special_tokens=False)
+
+    @staticmethod
+    def _first_index_at_or_after(values: List[float], target: float) -> int:
+        for idx, value in enumerate(values):
+            if value >= target:
+                return idx
+        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 get_qwen_template_labels(
         self,
         hf_messages,
@@ -406,6 +596,7 @@ def get_qwen_template_labels(
         video_metadata: List[dict],
         video_grid_thw=None,
         audio_per_chunk_per_video: Optional[List[List[int]]] = None,
+        timed_user_segments: Optional[List[Dict]] = None,
         system_message: str = "You are a helpful assistant",
         add_system_prompt: bool = True,
         add_generation_prompt: bool = False,
@@ -445,6 +636,7 @@ def get_qwen_template_labels(
                     curr_timestamp,
                     video_grid_thw,
                     audio_per_chunk_per_video=audio_per_chunk_per_video,
+                    timed_user_segments=timed_user_segments,
                 )
                 video_start_from += used_video
 
@@ -487,6 +679,7 @@ def _expand_encode_id_video_tokens(
         curr_timestamp: List[float] = None,
         video_grid_thw=None,
         audio_per_chunk_per_video: Optional[List[List[int]]] = None,
+        timed_user_segments: Optional[List[Dict]] = None,
     ):
         """Expand ``<|video_pad|>`` placeholders.
 
@@ -510,6 +703,7 @@ def _expand_encode_id_video_tokens(
         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)
         temporal_patch_size = getattr(self.processor.video_processor, "temporal_patch_size", 2)
+        timed_user_segments = timed_user_segments or []
 
         video_pos = [i for i, x in enumerate(encode_id) if x == self.video_token_id]
         expanded_encode_id = []
@@ -525,14 +719,26 @@ def _expand_encode_id_video_tokens(
             # them directly for the chunk start times.
             audio_per_chunk = audio_per_chunk_per_video[v_global]
             assert len(audio_per_chunk) == grid_t, f"audio_per_chunk len {len(audio_per_chunk)} != grid_t {grid_t}"
+            chunk_times = [
+                curr_timestamp[t] if t < len(curr_timestamp) else (t * temporal_patch_size) for t in range(grid_t)
+            ]
+            user_by_chunk = [[] for _ in range(grid_t)]
+            for segment in timed_user_segments:
+                chunk_idx = self._first_index_at_or_after(chunk_times, float(segment["time"]))
+                if chunk_idx >= grid_t:
+                    chunk_idx = grid_t - 1
+                user_by_chunk[chunk_idx].append(segment["text"])
 
             # Strip surrounding <|vision_start|> / <|vision_end|> from the
             # template (positions pos-1 and pos+1) -- we will emit our own.
             expanded_encode_id.extend(encode_id[prev : pos - 1])
 
             for t in range(grid_t):
+                for user_text in user_by_chunk[t]:
+                    expanded_encode_id.extend(self._encode_realtime_text(user_text))
+
                 # Per-frame timestamp (seconds) from the video metadata
-                t_sec = curr_timestamp[t] if t < len(curr_timestamp) else (t * temporal_patch_size)
+                t_sec = chunk_times[t]
                 timestamp_token_ids = self.processor.tokenizer.encode(f"<{t_sec:.1f} seconds>")
                 n_audio_t = audio_per_chunk[t]
                 expanded_encode_id.extend(timestamp_token_ids)