From 714abcf802c31a67c572d7e3ae57b001d86d1c4d Mon Sep 17 00:00:00 2001
From: jingyaogong <gongjy.cs@qq.com>
Date: Sat, 17 Jan 2026 12:00:17 +0800
Subject: [PATCH] [update] pretrain load

---
 dataset/lm_dataset.py | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/dataset/lm_dataset.py b/dataset/lm_dataset.py
index 51db696..1a57ce9 100644
--- a/dataset/lm_dataset.py
+++ b/dataset/lm_dataset.py
@@ -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