/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
applications/ColossalChat/benchmarks/dummy_dataset.py
30 строк
799 B
YeAnbang
refactor evaluation
22 июл 2024, 08:57
22 июл 2024, 08:57
12fe8b5
Код
Авторство
О чём код?
from typing import Callable from torch.utils.data import Dataset class DummyLLMDataset(Dataset): def __init__(self, keys, seq_len, size=500, gen_fn={}): self.keys = keys self.gen_fn = gen_fn self.seq_len = seq_len self.data = self._generate_data() self.size = size def _generate_data(self): data = {} for key in self.keys: if key in self.gen_fn: data[key] = self.gen_fn[key] else: data[key] = [1] * self.seq_len return data def __len__(self): return self.size def __getitem__(self, idx): return { key: self.data[key] if not isinstance(self.data[key], Callable) else self.data[key](idx) for key in self.keys }