bend.models package

Submodules

bend.models.dilated_cnn module

dilated_cnn.py

A ResNet with dilated convolutions masked language model. code from https://github.com/songlab-cal/gpn

class bend.models.dilated_cnn.ConvLayer(hidden_size=None, **kwargs)[source]

Bases: Module

A layer that performs a convolution.

Build a convolutional layer.

Parameters:

hidden_size (int) – Size of the hidden state.
**kwargs – Additional arguments passed to nn.Conv1d.

forward(x)[source]

Perform a convolution.

Parameters:: x (torch.Tensor) – Input tensor.
Return type:: torch.Tensor

class bend.models.dilated_cnn.ConvNetConfig(vocab_size=7, hidden_size=512, n_layers=30, kernel_size=9, dilation_double_every=1, dilation_max=32, dilation_cycle=6, initializer_range=0.02, **kwargs)[source]

Bases: PretrainedConfig

Configuration of a ResNet with dilated convolutions.

Build the configuration of a ResNet with dilated convolutions.

Parameters:

vocab_size (int) – Size of the vocabulary.
hidden_size (int) – Size of the hidden state.
n_layers (int) – Number of layers.
kernel_size (int) – Size of the kernel.
dilation_double_every (int) – Number of layers after which the dilation is doubled.
dilation_max (int) – Maximum dilation.
dilation_cycle (int) – Number of layers after which the dilation is reset.
initializer_range (float) – Range of the initializer.

model_type: str = 'ConvNet'

class bend.models.dilated_cnn.ConvNetForMaskedLM(config, **kwargs)[source]

Bases: ConvNetPreTrainedModel

A ResNet with dilated convolutions and a head for masked language modeling.

Build a ResNet with dilated convolutions and a head for masked language modeling.

Parameters:: config (ConvNetConfig) – Configuration for the model.

forward(input_ids=None, labels=None, return_last_hidden_state=True, **kwargs)[source]

Perform a forward pass through the model.

Parameters:

input_ids (torch.Tensor) – Input tensor of nucleotide tokens with mask tokens.
labels (torch.Tensor) – Input tensor of nucleotide tokens without mask tokens.
return_last_hidden_state (bool) – Whether to return the last hidden state of the model.

Returns:

Output of the model.

Return type:

transformers.modeling_outputs.MaskedLMOutput

class bend.models.dilated_cnn.ConvNetModel(config, **kwargs)[source]

Bases: ConvNetPreTrainedModel

A ResNet with dilated convolutions.

Build a ResNet with dilated convolutions.

Parameters:: config (ConvNetConfig) – Configuration for the model.

forward(input_ids=None, **kwargs)[source]

Perform a forward pass through the model.

Parameters:: input_ids (torch.Tensor) – Input tensor of nucleotide tokens.

class bend.models.dilated_cnn.ConvNetOnlyMLMHead(config)[source]

Bases: Module

A head for masked language modeling.

Build a head for masked language modeling. This is Linear -> GELU -> LayerNorm -> Linear decoder that takes the hidden state of the model as input.

Parameters:: config (ConvNetConfig) – Configuration for the model.

forward(hidden_state)[source]

Perform a forward pass through the head.

Parameters:: hidden_state (torch.Tensor) – Hidden state of the model.
Returns:: Logits for each token in the vocabulary.
Return type:: torch.Tensor

class bend.models.dilated_cnn.ConvNetPreTrainedModel(config: PretrainedConfig, *inputs, **kwargs)[source]

Bases: PreTrainedModel

Base class for a ResNet with dilated convolutions. Hanndles the initialization, loading and saving of the model.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

base_model_prefix = 'model'

config_class: alias of ConvNetConfig

class bend.models.dilated_cnn.OneHotEmbedding(hidden_size=None)[source]

Bases: Module

A layer that performs a one-hot embedding.

Build a one-hot embedding layer.

Parameters:: hidden_size (int) – Size of the hidden state - this is equal to the size of the vocabulary.

extra_repr()[source]

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(x)[source]

Perform a one-hot embedding. If the input is already one-hot embedded (it has two dimensions), then it is returned as is.

Parameters:: x (torch.Tensor) – Input tensor.
Return type:: torch.Tensor

class bend.models.dilated_cnn.TransposeLayer[source]

Bases: Module

A layer that transposes the input.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Transpose the input.

Parameters:: x (torch.Tensor) – Input tensor.
Returns:: Transposed tensor.
Return type:: torch.Tensor

bend.models.downstream module

downstream.py

This module contains the implementations of the supervised models used in the paper.

ConvNetForSupervised: a ResNet that we train as baseline model on one-hot encodings, if no dedicated baseline architecture is available for a task.
CNN: a two-layer CNN used for all downstream tasks.

