bio-attention.embed Reference
Here be modules that will help embed inputs in different ways
- class bio_attention.embed.DiscreteEmbedding(num_embeddings: int, embedding_dim: int, cls: bool = False, cls_axis: int = -2, init_cls_as_mask: bool = False)
Embedding module for discrete data. Uses nn.Embedding with added support for masking tokens and cls tokens. Masking tokens are recognized as torch.nans. Therefore, input should be float dtype. WARNING: Because this module expects a float dtype for what should be integers, there is a limit to the vocab size you can effectively accurately represent. In practice, vocab sizes up to 1 million in size should be faithfully represented with float32. For float16, expect issues.
- Parameters:
num_embeddings (int) – number of discrete classes.
embedding_dim (int) – umber of hidden dimensions of the embeddings
cls (bool, optional) – whether to include cls token, by default False
cls_axis (int, optional) – which axis in the embedded input to add cls token to, by default -2
init_cls_as_mask (bool, optional) – initialize the cls token to be equal to the masking token, by default False If False, cls will be initialized as N(0,1), just like all other tokens.
- class bio_attention.embed.ContinuousEmbedding(embedding_dim: int, depth: int = 1, norm: bool = False, cls: bool = False, cls_axis: int = -2, init_cls_as_bias: bool = False, init_mask_as_bias: bool = False, init_cls_as_mask: bool = False)
Embedding module for continuous data. Uses nn.Linear with added support for masking tokens cls tokens. Masking tokens are recognized as torch.nans. Therefore, input should be float dtype.
- Parameters:
embedding_dim (int) – number of hidden dimensions of the embeddings
depth (int, optional) – depth of the embedding module. Will follow this structure: Linear -> (Norm if norm) { -> Swish -> Linear -> (Norm if norm) } * (depth-1) By default 1, for a linear embedding.
norm (bool, optional) – Whether to use LayerNorms in the embedding net, by default False
cls (bool, optional) – whether to include cls token, by default False
cls_axis (int, optional) – which axis in the embedded input to add cls token to, by default -2
init_cls_as_bias (bool, optional) – initialize the cls token to be equal to the bias in the embedding network, by default False Mutually excluse with init_cls_as_mask If a multi-layer embedding net is taken, the bias is still defined as the output of said network with a zero input. If both this and init_cls_as_mask are False, cls will be initialized as U(-1,1), just like the bias by default.
init_mask_as_bias (bool, optional) – initialize the mask token to be equal to the bias in the embedding network, by default False If a multi-layer embedding net is taken, the bias is still defined as the output of said network with a zero input. If False, mask will be initialized as U(-1,1), just like the bias by default.
init_cls_as_mask (bool, optional) – initialize the cls token to be equal to the masking token, by default False. Mutually excluse with init_cls_as_bias If both this and init_cls_as_bias are False, cls will be initialized as U(-1,1), just like the bias by default.
- class bio_attention.embed.BinaryEmbedding(embedding_dim: int, lambda_fn: ~typing.Callable = <function BinaryEmbedding.<lambda>>, cls: bool = False, cls_axis: int = -2, init_cls_as_mask: bool = False)
Embedding module that binarizes input data and then embeds them using DiscreteEmbedding. Supports masking tokens as NaNs. For efficiency, consider tokenizing data before training using this module and using DiscreteEmbedding in the actual model.
- Parameters:
embedding_dim (int) – number of hidden dimensions of the embeddings
lambda_fn (Callable, optional) – function that binarizes data, by default lambdax:x==0
cls (bool, optional) – whether to include cls token, by default False
cls_axis (int, optional) – which axis in the embedded input to add cls token to, by default -2
init_cls_as_mask (bool, optional) – initialize the cls token to be equal to the masking token, by default False If False, cls will be initialized as N(0,1), just like all other tokens.
- class bio_attention.embed.BinEmbedding(embedding_dim: int, bins: Tensor, left_closed: bool = True, cls: bool = False, cls_axis: int = -2, init_cls_as_mask: bool = False)
Embedding module that bins input data and then embeds them using DiscreteEmbedding. Supports masking tokens as NaNs. For efficiency, consider tokenizing data before training using this module and using DiscreteEmbedding in the actual model.
- Parameters:
embedding_dim (int) – number of hidden dimensions of the embeddings
bins (torch.Tensor) – 1-D tensor that delineates bins. Should be monotonically increasing.
left_closed (bool, optional) – Whether the bins should start including the lower bound or not, by default True
cls (bool, optional) – whether to include cls token, by default False
cls_axis (int, optional) – which axis in the embedded input to add cls token to, by default -2
init_cls_as_mask (bool, optional) – initialize the cls token to be equal to the masking token, by default False If False, cls will be initialized as N(0,1), just like all other tokens.
- class bio_attention.embed.CLS(dim: int, cls_axis: int = -2, init: Tensor | None = None)
Module for adding a CLS token to an input
- Parameters:
dim (int) – number of hidden dimensions of the embeddings
cls_axis (int, optional) – which axis in the embedded input to add cls token to, by default -2
init (Optional[torch.Tensor], optional) – initial CLS token weights, by default None, for uniform [-1, 1] initialization.