terra.embed_dataset

terra.embed_dataset(dataset, model_folder_path, emb_layer=None, agg_excluded_genes=None, top_k=None, batch_size=128, pin_memory=False, num_workers=12, include_spatial_cell_emb=True, return_token_embeddings=False, ignore_spc_tokens=True, agg_type='avg')

Embed a tokenized dataset, returning cell-, neighborhood- and gene-level embeddings.

Parameters:

• dataset (Dataset) – Tokenized Hugging Face dataset.

• model_folder_path (str) – Path to the folder containing the model config, token dictionary, and normalization factors.

• emb_layer (Optional[int] (default: None)) – Layer for which to retrieve the embedding.

• agg_excluded_genes (Optional[list[str]] (default: None)) – List of gene ensembl IDs to be excluded from the aggregation.

• top_k (Optional[int] (default: None)) – Include only top_k genes in aggregation.

• batch_size (int (default: 128)) – Dataloader param.

• pin_memory (bool (default: False)) – Dataloader param.

• num_workers (int (default: 12)) – Number of workers used.

• include_spatial_cell_emb (bool (default: True)) – If True, also return a spatially contextualized cell embedding that attends to the neighborhood.

• return_token_embeddings (bool (default: False)) – If True, also return per-token embeddings for each sequence position (cell and neighborhood tokens; special tokens are excluded).

Return type:

dict

Returns:

output_embed: Dictionary with the cell, cell gene, neighborhood, and neighborhood gene embeddings.