terra.gene_embed_dataset

terra.gene_embed_dataset#

terra.gene_embed_dataset(dataset, model_folder_path, emb_layer=None, cell_gene_ids=[], neighborhood_gene_ids=[], batch_size=128, pin_memory=False, num_workers=12, return_gene=False, return_gene_per_data=False, compute_cosine_with_list=[], return_distance=False, return_cosine_sim=False, return_receptor_average=False, include_spatial_cell_emb=True, description='', ignore_spc_tokens=True)#

Compute per-gene embeddings or co-occurrence statistics over a tokenized dataset.

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.

  • cell_gene_ids (list (default: [])) – List with gene IDs for which cell gene embeddings will be retrieved.

  • neighborhood_gene_ids (list (default: [])) – List with gene IDs for which neighborhood gene embeddings will be retrived.

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

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

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

  • return_gene (bool (default: False)) – If True will return per-gene cell and neighborhood embeddings (and, when include_spatial_cell_emb is True, spatially contextualized cell gene embeddings).

  • return_gene_per_data (bool (default: False)) – If True will return gene_embedding for each gene per dataset.

  • compute_cosine_with_list (list[str] (default: [])) – A list that defines the items with which we want to compute cosine similarity. It could have value of cell or/and neighborhood.

  • return_distance (bool (default: False)) – If True will compute and return distance between cosine sim of cell_neb and cell_cell matrix.

  • return_cosine_sim (bool (default: False)) – If True will compute and return cosine_sim matrix.

  • return_receptor_average (bool (default: False)) – If True will compute and return receptor average embeddings for cell-neighborhood gene pairs.

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

  • description (str (default: '')) – Description for task that is currently using this function.

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

Returns:

  • output (tuple | dict | None) – The return value depends on which flag is set (checked in this order): - If return_gene is True: a tuple of per-gene embedding dictionaries keyed by gene ID. With include_spatial_cell_emb it is (cell_gene_emb, neighborhood_gene_emb, spatial_cell_gene_emb); otherwise (cell_gene_emb, neighborhood_gene_emb). - If return_gene_per_data is True: a tuple of per-gene (sum, count) accumulator dictionaries. With include_spatial_cell_emb it is (cell_gene_emb_per_data, neighborhood_gene_emb_per_data, spatial_cell_gene_emb_per_data); otherwise the first two. - If return_cosine_sim is True: a single dict mapping each entry of compute_cosine_with_list to (sum_cos_sim, pair_count, cell_count). - If return_distance is True: a 2-tuple of lists (emd_list, emd_matrix_list). - If return_receptor_average is True: a single dict keyed by (cell_gene_id, neighborhood_gene_id, 'present'|'absent') mapping to [sum, count]. - If none of the above flags is set: None.

Contents