Craft Module
CraftModel
Section titled “CraftModel”class CraftModel(config: CraftConfig)Bases: PreTrainedModel
CRAFT Model with a masked language modeling head.
Attributes
Section titled “Attributes”config_class=CraftConfigbase_model_prefix='craft'configgeneformer_configatacformer_configgene_encoderatac_encoderprojection_dimatac_embed_dimgene_embed_dimgene_projectionatac_projectionlogit_scale
Methods
Section titled “Methods”forward
Section titled “forward”def forward(gene_input_ids: torch.Tensor = None, gene_attention_mask: torch.Tensor = None, gene_token_type_ids: torch.Tensor = None, atac_input_ids: torch.Tensor = None, atac_attention_mask: torch.Tensor = None, return_dict: bool = None, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor], CraftOutput]Forward pass through the CRAFT model.
Args: gene_input_ids (torch.Tensor): Input IDs for the gene encoder. gene_attention_mask (torch.Tensor): Attention mask for the gene encoder. gene_token_type_ids (torch.Tensor): Token type IDs for the gene encoder. atac_input_ids (torch.Tensor): Input IDs for the ATAC encoder. atac_attention_mask (torch.Tensor): Attention mask for the ATAC encoder. return_dict (bool): Whether to return a dictionary or tuple. Returns: torch.Tensor: The logits from the CRAFT model.
CraftForContrastiveLearning
Section titled “CraftForContrastiveLearning”class CraftForContrastiveLearning(config: CraftConfig)Bases: PreTrainedModel
CRAFT model for contrastive learning between gene and ATAC embeddings. While this looks redudant with the CraftModel, it makes it easier to use the model for further tasks like gene activity prediction without needing to instantiate the CraftModel directly.
Mostly used for pre-training tasks
Attributes
Section titled “Attributes”Methods
Section titled “Methods”forward
Section titled “forward”def forward(gene_input_ids: torch.Tensor, gene_attention_mask: torch.Tensor, gene_token_type_ids: torch.Tensor, atac_input_ids: torch.Tensor, atac_attention_mask: torch.Tensor) -> CraftOutputForward pass through the model.
Args: gene_input_ids (torch.Tensor): Input IDs for the gene encoder. gene_attention_mask (torch.Tensor): Attention mask for the gene encoder. gene_token_type_ids (torch.Tensor): Token type IDs for the gene encoder. atac_input_ids (torch.Tensor): Input IDs for the ATAC encoder. atac_attention_mask (torch.Tensor): Attention mask for the ATAC encoder.
Returns: CraftOutput: The output of the CRAFT model containing loss and logits.
CraftForGeneActivityPrediction
Section titled “CraftForGeneActivityPrediction”class CraftForGeneActivityPrediction(config: CraftConfig)Bases: PreTrainedModel
CRAFT model for gene activity prediction.
Attributes
Section titled “Attributes”config_class=CraftConfigbase_model_prefix='craft_for_gene_activity_prediction'craftgene_activity_head
Methods
Section titled “Methods”forward
Section titled “forward”def forward(atac_input_ids: torch.Tensor, atac_attention_mask: torch.Tensor, gene_activity: Optional[torch.Tensor] = None, return_dict: bool = True) -> Union[Tuple[torch.Tensor | None, Any], CraftGeneActivityOutput]Forward pass through the model.
Args: atac_input_ids (torch.Tensor): Input IDs for the ATAC encoder. atac_attention_mask (torch.Tensor): Attention mask for the ATAC encoder. gene_activity (Optional[torch.Tensor]): Optional gene activity scores for computing loss. Returns: torch.Tensor: The predicted gene activity scores.
CraftConfig
Section titled “CraftConfig”class CraftConfig(geneformer_config: GeneformerConfig = None, atacformer_config: AtacformerConfig = None, projection_dim: int = 512, logit_scale_init_value: float = 2.6592, **kwargs)Bases: PretrainedConfig
Configuration for the CRAFT model, a contrastive RNA-ATAC transformer that attempts to learn leverage Geneformer and Atacformer to learn a joint representation of RNA and ATAC data.
Joint configuration for the CRAFT model.
Args: geneformer_config (GeneformerConfig): Configuration for the Geneformer model. atacformer_config (AtacformerConfig): Configuration for the Atacformer model. projection_dim (int): Dimension of the projection layer. logit_scale_init_value (float): Initial value for the logit scale parameter.
Attributes
Section titled “Attributes”DataCollatorForCraft
Section titled “DataCollatorForCraft”class DataCollatorForCraft(gene_pad: int, atac_pad: int, gene_max_len: int | None = None, atac_max_len: int | None = None)Pads + builds masks for gene/ATAC pairs used by CraftModel.
DataCollatorForCraftGeneActivityPrediction
Section titled “DataCollatorForCraftGeneActivityPrediction”class DataCollatorForCraftGeneActivityPrediction(atac_pad: int, atac_max_len: int | None = None)Pads + builds masks for ATAC-gene pairs used by CraftForGeneActivityPrediction.
Gene activity is always the same shape for everything, its just a set of floats representing the activity of each gene in the genome, so we don't pad it.