scEmbed Module

ScEmbed

class ScEmbed(model_path: str = None, tokenizer: Tokenizer = None, device: str = None, pooling_method: POOLING_TYPES = 'mean', **kwargs)

Initialize ScEmbed.

Args: model_path (str): Path to the pre-trained model on huggingface. tokenizer (Tokenizer): Tokenizer instance for tokenization. device (str): Device to use for computation. pooling_method (POOLING_TYPES): Pooling method for embeddings. **kwargs: Additional keyword arguments to pass to the model.

Attributes

• model_path (str)
• tokenizer (Tokenizer)
• trained (bool)
• pooling_method (POOLING_TYPES)

Properties

model

@property
def model

Get the core Region2Vec model.

Returns: Region2Vec: The core Region2Vec model.

Class methods

from_pretrained

@classmethod
def from_pretrained(path_to_files: str, model_file_name: str = MODEL_FILE_NAME, universe_file_name: str = UNIVERSE_FILE_NAME, config_file_name: str = CONFIG_FILE_NAME) -> ScEmbed

Load the model from a set of files that were exported using the export function.

Args: path_to_files (str): Path to the directory containing the files. model_file_name (str): Name of the model file. universe_file_name (str): Name of the universe file. config_file_name (str): Name of the config file.

Returns: ScEmbed: The loaded model.

Methods

add_tokenizer

def add_tokenizer(tokenizer: Tokenizer, **kwargs)

Add a tokenizer to the model. This should be use when the model is not initialized with a tokenizer.

Args: tokenizer (Tokenizer): Tokenizer to add to the model. **kwargs: Additional keyword arguments to pass to the model.

train

def train(dataset: Region2VecDataset, window_size: int = DEFAULT_WINDOW_SIZE, epochs: int = DEFAULT_EPOCHS, min_count: int = DEFAULT_MIN_COUNT, num_cpus: int = 1, seed: int = 42, save_checkpoint_path: str = None, gensim_params: dict = {}, load_from_checkpoint: str = None) -> bool

Train the model.

Args: dataset (Region2VecDataset): Data to train on. This is a dataset of tokens. window_size (int): Window size for the model. epochs (int): Number of epochs to train for. min_count (int): Minimum count for a region to be included in the vocabulary. num_cpus (int): Number of cpus to use for training. seed (int): Seed to use for training. save_checkpoint_path (str): Path to save the model checkpoints to. gensim_params (dict): Additional parameters to pass to the gensim model. load_from_checkpoint (str): Path to a checkpoint to load from.

Returns: bool: Whether or not the model was trained.

export

def export(path: str, checkpoint_file: str = MODEL_FILE_NAME, universe_file: str = UNIVERSE_FILE_NAME, config_file: str = CONFIG_FILE_NAME)

Function to facilitate exporting the model in a way that can be directly uploaded to huggingface.

This exports the model weights and the vocabulary.

Args: path (str): Path to export the model to. checkpoint_file (str): Name of the checkpoint file. universe_file (str): Name of the universe file. config_file (str): Name of the config file.

encode

def encode(regions: Union[sc.AnnData, str], pooling: POOLING_TYPES = None) -> np.ndarray

Get the vector for a region.

Args: regions (Union[sc.AnnData, str]): Region to get the vector for. pooling (POOLING_TYPES): Pooling type to use.

Returns: np.ndarray: Vector for the region.

AnnotationServer

class AnnotationServer(location: str = None, url: str = None, port: int = None, api_key: str = None, collection_name: str = None, timeout: float = 10)

Bases: QdrantClient

A class for querying a Qdrant server for cell type predictions.

This class requires that you have a Qdrant server running with a collection of cell type embeddings. You can create this collection using the geniml/examples/scembed/load_qdrant.ipynb script.

Args: collection_name (str): The name of the collection to query. location (str): The location of the Qdrant server. This should be in the format host:port. url (str): The URL of the Qdrant server. port (int): The port of the Qdrant server. api_key (str): The API key for the Qdrant server. timeout (float): The timeout for the Qdrant server.

Attributes

• collection_name

Annotator

class Annotator(collection_name: str = None, location: str = None, url: str = None, port: int = None, timeout: float = 10)

A class for annotating single cell data with cell type predictions.

This class requires that you have a Qdrant server running with a collection of cell type embeddings. You can create this collection using the geniml/examples/scembed/load_qdrant.ipynb script.

Args: collection_name (str): The name of the collection to query. location (str): The location of the Qdrant server. This should be in the format host:port. url (str): The URL of the Qdrant server. port (int): The port of the Qdrant server. timeout (float): The timeout for the Qdrant server.

Attributes

• collection_name
• url
• port
• timeout

Methods

annotate

def annotate(adata: sc.AnnData, embedding_key: str = 'embedding', key_added: str = 'pred_celltype', cluster_key: str = 'leiden', knn: int = 3, score_threshold: float = 0.5)

Annotate a sc.AnnData object with cell type predictions for each cluster.

This function requires that you have 1) clustered your data and 2) embedded your data using some pretrained model (databio/multiome). It is imperative that the model used to embed your single cells is the same model used to produce the embeddings in the database. Otherwise, the predictions will not be accurate; in fact they will be meaningless.

Args: adata (sc.AnnData): The annotated data. embedding_key (str): The key in adata.obsm where the embeddings are stored. key_added (str): The key in adata.obs where the cell type predictions will be stored. cluster_key (str): The key in adata.obs where the cluster labels are stored. knn (int): The number of nearest neighbors to use when querying the database. score_threshold (float): The score threshold to use when querying the database.