Atacformer Module

TrainingTokenizer

class TrainingTokenizer(*args, **kwargs)

Bases: Tokenizer, PreTrainedTokenizerBase

A special training tokenizer. This class is a subclass of both our Tokenizer and PreTrainedTokenizerBase. This is because the data collator requires a collator that is both a Tokenizer and a PreTrainedTokenizerBase. This is a workaround to make the code work with our Tokenizer.

Properties

added_tokens_decoder

@property
def added_tokens_decoder

added_tokens_encoder

@property
def added_tokens_encoder

all_special_ids

@property
def all_special_ids

Returns a list of all special token ids.

Methods

num_special_tokens_to_add

def num_special_tokens_to_add(pair=False)

save_vocabulary

def save_vocabulary(save_directory, filename_prefix=None)

convert_tokens_to_string

def convert_tokens_to_string(tokens: List[str]) -> str

AtacformerModel

class AtacformerModel(config: AtacformerConfig)

Bases: AtacformerPreTrainedModel

atacformer model with a simple embedding layer that skips positional encoding.

Attributes

• config_class = AtacformerConfig
• base_model_prefix = 'atacformer'
• embeddings
• encoder

Methods

get_input_embeddings

def get_input_embeddings()

set_input_embeddings

def set_input_embeddings(value)

forward

def forward(input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, return_dict: Optional[bool] = None, **kwargs: Optional[dict]) -> torch.Tensor

Forward pass of the model.

Args: input_ids (torch.LongTensor): Input tensor of shape (batch_size, sequence_length). attention_mask (torch.Tensor, optional): Mask to avoid performing attention on padding token indices. return_dict (bool, optional): Whether to return the outputs as a dict or a tuple. Raises: ValueError: If input_ids is not provided. Returns: torch.Tensor: The output of the model. It will be a tensor of shape (batch_size, sequence_length, hidden_size).

AtacformerForMaskedLM

class AtacformerForMaskedLM(config: AtacformerConfig)

Bases: EncodeTokenizedCellsMixin, AtacformerPreTrainedModel

Attributes

• atacformer
• lm_head
• loss_fct

Methods

get_input_embeddings

def get_input_embeddings()

get_output_embeddings

def get_output_embeddings()

forward

def forward(input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.LongTensor] = None, return_dict: Optional[bool] = None) -> Union[Tuple[torch.Tensor], MaskedLMOutput]

Forward pass of the model.

Args: input_ids (torch.LongTensor): Input tensor of shape (batch_size, sequence_length). attention_mask (torch.Tensor, optional): Mask to avoid performing attention on padding token indices. labels (torch.LongTensor, optional): Labels for masked language modeling. return_dict (bool, optional): Whether to return the outputs as a dict or a tuple. Returns: MaskedLMOutput: The output of the model.

AtacformerForReplacedTokenDetection

class AtacformerForReplacedTokenDetection(config: AtacformerConfig)

Bases: EncodeTokenizedCellsMixin, AtacformerPreTrainedModel

Atacformer model for replaced token detection. This model uses the ELECTRA framework to train a discriminator (this model) to detect replaced tokens.

https://arxiv.org/abs/2003.10555

Attributes

• atacformer
• discriminator

Methods

get_input_embeddings

def get_input_embeddings()

forward

def forward(input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.LongTensor] = None, return_dict: Optional[bool] = None) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]

Forward pass of the model.

Args: input_ids (torch.LongTensor): Input tensor of shape (batch_size, sequence_length). attention_mask (torch.Tensor, optional): Mask to avoid performing attention on padding token indices. labels (torch.LongTensor, optional): Labels for masked language modeling. return_dict (bool, optional): Whether to return the outputs as a dict or a tuple. Returns: TokenClassifierOutput: The output of the model.

AtacformerForCellClustering

class AtacformerForCellClustering(config: AtacformerConfig)

Bases: EncodeTokenizedCellsMixin, AtacformerPreTrainedModel

Atacformer model for cell clustering. It follows a similar learning framework to SentenceBERT (SBERT), where the model is trained to minimize the distance between positive pairs and maximize the distance between negative pairs.

Attributes

• atacformer
• triplet_loss

Methods

get_input_embeddings

def get_input_embeddings()

forward

def forward(input_ids_anchor: torch.LongTensor, attention_mask_anchor: Optional[torch.Tensor] = None, input_ids_positive: torch.LongTensor = None, attention_mask_positive: Optional[torch.Tensor] = None, input_ids_negative: torch.LongTensor = None, attention_mask_negative: Optional[torch.Tensor] = None, return_dict: Optional[bool] = None) -> Union[Tuple[torch.Tensor], BaseModelOutput]

