Streaming MLM Dataset

Streaming MLM dataset with fault-tolerant parallel processing.

Provides StreamingMLMDataset, a PyTorch Dataset that streams text from HuggingFace datasets or local Parquet files, tokenizes and applies MLM masking in parallel via process-pool workers, and caches processed batches to disk for fault-tolerant resumption.

class src.training.streaming_mlm_dataset.StreamingMLMDataset(*args: Any, **kwargs: Any)[source]

Bases: Dataset

Streaming MLM dataset with fault-tolerant parallel processing.

Streams raw text from HuggingFace datasets or local Parquet files, tokenizes and applies MLM masking in parallel using process-pool workers, and caches processed batches to disk as pickle files. Supports resumption from partially completed caches and optional context-window joining for longer contiguous sequences.

tokenizer

The tokenizer instance.

max_length

Maximum sequence length for tokenization.

mlm_probability

Fraction of tokens to mask for MLM.

max_samples

Target maximum number of processed examples.

batch_size

Number of examples per cached batch file.

num_workers

Number of parallel process-pool workers.

total_examples

Total indexed examples available for __getitem__.

cache_dir

Active cache directory path.

data_source

Label describing the data source used.

__init__(tokenizer: Any, max_length: int = 512, mlm_probability: float = 0.15, max_samples: int = 1000000, batch_size: int = 5000, num_workers: int | None = None, cache_dir: str | None = None, max_batch_in_memory: int | None = None, parallel_chunksize: int = 32, local_parquet_dir: str | None = None, prefer_local_cache: bool = True, stream_local_parquet: bool = True, join_temp_data_context_window: int = 0, join_temp_data_min_remainder_tokens: int = 128, apply_mlm_mask: bool = True)[source]

Initialize the streaming MLM dataset.

Parameters:

  • tokenizer – Tokenizer instance (SPM or HuggingFace).

  • max_length – Maximum sequence length after tokenization.

  • mlm_probability – Fraction of eligible tokens to mask.

  • max_samples – Target maximum number of examples to produce.

  • batch_size – Number of examples per cached batch file.

  • num_workers – Number of parallel workers. Defaults to min(8, cpu_count // 2).

  • cache_dir – Directory for cached batch files. Defaults to ./temp_data/dataset_cache.

  • max_batch_in_memory – Maximum in-memory batch size during processing. Defaults to min(batch_size, 2000).

  • parallel_chunksize – Chunksize for ProcessPoolExecutor.map.

  • local_parquet_dir – Optional path to local Parquet files.

  • prefer_local_cache – If True, prefer local Parquet over streaming from HuggingFace.

  • stream_local_parquet – If True, stream local Parquet files instead of loading them entirely into memory.

  • join_temp_data_context_window – If > 0, join cached examples into contiguous sequences of this length.

  • join_temp_data_min_remainder_tokens – Minimum remaining tokens to form a final joined sequence.

  • apply_mlm_mask – If True (default), apply BERT-style MLM masking at cache-build time. If False (e.g. for causal language modeling), sequences are stored unmasked and labels == input_ids so the trainer can compute shifted next-token loss.

cleanup_cache()[source]
get_stats() → Dict[str, Any][source]

Return dataset statistics and configuration summary.

Returns:

Dictionary with keys including total_examples, completed_batches, total_samples_processed, batch_size, num_workers, cache_dir, data_source, and tokenizer backend.