Spaces:
Running on Zero
Running on Zero
File size: 2,669 Bytes
3e936b2 | 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 | from torch.utils.data import Dataset
import numpy as np
import torch
import lmdb
import json
from pathlib import Path
from PIL import Image
import os
import datasets
class TextDataset(Dataset):
def __init__(self, prompt_path, extended_prompt_path=None):
with open(prompt_path, encoding='utf-8') as f:
self.prompt_list = [line.rstrip() for line in f]
if extended_prompt_path is not None:
with open(extended_prompt_path, encoding='utf-8') as f:
self.extended_prompt_list = [line.rstrip() for line in f]
assert len(self.extended_prompt_list) == len(self.prompt_list)
else:
self.extended_prompt_list = None
def __len__(self):
return len(self.prompt_list)
def __getitem__(self, idx):
batch = {'prompts': self.prompt_list[idx], 'idx': idx}
if self.extended_prompt_list is not None:
batch['extended_prompts'] = self.extended_prompt_list[idx]
return batch
class TwoTextDataset(Dataset):
def __init__(self, prompt_path: str, switch_prompt_path: str):
with open(prompt_path, encoding='utf-8') as f:
self.prompt_list = [line.rstrip() for line in f]
with open(switch_prompt_path, encoding='utf-8') as f:
self.switch_prompt_list = [line.rstrip() for line in f]
assert len(self.switch_prompt_list) == len(self.prompt_list), 'The two prompt files must contain the same number of lines so that each first-segment prompt is paired with exactly one second-segment prompt.'
def __len__(self):
return len(self.prompt_list)
def __getitem__(self, idx):
return {'prompts': self.prompt_list[idx], 'switch_prompts': self.switch_prompt_list[idx], 'idx': idx}
class MultiTextDataset(Dataset):
def __init__(self, prompt_path: str, field: str='prompts', cache_dir: str | None=None):
self.ds = datasets.load_dataset('json', data_files=prompt_path, split='train', cache_dir=cache_dir, streaming=False)
assert len(self.ds) > 0, 'JSONL is empty'
assert field in self.ds.column_names, f"Missing field '{field}'"
seg_len = len(self.ds[0][field])
for i, ex in enumerate(self.ds):
val = ex[field]
assert isinstance(val, list), f"Line {i} field '{field}' is not a list"
assert len(val) == seg_len, f'Line {i} list length mismatch'
self.field = field
def __len__(self):
return len(self.ds)
def __getitem__(self, idx: int):
return {'idx': idx, 'prompts_list': self.ds[idx][self.field]}
def cycle(dl):
while True:
for data in dl:
yield data
|