PebbleLM-117M / src /data /dataloader.py

Add model architecture code

27871e7 9 months ago

8.73 kB

"""
DataLoader utilities for SLM training.
Provides efficient batching and data loading for training.
"""

import os
from typing import Dict, Optional, List

import torch
from torch.utils.data import DataLoader, Dataset, DistributedSampler

from .dataset import ConversationalDataset, StreamingTextDataset, PackedDataset
from .tokenizer import SLMTokenizer


def create_dataloader(
    dataset: Dataset,
    batch_size: int,
    shuffle: bool = True,
    num_workers: int = 4,
    pin_memory: bool = None,  # Auto-detect based on device
    drop_last: bool = True,
    distributed: bool = False,
    world_size: int = 1,
    rank: int = 0,
) -> DataLoader:
    """Create a DataLoader with optimal settings.
    Args:
        dataset: The dataset to load from
        batch_size: Batch size per device
        shuffle: Whether to shuffle data
        num_workers: Number of data loading workers
        pin_memory: Pin memory for faster GPU transfer
        drop_last: Drop last incomplete batch
        distributed: Whether using distributed training
        world_size: Number of distributed processes
        rank: Current process rank
    Returns:
        Configured DataLoader
    """
    sampler = None
    if distributed:
        sampler = DistributedSampler(
            dataset,
            num_replicas=world_size,
            rank=rank,
            shuffle=shuffle,
        )
        shuffle = False  # Sampler handles shuffling

    # Auto-detect pin_memory: disable for MPS (not supported)
    if pin_memory is None:
        import torch
        pin_memory = torch.cuda.is_available()  # Only True for CUDA

    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle if sampler is None else False,
        sampler=sampler,
        num_workers=num_workers,
        pin_memory=pin_memory,
        drop_last=drop_last,
        collate_fn=default_collate_fn,
    )


def default_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
    """Collate function for batching samples.
    Args:
        batch: List of sample dictionaries
    Returns:
        Batched dictionary with stacked tensors
    """
    return {
        "input_ids": torch.stack([s["input_ids"] for s in batch]),
        "attention_mask": torch.stack([s["attention_mask"] for s in batch]),
        "labels": torch.stack([s["labels"] for s in batch]),
    }


class DataModule:
    """Data module for managing train/val dataloaders.
    Provides a unified interface for data loading during training.
    """

    def __init__(
        self,
        data_dir: str,
        tokenizer_path: str,
        max_length: int = 1024,
        batch_size: int = 32,
        num_workers: int = 4,
        val_batch_size: Optional[int] = None,
    ):
        """Initialize data module.
        Args:
            data_dir: Directory containing processed data
            tokenizer_path: Path to tokenizer.json
            max_length: Maximum sequence length
            batch_size: Training batch size
            num_workers: Number of data loading workers
            val_batch_size: Validation batch size (defaults to batch_size)
        """
        self.data_dir = data_dir
        self.max_length = max_length
        self.batch_size = batch_size
        self.val_batch_size = val_batch_size or batch_size
        self.num_workers = num_workers

        # Load tokenizer
        self.tokenizer = SLMTokenizer.from_file(tokenizer_path)

        # Datasets (created on first access)
        self._train_dataset = None
        self._val_dataset = None

    @property
    def train_dataset(self) -> Dataset:
        """Get or create training dataset."""
        if self._train_dataset is None:
            self._train_dataset = ConversationalDataset(
                data_path=self.data_dir,
                tokenizer=self.tokenizer,
                max_length=self.max_length,
                split="train",
            )
        return self._train_dataset

    @property
    def val_dataset(self) -> Dataset:
        """Get or create validation dataset."""
        if self._val_dataset is None:
            self._val_dataset = ConversationalDataset(
                data_path=self.data_dir,
                tokenizer=self.tokenizer,
                max_length=self.max_length,
                split="val",
            )
        return self._val_dataset

    def train_dataloader(
        self,
        distributed: bool = False,
        world_size: int = 1,
        rank: int = 0,
    ) -> DataLoader:
        """Get training dataloader."""
        return create_dataloader(
            self.train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.num_workers,
            drop_last=True,
            distributed=distributed,
            world_size=world_size,
            rank=rank,
        )

    def val_dataloader(self) -> DataLoader:
        """Get validation dataloader."""
        return create_dataloader(
            self.val_dataset,
            batch_size=self.val_batch_size,
            shuffle=False,
            num_workers=self.num_workers,
            drop_last=False,
        )


class StreamingDataModule:
    """Data module for streaming large datasets.
    Memory-efficient loading for large text corpora.
    """

    def __init__(
        self,
        data_files: List[str],
        tokenizer_path: str,
        max_length: int = 1024,
        batch_size: int = 32,
        num_workers: int = 4,
    ):
        """Initialize streaming data module.
        Args:
            data_files: List of text file paths
            tokenizer_path: Path to tokenizer.json
            max_length: Maximum sequence length
            batch_size: Batch size
            num_workers: Number of data loading workers
        """
        self.data_files = data_files
        self.max_length = max_length
        self.batch_size = batch_size
        self.num_workers = num_workers

        # Load tokenizer
        self.tokenizer = SLMTokenizer.from_file(tokenizer_path)

    def train_dataloader(self) -> DataLoader:
        """Get training dataloader for streaming data."""
        dataset = StreamingTextDataset(
            data_files=self.data_files,
            tokenizer=self.tokenizer,
            max_length=self.max_length,
            shuffle=True,
        )

        return DataLoader(
            dataset,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=True,
            collate_fn=default_collate_fn,
        )


