fix cell state gene embeddings bug (#345)

- added quality of life improvements; fixed gene similarities with cell_states_to_model (4b4547f0634eed07560e599766c30326138b7a32)
- reinstate save_to_disk patch (344f263c6173a6bbe96eabcc5ac65e45fa4756e7)

geneformer/__init__.py

@@ -11,7 +11,7 @@ from .collator_for_classification import (
     DataCollatorForCellClassification,
     DataCollatorForGeneClassification,
 )
-from .emb_extractor import EmbExtractor
+from .emb_extractor import EmbExtractor, get_embs
 from .in_silico_perturber import InSilicoPerturber
 from .in_silico_perturber_stats import InSilicoPerturberStats
 from .pretrainer import GeneformerPretrainer

geneformer/in_silico_perturber.py

@@ -39,6 +39,7 @@ import os
 import pickle
 from collections import defaultdict
 from typing import List
+from multiprocess import set_start_method
 
 import seaborn as sns
 import torch
@@ -47,7 +48,8 @@ from tqdm.auto import trange
 
 from . import perturber_utils as pu
 from .emb_extractor import get_embs
-from .tokenizer import TOKEN_DICTIONARY_FILE
+from .perturber_utils import TOKEN_DICTIONARY_FILE
+
 
 sns.set()
 
@@ -185,6 +187,10 @@ class InSilicoPerturber:
         token_dictionary_file : Path
             | Path to pickle file containing token dictionary (Ensembl ID:token).
         """
+        try:
+            set_start_method("spawn")
+        except RuntimeError:
+            pass
 
         self.perturb_type = perturb_type
         self.perturb_rank_shift = perturb_rank_shift
@@ -422,6 +428,7 @@ class InSilicoPerturber:
         self.max_len = pu.get_model_input_size(model)
         layer_to_quant = pu.quant_layers(model) + self.emb_layer
 
+
         ### filter input data ###
         # general filtering of input data based on filter_data argument
         filtered_input_data = pu.load_and_filter(
@@ -520,6 +527,7 @@ class InSilicoPerturber:
         perturbed_data = filtered_input_data.map(
             make_group_perturbation_batch, num_proc=self.nproc
         )
+    
         if self.perturb_type == "overexpress":
             filtered_input_data = filtered_input_data.add_column(
                 "n_overflow", perturbed_data["n_overflow"]

geneformer/perturber_utils.py

@@ -4,6 +4,8 @@ import pickle
 import re
 from collections import defaultdict
 from typing import List
+from pathlib import Path
+
 
 import numpy as np
 import pandas as pd
@@ -16,6 +18,11 @@ from transformers import (
     BertForTokenClassification,
 )
 
+GENE_MEDIAN_FILE = Path(__file__).parent / "gene_median_dictionary.pkl"
+TOKEN_DICTIONARY_FILE = Path(__file__).parent / "token_dictionary.pkl"
+ENSEMBL_DICTIONARY_FILE = Path(__file__).parent / "gene_name_id_dict.pkl"
+
+
 sns.set()
 
 logger = logging.getLogger(__name__)
@@ -581,9 +588,11 @@ def quant_cos_sims(
     elif emb_mode == "cell":
         cos = torch.nn.CosineSimilarity(dim=1)
 
-    if cell_states_to_model is None:
+    # if emb_mode == "gene", can only calculate gene cos sims
+    # against original cell anyways
+    if cell_states_to_model is None or emb_mode == "gene":
         cos_sims = cos(perturbation_emb, original_emb).to("cuda")
-    else:
+    elif cell_states_to_model is not None and emb_mode == "cell":
         possible_states = get_possible_states(cell_states_to_model)
         cos_sims = dict(zip(possible_states, [[] for _ in range(len(possible_states))]))
         for state in possible_states:
@@ -705,3 +714,48 @@ def validate_cell_states_to_model(cell_states_to_model):
                 "'alt_states': ['hcm', 'other1', 'other2']}"
             )
             raise
+
+class GeneIdHandler:
+    def __init__(self, raise_errors=False):
+        def invert_dict(dict_obj):
+            return {v:k for k,v in dict_obj.items()}
+        
+        self.raise_errors = raise_errors
+        
+        with open(TOKEN_DICTIONARY_FILE, 'rb') as f:
+            self.gene_token_dict = pickle.load(f)
+            self.token_gene_dict = invert_dict(self.gene_token_dict)
+
+        with open(ENSEMBL_DICTIONARY_FILE, 'rb') as f:
+            self.id_gene_dict = pickle.load(f)
+            self.gene_id_dict = invert_dict(self.id_gene_dict)
+            
+    def ens_to_token(self, ens_id):
+        if not self.raise_errors:
+            return self.gene_token_dict.get(ens_id, ens_id)
+        else:
+            return self.gene_token_dict[ens_id]
+    
+    def token_to_ens(self, token):
+        if not self.raise_errors:
+            return self.token_gene_dict.get(token, token)
+        else:
+            return self.token_gene_dict[token]
+
+    def ens_to_symbol(self, ens_id):
+        if not self.raise_errors:
+            return self.gene_id_dict.get(ens_id, ens_id)
+        else:
+            return self.gene_id_dict[ens_id]
+    
+    def symbol_to_ens(self, symbol):
+        if not self.raise_errors:
+            return self.id_gene_dict.get(symbol, symbol)
+        else:
+            return self.id_gene_dict[symbol]
+    
+    def token_to_symbol(self, token):
+        return self.ens_to_symbol(self.token_to_ens(token))
+    
+    def symbol_to_token(self, symbol):
+        return self.ens_to_token(self.symbol_to_ens(symbol))

geneformer/tokenizer.py

@@ -52,8 +52,8 @@ import loompy as lp  # noqa
 
 logger = logging.getLogger(__name__)
 
-GENE_MEDIAN_FILE = Path(__file__).parent / "gene_median_dictionary.pkl"
-TOKEN_DICTIONARY_FILE = Path(__file__).parent / "token_dictionary.pkl"
+from .perturber_utils import GENE_MEDIAN_FILE, TOKEN_DICTIONARY_FILE
+
 
 def rank_genes(gene_vector, gene_tokens):
     """