trainers

Trainer objects help train a model on a dataset within a DataDreamer session. All trainers derive from the Trainer base class. Trainers typically have a train() method that takes OutputDatasetColumn or OutputIterableDatasetColumn dataset columns as inputs and outputs for training. Dataset columns can be accessed from the output produced by steps.

Resuming Training

For trainers that support resumability, trainers will resume training from the last checkpoint saved automatically if interrupted.

Efficient Training Techniques

Training on Multiple GPUs

DataDreamer makes training on multiple GPUs with extremely simple and straightforward. All you need to do is pass in a list of devices to the device parameter of Trainer at construction instead of a single device. That’s it. For more advanced configuration, see the Training on Multiple GPUs page.

Quantization

See the Quantization page.

Parameter-Efficient Training

See the Parameter-Efficient Training page.

Configuring Training

You can make use of **kwargs to pass in additional keyword arguments to the underlying model’s training method.

Configuring Early Stopping

To override the default early stopping, pass in the early_stopping_patience parameter and the early_stopping_threshold parameter. To disable early stopping, set early_stopping_patience to None.

See EarlyStoppingCallback for more details.

Advanced Configuration

Most configuration you need can be done by passing **kwargs. However, if you need more advanced configuration, you can subclass the trainer class and override as needed or create your own trainer.

Model Card Generation

An automatically generated model card can be viewed by calling model_card(). The model card can be helpful for reproducibility and for sharing your work with others when published alongside your code. When publishing models, the model card will be published alongside the model.

The model card traces what steps were run to produce the model’s training dataset, what models were used, what paper citations and software licenses may apply, among other useful information. Reproducibility information such as the versions of packages used and a fingerprint hash (a signature of all steps chained together to produce the final training dataset) is also included.

Exporting and Publishing Models

You can export a trained model to disk by calling export_to_disk(). You can publish a trained model to the Hugging Face Hub by calling publish_to_hf_hub(). DataDreamer will also helpfully include instructions with a code snippet on how to load and use your model and setup the demonstration widget on the Hugging Face Hub model page with some examples.

class datadreamer.trainers.Trainer(name, force=False, verbose=None, log_level=None)

Bases: ABC

Base class for all trainers.

Parameters:

• name (str) – The name of the trainer.

• force (bool, default: False) – Whether to force run the trainer (ignore saved results).

• verbose (Optional[bool], default: None) – Whether or not to print verbose logs.

• log_level (Optional[int], default: None) – The logging level to use (DEBUG, INFO, etc.).

abstract train()

Train the model.

Return type:

Trainer

property model

An instance of the trained model after training.

abstract property model_path: str

The path to the trained model after training.

model_card()

Print the data card for the step.

unload_model()

Unloads the trained model from memory.

class datadreamer.trainers.TrainOpenAIFineTune(name, model_name, system_prompt=None, organization=None, api_key=None, base_url=None, api_version=None, retry_on_fail=True, force=False, verbose=None, log_level=None, **kwargs)

Bases: Trainer

Base class for all trainers.

Parameters:

• name (str) – The name of the trainer.

• force (bool, default: False) – Whether to force run the trainer (ignore saved results).

• verbose (Optional[bool], default: None) – Whether or not to print verbose logs.

• log_level (Optional[int], default: None) – The logging level to use (DEBUG, INFO, etc.).

train(train_input, train_output, validation_input, validation_output, truncate=True, epochs=AUTO, batch_size=AUTO, learning_rate_multiplier=AUTO, **kwargs)

Train the model.

Return type:

TrainOpenAIFineTune

property model

The name of the trained model after training.

property model_path: str

The name of the trained model after training.

class datadreamer.trainers.TrainHFClassifier(name, model_name, revision=None, trust_remote_code=False, device=None, dtype=None, quantization_config=None, peft_config=None, distributed_config=AUTO, fsdp=AUTO, fsdp_config=AUTO, force=False, verbose=None, log_level=None, **kwargs)

Bases: _TrainHFBase

Base class for all trainers.

Parameters:

• name (str) – The name of the trainer.

• force (bool, default: False) – Whether to force run the trainer (ignore saved results).

• verbose (Optional[bool], default: None) – Whether or not to print verbose logs.

• log_level (Optional[int], default: None) – The logging level to use (DEBUG, INFO, etc.).

