obscura/minimind
1
0
Fork 0
mirror of https://github.com/jingyaogong/minimind.git synced 2026-10-02 15:47:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

[update] pretrain load

Browse Source
This commit is contained in:
jingyaogong
2026-01-17 12:00:17 +08:00
parent aa539a824a
commit 714abcf802
1 changed files with 4 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset/lm_dataset.py
+4 -8
View File
@@ -16,14 +16,10 @@ class PretrainDataset(Dataset):
def __getitem__(self, index):
sample = self.samples[index]
encoding = self.tokenizer(
str(sample['text']),
max_length=self.max_length,
padding='max_length',
truncation=True,
return_tensors='pt'
)
input_ids = encoding.input_ids.squeeze()
tokens = self.tokenizer(str(sample['text']), add_special_tokens=False).input_ids
tokens = [self.tokenizer.bos_token_id] + tokens[:self.max_length - 2] + [self.tokenizer.eos_token_id]
input_ids = tokens + [self.tokenizer.pad_token_id] * (self.max_length - len(tokens))
input_ids = torch.tensor(input_ids, dtype=torch.long)
labels = input_ids.clone()
labels[input_ids == self.tokenizer.pad_token_id] = -100
return input_ids, labels