Tasks

Note: Adapters and Backbones are typed as Callables, since some args are reserved to be automatically configured within the task. As a general rule, positional arguments are reserved while keyword arguments are free to use. For example, the backbone, adapter, optimizer, and lr_scheduler can be configured as

# My SequenceClassification task
model:
  class_path: SequenceClassification
  init_args:
    backbone:
      class_path: aido_dna_dummy
      init_args:
        use_peft: true
        lora_r: 16
        lora_alpha: 32
        lora_dropout: 0.1  
    adapter:
      class_path: modelgenerator.adapters.MLPPoolAdapter
      init_args:
        pooling: mean_pooling
        hidden_sizes: 
        - 512
        - 256
        bias: true
        dropout: 0.1
        dropout_in_middle: false
    optimizer:
      class_path: torch.optim.AdamW
      init_args:
        lr: 1e-4
    lr_scheduler:
      class_path: torch.optim.lr_scheduler.StepLR
      init_args:
        step_size: 1
        gamma: 0.1

modelgenerator.tasks.SequenceClassification

Bases: TaskInterface

Task for fine-tuning a model on a sequence classification task. Inherits from TaskInterface.

Note

Supports binary, multiclass, and multi-label classification tasks.

Parameters:

Name	Type	Description	Default
backbone	BackboneCallable	The callable that returns a backbone. Defaults to aido_dna_dummy.	aido_dna_dummy
adapter	Callable[[int, int], SequenceAdapter]	The callable that returns an adapter. Defaults to LinearCLSAdapter.	LinearCLSAdapter
n_classes	int	The number of classes in the classification task. Defaults to 2.	2
multilabel	bool	Indicate whether it is a multilabel classification task. If True, the n_classes should be set to the number of targets. Defaults to False.	False
optimizer	OptimizerCallable	The optimizer to use for training. Defaults to torch.optim.AdamW.	required
lr_scheduler	LRSchedulerCallable	The learning rate scheduler to use for training. Defaults to None.	required
batch_size	int	The batch size to use for training. Defaults to None.	required
use_legacy_adapter	bool	Whether to use the adapter from the backbone. Defaults to False. Will be deprecated once the new adapter API is fully verified.	required
strict_loading	bool	Whether to strictly load the model. Defaults to True. Set it to False if you want to replace the adapter (e.g. for continue pretraining)	required
reset_optimizer_states	bool	Whether to reset the optimizer states. Defaults to False. Set it to True if you want to replace the adapter (e.g. for continue pretraining).	required

modelgenerator.tasks.TokenClassification

Bases: SequenceClassification

Task for fine-tuning a model on a token classification task.

Parameters:

Name	Type	Description	Default
backbone	BackboneCallable	The callable that returns a backbone. Defaults to aido_dna_dummy.	required
adapter	Callable[[int, int], TokenAdapter]	The callable that returns an adapter. Defaults to LinearAdapter.	LinearAdapter
n_classes	int	The number of classes in the classification task. Defaults to 2.	required
optimizer	OptimizerCallable	The optimizer to use for training. Defaults to torch.optim.AdamW.	required
lr_scheduler	LRSchedulerCallable	The learning rate scheduler to use for training. Defaults to None.	required
batch_size	int	The batch size to use for training. Defaults to None.	required
use_legacy_adapter	bool	Whether to use the adapter from the backbone. Defaults to False. Will be deprecated once the new adapter API is fully verified.	required
strict_loading	bool	Whether to strictly load the model. Defaults to True. Set it to False if you want to replace the adapter (e.g. for continue pretraining)	required
reset_optimizer_states	bool	Whether to reset the optimizer states. Defaults to False. Set it to True if you want to replace the adapter (e.g. for continue pretraining).	required

modelgenerator.tasks.PairwiseTokenClassification

Bases: SequenceClassification

Task for fine-tuning a model on a pairwise token classification task.

Parameters:

