| import json |
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
| from video_inputs import build_video_item, prepare |
| from modeling_skivlm import POSE_START, POSE_PATCH, POSE_END |
|
|
| IGNORE = -100 |
| QUESTION = "Describe what the person is doing in this video." |
|
|
|
|
| def pose_placeholder(num_pose_tokens): |
| return POSE_START + POSE_PATCH * num_pose_tokens + POSE_END |
|
|
|
|
| class MotionXSkiDataset(Dataset): |
| def __init__(self, manifest_path, pose_path, processor, use_skeleton=True, |
| num_pose_tokens=8, video_fps=1.0, max_frames=8, |
| max_pixels=200704, limit=0, seed=0): |
| items = json.load(open(manifest_path)) |
| if limit and limit < len(items): |
| rng = np.random.default_rng(seed) |
| idx = rng.choice(len(items), size=limit, replace=False) |
| items = [items[k] for k in sorted(idx)] |
| self.items = items |
| self.pose = np.load(pose_path, mmap_mode="r") |
| self.processor = processor |
| self.use_skeleton = use_skeleton |
| self.num_pose_tokens = num_pose_tokens |
| self.video_fps = video_fps |
| self.max_frames = max_frames |
| self.max_pixels = max_pixels |
|
|
| def __len__(self): |
| return len(self.items) |
|
|
| def _messages(self, m): |
| text = QUESTION |
| if self.use_skeleton: |
| text = pose_placeholder(self.num_pose_tokens) + "\n" + QUESTION |
| return [{ |
| "role": "user", |
| "content": [build_video_item(m["video"], self.video_fps, |
| self.max_frames, self.max_pixels), |
| {"type": "text", "text": text}], |
| }] |
|
|
| def __getitem__(self, i): |
| for attempt in range(8): |
| try: |
| return self._build(i) |
| except Exception as e: |
| print("[dataset] skipping", self.items[i]["stem"], type(e).__name__, flush=True) |
| i = int(np.random.default_rng(i + attempt).integers(len(self.items))) |
| raise RuntimeError("too many unreadable videos") |
|
|
| def _build(self, i): |
| m = self.items[i] |
| messages = self._messages(m) |
| prompt = self.processor.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True) |
| answer = m["text"] + self.processor.tokenizer.eos_token |
| video_inputs, vkw = prepare(messages, m["video"]) |
|
|
| enc = self.processor(text=[prompt + answer], videos=video_inputs, |
| return_tensors="pt", **vkw) |
| prompt_only = self.processor(text=[prompt], videos=video_inputs, |
| return_tensors="pt", **vkw) |
| n_prompt = prompt_only["input_ids"].shape[1] |
| input_ids = enc["input_ids"][0] |
| labels = input_ids.clone() |
| labels[:n_prompt] = IGNORE |
|
|
| x = torch.from_numpy(np.asarray(self.pose[m["i"]], dtype=np.float32)) |
| x = x.permute(2, 0, 1).unsqueeze(-1) |
| return { |
| "input_ids": input_ids, |
| "attention_mask": enc["attention_mask"][0], |
| "labels": labels, |
| "pixel_values_videos": enc.get("pixel_values_videos"), |
| "video_grid_thw": enc.get("video_grid_thw"), |
| "pose": x, |
| } |
|
|
|
|
| def make_collator(pad_token_id): |
| def collate(batch): |
| maxlen = max(b["input_ids"].shape[0] for b in batch) |
|
|
| def pad(x, val): |
| return torch.cat([x, x.new_full((maxlen - x.shape[0],), val)]) |
|
|
| out = { |
| "input_ids": torch.stack([pad(b["input_ids"], pad_token_id) for b in batch]), |
| "attention_mask": torch.stack([pad(b["attention_mask"], 0) for b in batch]), |
| "labels": torch.stack([pad(b["labels"], IGNORE) for b in batch]), |
| "pose": torch.stack([b["pose"] for b in batch]), |
| } |
| if batch[0]["pixel_values_videos"] is not None: |
| out["pixel_values_videos"] = torch.cat([b["pixel_values_videos"] for b in batch], dim=0) |
| out["video_grid_thw"] = torch.cat([b["video_grid_thw"] for b in batch], dim=0) |
| return out |
|
|
| return collate |
|
|