train(train_input, train_output, validation_input, validation_output, truncate=True, epochs=3.0, batch_size=8, optim='adamw_torch', learning_rate=0.001, weight_decay=0.01, lr_scheduler_type='linear', warmup_steps=0, neftune_noise_alpha=None, seed=42, **kwargs)

Train the model.

Return type:

TrainHFClassifier

export_to_disk(path, adapter_only=False)

Export the trained model to disk.

Parameters:

• path (str) – The path to export the model to.

• adapter_only (bool, default: False) – Whether to export only the adapter.

Return type:

PreTrainedModel

Returns:

The exported model.

publish_to_hf_hub(repo_id, branch=None, private=False, token=None, adapter_only=False, is_synthetic=True, **kwargs)

Publish the model to the Hugging Face Hub.

Parameters:

• repo_id (str) – The repository ID to publish the model to.

• branch (Optional[str], default: None) – The branch to push the model to.

• private (bool, default: False) – Whether to make the model private.

• token (Optional[str], default: None) – The Hugging Face API token to use for authentication.

• adapter_only (bool, default: False) – Whether to publish only the adapter.

• is_synthetic (bool, default: True) – Whether the dataset is synthetic (applies certain metadata when publishing).

• **kwargs – Additional keyword arguments to pass to push_to_hub().

Return type:

str

Returns:

The URL to the published model.

property model: PreTrainedModel

An instance of the trained model after training.

property model_path: str

The path to the trained model after training.

property tokenizer: PreTrainedTokenizer

The tokenizer instance being used.

class datadreamer.trainers.TrainHFFineTune(name, model_name, chat_prompt_template=AUTO, system_prompt=AUTO, revision=None, trust_remote_code=False, device=None, dtype=None, quantization_config=None, peft_config=None, distributed_config=AUTO, fsdp=AUTO, fsdp_config=AUTO, force=False, verbose=None, log_level=None, **kwargs)

Bases: _TrainHFBase

Base class for all trainers.

Parameters:

• name (str) – The name of the trainer.

• force (bool, default: False) – Whether to force run the trainer (ignore saved results).

• verbose (Optional[bool], default: None) – Whether or not to print verbose logs.

• log_level (Optional[int], default: None) – The logging level to use (DEBUG, INFO, etc.).

train(train_input, train_output, validation_input, validation_output, truncate=True, epochs=3.0, batch_size=8, optim='adamw_torch', learning_rate=0.001, weight_decay=0.01, lr_scheduler_type='linear', warmup_steps=0, neftune_noise_alpha=None, seed=42, **kwargs)

Train the model.

Return type:

TrainHFFineTune

export_to_disk(path, adapter_only=False)

Export the trained model to disk.

Parameters:

• path (str) – The path to export the model to.

• adapter_only (bool, default: False) – Whether to export only the adapter.

Return type:

PreTrainedModel

Returns:

The exported model.

publish_to_hf_hub(repo_id, branch=None, private=False, token=None, adapter_only=False, is_synthetic=True, **kwargs)

Publish the model to the Hugging Face Hub.

Parameters:

• repo_id (str) – The repository ID to publish the model to.

• branch (Optional[str], default: None) – The branch to push the model to.

• private (bool, default: False) – Whether to make the model private.

• token (Optional[str], default: None) – The Hugging Face API token to use for authentication.

• adapter_only (bool, default: False) – Whether to publish only the adapter.

• is_synthetic (bool, default: True) – Whether the dataset is synthetic (applies certain metadata when publishing).

• **kwargs – Additional keyword arguments to pass to push_to_hub().

Return type:

str

Returns:

The URL to the published model.

property model: PreTrainedModel

An instance of the trained model after training.

property model_path: str

The path to the trained model after training.

property tokenizer: PreTrainedTokenizer

The tokenizer instance being used.

class datadreamer.trainers.TrainSentenceTransformer(name, model_name, trust_remote_code=False, device=None, dtype=None, peft_config=None, distributed_config=AUTO, fsdp=AUTO, fsdp_config=AUTO, force=False, verbose=None, log_level=None, **kwargs)

Bases: _TrainHFBase

Base class for all trainers.

Parameters:

• name (str) – The name of the trainer.

