terra.embed_dataset

Contents

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).

  • ignore_spc_tokens (bool (default: True)) – Whether to ignore special tokens when retrieving layer embeddings.

  • agg_type (Literal['avg', 'softmax'] (default: 'avg')) – How gene embeddings are aggregated into cell and neighborhood embeddings ('avg' for an unweighted mean or 'softmax' for a count-softmax-weighted mean).

Returns:

  • output_embed (dict) – Dictionary with one numpy array per embedding type. Always contains cell_emb (cell embedding) and neighborhood_emb (neighborhood embedding). If include_spatial_cell_emb is True, also contains spatial_cell_emb (spatially contextualized cell embedding). If return_token_embeddings is True, also contains token_emb (per-token embeddings).

Contents