File size: 4,079 Bytes
723126c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | 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
|