• force (bool, default: False) – Whether to force run the trainer (ignore saved results).

• verbose (Optional[bool], default: None) – Whether or not to print verbose logs.

• log_level (Optional[int], default: None) – The logging level to use (DEBUG, INFO, etc.).

train_with_triplets(train_anchors, train_positives, train_negatives, validation_anchors, validation_positives, validation_negatives, truncate=True, margin=DEFAULT, epochs=3.0, batch_size=8, loss=AUTO, optim='adamw_torch', learning_rate=0.001, weight_decay=0.01, lr_scheduler_type='linear', warmup_steps=0, neftune_noise_alpha=None, seed=42, **kwargs)

Return type:

TrainSentenceTransformer

train_with_positive_pairs(train_anchors, train_positives, validation_anchors, validation_positives, truncate=True, epochs=3.0, batch_size=8, loss=AUTO, optim='adamw_torch', learning_rate=0.001, weight_decay=0.01, lr_scheduler_type='linear', warmup_steps=0, neftune_noise_alpha=None, seed=42, **kwargs)

Return type:

TrainSentenceTransformer

train_with_labeled_pairs(train_anchors, train_others, train_labels, validation_anchors, validation_others, validation_labels, truncate=True, epochs=3.0, batch_size=8, loss=AUTO, optim='adamw_torch', learning_rate=0.001, weight_decay=0.01, lr_scheduler_type='linear', warmup_steps=0, neftune_noise_alpha=None, seed=42, **kwargs)

Return type:

TrainSentenceTransformer

export_to_disk(path, adapter_only=False)

Export the trained model to disk.

Parameters:

• path (str) – The path to export the model to.

• adapter_only (bool, default: False) – Whether to export only the adapter.

Return type:

SentenceTransformer

Returns:

The exported model.

publish_to_hf_hub(repo_id, private=False, token=None, adapter_only=False, is_synthetic=True, **kwargs)

Publish the model to the Hugging Face Hub.

Parameters:

• repo_id (str) – The repository ID to publish the model to.

• branch – The branch to push the model to.

• private (bool, default: False) – Whether to make the model private.

• token (Optional[str], default: None) – The Hugging Face API token to use for authentication.

• adapter_only (bool, default: False) – Whether to publish only the adapter.

• is_synthetic (bool, default: True) – Whether the dataset is synthetic (applies certain metadata when publishing).

• **kwargs – Additional keyword arguments to pass to push_to_hub().

Return type:

str

Returns:

The URL to the published model.

property model: SentenceTransformer

An instance of the trained model after training.

property model_path: str

The path to the trained model after training.

property tokenizer: PreTrainedTokenizer

The tokenizer instance being used.

class datadreamer.trainers.TrainHFDPO(name, model_name, chat_prompt_template=AUTO, system_prompt=AUTO, revision=None, trust_remote_code=False, device=None, dtype=None, quantization_config=None, peft_config=None, distributed_config=AUTO, fsdp=AUTO, fsdp_config=AUTO, force=False, verbose=None, log_level=None, **kwargs)

Bases: TrainHFFineTune

Base class for all trainers.

Parameters:

• name (str) – The name of the trainer.

• force (bool, default: False) – Whether to force run the trainer (ignore saved results).

• verbose (Optional[bool], default: None) – Whether or not to print verbose logs.

• log_level (Optional[int], default: None) – The logging level to use (DEBUG, INFO, etc.).

train(train_prompts, train_chosen, train_rejected, validation_prompts, validation_chosen, validation_rejected, truncate=True, epochs=3.0, batch_size=8, optim='adamw_torch', learning_rate=0.001, weight_decay=0.01, lr_scheduler_type='linear', warmup_steps=0, neftune_noise_alpha=None, dpo_beta=0.1, loss_type='kto_pair', disable_dropout=True, seed=42, **kwargs)

Train the model.

Return type:

TrainHFDPO

export_to_disk(path, adapter_only=False)

Export the trained model to disk.

Parameters:

• path (str) – The path to export the model to.

• adapter_only (bool, default: False) – Whether to export only the adapter.

Return type:

PreTrainedModel

Returns:

The exported model.

publish_to_hf_hub(repo_id, branch=None, private=False, token=None, adapter_only=False, is_synthetic=True, **kwargs)

