初始化项目,由ModelHub XC社区提供模型
Model: divakar-yadav/transformer-1b-chat Source: Original Platform
This commit is contained in:
79
training_code/model/data.py
Normal file
79
training_code/model/data.py
Normal file
@@ -0,0 +1,79 @@
|
||||
"""
|
||||
Data pipeline: streams and tokenizes OpenWebText for pretraining.
|
||||
Packs sequences to max_seq_len for efficiency (no padding waste).
|
||||
"""
|
||||
|
||||
import os
|
||||
import torch
|
||||
from torch.utils.data import IterableDataset, DataLoader
|
||||
from datasets import load_dataset
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
|
||||
def get_tokenizer(name: str = "mistralai/Mistral-7B-v0.1"):
|
||||
"""Use Mistral's tokenizer — 32k vocab, BPE, well-trained on diverse data."""
|
||||
tok = AutoTokenizer.from_pretrained(name, use_fast=True)
|
||||
if tok.pad_token is None:
|
||||
tok.pad_token = tok.eos_token
|
||||
return tok
|
||||
|
||||
|
||||
class PackedPretrainDataset(IterableDataset):
|
||||
"""
|
||||
Streams text from HuggingFace dataset, tokenizes on the fly,
|
||||
and packs into fixed-length sequences for maximum GPU utilization.
|
||||
"""
|
||||
|
||||
def __init__(self, tokenizer, max_seq_len: int, split: str = "train", cache_dir: str = None, seed: int = 42):
|
||||
self.tokenizer = tokenizer
|
||||
self.max_seq_len = max_seq_len
|
||||
self.split = split
|
||||
self.cache_dir = cache_dir
|
||||
self.seed = seed
|
||||
self.eos_id = tokenizer.eos_token_id
|
||||
|
||||
def _token_stream(self):
|
||||
ds = load_dataset(
|
||||
"HuggingFaceFW/fineweb-edu",
|
||||
name="sample-10BT",
|
||||
split=self.split,
|
||||
streaming=True,
|
||||
cache_dir=self.cache_dir,
|
||||
)
|
||||
ds = ds.shuffle(seed=self.seed, buffer_size=10_000)
|
||||
|
||||
for example in ds:
|
||||
text = example.get("text", "")
|
||||
if len(text.strip()) < 50:
|
||||
continue
|
||||
token_ids = self.tokenizer.encode(text, add_special_tokens=False)
|
||||
yield from token_ids
|
||||
yield self.eos_id
|
||||
|
||||
def __iter__(self):
|
||||
buffer = []
|
||||
for token_id in self._token_stream():
|
||||
buffer.append(token_id)
|
||||
if len(buffer) == self.max_seq_len + 1:
|
||||
input_ids = torch.tensor(buffer[:-1], dtype=torch.long)
|
||||
labels = torch.tensor(buffer[1:], dtype=torch.long)
|
||||
yield input_ids, labels
|
||||
buffer = []
|
||||
|
||||
|
||||
def create_dataloader(tokenizer, config, rank: int = 0, world_size: int = 1, seed_override: int = None):
|
||||
seed = seed_override if seed_override is not None else config.seed
|
||||
dataset = PackedPretrainDataset(
|
||||
tokenizer=tokenizer,
|
||||
max_seq_len=config.max_seq_len,
|
||||
split="train",
|
||||
cache_dir=config.data_cache_dir,
|
||||
seed=seed + rank,
|
||||
)
|
||||
return DataLoader(
|
||||
dataset,
|
||||
batch_size=config.batch_size_per_gpu,
|
||||
num_workers=config.num_workers,
|
||||
pin_memory=True,
|
||||
prefetch_factor=4,
|
||||
)
|
||||
Reference in New Issue
Block a user