Name	Type	Description	Default
backbone	BackboneCallable	The callable that returns a backbone. Defaults to aido_dna_dummy.	required
adapter	Callable[[int, int], TokenAdapter]	The callable that returns an adapter. Defaults to LinearAdapter.	LinearAdapter
n_classes	int	The number of classes in the classification task. Defaults to 2.	2
optimizer	OptimizerCallable	The optimizer to use for training. Defaults to torch.optim.AdamW.	required
lr_scheduler	LRSchedulerCallable	The learning rate scheduler to use for training. Defaults to None.	required
batch_size	int	The batch size to use for training. Defaults to None.	required
use_legacy_adapter	bool	Whether to use the adapter from the backbone. Defaults to False. Will be deprecated once the new adapter API is fully verified.	required
strict_loading	bool	Whether to strictly load the model. Defaults to True. Set it to False if you want to replace the adapter (e.g. for continue pretraining)	required
reset_optimizer_states	bool	Whether to reset the optimizer states. Defaults to False. Set it to True if you want to replace the adapter (e.g. for continue pretraining).	required

modelgenerator.tasks.Diffusion

Bases: TaskInterface

Task Masked Diffusion Language Modeling training and denoising on sequences (https://arxiv.org/abs/2406.07524). Inherits from TaskInterface.

Parameters:

Name	Type	Description	Default
backbone	BackboneCallable	The callable that returns a backbone. Defaults to aido_dna_dummy.	aido_dna_dummy
adapter	Callable[[int, int], TokenAdapter]	The callable that returns an adapter. Defaults to None.	None
use_legacy_adapter	bool	Whether to use the legacy adapter. Defaults to True. Will be deprecated once the new adapter API is fully verified.	True
sample_seq	bool	Whether to sample tokens during denoising. Defaults to False.	False
num_denoise_steps	int	The number of denoising steps to take. Defaults to 4.	4
sampling_temperature	float	The temperature for sampling tokens. Defaults to 1.0.	1.0
normalize_posterior_weights	bool	Whether to normalize posterior weights. Defaults to False.	False
optimizer	OptimizerCallable	The optimizer to use for training. Defaults to torch.optim.AdamW.	required
lr_scheduler	LRSchedulerCallable	The learning rate scheduler to use for training. Defaults to None.	required
batch_size	int	The batch size to use for training. Defaults to None.	required
use_legacy_adapter	bool	Whether to use the adapter from the backbone. Defaults to False. Will be deprecated once the new adapter API is fully verified.	True
strict_loading	bool	Whether to strictly load the model. Defaults to True. Set it to False if you want to replace the adapter (e.g. for continue pretraining)	required
reset_optimizer_states	bool	Whether to reset the optimizer states. Defaults to False. Set it to True if you want to replace the adapter (e.g. for continue pretraining).	required

modelgenerator.tasks.ConditionalDiffusion

Bases: Diffusion

Task for Conditional Diffusion Language Modeling training and denoising on sequences (https://arxiv.org/abs/2406.07524). Inherits from Diffusion.

Parameters:

Name	Type	Description	Default
backbone	BackboneCallable	The callable that returns a backbone. Defaults to aido_dna_dummy.	required
adapter	Callable[[int, int, int, Module], ConditionalGenerationAdapter]	The callable that returns an adapter. Defaults to ConditionalLMAdapter.	ConditionalLMAdapter
use_legacy_adapter	bool	Whether to use the pre-trained legacy adapter within the conditional decoder. Defaults to True.	True
condition_dim	int	The dimension of the condition. Defaults to 1.	1
sample_seq	bool	Whether to sample tokens during denoising. Defaults to False.	required
num_denoise_steps	int	The number of denoising steps to take. Defaults to 4.	required
sampling_temperature	float	The temperature for sampling tokens. Defaults to 1.0.	required
normalize_posterior_weights	bool	Whether to normalize posterior weights. Defaults to False.	required
optimizer	OptimizerCallable	The optimizer to use for training. Defaults to torch.optim.AdamW.	required
lr_scheduler	LRSchedulerCallable	The learning rate scheduler to use for training. Defaults to None.	required
batch_size	int	The batch size to use for training. Defaults to None.	required
use_legacy_adapter	bool	Whether to use the adapter from the backbone. Defaults to False. Will be deprecated once the new adapter API is fully verified.	True
strict_loading	bool	Whether to strictly load the model. Defaults to True. Set it to False if you want to replace the adapter (e.g. for continue pretraining)	required
reset_optimizer_states	bool	Whether to reset the optimizer states. Defaults to False. Set it to True if you want to replace the adapter (e.g. for continue pretraining).	required

modelgenerator.tasks.ConditionalMLM

Bases: TaskInterface

Task for Conditional Masked Language Modeling training and denoising on sequences. Inherits from TaskInterface.

Parameters:

Name	Type	Description	Default
backbone	BackboneCallable	The callable that returns a backbone. Defaults to aido_dna_dummy.	aido_dna_dummy
adapter	Callable[[int, int, int, Module], ConditionalGenerationAdapter]	The callable that returns an adapter. Defaults to ConditionalLMAdapter.	ConditionalLMAdapter
use_legacy_adapter	bool	Whether to use the pre-trained legacy adapter within the conditional decoder. Defaults to True.	True
condition_dim	int	The dimension of the condition. Defaults to 1.	1
use_pretrained_decoder_head	bool	Whether to use a pretrained decoder head in the condition adapter. Defaults to True.	required
sample_seq	bool	Whether to sample tokens during denoising. Defaults to False.	required
num_denoise_steps	int	The number of denoising steps to take. Defaults to 4.	required
sampling_temperature	float	The temperature for sampling tokens. Defaults to 1.0.	required
optimizer	OptimizerCallable	The optimizer to use for training. Defaults to torch.optim.AdamW.	required
lr_scheduler	LRSchedulerCallable	The learning rate scheduler to use for training. Defaults to None.	required
batch_size	int	The batch size to use for training. Defaults to None.	required
strict_loading	bool	Whether to strictly load the model. Defaults to True. Set it to False if you want to replace the adapter (e.g. for continue pretraining)	required
reset_optimizer_states	bool	Whether to reset the optimizer states. Defaults to False. Set it to True if you want to replace the adapter (e.g. for continue pretraining).	required

modelgenerator.tasks.Embed

Bases: TaskInterface

Task for getting embeddings from a backbone. This task is used only for inference.

Note

Must be used with modelgenerator.callbacks.PredictionWriter. Embeddings are stored under "predictions".

Parameters:

Name	Type	Description	Default
backbone	BackboneCallable	The callable that returns a backbone. Defaults to aido_dna_dummy.	aido_dna_dummy
batch_size	int	The batch size to use for training. Defaults to None.	required
strict_loading	bool	Whether to strictly load the model. Defaults to True. Set it to False if you want to replace the adapter (e.g. for continue pretraining)	required

modelgenerator.tasks.Inference

Bases: MLM

Task for performing inference of token probabilities with a pre-trained backbone

Parameters:

Name	Type	Description	Default
backbone	BackboneCallable	The callable that returns a backbone. Defaults to aido_dna_dummy.	aido_dna_dummy
optimizer	OptimizerCallable	The optimizer to use for training. Defaults to torch.optim.AdamW.	required
lr_scheduler	LRSchedulerCallable	The learning rate scheduler to use for training. Defaults to None.	required
batch_size	int	The batch size to use for training. Defaults to None.	required
use_legacy_adapter	bool	Whether to use the adapter from the backbone. Defaults to False. Will be deprecated once the new adapter API is fully verified.	True
strict_loading	bool	Whether to strictly load the model. Defaults to True. Set it to False if you want to replace the adapter (e.g. for continue pretraining)	required
reset_optimizer_states	bool	Whether to reset the optimizer states. Defaults to False. Set it to True if you want to replace the adapter (e.g. for continue pretraining).	required

modelgenerator.tasks.MLM

Bases: TaskInterface

Task for continuing pretraining on a masked language model. This task is used to fine-tune a model on a downstream task by continuing pretraining on a dataset with masked sequences.

Parameters:

Name	Type	Description	Default
backbone	BackboneCallable	The callable that returns a backbone. Defaults to aido_dna_dummy.	aido_dna_dummy
optimizer	OptimizerCallable	The optimizer to use for training. Defaults to torch.optim.AdamW.	required
lr_scheduler	LRSchedulerCallable	The learning rate scheduler to use for training. Defaults to None.	required
batch_size	int	The batch size to use for training. Defaults to None.	required
use_legacy_adapter	bool	Whether to use the adapter from the backbone. Defaults to False. Will be deprecated once the new adapter API is fully verified.	True
strict_loading	bool	Whether to strictly load the model. Defaults to True. Set it to False if you want to replace the adapter (e.g. for continue pretraining)	required
reset_optimizer_states	bool	Whether to reset the optimizer states. Defaults to False. Set it to True if you want to replace the adapter (e.g. for continue pretraining).	required