Publish the model to the Hugging Face Hub.

Parameters:

• repo_id (str) – The repository ID to publish the model to.

• branch (Optional[str], default: None) – The branch to push the model to.

• private (bool, default: False) – Whether to make the model private.

• token (Optional[str], default: None) – The Hugging Face API token to use for authentication.

• adapter_only (bool, default: False) – Whether to publish only the adapter.

• is_synthetic (bool, default: True) – Whether the dataset is synthetic (applies certain metadata when publishing).

• **kwargs – Additional keyword arguments to pass to push_to_hub().

Return type:

str

Returns:

The URL to the published model.

property model: PreTrainedModel

An instance of the trained model after training.

property model_path: str

The path to the trained model after training.

property tokenizer: PreTrainedTokenizer

The tokenizer instance being used.

class datadreamer.trainers.TrainHFRewardModel(name, model_name, chat_prompt_template=AUTO, system_prompt=AUTO, revision=None, trust_remote_code=False, device=None, dtype=None, quantization_config=None, peft_config=None, distributed_config=AUTO, fsdp=AUTO, fsdp_config=AUTO, force=False, verbose=None, log_level=None, **kwargs)

Bases: TrainHFClassifier

Base class for all trainers.

Parameters:

• name (str) – The name of the trainer.

• force (bool, default: False) – Whether to force run the trainer (ignore saved results).

• verbose (Optional[bool], default: None) – Whether or not to print verbose logs.

• log_level (Optional[int], default: None) – The logging level to use (DEBUG, INFO, etc.).

train_with_pairs(train_prompts, train_chosen, train_rejected, validation_prompts, validation_chosen, validation_rejected, truncate=True, epochs=3.0, batch_size=8, optim='adamw_torch', learning_rate=0.001, weight_decay=0.01, lr_scheduler_type='linear', warmup_steps=0, neftune_noise_alpha=None, seed=42, **kwargs)

Return type:

TrainHFRewardModel

train_with_pairs_and_scores(train_prompts, train_chosen, train_chosen_scores, train_rejected, train_rejected_scores, validation_prompts, validation_chosen, validation_chosen_scores, validation_rejected, validation_rejected_scores, truncate=True, epochs=3.0, batch_size=8, optim='adamw_torch', learning_rate=0.001, weight_decay=0.01, lr_scheduler_type='linear', warmup_steps=0, neftune_noise_alpha=None, seed=42, **kwargs)

Return type:

TrainHFRewardModel

train_with_scores(train_prompts, train_generations, train_scores, validation_prompts, validation_generations, validation_scores, truncate=True, epochs=3.0, batch_size=8, optim='adamw_torch', learning_rate=0.001, weight_decay=0.01, lr_scheduler_type='linear', warmup_steps=0, neftune_noise_alpha=None, seed=42, **kwargs)

Return type:

TrainHFRewardModel

export_to_disk(path, adapter_only=False)

Export the trained model to disk.

Parameters:

• path (str) – The path to export the model to.

• adapter_only (bool, default: False) – Whether to export only the adapter.

Return type:

PreTrainedModel

Returns:

The exported model.

publish_to_hf_hub(repo_id, branch=None, private=False, token=None, adapter_only=False, is_synthetic=True, **kwargs)

Publish the model to the Hugging Face Hub.

Parameters:

• repo_id (str) – The repository ID to publish the model to.

• branch (Optional[str], default: None) – The branch to push the model to.

• private (bool, default: False) – Whether to make the model private.

• token (Optional[str], default: None) – The Hugging Face API token to use for authentication.

• adapter_only (bool, default: False) – Whether to publish only the adapter.

• is_synthetic (bool, default: True) – Whether the dataset is synthetic (applies certain metadata when publishing).

• **kwargs – Additional keyword arguments to pass to push_to_hub().

Return type:

str

Returns:

The URL to the published model.

property model: PreTrainedModel

An instance of the trained model after training.

property model_path: str

The path to the trained model after training.

property tokenizer: PreTrainedTokenizer

The tokenizer instance being used.

class datadreamer.trainers.TrainHFPPO(name, model_name, chat_prompt_template=AUTO, system_prompt=AUTO, revision=None, trust_remote_code=False, device=None, dtype=None, quantization_config=None, peft_config=None, distributed_config=AUTO, fsdp=AUTO, fsdp_config=AUTO, force=False, verbose=None, log_level=None, **kwargs)

