File size: 2,854 Bytes
fd448dd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | """Pretrain dataset for HuggingFace datasets."""
from typing import Dict
import torch
from taoTrain.config import TrainingConfig
from taoTrain.data.hf_base import BaseHFDataset
class PretrainDataset(BaseHFDataset):
"""Dataset for pretraining with raw text."""
def _preprocess(self):
"""Tokenize text data."""
dataset_config = self.config.dataset
text_column = dataset_config.text_column
def tokenize_function(examples):
# Concatenate all texts
concatenated_examples = {
k: sum(examples[k], []) for k in examples.keys()
}
total_length = len(concatenated_examples[text_column])
# We'll use max_seq_length for training
total_length = (total_length // self.config.model.max_seq_length) * self.config.model.max_seq_length
# Tokenize
tokenized = self.tokenizer(
concatenated_examples[text_column],
truncation=False, # We'll chunk below
return_special_tokens_mask=False,
)
# Chunk tokenized text
result = {
"input_ids": [],
"attention_mask": [],
}
for i in range(0, total_length, self.config.model.max_seq_length):
result["input_ids"].append(
tokenized["input_ids"][i:i + self.config.model.max_seq_length]
)
result["attention_mask"].append(
tokenized["attention_mask"][i:i + self.config.model.max_seq_length]
)
return result
# Preprocess in batches
self.data = self.data.map(
tokenize_function,
batched=True,
batch_size=100,
remove_columns=self.data.column_names,
desc="Tokenizing...",
)
def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
"""Get preprocessed sample."""
item = self.data[idx]
input_ids = torch.tensor(item["input_ids"], dtype=torch.long)
attention_mask = torch.tensor(item["attention_mask"], dtype=torch.long)
# For pretrain, labels = input_ids shifted by 1 (next token prediction)
# Position i predicts token at position i+1
labels = input_ids[1:].clone()
labels = torch.cat([labels, torch.tensor([-100])], dim=0)
# Mark padding tokens as -100 to ignore in loss computation
labels[attention_mask == 0] = -100
return {
"input_ids": input_ids,
"attention_mask": attention_mask,
"labels": labels,
}
|