github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 18:40:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse Source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot]
2024-07-10 10:44:30 +00:00
parent 33f15203d3
commit 8a9721bafe
8 changed files with 29 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
applications/ColossalChat/benchmarks/dummy_dataset.py
View File

@@ -1,5 +1,6 @@
import torch
from torch.utils.data import Dataset, DataLoader
from torch.utils.data import Dataset
class DummyLLMDataset(Dataset):
def __init__(self, keys, seq_len, size=500):
@@ -11,11 +12,11 @@ class DummyLLMDataset(Dataset):
def _generate_data(self):
data = {}
for key in self.keys:
data[key] = torch.ones(self.seq_len, dtype = torch.long)
data[key] = torch.ones(self.seq_len, dtype=torch.long)
return data
def __len__(self):
return self.size
def __getitem__(self, idx):
return {key: self.data[key] for key in self.keys}
return {key: self.data[key] for key in self.keys}