Bases: TrainHFFineTune

Base class for all trainers.

Parameters:

• name (str) – The name of the trainer.

• force (bool, default: False) – Whether to force run the trainer (ignore saved results).

• verbose (Optional[bool], default: None) – Whether or not to print verbose logs.

• log_level (Optional[int], default: None) – The logging level to use (DEBUG, INFO, etc.).

train(train_prompts, validation_prompts, reward_model, reward_model_tokenizer=None, max_new_tokens=None, temperature=1.0, top_p=1.0, generation_kwargs=None, truncate=True, epochs=3.0, batch_size=8, optimizer=AUTO, learning_rate=1e-05, weight_decay=0.01, lr_scheduler=None, seed=42, length_sampler=None, init_kl_coef=0.2, adap_kl_ctrl=True, adap_kl_target=6.0, max_kl=AUTO, **kwargs)

Train the model.

Return type:

TrainHFPPO

export_to_disk(path, adapter_only=False)

Export the trained model to disk.

Parameters:

• path (str) – The path to export the model to.

• adapter_only (bool, default: False) – Whether to export only the adapter.

Return type:

PreTrainedModel

Returns:

The exported model.

publish_to_hf_hub(repo_id, branch=None, private=False, token=None, adapter_only=False, is_synthetic=True, **kwargs)

Publish the model to the Hugging Face Hub.

Parameters:

• repo_id (str) – The repository ID to publish the model to.

• branch (Optional[str], default: None) – The branch to push the model to.

• private (bool, default: False) – Whether to make the model private.

• token (Optional[str], default: None) – The Hugging Face API token to use for authentication.

• adapter_only (bool, default: False) – Whether to publish only the adapter.

• is_synthetic (bool, default: True) – Whether the dataset is synthetic (applies certain metadata when publishing).

• **kwargs – Additional keyword arguments to pass to push_to_hub().

Return type:

str

Returns:

The URL to the published model.

property model: PreTrainedModel

An instance of the trained model after training.

property model_path: str

The path to the trained model after training.

property tokenizer: PreTrainedTokenizer

The tokenizer instance being used.

class datadreamer.trainers.TrainSetFitClassifier(name, model_name, multi_target_strategy=AUTO, device=None, dtype=None, peft_config=None, force=False, verbose=None, log_level=None, **kwargs)

Bases: TrainHFClassifier

Base class for all trainers.

Parameters:

• name (str) – The name of the trainer.

• force (bool, default: False) – Whether to force run the trainer (ignore saved results).

• verbose (Optional[bool], default: None) – Whether or not to print verbose logs.

• log_level (Optional[int], default: None) – The logging level to use (DEBUG, INFO, etc.).

train(train_input, train_output, validation_input, validation_output, truncate=True, sampling_strategy='oversampling', end_to_end=False, epochs=3.0, batch_size=8, body_learning_rate=0.001, head_learning_rate=0.001, weight_decay=0.01, warmup_steps=0, seed=42, **kwargs)

Train the model.

Return type:

TrainSetFitClassifier

export_to_disk(path)

Export the trained model to disk.

Parameters:

• path (str) – The path to export the model to.

• adapter_only – Whether to export only the adapter.

Return type:

PreTrainedModel

Returns:

The exported model.

publish_to_hf_hub(repo_id, branch=None, private=False, token=None, is_synthetic=True, **kwargs)

Publish the model to the Hugging Face Hub.

Parameters:

• repo_id (str) – The repository ID to publish the model to.

• branch (Optional[str], default: None) – The branch to push the model to.

• private (bool, default: False) – Whether to make the model private.

• token (Optional[str], default: None) – The Hugging Face API token to use for authentication.

• adapter_only – Whether to publish only the adapter.

• is_synthetic (bool, default: True) – Whether the dataset is synthetic (applies certain metadata when publishing).

• **kwargs – Additional keyword arguments to pass to push_to_hub().

Return type:

str

Returns:

The URL to the published model.

property model: SetFitModel

An instance of the trained model after training.

property model_path: str

The path to the trained model after training.

property tokenizer: PreTrainedTokenizer

The tokenizer instance being used.