262 lines
7.2 KiB
Python
262 lines
7.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Text prompt datasets and samplers for RL training."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
|
|
import torch
|
|
from torch.utils.data import DataLoader, Dataset, Sampler
|
|
|
|
from fastvideo.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class TextPromptDataset(Dataset):
|
|
"""Load plain text prompts from train.txt / test.txt."""
|
|
|
|
def __init__(self, dataset: str, split: str = "train"):
|
|
self.file_path = os.path.join(
|
|
dataset, f"{split}.txt"
|
|
)
|
|
with open(self.file_path) as f:
|
|
self.prompts = [
|
|
line.strip() for line in f.readlines()
|
|
]
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.prompts)
|
|
|
|
def __getitem__(
|
|
self, idx: int | tuple[int, int]
|
|
) -> dict:
|
|
epoch_tag = None
|
|
if isinstance(idx, tuple):
|
|
epoch_tag, idx = idx
|
|
return {
|
|
"epoch": epoch_tag,
|
|
"prompt": self.prompts[idx],
|
|
"metadata": {},
|
|
}
|
|
|
|
@staticmethod
|
|
def collate_fn(
|
|
examples: list[dict],
|
|
) -> tuple[int | None, list[str], list[dict]]:
|
|
epoch_tags = [
|
|
example.get("epoch") for example in examples
|
|
]
|
|
epoch_tag = (
|
|
epoch_tags[0]
|
|
if all(tag == epoch_tags[0] for tag in epoch_tags)
|
|
else None
|
|
)
|
|
prompts = [
|
|
example["prompt"] for example in examples
|
|
]
|
|
metadatas = [
|
|
example["metadata"] for example in examples
|
|
]
|
|
return epoch_tag, prompts, metadatas
|
|
|
|
|
|
class JsonPromptDataset(Dataset):
|
|
"""Load prompts from JSONL files."""
|
|
|
|
def __init__(self, dataset: str, split: str = "train"):
|
|
self.file_path = os.path.join(
|
|
dataset, f"{split}.json"
|
|
)
|
|
self._prompts: list[str] = []
|
|
self._metadatas: list[dict] = []
|
|
self._load_all_prompts()
|
|
|
|
def _load_all_prompts(self):
|
|
with open(self.file_path, encoding="utf-8") as f:
|
|
for raw_line in f:
|
|
line = raw_line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
item = json.loads(line)
|
|
prompt = item.get("prompt", "")
|
|
if prompt:
|
|
self._prompts.append(prompt)
|
|
metadata = {
|
|
k: v
|
|
for k, v in item.items()
|
|
if k != "prompt"
|
|
}
|
|
self._metadatas.append(metadata)
|
|
except json.JSONDecodeError as e:
|
|
logger.warning(
|
|
"Skipping invalid JSON line: %s",
|
|
e,
|
|
)
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._prompts)
|
|
|
|
def __getitem__(
|
|
self, idx: int | tuple[int, int]
|
|
) -> dict:
|
|
epoch_tag = None
|
|
if isinstance(idx, tuple):
|
|
epoch_tag, idx = idx
|
|
return {
|
|
"epoch": epoch_tag,
|
|
"prompt": self._prompts[idx],
|
|
"metadata": (
|
|
self._metadatas[idx]
|
|
if self._metadatas
|
|
else {}
|
|
),
|
|
}
|
|
|
|
@staticmethod
|
|
def collate_fn(
|
|
examples: list[dict],
|
|
) -> tuple[int | None, list[str], list[dict]]:
|
|
epoch_tags = [
|
|
example.get("epoch") for example in examples
|
|
]
|
|
epoch_tag = (
|
|
epoch_tags[0]
|
|
if all(tag == epoch_tags[0] for tag in epoch_tags)
|
|
else None
|
|
)
|
|
prompts = [
|
|
example["prompt"] for example in examples
|
|
]
|
|
metadatas = [
|
|
example["metadata"] for example in examples
|
|
]
|
|
return epoch_tag, prompts, metadatas
|
|
|
|
|
|
class DistributedKRepeatSampler(Sampler):
|
|
|
|
def __init__(
|
|
self,
|
|
dataset: Dataset,
|
|
batch_size: int,
|
|
k: int,
|
|
num_replicas: int,
|
|
rank: int,
|
|
seed: int = 0,
|
|
):
|
|
self.dataset = dataset
|
|
self.batch_size = batch_size
|
|
self.k = k
|
|
self.num_replicas = num_replicas
|
|
self.rank = rank
|
|
self.seed = seed
|
|
self.total_samples = num_replicas * batch_size
|
|
assert self.total_samples % self.k == 0, (
|
|
f"k cannot divide n*b: k={k} "
|
|
f"num_replicas={num_replicas} "
|
|
f"batch_size={batch_size}"
|
|
)
|
|
self.m = self.total_samples // self.k
|
|
self.epoch = 0
|
|
|
|
def __iter__(self):
|
|
while True:
|
|
g = torch.Generator()
|
|
g.manual_seed(self.seed + self.epoch)
|
|
indices = torch.randperm(
|
|
len(self.dataset), generator=g
|
|
)[: self.m].tolist()
|
|
repeated = [
|
|
idx
|
|
for idx in indices
|
|
for _ in range(self.k)
|
|
]
|
|
shuffled_idx = torch.randperm(
|
|
len(repeated), generator=g
|
|
).tolist()
|
|
shuffled = [
|
|
repeated[i] for i in shuffled_idx
|
|
]
|
|
per_card = []
|
|
for i in range(self.num_replicas):
|
|
start = i * self.batch_size
|
|
end = start + self.batch_size
|
|
per_card.append(
|
|
[
|
|
(self.epoch, idx)
|
|
for idx in shuffled[start:end]
|
|
]
|
|
)
|
|
yield per_card[self.rank]
|
|
|
|
def set_epoch(self, epoch: int):
|
|
self.epoch = epoch
|
|
|
|
|
|
def build_prompt_dataloaders(
|
|
prompt_dataset_path: str,
|
|
prompt_fn: str,
|
|
sample_batch_size: int,
|
|
eval_batch_size: int,
|
|
num_video_per_prompt: int,
|
|
num_processes: int,
|
|
process_index: int,
|
|
seed: int,
|
|
) -> tuple[DataLoader, DataLoader, DistributedKRepeatSampler]:
|
|
"""Build train/eval prompt dataloaders.
|
|
|
|
Returns:
|
|
(train_dataloader, test_dataloader, train_sampler)
|
|
"""
|
|
if prompt_fn == "general_ocr":
|
|
train_ds = TextPromptDataset(
|
|
prompt_dataset_path, "train"
|
|
)
|
|
test_ds = TextPromptDataset(
|
|
prompt_dataset_path, "test"
|
|
)
|
|
collate = TextPromptDataset.collate_fn
|
|
elif prompt_fn == "filtered_prompts":
|
|
train_ds = JsonPromptDataset(
|
|
prompt_dataset_path, "train"
|
|
)
|
|
test_ds = JsonPromptDataset(
|
|
prompt_dataset_path, "test"
|
|
)
|
|
collate = JsonPromptDataset.collate_fn
|
|
else:
|
|
msg = (
|
|
f"Unsupported prompt_fn: {prompt_fn}. "
|
|
"Use 'general_ocr' or 'filtered_prompts'."
|
|
)
|
|
raise NotImplementedError(msg)
|
|
|
|
train_sampler = DistributedKRepeatSampler(
|
|
dataset=train_ds,
|
|
batch_size=sample_batch_size,
|
|
k=num_video_per_prompt,
|
|
num_replicas=num_processes,
|
|
rank=process_index,
|
|
seed=seed,
|
|
)
|
|
|
|
train_dl = DataLoader(
|
|
train_ds,
|
|
batch_sampler=train_sampler,
|
|
num_workers=1,
|
|
collate_fn=collate,
|
|
prefetch_factor=1,
|
|
persistent_workers=False,
|
|
)
|
|
test_dl = DataLoader(
|
|
test_ds,
|
|
batch_size=eval_batch_size,
|
|
collate_fn=collate,
|
|
shuffle=False,
|
|
num_workers=8,
|
|
)
|
|
return train_dl, test_dl, train_sampler
|