def estimate_dataset_tokens(data_dir: str, tokenizer_path: str) -> Dict[str, int]:
    """Estimate total tokens in a dataset.
    Args:
        data_dir: Directory containing data files
        tokenizer_path: Path to tokenizer
    Returns:
        Dictionary with token counts
    """
    import json
    from pathlib import Path

    tokenizer = SLMTokenizer.from_file(tokenizer_path)

    total_tokens = 0
    total_samples = 0

    for file_path in Path(data_dir).glob("*.json*"):
        with open(file_path, "r") as f:
            if file_path.suffix == ".jsonl":
                samples = [json.loads(line) for line in f if line.strip()]
            else:
                samples = json.load(f)
                if not isinstance(samples, list):
                    samples = [samples]

        for sample in samples:
            if "user" in sample and "assistant" in sample:
                tokens = tokenizer.encode_conversation(
                    sample["user"], sample["assistant"]
                )
            elif "text" in sample:
                tokens = tokenizer.encode(sample["text"])
            else:
                continue

            total_tokens += len(tokens)
            total_samples += 1

    return {
        "total_tokens": total_tokens,
        "total_samples": total_samples,
        "avg_tokens_per_sample": total_tokens / max(total_samples, 1),
    }


def get_dataloader_stats(dataloader: DataLoader) -> Dict[str, float]:
    """Get statistics from a dataloader.
    Args:
        dataloader: The dataloader to analyze
    Returns:
        Dictionary with statistics
    """
    total_batches = 0
    total_tokens = 0
    total_non_pad_tokens = 0

    for batch in dataloader:
        total_batches += 1
        total_tokens += batch["input_ids"].numel()
        total_non_pad_tokens += batch["attention_mask"].sum().item()

        # Only sample first 100 batches
        if total_batches >= 100:
            break

    return {
        "batches_sampled": total_batches,
        "tokens_per_batch": total_tokens / max(total_batches, 1),
        "non_pad_ratio": total_non_pad_tokens / max(total_tokens, 1),
    }

	"""
	DataLoader utilities for SLM training.

	Provides efficient batching and data loading for training.
	"""

	import os
	from typing import Dict, Optional, List

	import torch
	from torch.utils.data import DataLoader, Dataset, DistributedSampler

	from .dataset import ConversationalDataset, StreamingTextDataset, PackedDataset
	from .tokenizer import SLMTokenizer


	def create_dataloader(
	dataset: Dataset,
	batch_size: int,
	shuffle: bool = True,
	num_workers: int = 4,
	pin_memory: bool = None, # Auto-detect based on device
	drop_last: bool = True,
	distributed: bool = False,
	world_size: int = 1,
	rank: int = 0,
	) -> DataLoader:
	"""Create a DataLoader with optimal settings.

	Args:
	dataset: The dataset to load from
	batch_size: Batch size per device
	shuffle: Whether to shuffle data
	num_workers: Number of data loading workers
	pin_memory: Pin memory for faster GPU transfer
	drop_last: Drop last incomplete batch
	distributed: Whether using distributed training
	world_size: Number of distributed processes
	rank: Current process rank

	Returns:
	Configured DataLoader
	"""
	sampler = None
	if distributed:
	sampler = DistributedSampler(
	dataset,
	num_replicas=world_size,
	rank=rank,
	shuffle=shuffle,
	)
	shuffle = False # Sampler handles shuffling

	# Auto-detect pin_memory: disable for MPS (not supported)
	if pin_memory is None:
	import torch
	pin_memory = torch.cuda.is_available() # Only True for CUDA

	return DataLoader(
	dataset,
	batch_size=batch_size,
	shuffle=shuffle if sampler is None else False,
	sampler=sampler,
	num_workers=num_workers,
	pin_memory=pin_memory,
	drop_last=drop_last,
	collate_fn=default_collate_fn,
	)


	def default_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
	"""Collate function for batching samples.

	Args:
	batch: List of sample dictionaries

	Returns:
	Batched dictionary with stacked tensors
	"""
	return {
	"input_ids": torch.stack([s["input_ids"] for s in batch]),
	"attention_mask": torch.stack([s["attention_mask"] for s in batch]),
	"labels": torch.stack([s["labels"] for s in batch]),
	}


	class DataModule:
	"""Data module for managing train/val dataloaders.

	Provides a unified interface for data loading during training.
	"""

	def __init__(
	self,
	data_dir: str,
	tokenizer_path: str,
	max_length: int = 1024,
	batch_size: int = 32,
	num_workers: int = 4,
	val_batch_size: Optional[int] = None,
	):
	"""Initialize data module.

	Args:
	data_dir: Directory containing processed data
	tokenizer_path: Path to tokenizer.json
	max_length: Maximum sequence length
	batch_size: Training batch size
	num_workers: Number of data loading workers
	val_batch_size: Validation batch size (defaults to batch_size)
	"""
	self.data_dir = data_dir
	self.max_length = max_length
	self.batch_size = batch_size
	self.val_batch_size = val_batch_size or batch_size
	self.num_workers = num_workers

	# Load tokenizer
	self.tokenizer = SLMTokenizer.from_file(tokenizer_path)

	# Datasets (created on first access)
	self._train_dataset = None
	self._val_dataset = None

	@property
	def train_dataset(self) -> Dataset:
	"""Get or create training dataset."""
	if self._train_dataset is None:
	self._train_dataset = ConversationalDataset(
	data_path=self.data_dir,
	tokenizer=self.tokenizer,
	max_length=self.max_length,
	split="train",
	)
	return self._train_dataset

	@property
	def val_dataset(self) -> Dataset:
	"""Get or create validation dataset."""
	if self._val_dataset is None:
	self._val_dataset = ConversationalDataset(
	data_path=self.data_dir,
	tokenizer=self.tokenizer,
	max_length=self.max_length,
	split="val",
	)
	return self._val_dataset

	def train_dataloader(
	self,
	distributed: bool = False,
	world_size: int = 1,
	rank: int = 0,
	) -> DataLoader:
	"""Get training dataloader."""
	return create_dataloader(
	self.train_dataset,
	batch_size=self.batch_size,
	shuffle=True,
	num_workers=self.num_workers,
	drop_last=True,
	distributed=distributed,
	world_size=world_size,
	rank=rank,
	)

	def val_dataloader(self) -> DataLoader:
	"""Get validation dataloader."""
	return create_dataloader(
	self.val_dataset,
	batch_size=self.val_batch_size,
	shuffle=False,
	num_workers=self.num_workers,
	drop_last=False,
	)


	class StreamingDataModule:
	"""Data module for streaming large datasets.

	Memory-efficient loading for large text corpora.
	"""

	def __init__(
	self,
	data_files: List[str],
	tokenizer_path: str,
	max_length: int = 1024,
	batch_size: int = 32,
	num_workers: int = 4,
	):
	"""Initialize streaming data module.

	Args:
	data_files: List of text file paths
	tokenizer_path: Path to tokenizer.json
	max_length: Maximum sequence length
	batch_size: Batch size
	num_workers: Number of data loading workers
	"""
	self.data_files = data_files
	self.max_length = max_length
	self.batch_size = batch_size
	self.num_workers = num_workers

	# Load tokenizer
	self.tokenizer = SLMTokenizer.from_file(tokenizer_path)

	def train_dataloader(self) -> DataLoader:
	"""Get training dataloader for streaming data."""
	dataset = StreamingTextDataset(
	data_files=self.data_files,
	tokenizer=self.tokenizer,
	max_length=self.max_length,
	shuffle=True,
	)

	return DataLoader(
	dataset,
	batch_size=self.batch_size,
	num_workers=self.num_workers,
	pin_memory=True,
	collate_fn=default_collate_fn,
	)


	def estimate_dataset_tokens(data_dir: str, tokenizer_path: str) -> Dict[str, int]:
	"""Estimate total tokens in a dataset.

	Args:
	data_dir: Directory containing data files
	tokenizer_path: Path to tokenizer

	Returns:
	Dictionary with token counts
	"""
	import json
	from pathlib import Path

	tokenizer = SLMTokenizer.from_file(tokenizer_path)

	total_tokens = 0
	total_samples = 0

	for file_path in Path(data_dir).glob(".json"):
	with open(file_path, "r") as f:
	if file_path.suffix == ".jsonl":
	samples = [json.loads(line) for line in f if line.strip()]
	else:
	samples = json.load(f)
	if not isinstance(samples, list):
	samples = [samples]

	for sample in samples:
	if "user" in sample and "assistant" in sample:
	tokens = tokenizer.encode_conversation(
	sample["user"], sample["assistant"]
	)
	elif "text" in sample:
	tokens = tokenizer.encode(sample["text"])
	else:
	continue

	total_tokens += len(tokens)
	total_samples += 1

	return {
	"total_tokens": total_tokens,
	"total_samples": total_samples,
	"avg_tokens_per_sample": total_tokens / max(total_samples, 1),
	}


	def get_dataloader_stats(dataloader: DataLoader) -> Dict[str, float]:
	"""Get statistics from a dataloader.

	Args:
	dataloader: The dataloader to analyze

	Returns:
	Dictionary with statistics
	"""
	total_batches = 0
	total_tokens = 0
	total_non_pad_tokens = 0

	for batch in dataloader:
	total_batches += 1
	total_tokens += batch["input_ids"].numel()
	total_non_pad_tokens += batch["attention_mask"].sum().item()

	# Only sample first 100 batches
	if total_batches >= 100:
	break

	return {
	"batches_sampled": total_batches,
	"tokens_per_batch": total_tokens / max(total_batches, 1),
	"non_pad_ratio": total_non_pad_tokens / max(total_tokens, 1),
	}