AtacformerForUnsupervisedBatchCorrection

class AtacformerForUnsupervisedBatchCorrection(config: AtacformerConfig)

Bases: EncodeTokenizedCellsMixin, AtacformerPreTrainedModel

Atacformer model for batch correction. It follows a similar learning framework to the one used in domain adaptation, where the model is trained to correct batch effects in the embeddings.

Attributes

• atacformer
• discriminator
• grl
• batch_prediction_head
• lambda_adversarial

Methods

get_input_embeddings

def get_input_embeddings()

forward

def forward(input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, batch_labels: Optional[torch.Tensor] = None, return_dict: Optional[bool] = None) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]

Forward pass of the model.

Args: input_ids (torch.LongTensor): Input tensor of shape (batch_size, sequence_length). attention_mask (torch.Tensor, optional): Mask to avoid performing attention on padding token indices. labels (torch.Tensor, optional): Labels for masked language modeling (ELECTRA). batch_labels (torch.Tensor, optional): Labels for batch prediction. Should be of shape (batch_size,). return_dict (bool, optional): Whether to return the outputs as a dict or a tuple. Returns: BaseModelOutput: The output of the model.

AtacformerConfig

class AtacformerConfig(use_pos_embeddings: bool = True, vocab_size: int = 890711, max_position_embeddings: int = 8192, hidden_size: int = 384, intermediate_size: int = 1536, num_hidden_layers: int = 6, num_attention_heads: int = 8, pad_token_id: int = 890705, eos_token_id: int = 890708, bos_token_id: int = 890709, cls_token_id: int = 890707, sep_token_id: int = 890710, sparse_prediction: bool = True, norm_eps: float = 1e-05, embedding_dropout: float = 0.0, initializer_range: float = 0.02, initializer_cutoff_factor: float = 2.0, tie_word_embeddings: bool = True, num_batches: int = None, lambda_adv: float = 1.0, grl_alpha: float = 1.0, bc_unfreeze_last_n_layers: int = 2, **kwargs)

Bases: PretrainedConfig

This is the configuration class to store the configuration of an AtacformerModel. it inherits from [ModernBertConfig] and expands it for Atacformer specific settings. instantiating a configuration with the defaults will yield a similar configuration to that of the modernbert base configuration.

Args: use_pos_embeddings (bool, optional, defaults to True): whether to use positional embeddings. vocab_size (int, optional, defaults to 890711): vocabulary size tailored for genomic regions. max_position_embeddings (int, optional, defaults to 8192): the maximum sequence length that this model might ever be used with. hidden_size (int, optional, defaults to 384): the size of the encoder layers and the pooler layer. intermediate_size (int, optional, defaults to 1536): the size of the "intermediate" (often named feed-forward) layer in the transformer. num_hidden_layers (int, optional, defaults to 6): the number of hidden layers in the Transformer encoder. num_attention_heads (int, optional, defaults to 8): the number of attention heads in each attention layer. pad_token_id (int, optional, defaults to 890705): the id of the token used for padding. eos_token_id (int, optional, defaults to 890708): the id of the token used for the end of a sequence. bos_token_id (int, optional, defaults to 890709): the id of the token used for the beginning of a sequence. cls_token_id (int, optional, defaults to 890707): the id of the token used for classification tasks. sep_token_id (int, optional, defaults to 890710): the id of the token used to separate segments in a sequence. sparse_prediction (bool, optional, defaults to True): whether to use sparse prediction for the output layer. norm_eps (float, optional, defaults to 1e-5): the epsilon value used for layer normalization. embedding_dropout (float, optional, defaults to 0.0): the dropout probability for the embedding layer. initializer_range (float, optional, defaults to 0.02): the standard deviation of the truncated_normal_initializer for initializing all weight matrices. initializer_cutoff_factor (float, optional, defaults to 2.0): the cutoff factor for the truncated normal initializer. tie_word_embeddings (bool, optional, defaults to True): whether to tie the word embeddings with the output layer. num_batches (int, optional, defaults to 1): the number of batches when doing batch correction training. lambda_adv: float = 1.0, the weight for the adversarial loss. grl_alpha: float = 1.0, the alpha value for the gradient reversal layer. bc_unfreeze_last_n_layers (int, optional, defaults to 0): the number of last layers to unfreeze during training for batch correction. **kwargs: (additional keyword arguments, optional): additional configuration parameters.

