kempnerforge.config.data

Data pipeline configuration.

Functions

check_token_dtype(value, key)

Reject a token_dtype that is neither "" nor one of TOKEN_DTYPES.

Classes

DataConfig

Data pipeline settings.

DatasetSource

A single data source in a multi-dataset mixture.

TrainingPhase

A training phase with custom dataset weights and LR scaling.
kempnerforge.config.data.check_token_dtype(value, key)[source]

Reject a token_dtype that is neither "" nor one of TOKEN_DTYPES.

Parameters:
Return type:

None

class kempnerforge.config.data.DatasetSource[source]

Bases: object

A single data source in a multi-dataset mixture.

Either path (pre-tokenized) or hf_name (HuggingFace) must be set. weight controls the relative sampling probability. token_dtype overrides data.token_dtype for the .bin shards under path.

path: str = ''
weight: float = 1.0
name: str = ''
hf_name: str = ''
hf_config: str = ''
token_dtype: str = ''
__init__(path='', weight=1.0, name='', hf_name='', hf_config='', token_dtype='')
Parameters:
Return type:

None

class kempnerforge.config.data.TrainingPhase[source]

Bases: object

A training phase with custom dataset weights and LR scaling.

Used for data annealing: at start_step, the mixture sampler switches to dataset_weights and the learning rate is multiplied by lr_scale.

start_step: int = 0
dataset_weights: dict[str, float]
lr_scale: float = 1.0
__init__(start_step=0, dataset_weights=<factory>, lr_scale=1.0)
Parameters:
Return type:

None

class kempnerforge.config.data.DataConfig[source]

Bases: object

Data pipeline settings.

token_dtype is the integer type of the raw .bin token shards, "uint16" or "uint32", for dataset_path and every datasets source that does not set its own. Left empty, the metadata.yaml next to the shards decides, as MemoryMappedDataset describes. .npy shards carry their dtype in their header and ignore it.

dataset_path: str = ''
file_pattern: str = '*.npy'
token_dtype: str = ''
tokenizer_path: str = ''
num_workers: int = 4
pin_memory: bool = True
prefetch_factor: int = 2
hf_dataset_name: str | None = None
hf_dataset_config: str | None = None
hf_dataset_split: str = 'train'
hf_dataset_text_field: str = 'text'
hf_dataset_image_field: str = 'image'
hf_dataset_prompt_field: str = ''
hf_image_size: int = 224
hf_streaming: bool = False
pack_sequences: bool = False
datasets: list[DatasetSource]
mix_temperature: float = 1.0
phases: list[TrainingPhase]
anneal_start_step: int = 0
anneal_weights: dict[str, float]
__init__(dataset_path='', file_pattern='*.npy', token_dtype='', tokenizer_path='', num_workers=4, pin_memory=True, prefetch_factor=2, hf_dataset_name=None, hf_dataset_config=None, hf_dataset_split='train', hf_dataset_text_field='text', hf_dataset_image_field='image', hf_dataset_prompt_field='', hf_image_size=224, hf_streaming=False, pack_sequences=False, datasets=<factory>, mix_temperature=1.0, phases=<factory>, anneal_start_step=0, anneal_weights=<factory>)
Parameters:
Return type:

None