MCPcopy Create free account
hub / github.com/THUDM/LongWriter / LMSortDataset

Class LMSortDataset

train/dataset.py:27–45  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

25 return self.input_ids.size(0)
26
27class LMSortDataset(torch.utils.data.Dataset):
28 def __init__(self, filepath):
29 self.input_ids, self.labels = self.process_data(filepath)
30 self.input_ids = self.input_ids
31 self.labels = self.labels
32
33 def process_data(self, filepath):
34 input_ids = torch.from_numpy(np.load(os.path.join(filepath, 'inputs_sort.npy')))
35 labels = torch.from_numpy(np.load(os.path.join(filepath, 'labels_sort.npy')))
36 return input_ids, labels
37
38 def __getitem__(self, idx):
39 return {
40 'input_ids': self.input_ids[idx],
41 'labels': self.labels[idx]
42 }
43
44 def __len__(self):
45 return self.input_ids.size(0)
46
47class LMPackDataset(torch.utils.data.Dataset):
48 def __init__(self, filepath):

Callers 1

Calls

no outgoing calls

Tested by

no test coverage detected