import json
import math
import random
import logging
import time
from dataclasses import dataclass
from typing import Dict, Optional, Sequence, List, Tuple, Any
from pathlib import Path
import numpy as np
import torch
from torch.utils.data import Dataset
import torch.distributed as dist
from qwen_vl_utils import process_vision_info
import qwen_vl_utils.vision_process as qwen_vision_process
import transformers
from torchcodec.decoders import VideoDecoder
from . import data_list
from .rope2d import ROPE_INDEX_FN
def _get_video_pixels(processor):
"""Return (min_pixels, max_pixels) from the video processor.
Qwen2/2.5 VL exposes ``min_pixels`` / ``max_pixels`` directly,
while Qwen3 VL only provides ``size["shortest_edge"]`` /
``size["longest_edge"]``.
"""
vp = processor.video_processor
min_px = getattr(vp, "min_pixels", None)
max_px = getattr(vp, "max_pixels", None)
if min_px is None or max_px is None:
size = getattr(vp, "size", {})
min_px = min_px or size.get("shortest_edge")
max_px = max_px or size.get("longest_edge")
if min_px is None or max_px is None:
raise ValueError(
"Cannot resolve video pixel limits from processor.video_processor: "
f"min_pixels={min_px}, max_pixels={max_px}, size={getattr(vp, 'size', 'N/A')}"
)
return min_px, max_px
def _resolve_vit_patch_size(processor) -> int:
"""Extract ViT patch size from a HuggingFace processor."""
return getattr(
processor.video_processor,
"patch_size",
getattr(processor.image_processor, "patch_size", 14),
)
def load_video_frames(
video_path: str,
video_start: float,
video_end: float,
total_nframes: int,
frames_per_chunk: int,
num_chunks: int,
*,
min_pixels: Optional[int] = None,
max_pixels: Optional[int] = None,
processor=None,
vit_patch_size: Optional[int] = None,
model_type: str,
) -> Tuple[List[torch.Tensor], dict, Optional[List[dict]]]:
"""Load video frames via a ghost message and split into per-chunk tensors.
This is the single entry-point for the "ghost message" video loading
pattern shared by SFT (``process_messages_to_model_inputs``), GRPO
(``build_grpo_inputs``), and inference (``streaming_video_chat``).
Pixel limits are resolved in order of priority:
1. Explicit *min_pixels* / *max_pixels* arguments (inference path).
2. Extracted from *processor* via ``_get_video_pixels`` (training path).
Args:
video_path: Absolute filesystem path to the video.
video_start: Start time in seconds.
video_end: End time in seconds.
total_nframes: Total number of frames to extract.
frames_per_chunk: Frames per temporal chunk.
num_chunks: Expected number of chunks.
min_pixels: Minimum pixel budget (overrides processor).
max_pixels: Maximum pixel budget (overrides processor).
processor: HuggingFace processor (required when *min_pixels*
or *max_pixels* is ``None``).
vit_patch_size: ViT patch size. When provided, *processor* is not
needed for the ``image_patch_size`` kwarg.
model_type: ``"qwen2.5vl"`` or ``"qwen3vl"``.
Returns:
``(split_videos, video_kwargs, chunk_metadatas)``
- *split_videos*: ``List[Tensor]`` of length *num_chunks*.
- *video_kwargs*: Extra kwargs from ``process_vision_info``
(e.g. ``fps``). May be empty.
- *chunk_metadatas*: Per-chunk Qwen3VL metadata list, or ``None``.
"""
is_qwen3vl = model_type == "qwen3vl"
# DataLoader uses the spawn start method, so every worker imports
# qwen-vl-utils afresh. Align its module global with the worker's processor
# immediately before smart_nframes runs.
if processor is not None:
temporal_patch_size = int(processor.video_processor.temporal_patch_size)
qwen_vision_process.FRAME_FACTOR = temporal_patch_size
# --- resolve pixel limits ---
if min_pixels is None or max_pixels is None:
if processor is None:
raise ValueError(
"Either (min_pixels, max_pixels) or processor must be provided."
)
_min, _max = _get_video_pixels(processor)
min_pixels = min_pixels or _min
max_pixels = max_pixels or _max
# --- guard: video_start == video_end ---
if video_end <= video_start:
video_end = video_start + 1e-3
# --- ghost message ---
ghost_message = [
{
"role": "user",
"content": [
{
"type": "video",
"video": video_path,
"video_start": video_start,
"video_end": video_end,
"nframes": total_nframes,
"min_pixels": min_pixels,
"max_pixels": max_pixels,
}
],
}
]
pvi_kwargs: dict = dict(return_video_kwargs=True)
if vit_patch_size is not None:
pvi_kwargs["image_patch_size"] = vit_patch_size
elif processor is not None:
pvi_kwargs["image_patch_size"] = _resolve_vit_patch_size(processor)
if is_qwen3vl:
pvi_kwargs["return_video_metadata"] = True
_, video_inputs_list, video_kwargs = process_vision_info(
ghost_message,
**pvi_kwargs,
)
# --- unpack ---
if is_qwen3vl:
big_video_tensor, video_metadata = video_inputs_list[0]
else:
big_video_tensor = video_inputs_list[0]
video_metadata = None
# --- split ---
split_videos = list(torch.split(big_video_tensor, frames_per_chunk, dim=0))
if len(split_videos) > num_chunks:
split_videos = split_videos[:num_chunks]
elif len(split_videos) < num_chunks:
raise RuntimeError(
f"Video split mismatch: expected {num_chunks}, got {len(split_videos)}"
)
# --- per-chunk metadata (Qwen3VL) or replicate fps (Qwen2/2.5VL) ---
chunk_metadatas: Optional[List[dict]] = None
if is_qwen3vl and video_metadata is not None:
all_indices = video_metadata["frames_indices"]
if isinstance(all_indices, torch.Tensor):
chunk_idx_splits = list(torch.split(all_indices, frames_per_chunk))
else:
chunk_idx_splits = [
all_indices[i : i + frames_per_chunk]
for i in range(0, len(all_indices), frames_per_chunk)
]
chunk_idx_splits = chunk_idx_splits[: len(split_videos)]
chunk_metadatas = [
{**video_metadata, "frames_indices": ci} for ci in chunk_idx_splits
]
else:
if "fps" in video_kwargs and isinstance(video_kwargs["fps"], list):
video_kwargs["fps"] = [video_kwargs["fps"][0]] * len(split_videos)
return split_videos, video_kwargs, chunk_metadatas
def preload_video(
video_path: str,
*,
video_start: float = 0.0,
video_end: Optional[float] = None,
frames_per_chunk: int = 2,
max_chunks: int = 120,
min_pixels: Optional[int] = None,
max_pixels: Optional[int] = None,
processor=None,
vit_patch_size: Optional[int] = None,
model_type: str,
) -> Dict[str, Any]:
"""Pre-load video metadata and decoded frames in one shot.
This combines the ``VideoDecoder`` metadata probe, chunking strategy
computation, and ``load_video_frames`` call that are otherwise spread
across ``streaming_video_chat`` and dataset code. By calling this in
``Dataset.__getitem__``, the heavy video I/O can be parallelised via
DataLoader ``num_workers``.
Returns a dict suitable for passing as ``preloaded_video`` to
``streaming_video_chat`` or extracting ``(split_videos, video_kwargs,
chunk_metadatas)`` for ``process_messages_to_model_inputs``.
"""
decoder = VideoDecoder(video_path)
video_fps = decoder.metadata.average_fps
total_video_frames = decoder.metadata.num_frames
if video_end is None:
video_end = decoder.metadata.duration_seconds
del decoder
# When start == end, take a window of 4 frames centered at that time, clamped to video bounds.
if video_end <= video_start:
t, max_sec = video_start, (total_video_frames - 1) / video_fps
half_span = 2.0 / video_fps
video_start = max(0.0, t - half_span)
video_end = min(max_sec, t + half_span)
video_duration = video_end - video_start
video_chunk_size = max(1.0, float(math.ceil(video_duration / max_chunks)))
num_iterations = max(1, int(video_duration // video_chunk_size))
frames_per_chunk = int(frames_per_chunk)
total_nframes = num_iterations * frames_per_chunk
# Clamp for low-fps videos (mirrors streaming_video_chat logic).
_start_frame = math.ceil(max(0.0, video_start) * video_fps)
_end_frame = min(math.floor(video_end * video_fps), total_video_frames - 1)
available_frames = max(_end_frame - _start_frame + 1, frames_per_chunk)
if total_nframes > available_frames:
num_iterations = max(1, available_frames // frames_per_chunk)
total_nframes = num_iterations * frames_per_chunk
video_chunk_size = video_duration / num_iterations
split_videos, video_kwargs, chunk_metadatas = load_video_frames(
video_path=video_path,
video_start=video_start,
video_end=video_end,
total_nframes=total_nframes,
frames_per_chunk=frames_per_chunk,
num_chunks=num_iterations,
min_pixels=min_pixels,
max_pixels=max_pixels,
processor=processor,
vit_patch_size=vit_patch_size,
model_type=model_type,
)
return {
"video_path": video_path,
"video_start": video_start,
"video_end": video_end,
"video_fps": video_fps,
"total_video_frames": total_video_frames,
"video_duration": video_duration,
"video_chunk_size": video_chunk_size,
"num_iterations": num_iterations,
"frames_per_chunk": frames_per_chunk,
"total_nframes": total_nframes,
"split_videos": split_videos,
"video_kwargs": video_kwargs,
"chunk_metadatas": chunk_metadatas,
}
IGNORE_INDEX = -100
FRAMES_PER_CHUNK = 1
# FRAMES_PER_CHUNK = 2
DEFAULT_MAX_CHUNKS = 120
DEFAULT_INFERENCE_MIN_PIXELS = 100352 * 2
DEFAULT_INFERENCE_MAX_PIXELS = 100352 * 4
def build_video_meta(
abs_path: str,
total_start: float,
total_end: float,
num_chunks: int,
frames_per_chunk: int = FRAMES_PER_CHUNK,
) -> Dict[str, Any]:
"""Build the canonical ``video_meta`` dict consumed by ``load_video_frames``
and ``process_messages_to_model_inputs``."""
return {
"abs_path": abs_path,
"total_start": total_start,
"total_end": total_end,
"num_chunks": num_chunks,
"frames_per_chunk": frames_per_chunk,
"total_nframes": num_chunks * frames_per_chunk,
}
def rank0_print(*args):
if dist.is_available() and dist.is_initialized():
if dist.get_rank() == 0:
print(*args)
else:
print(*args)
def read_jsonl(path):
with open(path, "r") as f:
return [json.loads(line) for line in f]
def _make_abs_paths(base: Path, files: str) -> str:
return f"{(base / files).resolve()}"
def update_processor_pixels(processor, data_args):
# --- Image Processor ---
ip = processor.image_processor
rank0_print("=== BEFORE IMAGE PROCESSOR PARAMETERS ===")
rank0_print(f"Image min_pixels: {getattr(ip, 'min_pixels', 'N/A')}")
rank0_print(f"Image max_pixels: {getattr(ip, 'max_pixels', 'N/A')}")
rank0_print(f"ip.size: {ip.size}")
rank0_print(f"Image size (shortest_edge): {ip.size.get('shortest_edge', 'N/A')}")
rank0_print(f"Image size (longest_edge): {ip.size.get('longest_edge', 'N/A')}")
ip.min_pixels = data_args.min_pixels
ip.max_pixels = data_args.max_pixels
rank0_print(f"✅ Updated image_processor min_pixels to {data_args.min_pixels}")
rank0_print(f"✅ Updated image_processor max_pixels to {data_args.max_pixels}")
if hasattr(ip, "size") and isinstance(ip.size, dict):
ip.size["shortest_edge"] = data_args.min_pixels
ip.size["longest_edge"] = data_args.max_pixels
rank0_print(
f"✅ Updated image_processor size['shortest_edge'] to {data_args.min_pixels}"
)
rank0_print(
f"✅ Updated image_processor size['longest_edge'] to {data_args.max_pixels}"
)
rank0_print("=== AFTER IMAGE PROCESSOR PARAMETERS ===")
rank0_print(f"Image min_pixels: {getattr(ip, 'min_pixels', 'N/A')}")
rank0_print(f"Image max_pixels: {getattr(ip, 'max_pixels', 'N/A')}")
rank0_print(f"Image size (shortest_edge): {ip.size.get('shortest_edge', 'N/A')}")
rank0_print(f"Image size (longest_edge): {ip.size.get('longest_edge', 'N/A')}")
# --- Video Processor ---
if hasattr(processor, "video_processor") and processor.video_processor is not None:
vp = processor.video_processor
rank0_print("\n=== BEFORE VIDEO PROCESSOR PARAMETERS ===")
rank0_print(f"Video min_pixels: {getattr(vp, 'min_pixels', 'N/A')}")
rank0_print(f"Video max_pixels: {getattr(vp, 'max_pixels', 'N/A')}")
rank0_print(f"Video min_frames: {getattr(vp, 'min_frames', 'N/A')}")
rank0_print(f"Video max_frames: {getattr(vp, 'max_frames', 'N/A')}")
rank0_print(f"Video fps: {getattr(vp, 'fps', 'N/A')}")
rank0_print(
f"Video size (shortest_edge): {vp.size.get('shortest_edge', 'N/A')}"
)
rank0_print(f"Video size (longest_edge): {vp.size.get('longest_edge', 'N/A')}")
vp.min_pixels = data_args.video_min_pixels
vp.max_pixels = data_args.video_max_pixels
rank0_print(
f"✅ Updated video_processor min_pixels to {data_args.video_min_pixels}"
)
rank0_print(
f"✅ Updated video_processor max_pixels to {data_args.video_max_pixels}"
)
if hasattr(vp, "min_frames") and hasattr(vp, "max_frames"):
vp.min_frames = data_args.video_min_frames
vp.max_frames = data_args.video_max_frames
rank0_print(
f"✅ Updated video_processor min_frames to {data_args.video_min_frames}"
)
rank0_print(
f"✅ Updated video_processor max_frames to {data_args.video_max_frames}"
)
if hasattr(vp, "fps"):
vp.fps = data_args.video_fps
rank0_print(f"✅ Updated video_processor fps to {data_args.video_fps}")
if hasattr(vp, "size") and isinstance(vp.size, dict):
vp.size["shortest_edge"] = data_args.video_min_pixels
vp.size["longest_edge"] = data_args.video_max_pixels
rank0_print(
f"✅ Updated Video size (shortest_edge): {vp.size.get('shortest_edge', 'N/A')}"
)
rank0_print(
f"✅ Updated Video size (longest_edge): {vp.size.get('longest_edge', 'N/A')}"
)
rank0_print("=== AFTER VIDEO PROCESSOR PARAMETERS ===")
rank0_print(f"Video min_pixels: {getattr(vp, 'min_pixels', 'N/A')}")
rank0_print(f"Video max_pixels: {getattr(vp, 'max_pixels', 'N/A')}")
rank0_print(f"Video min_frames: {getattr(vp, 'min_frames', 'N/A')}")
rank0_print(f"Video max_frames: {getattr(vp, 'max_frames', 'N/A')}")
rank0_print(f"Video fps: {getattr(vp, 'fps', 'N/A')}")
rank0_print(
f"Video size (shortest_edge): {vp.size.get('shortest_edge', 'N/A')}"
)
rank0_print(f"Video size (longest_edge): {vp.size.get('longest_edge', 'N/A')}")
# Vision architecture params that affect visual token count
ip = processor.image_processor
vp = getattr(processor, "video_processor", None)
vit_patch = getattr(vp, "patch_size", getattr(ip, "patch_size", "N/A"))
vit_merge = getattr(ip, "merge_size", "N/A")
vit_temporal = getattr(vp, "temporal_patch_size", "N/A") if vp else "N/A"
factor = (
(vit_patch * vit_merge)
if isinstance(vit_patch, int) and isinstance(vit_merge, int)
else "N/A"
)
rank0_print(f"\n=== VISION ARCHITECTURE ===")
rank0_print(
f"patch_size={vit_patch}, merge_size={vit_merge}, temporal_patch_size={vit_temporal}, factor={factor}"
)
if hasattr(data_args, "model_max_length"):
rank0_print(f"Before: max_length: ", processor.tokenizer.model_max_length)
processor.tokenizer.model_max_length = data_args.model_max_length
rank0_print(f"After: max_length: ", processor.tokenizer.model_max_length)
return processor
def _get_duration(path: str) -> float:
return VideoDecoder(path).metadata.duration_seconds
SYSTEM_PROMPT = (
"You are a helpful assistant. "
"You will see a continuous stream of video chunks. "
"Based on the user's query and the video content, first output your internal reasoning enclosed in <think>...</think> tags. "
"Then, if you determine that a response is needed at this moment, output <response> followed by the content. "
"If no response is needed, output <silent>. "
"Your generated thoughts and responses should be continuous and fluent across the video chunks."
)
# Chat template without system prompt (for subsequent streaming chunks)
QWEN_TEMPLATE_WO_SYSTEM = (
"\n{% set image_count = namespace(value=0) %}"
"{% set video_count = namespace(value=0) %}"
"{% for message in messages %}"
"<|im_start|>{{ message['role'] }}\n"
"{% if message['content'] is string %}"
"{{ message['content'] }}<|im_end|>\n"
"{% else %}"
"{% for content in message['content'] %}"
"{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}"
"{% set image_count.value = image_count.value + 1 %}"
"{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}"
"<|vision_start|><|image_pad|><|vision_end|>"
"{% elif content['type'] == 'video' or 'video' in content %}"
"{% set video_count.value = video_count.value + 1 %}"
"{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}"
"<|vision_start|><|video_pad|><|vision_end|>"
"{% elif 'text' in content %}"
"{{ content['text'] }}"
"{% endif %}"
"{% endfor %}<|im_end|>\n"
"{% endif %}"
"{% endfor %}"
"{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
)
def _build_messages(
item: Dict[str, Any], base_path: Path, remaining_video_chunks: int = 3
) -> Dict[str, Any]:
video_path = item.get("video_path")
if not video_path:
raise ValueError("video_path is required")
abs_video_path = str(_make_abs_paths(base_path, video_path))
video_duration = _get_duration(abs_video_path)
video_chunk_size = float(math.ceil(video_duration / DEFAULT_MAX_CHUNKS))
user_queue = sorted(
[
(float(c.get("timestamp", 0.0)), c.get("content", ""))
for c in item.get("conversations", [])
if c.get("role") == "user"
],
key=lambda x: x[0],
)
assistant_queue = sorted(
[
(float(c.get("timestamp", 0.0)), c.get("content", ""))
for c in item.get("conversations", [])
if c.get("role") == "assistant"
],
key=lambda x: x[0],
)
# [新增] 读取并排序 thoughts 队列
thoughts_queue = sorted(
[
(float(t.get("timestamp", 0.0)), t.get("think", ""))
for t in item.get("thoughts", [])
],
key=lambda x: x[0],
)
max_video_chunks = int(video_duration // video_chunk_size)
# [修改点 1] 保持原逻辑：只根据 assistant 的最后一条消息来决定视频切分长度
last_assist_ts = assistant_queue[-1][0] if assistant_queue else 0.0
last_assist_chunk_idx = int(np.floor(last_assist_ts / video_chunk_size))
target_stop_idx = min(
last_assist_chunk_idx + remaining_video_chunks, max_video_chunks - 1
)
# [修改点 2] 插入 System Prompt 到 messages 开头
messages = [{"role": "system", "content": SYSTEM_PROMPT}]
video_chunks_count = 0
for chunk_idx in range(target_stop_idx + 1):
window_start = chunk_idx * video_chunk_size
window_end = (chunk_idx + 1) * video_chunk_size
is_last_chunk = chunk_idx == max_video_chunks - 1
# --- User Content 构建 ---
user_content = []
user_content.append(
{
"type": "video",
"video": abs_video_path,
"video_start": window_start,
"video_end": window_end,
}
)
video_chunks_count += 1
while user_queue and user_queue[0][0] < window_end:
ts, text = user_queue[0]
if ts >= window_start:
user_queue.pop(0)
user_content.append({"type": "text", "text": "\n" + text})
else:
user_queue.pop(0) # 丢弃过期
messages.append({"role": "user", "content": user_content})
# --- Assistant Content 构建 (Thoughts + Response) ---
# 1. 处理 Thoughts (与 assistant 逻辑一致)
current_thought_segments = []
while thoughts_queue:
ts, content = thoughts_queue[0]
if not is_last_chunk and ts >= window_end:
break
if ts >= window_start:
thoughts_queue.pop(0)
current_thought_segments.append(content)
else:
thoughts_queue.pop(0) # 丢弃过期
# 2. 处理 Response
current_text_segments = []
while assistant_queue:
ts, content = assistant_queue[0]
if not is_last_chunk and ts >= window_end:
break
if ts >= window_start:
assistant_queue.pop(0)
current_text_segments.append(content)
else:
assistant_queue.pop(0)
# [修改点 3] 拼接逻辑：<think>...</think> + <response>/<silent>
# 即使没有思考内容，也要输出空的 <think></think>
thought_part = "<think>" + " ".join(current_thought_segments) + "</think>"
if current_text_segments:
response_part = "<response>" + " ".join(current_text_segments)
else:
response_part = "<silent>"
final_content = thought_part + response_part
assistant_content = [{"type": "text", "text": final_content}]
messages.append({"role": "assistant", "content": assistant_content})
total_covered_duration = min(video_chunks_count * video_chunk_size, video_duration)
video_meta = build_video_meta(
abs_path=abs_video_path,
total_start=0.0,
total_end=total_covered_duration,
num_chunks=video_chunks_count,
)
return {
"messages": messages,
"video_meta": video_meta,
"video_chunk_size": video_chunk_size,
}
def process_messages_to_model_inputs(
messages: List[Dict[str, Any]],
video_meta: Dict[str, Any],
video_chunk_size: float,
processor,
model_type: str,
add_generation_prompt: bool = False,
preloaded_frames: Optional[Tuple[List, dict, Optional[List]]] = None,
) -> Dict:
"""
Common function: takes pre-built chat messages + video metadata, loads video
frames, tokenises, and returns model-forward-ready tensors.
This is used by both SFT (``preprocess_qwen_visual``) and GRPO (after
rollout, to build training inputs from generated completions).
Video pixel limits are read from the processor's ``video_processor``
configuration. Both ``LazySupervisedDataset`` (SFT) and
``LazyRawDataset`` (GRPO) call ``update_processor_pixels`` at init
time, so the processor always carries the correct limits.
Args:
messages: Full chat-format message list (system + alternating user/assistant).
User messages may contain ``{"type": "video", ...}`` entries.
video_meta: Dict with keys:
- abs_path: str, absolute path to the video file
- total_start: float, start time (seconds)
- total_end: float, end time (seconds)
- num_chunks: int, number of video chunks
- frames_per_chunk: int, frames per chunk
- total_nframes: int, total number of frames
video_chunk_size: float, seconds per video chunk (needed for RoPE
computation downstream).
processor: HuggingFace processor (e.g. ``AutoProcessor``).
add_generation_prompt: Whether to append the generation prompt
(``<|im_start|>assistant\\n``) at the end.
preloaded_frames: Optional pre-loaded ``(split_videos, video_kwargs,
chunk_metadatas)`` tuple. When provided, ``load_video_frames``
is skipped entirely.
Returns:
Dict containing at least:
- input_ids: [1, L]
- pixel_values_videos: video pixel tensors
- video_grid_thw: [num_chunks, 3]
- video_mask: [1, L] (True where video tokens appear)
- video_chunk_size: float (pass-through for RoPE)
"""
is_qwen3vl = model_type == "qwen3vl"
# 1. Load ALL video frames at once via unified loader
if preloaded_frames is not None:
split_videos, video_kwargs, chunk_metadatas = preloaded_frames
else:
split_videos, video_kwargs, chunk_metadatas = load_video_frames(
video_path=video_meta["abs_path"],
video_start=video_meta["total_start"],
video_end=video_meta["total_end"],
total_nframes=video_meta["total_nframes"],
frames_per_chunk=video_meta["frames_per_chunk"],
num_chunks=video_meta["num_chunks"],
processor=processor,
model_type=model_type,
)
# 2. Tokenise
text = processor.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=add_generation_prompt,
)
processor_call_kwargs = dict(
text=text,
images=None,
videos=split_videos,
return_tensors="pt",
**video_kwargs,
)
if is_qwen3vl:
processor_call_kwargs["do_resize"] = False
if chunk_metadatas is not None:
processor_call_kwargs["video_metadata"] = chunk_metadatas
full_result = processor(**processor_call_kwargs)
input_ids = full_result["input_ids"]
if isinstance(input_ids, list):
input_ids = torch.tensor(input_ids).unsqueeze(0)
full_result["input_ids"] = input_ids
# 4. Video mask
video_token_id = processor.tokenizer.convert_tokens_to_ids(["<|video_pad|>"])[0]
assert "video_mask" not in full_result, "Key Conflict!"
full_result["video_mask"] = input_ids == video_token_id
# 5. Pass-through for RoPE computation
full_result["video_chunk_size"] = video_chunk_size
return full_result
# ---------------------------------------------------------------------------
# Unified helpers (shared across SFT, GRPO, and Eval)
# ---------------------------------------------------------------------------
def find_assistant_spans(
input_ids_1d: List[int],
tokenizer,
) -> List[Tuple[int, int]]:
"""Locate all assistant-turn token spans in a flat token sequence.
Scans *input_ids_1d* for the ``assistant`` token. For each occurrence the
assistant content starts 2 positions later (skipping ``assistant\\n``) and
extends until ``<|im_end|>`` (inclusive) plus one trailing token (the
``\\n`` that follows ``<|im_end|>``).
Returns:
A list of ``(start, end_exclusive)`` tuples suitable for slicing, e.g.
``input_ids[start:end_exclusive]``.
"""
assistant_id, im_end_id = tokenizer.convert_tokens_to_ids(
["assistant", "<|im_end|>"]
)
spans: List[Tuple[int, int]] = []
L = len(input_ids_1d)
pos = 0
while pos < L:
if input_ids_1d[pos] == assistant_id:
ans_start = pos + 2
ans_end = ans_start
while ans_end < L and input_ids_1d[ans_end] != im_end_id:
ans_end += 1
if ans_end < L:
spans.append((ans_start, ans_end + 2))
pos = ans_end
pos += 1
return spans
def compute_position_ids(
processor_output: Dict[str, Any],
processor,
model_type: str,
) -> torch.Tensor:
"""Compute MROPE position IDs – the **single** path for all workflows.
Used by SFT, GRPO training reconstruction, **and** inference
(``streaming_video_chat``). Callers must ensure ``video_chunk_size``
is present in *processor_output* so that ``second_per_grid_ts`` is
computed identically everywhere.
Handles ``image_grid_thw`` / ``video_grid_thw`` normalisation,
``second_per_grid_ts`` computation from ``video_chunk_size``, and
dispatches to the correct RoPE function for the given *model_type*.
The ``video_chunk_size`` key is **consumed** (popped) from
*processor_output* during this call.
Returns:
``position_ids`` tensor of shape ``(3, B, L)`` for MROPE models.
"""
input_ids = processor_output["input_ids"]
# --- image_grid_thw ---
if "image_grid_thw" in processor_output:
image_grid_thw = processor_output["image_grid_thw"]
if not isinstance(image_grid_thw, (list, tuple)):
image_grid_thw = [image_grid_thw]
image_grid_thw = torch.cat(image_grid_thw, dim=0)
else:
image_grid_thw = None
# --- video_grid_thw + second_per_grid_ts ---
if "video_grid_thw" in processor_output:
video_grid_thw = processor_output["video_grid_thw"]
if not isinstance(video_grid_thw, (list, tuple)):
video_grid_thw = [video_grid_thw]
video_grid_thw = torch.cat(video_grid_thw, dim=0)
second_per_grid_ts = [
processor_output.pop("video_chunk_size", 1)
* processor.video_processor.temporal_patch_size
/ processor.video_processor.fps
] * len(video_grid_thw)
else:
video_grid_thw = None
second_per_grid_ts = None
merge_size = getattr(processor.image_processor, "merge_size", 2)
rope_fn = ROPE_INDEX_FN[model_type]
position_ids, _ = rope_fn(
merge_size,
input_ids,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
second_per_grid_ts=second_per_grid_ts,
)
return position_ids
def preprocess_qwen_visual(sources, processor, model_type: str) -> Dict:
if len(sources) != 1:
raise ValueError(f"Expected 1 source, got {len(sources)}")
source = sources[0]
base_path = Path(source.get("data_path", ""))
build_result = _build_messages(source, base_path)
messages = build_result["messages"]
video_meta = build_result["video_meta"]
video_chunk_size = build_result["video_chunk_size"]
# Use common function for video loading + tokenisation
full_result = process_messages_to_model_inputs(
messages=messages,
video_meta=video_meta,
video_chunk_size=video_chunk_size,
processor=processor,
model_type=model_type,
add_generation_prompt=False,
)
# SFT-specific: create labels (only train on assistant turns)
input_ids = full_result["input_ids"]
labels = torch.full_like(input_ids, IGNORE_INDEX)
for start, end in find_assistant_spans(input_ids[0].tolist(), processor.tokenizer):
labels[0, start:end] = input_ids[0, start:end]
full_result["labels"] = labels
return full_result
class LazySupervisedDataset(Dataset):
"""Dataset for supervised fine-tuning."""
def __init__(self, processor, data_args):
super(LazySupervisedDataset, self).__init__()
dataset = data_args.dataset_use.split(",")
dataset_list = data_list(dataset)
rank0_print(f"Loading datasets: {dataset_list}")
self.video_max_total_pixels = getattr(
data_args, "video_max_total_pixels", 1664 * 28 * 28
)
self.video_min_total_pixels = getattr(
data_args, "video_min_total_pixels", 256 * 28 * 28
)
self.model_type = data_args.model_type
if data_args.model_type not in ROPE_INDEX_FN:
raise ValueError(
f"model_type: {data_args.model_type} not supported. "
f"Choose from {list(ROPE_INDEX_FN.keys())}"
)
self.get_rope_index = ROPE_INDEX_FN[data_args.model_type]
list_data_dict = []
for data in dataset_list:
file_format = data["annotation_path"].split(".")[-1]
if file_format == "jsonl":
annotations = read_jsonl(data["annotation_path"])
else:
annotations = json.load(open(data["annotation_path"], "r"))
sampling_rate = data.get("sampling_rate", 1.0)
if sampling_rate < 1.0:
annotations = random.sample(
annotations, int(len(annotations) * sampling_rate)
)
rank0_print(f"sampling {len(annotations)} examples from dataset {data}")
else:
rank0_print(f"dataset name: {data}")
for ann in annotations:
if isinstance(ann, list):
for sub_ann in ann:
sub_ann["data_path"] = data["data_path"]
else:
ann["data_path"] = data["data_path"]
list_data_dict += annotations
rank0_print(f"Total training samples: {len(list_data_dict)}")
random.shuffle(list_data_dict) # Randomly shuffle the data for training
rank0_print("Formatting inputs...Skip in lazy mode")
processor = update_processor_pixels(processor, data_args)
self.processor = processor
self.tokenizer = processor.tokenizer
self.data_args = data_args
self.merge_size = getattr(processor.image_processor, "merge_size", 2)
self.list_data_dict = list_data_dict
self.item_fn = self._get_item
def __len__(self):
return len(self.list_data_dict)
@property
def lengths(self):
length_list = []
for sample in self.list_data_dict:
img_tokens = 128 if "image" in sample else 0
length_list.append(
sum(len(conv["value"].split()) for conv in sample["conversations"])
+ img_tokens
)
return length_list
@property
def modality_lengths(self):
length_list = []
for sample in self.list_data_dict:
cur_len = sum(
len(conv["value"].split()) for conv in sample["conversations"]
)
cur_len = (
cur_len if ("image" in sample) or ("video" in sample) else -cur_len
)
length_list.append(cur_len)
return length_list
@property
def pre_calculated_length(self):
if "num_tokens" in self.list_data_dict[0]:
length_list = [sample["num_tokens"] for sample in self.list_data_dict]
return np.array(length_list)
else:
print("No pre-calculated length available.")
return np.array([1] * len(self.list_data_dict))
def __getitem__(self, i) -> Dict[str, torch.Tensor]:
num_base_retries = 3
num_final_retries = 30
for attempt_idx in range(num_base_retries):
try:
sources = self.list_data_dict[i]
if isinstance(sources, dict):
sources = [sources]
sample = self.item_fn(sources)
return sample
except Exception as e:
# sleep 1s in case it is a cloud disk issue
print(f"[Try #{attempt_idx}] Failed to fetch sample {i}. Exception:", e)
time.sleep(1)
for attempt_idx in range(num_final_retries):
try:
random_index = random.randint(0, len(self.list_data_dict) - 1)
if random_index == i:
continue
sources = self.list_data_dict[random_index]
if isinstance(sources, dict):
sources = [sources]
sample = self.item_fn(sources)
return sample
except Exception as e:
# no need to sleep
print(
f"[Try random #{attempt_idx}] Failed to fetch random sample {random_index}. Exception:",
e,
)
pass
try:
sources = self.list_data_dict[i]
if isinstance(sources, dict):
sources = [sources]
sample = self.item_fn(sources)
return sample
except Exception as e:
raise
def _get_item(self, sources) -> Dict[str, torch.Tensor]:
data_dict = preprocess_qwen_visual(
sources,
self.processor,
model_type=self.model_type,
)
seq_len = data_dict["input_ids"][0].size(0)
data_dict["position_ids"] = compute_position_ids(
data_dict,
self.processor,
self.model_type,
)
data_dict["attention_mask"] = [seq_len]
return data_dict
def pad_and_cat(tensor_list):
max_length = max(tensor.shape[2] for tensor in tensor_list)
padded_tensors = []
for tensor in tensor_list:
pad_length = max_length - tensor.shape[2]
padded_tensor = torch.nn.functional.pad(tensor, (0, pad_length), "constant", 1)
padded_tensors.append(padded_tensor)
stacked_tensor = torch.cat(padded_tensors, dim=1)
return stacked_tensor
@dataclass
class DataCollatorForSupervisedDataset:
"""Collate examples for supervised fine-tuning."""
tokenizer: transformers.PreTrainedTokenizer
vocab_size: int
def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
input_ids, labels, position_ids, video_masks = tuple(
[instance[key] for instance in instances]
for key in ("input_ids", "labels", "position_ids", "video_mask")
)
input_ids = [ids.squeeze(0) for ids in input_ids]
labels = [ids.squeeze(0) for ids in labels]
video_masks = [ids.squeeze(0) for ids in video_masks]
input_ids = torch.nn.utils.rnn.pad_sequence(
input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
)
labels = torch.nn.utils.rnn.pad_sequence(
labels, batch_first=True, padding_value=IGNORE_INDEX
)
video_masks = torch.nn.utils.rnn.pad_sequence(
video_masks, batch_first=True, padding_value=0
)
position_ids = pad_and_cat(position_ids)
input_ids = input_ids[:, : self.tokenizer.model_max_length]
labels = labels[:, : self.tokenizer.model_max_length]
position_ids = position_ids[:, :, : self.tokenizer.model_max_length]
video_masks = video_masks[:, : self.tokenizer.model_max_length]
batch = dict(
input_ids=input_ids,
labels=labels,
attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
video_mask=video_masks,
)
images = list(
instance["pixel_values"]
for instance in instances
if "pixel_values" in instance
)
videos = list(
instance["pixel_values_videos"]
for instance in instances
if "pixel_values_videos" in instance
)
if len(images) != 0:
concat_images = torch.cat([image for image in images], dim=0)
grid_thw = [
instance["image_grid_thw"]
for instance in instances
if "image_grid_thw" in instance
]
grid_thw = torch.cat(grid_thw, dim=0)
else:
concat_images = None
grid_thw = None
if len(videos) != 0:
concat_videos = torch.cat([video for video in videos], dim=0)
video_grid_thw = [
instance["video_grid_thw"]
for instance in instances
if "video_grid_thw" in instance
]
video_grid_thw = torch.cat(video_grid_thw, dim=0)
else:
concat_videos = None
video_grid_thw = None
batch["pixel_values"] = concat_images
batch["image_grid_thw"] = grid_thw
batch["pixel_values_videos"] = concat_videos
batch["video_grid_thw"] = video_grid_thw
batch["position_ids"] = position_ids
# CE Weight
response_id = self.tokenizer.convert_tokens_to_ids(["<response>"])[0]
silent_id = self.tokenizer.convert_tokens_to_ids(["<silent>"])[0]
n_response = (batch["labels"] == response_id).sum()
n_silent = (batch["labels"] == silent_id).sum()
total_n = n_response + n_silent
ce_weight = torch.ones(self.vocab_size)
eps = 1e-3
ce_weight[silent_id] = total_n / (2 * n_silent + eps)
ce_weight[response_id] = total_n / (2 * n_response + eps)
batch["ce_weight"] = torch.clamp(ce_weight, 0, 20)
print(input_ids.shape, batch["ce_weight"][[silent_id, response_id]])
return batch
def make_supervised_data_module(processor, data_args, vocab_size: int) -> Dict:
"""Make dataset and collator for supervised fine-tuning."""
train_dataset = LazySupervisedDataset(processor, data_args=data_args)
data_collator = DataCollatorForSupervisedDataset(
processor.tokenizer, vocab_size=vocab_size
)
return dict(
train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator
)
# ---------------------------------------------------------------------------
# Raw (unprocessed) Dataset for GRPO
# ---------------------------------------------------------------------------
class LazyRawDataset(Dataset):
"""
Dataset that returns raw JSON annotations with optional video preloading.
Used by the GRPO pipeline: the rollout stage receives raw data and
performs streaming inference to generate completions.
When *model_type* is provided, ``__getitem__`` also pre-loads video
frames via ``preload_video`` so that DataLoader ``num_workers`` can
parallelise the heavy video I/O.
Like ``LazySupervisedDataset``, the processor's pixel limits are updated
via ``update_processor_pixels`` at init time so that downstream callers
(``process_messages_to_model_inputs``, ``load_video_frames``) always see
the correct resolution config.
"""
def __init__(
self,
processor,
data_args,
*,
frames_per_chunk: int = FRAMES_PER_CHUNK,
max_chunks: int = DEFAULT_MAX_CHUNKS,
model_type: str = "",
):
super().__init__()
dataset = data_args.dataset_use.split(",")
dataset_list = data_list(dataset)
rank0_print(f"[LazyRawDataset] Loading datasets: {dataset_list}")
list_data_dict: List[Dict] = []
for data in dataset_list:
file_format = data["annotation_path"].split(".")[-1]
if file_format == "jsonl":
annotations = read_jsonl(data["annotation_path"])
else:
annotations = json.load(open(data["annotation_path"], "r"))
sampling_rate = data.get("sampling_rate", 1.0)
if sampling_rate < 1.0:
annotations = random.sample(
annotations, int(len(annotations) * sampling_rate)
)
rank0_print(
f" sampling {len(annotations)} examples from dataset {data}"
)
else:
rank0_print(f" dataset name: {data}")
for ann in annotations:
if isinstance(ann, list):
for sub_ann in ann:
sub_ann["data_path"] = data["data_path"]
else:
ann["data_path"] = data["data_path"]
list_data_dict += annotations
random.shuffle(list_data_dict)
rank0_print(f"[LazyRawDataset] Total raw samples: {len(list_data_dict)}")
update_processor_pixels(processor, data_args)
self.processor = processor
self.list_data_dict = list_data_dict
self._do_preload = bool(model_type)
self._model_type = model_type
self._frames_per_chunk = frames_per_chunk
self._max_chunks = max_chunks
min_px, max_px = _get_video_pixels(processor)
self._min_pixels = min_px
self._max_pixels = max_px
self._vit_patch_size = _resolve_vit_patch_size(processor)
def __len__(self):
return len(self.list_data_dict)
def __getitem__(self, i) -> Dict[str, Any]:
item = self.list_data_dict[i]
if not self._do_preload:
return item
data_path = item.get("data_path", "")
video_path = item.get("video_path", "")
abs_video_path = str(_make_abs_paths(Path(data_path), video_path))
preloaded = preload_video(
abs_video_path,
frames_per_chunk=self._frames_per_chunk,
max_chunks=self._max_chunks,
min_pixels=self._min_pixels,
max_pixels=self._max_pixels,
vit_patch_size=self._vit_patch_size,
model_type=self._model_type,
)
item = {**item, "_preloaded_video": preloaded}
return item
def raw_data_collate_fn(instances: Sequence[Dict]) -> List[Dict]:
"""Collator that passes through raw dicts without any processing."""
return list(instances)
def make_raw_data_module(
processor,
data_args,
*,
frames_per_chunk: int = FRAMES_PER_CHUNK,
max_chunks: int = DEFAULT_MAX_CHUNKS,
model_type: str = "",
) -> Dict:
"""Make dataset and collator that return raw JSON data (for GRPO)."""
train_dataset = LazyRawDataset(
processor,
data_args,
frames_per_chunk=frames_per_chunk,
max_chunks=max_chunks,
model_type=model_type,
)
return dict(
train_dataset=train_dataset,
eval_dataset=None,
data_collator=raw_data_collate_fn,
)
import json
import math
import random
import logging
import time
from dataclasses import dataclass
from typing import Dict, Optional, Sequence, List, Tuple, Any
from pathlib import Path
import numpy as np
import torch
from torch.utils.data import Dataset
import torch.distributed as dist
from qwen_vl_utils import process_vision_info
import qwen_vl_utils.vision_process as qwen_vision_process
import transformers
from torchcodec.decoders import VideoDecoder
from . import data_list
from .rope2d import ROPE_INDEX_FN
def _get_video_pixels(processor):
"""Return (min_pixels, max_pixels) from the video processor.
Qwen2/2.5 VL exposes ``min_pixels`` / ``max_pixels`` directly,
while Qwen3 VL only provides ``size["shortest_edge"]`` /
``size["longest_edge"]``.
"""
vp = processor.video_processor
min_px = getattr(vp, "min_pixels", None)
max_px = getattr(vp, "max_pixels", None)
if min_px is None or max_px is None:
size = getattr(vp, "size", {})
min_px = min_px or size.get("shortest_edge")
max_px = max_px or size.get("longest_edge")
if min_px is None or max_px is None:
raise ValueError(
"Cannot resolve video pixel limits from processor.video_processor: "
f"min_pixels={min_px}, max_pixels={max_px}, size={getattr(vp, 'size', 'N/A')}"
)
return min_px, max_px
def _resolve_vit_patch_size(processor) -> int:
"""Extract ViT patch size from a HuggingFace processor."""
return getattr(
processor.video_processor,
"patch_size",
getattr(processor.image_processor, "patch_size", 14),
)
def load_video_frames(
video_path: str,
video_start: float,
video_end: float,
total_nframes: int,
frames_per_chunk: int,
num_chunks: int,
*,
min_pixels: Optional[int] = None,
max_pixels: Optional[int] = None,
processor=None,
vit_patch_size: Optional[int] = None,
model_type: str,
) -> Tuple[List[torch.Tensor], dict, Optional[List[dict]]]:
"""Load video frames via a ghost message and split into per-chunk tensors.
This is the single entry-point for the "ghost message" video loading
pattern shared by SFT (``process_messages_to_model_inputs``), GRPO
(``build_grpo_inputs``), and inference (``streaming_video_chat``).
Pixel limits are resolved in order of priority:
1. Explicit *min_pixels* / *max_pixels* arguments (inference path).
2. Extracted from *processor* via ``_get_video_pixels`` (training path).
Args:
video_path: Absolute filesystem path to the video.
video_start: Start time in seconds.
video_end: End time in seconds.
total_nframes: Total number of frames to extract.
frames_per_chunk: Frames per temporal chunk.
num_chunks: Expected number of chunks.
min_pixels: Minimum pixel budget (overrides processor).
max_pixels: Maximum pixel budget (overrides processor).
processor: HuggingFace processor (required when *min_pixels*
or *max_pixels* is ``None``).
vit_patch_size: ViT patch size. When provided, *processor* is not
needed for the ``image_patch_size`` kwarg.
model_type: ``"qwen2.5vl"`` or ``"qwen3vl"``.
Returns:
``(split_videos, video_kwargs, chunk_metadatas)``
- *split_videos*: ``List[Tensor]`` of length *num_chunks*.
- *video_kwargs*: Extra kwargs from ``process_vision_info``
(e.g. ``fps``). May be empty.
- *chunk_metadatas*: Per-chunk Qwen3VL metadata list, or ``None``.
"""