github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[update] pretrain load

pull/825/head
jingyaogong 2026-01-17 12:00:17 +08:00
parent b832bd5253
commit 33b0dfbec5
1 changed files with 4 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
dataset/lm_dataset.py
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