class bend.models.downstream.CNN(input_size=5, output_size=2, hidden_size=64, kernel_size=3, upsample_factor: bool | int = False, output_downsample_window=None, encoder=None, *args, **kwargs)[source]

Bases: Module

A two-layer CNN with step size 1, ReLU activation, and a linear layer.

Build a two-layer CNN with step size 1, ReLU activation, and a linear layer.

Parameters:

input_size (int) – The embedding size of the input sequence.
output_size (int) – The size of the output sequence.
hidden_size (int) – The embedding size of the hidden layer.
kernel_size (int) – The kernel size of the convolutional layers.
upsample_factor (int) – The factor by which to upsample the input.
output_downsample_window (int) – The window size for downsampling the output along the sequence dimension. This is done by taking the average of the output values in the window.

forward(x, activation='none', length=None, **kwargs)[source]

Forward pass of the CNN.

Parameters:

x (torch.Tensor) – Input tensor. Should have shape (batch_size, length, embedding_size).
activation (str) – The activation function to use. Can be ‘softmax’, ‘softplus’, ‘sigmoid’, or ‘none’.
length (int) – The actual length (in nucleotides) of the input sequence. Only required when embedding upsampling is used.

Returns:

Output tensor. Has shape (batch_size, output_length, output_size). output_length is determined by the input length, the upsampling factor, and the output downsampling window.

Return type:

torch.Tensor

class bend.models.downstream.ConvNetForSupervised(hidden_size=256, vocab_size=7, n_layers=30, kernel_size=9, dilation_double_every=1, dilation_max=32, initializer_range=0.02, dilation_cycle=6, output_size=2, hidden_size_downstream=64, kernel_size_downstream=3, upsample_factor: bool | int = False, output_downsample_window=None, **kwargs)[source]

Bases: Module

A convolutional neural network for supervised learning. We use this as a baseline, when no dedicated supervised model for a particular task is available.

Build a convolutional neural network for supervised learning.

Parameters:

hidden_size (int) – The size of the hidden layers.
vocab_size (int) – The size of the input embeddings. This is called vocab_size because in the one-hot encoding case, the embedding size will be equal to the size of the vocabulary.
n_layers (int) – The number of convolutional layers.
kernel_size (int) – The kernel size of the convolutional layers.
dilation_double_every (int) – The number of layers after which to double the dilation rate.
dilation_max (int) – The maximum dilation rate.
dilation_cycle (int) – The number of layers after which to reset the dilation rate to 1.
output_size (int) – The size of the output sequence.
hidden_size_downstream (int) – The embedding size of the hidden layer in the downstream CNN.
kernel_size_downstream (int) – The kernel size of the convolutional layers in the downstream CNN.
upsample_factor (int) – The factor by which to upsample the input.
output_downsample_window (int) – The window size for downsampling the output along the sequence dimension. This is done by taking the average of the output values in the window.

forward(x, activation='none', **kwargs)[source]

Forward pass of the model.

Parameters:

x (torch.Tensor) – Input tensor. Should have shape (batch_size, length, vocab_size).
activation (str) – The activation function to use. Can be ‘softmax’, ‘softplus’, ‘sigmoid’, or ‘none’.

Returns:

Output tensor. Has shape (batch_size, output_length, output_size). output_length is determined by the input length, the upsampling factor, and the output downsampling window.

Return type:

torch.Tensor

class bend.models.downstream.CustomDataParallel(module: T, device_ids: Sequence[int | device] | None = None, output_device: int | device | None = None, dim: int = 0)[source]

Bases: DataParallel

A custom DataParallel class that allows for attribute access to the wrapped module.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

class bend.models.downstream.TransposeLayer[source]

Bases: Module

A layer that transposes the input.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Transpose the input.

Parameters:: x (torch.Tensor) – Input tensor.
Returns:: Transposed tensor.
Return type:: torch.Tensor

class bend.models.downstream.UpsampleLayer(scale_factor=6, input_size=2560)[source]

Bases: Module

A layer that upsamples the input along the sequence dimension. This is useful when a position in the input sequence corresponds to multiple positions in the output sequence. The one-to-n mapping needs to be a fixed factor.

Build an upsampling layer.

Parameters:

scale_factor (int) – The factor by which to upsample the input.
input_size (int) – The embedding size of the input sequence.

forward(x)[source]

Upsample the input.

Parameters:: x (torch.Tensor) – Input tensor. Should have shape (batch_size, length, embedding_size).
Returns:: Upsampled tensor. Has shape (batch_size, length * scale_factor, embedding_size).
Return type:: torch.Tensor

Module contents

bend.models

This module contains the implementations of the supervised models used in the paper.

ConvNetForSupervised: a ResNet that we train as baseline model on one-hot encodings, if no dedicated baseline architecture is available for a task.
CNN: a two-layer CNN used for all downstream tasks.