Attributes

• model_type = 'atacformer'
• use_pos_embeddings
• num_batches
• lambda_adv
• grl_alpha
• bc_unfreeze_last_n_layers

freeze_except_last_n

def freeze_except_last_n(model: AtacformerModel, n: int = 2)

Freeze all parameters except the last n layers of the encoder. Also keeps all layer norms trainable for stability.

Args: model (AtacformerModel): The model to freeze. n (int): The number of last layers to keep trainable.

patch_atacformer_model_for_mps

def patch_atacformer_model_for_mps(model: nn.Module)

Look for any TransformerEncoder layers in the model and patch them by setting enable_nested_tensor to False and setting use_nested_tensor to False.

Args: model (nn.Module): The model to patch.

DataCollatorForReplacedTokenDetection

class DataCollatorForReplacedTokenDetection(tokenizer: TrainingTokenizer, mlm_probability: float = 0.15, seed: int | None = None)

Bases: WandbMixin, DataCollatorForLanguageModeling

Like HF’s MLM collator but: • never uses [MASK] • picks replacement tokens from a user-supplied distribution • returns per-token 0/1 labels for ELECTRA-style discrimination

Simple data collator for ELECTRA-style token replacement detection. Args: tokenizer (TrainingTokenizer): The tokenizer to use. vocab_counts (torch.Tensor | None): 1-D tensor, size == vocab, log-probs OR probs mlm_probability (float): Probability of masking a token. seed (int | None): Random seed for reproducibility.

Attributes

• mlm_probability
• uniform_vocab_probs

Methods

torch_mask_tokens

def torch_mask_tokens(inputs: torch.Tensor, special_tokens_mask: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]

DataCollatorForTripletLoss

class DataCollatorForTripletLoss(tokenizer: TrainingTokenizer, max_position_embeddings: int = None)

A simple data collator for triplet loss to fine-tune Atacformer for cell-type clustering

Attributes

• tokenizer
• pad_token_id
• max_position_embeddings

DataCollatorForUnsupervisedBatchCorrection

class DataCollatorForUnsupervisedBatchCorrection(tokenizer: TrainingTokenizer, max_position_embeddings: int = None, mlm_probability: float = 0.15, seed: int | None = None)

Bases: DataCollatorForReplacedTokenDetection

Attributes

• max_position_embeddings

ModelParameterChangeCallback

class ModelParameterChangeCallback(initial_params: dict[str, torch.Tensor])

Bases: WandbMixin, TrainerCallback

A callback to log the changes in model parameters after training.

Attributes

• initial_params

Methods

on_log

def on_log(args, state, control, **kwargs)

Log the changes in model parameters after training.

AdjustedRandIndexCallback

class AdjustedRandIndexCallback(input_ids: List[List[int]], cell_type_labels: List[int], pad_token_id: int, batch_size: int = 128, log_every_n_steps: int = 500)

Bases: WandbMixin, TrainerCallback

A callback to log the adjusted Rand index (ARI) during training.

Attributes

• initial_labels
• num_classes
• input_ids
• pad_token_id
• batch_size
• log_every_n_steps

Methods

on_log

def on_log(args, state, control, **kwargs)

Log the adjusted Rand index (ARI) during training.

get_git_hash

def get_git_hash() -> str

Get the current git hash of the repository.

get_decaying_cosine_with_hard_restarts_schedule_with_warmup

def get_decaying_cosine_with_hard_restarts_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1)

Very similar to huggingfaces built-in cosine with restarts, however the amplitude slowly decreases so that the "kick ups" are less aggressive.

Create a schedule with a learning rate that decreases following the values of the cosine function between the initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases linearly between 0 and the initial lr set in the optimizer.

Args: optimizer ([~torch.optim.Optimizer]): The optimizer for which to schedule the learning rate. num_warmup_steps (int): The number of steps for the warmup phase. num_training_steps (int): The total number of training steps. num_cycles (int, optional, defaults to 1): The number of hard restarts to use. last_epoch (int, optional, defaults to -1): The index of the last epoch when resuming training.

Return: torch.optim.lr_scheduler.LambdaLR with the appropriate schedule.

tokenize_anndata

def tokenize_anndata(adata: sc.AnnData, tokenizer: Tokenizer)

Tokenize an AnnData object. This is more involved, so it gets its own function.

Args: adata (sc.AnnData): The AnnData object to tokenize. tokenizer (Tokenizer): The tokenizer to use.