wehub-resource-sync
108 lines
3.7 KiB
Python
108 lines
3.7 KiB
Python
import json
|
|
import os
|
|
|
|
import torch
|
|
from torch.utils.data import Dataset, Sampler
|
|
|
|
_LOCAL_DATASET_ROOT = os.path.join(os.path.dirname(__file__), "dataset")
|
|
|
|
|
|
def _candidate_dataset_dirs(dataset):
|
|
yield dataset
|
|
|
|
dataset_name = os.path.basename(os.path.normpath(dataset))
|
|
if not dataset_name:
|
|
return
|
|
|
|
local_dataset_dir = os.path.join(_LOCAL_DATASET_ROOT, dataset_name)
|
|
if os.path.normpath(local_dataset_dir) != os.path.normpath(dataset):
|
|
yield local_dataset_dir
|
|
|
|
|
|
def _resolve_dataset_file(dataset, filename):
|
|
checked_paths = []
|
|
for dataset_dir in _candidate_dataset_dirs(dataset):
|
|
file_path = os.path.join(dataset_dir, filename)
|
|
checked_paths.append(file_path)
|
|
if os.path.exists(file_path):
|
|
return file_path
|
|
|
|
checked_str = ", ".join(checked_paths)
|
|
raise FileNotFoundError(f"Could not find dataset file {filename!r}. Checked: {checked_str}")
|
|
|
|
|
|
class TextPromptDataset(Dataset):
|
|
def __init__(self, dataset, split="train"):
|
|
self.file_path = _resolve_dataset_file(dataset, f"{split}.txt")
|
|
with open(self.file_path, encoding="utf-8") as f:
|
|
self.prompts = [line.strip() for line in f.readlines()]
|
|
|
|
def __len__(self):
|
|
return len(self.prompts)
|
|
|
|
def __getitem__(self, idx):
|
|
return {"prompt": self.prompts[idx], "metadata": {}}
|
|
|
|
@staticmethod
|
|
def collate_fn(examples):
|
|
prompts = [example["prompt"] for example in examples]
|
|
metadatas = [example["metadata"] for example in examples]
|
|
return prompts, metadatas
|
|
|
|
|
|
class GenevalPromptDataset(Dataset):
|
|
def __init__(self, dataset, split="train"):
|
|
self.file_path = _resolve_dataset_file(dataset, f"{split}_metadata.jsonl")
|
|
with open(self.file_path, encoding="utf-8") as f:
|
|
self.metadatas = [json.loads(line) for line in f]
|
|
self.prompts = [item["prompt"] for item in self.metadatas]
|
|
|
|
def __len__(self):
|
|
return len(self.prompts)
|
|
|
|
def __getitem__(self, idx):
|
|
return {"prompt": self.prompts[idx], "metadata": self.metadatas[idx]}
|
|
|
|
@staticmethod
|
|
def collate_fn(examples):
|
|
prompts = [example["prompt"] for example in examples]
|
|
metadatas = [example["metadata"] for example in examples]
|
|
return prompts, metadatas
|
|
|
|
|
|
class DistributedKRepeatSampler(Sampler):
|
|
def __init__(self, dataset, batch_size, k, num_replicas, rank, seed=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 = self.num_replicas * self.batch_size
|
|
assert (
|
|
self.total_samples % self.k == 0
|
|
), f" total {self.total_samples} {self.k} k can not div n*b, k{k}-num_replicas{num_replicas}-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_indices = [idx for idx in indices for _ in range(self.k)]
|
|
|
|
shuffled_indices = torch.randperm(len(repeated_indices), generator=g).tolist()
|
|
shuffled_samples = [repeated_indices[i] for i in shuffled_indices]
|
|
|
|
per_card_samples = []
|
|
for i in range(self.num_replicas):
|
|
start = i * self.batch_size
|
|
end = start + self.batch_size
|
|
per_card_samples.append(shuffled_samples[start:end])
|
|
yield per_card_samples[self.rank]
|
|
|
|
def set_epoch(self, epoch):
|
|
self.epoch = epoch
|