move dicts to init

geneformer/__init__.py

@@ -1,4 +1,10 @@
 # ruff: noqa: F401
+from pathlib import Path
+
+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"
+
 from . import (
     collator_for_classification,
     emb_extractor,
@@ -18,4 +24,4 @@ from .pretrainer import GeneformerPretrainer
 from .tokenizer import TranscriptomeTokenizer
 
 from . import classifier  # noqa # isort:skip
-from .classifier import Classifier  # noqa # isort:skip
+from .classifier import Classifier  # noqa # isort:skip

geneformer/classifier.py

@@ -61,7 +61,7 @@ from . import DataCollatorForCellClassification, DataCollatorForGeneClassificati
 from . import classifier_utils as cu
 from . import evaluation_utils as eu
 from . import perturber_utils as pu
-from .tokenizer import TOKEN_DICTIONARY_FILE
+from . import TOKEN_DICTIONARY_FILE
 
 sns.set()
 

geneformer/collator_for_classification.py

@@ -4,6 +4,7 @@ Geneformer collator for gene and cell classification.
 Huggingface data collator modified to accommodate single-cell transcriptomics data for gene and cell classification.
 """
 import numpy as np
+import pickle
 import torch
 import warnings
 from enum import Enum
@@ -17,7 +18,11 @@ from transformers import (
 from transformers.utils import is_tf_available, is_torch_available, logging, to_py_obj
 from transformers.utils.generic import _is_tensorflow, _is_torch
 
-from .pretrainer import token_dictionary
+from . import TOKEN_DICTIONARY_FILE
+
+# load token dictionary (Ensembl IDs:token)
+with open(TOKEN_DICTIONARY_FILE, "rb") as f:
+    token_dictionary = pickle.load(f)
 
 EncodedInput = List[int]
 logger = logging.get_logger(__name__)

geneformer/emb_extractor.py

@@ -25,7 +25,7 @@ from tdigest import TDigest
 from tqdm.auto import trange
 
 from . import perturber_utils as pu
-from .tokenizer import TOKEN_DICTIONARY_FILE
+from . import TOKEN_DICTIONARY_FILE
 
 logger = logging.getLogger(__name__)
 

geneformer/evaluation_utils.py

@@ -21,7 +21,7 @@ from sklearn.metrics import (
 from tqdm.auto import trange
 
 from .emb_extractor import make_colorbar
-from .tokenizer import TOKEN_DICTIONARY_FILE
+from . import TOKEN_DICTIONARY_FILE
 
 logger = logging.getLogger(__name__)
 

geneformer/in_silico_perturber.py

@@ -38,21 +38,17 @@ import logging
 import os
 import pickle
 from collections import defaultdict
-from typing import List
 from multiprocess import set_start_method
 
-import seaborn as sns
 import torch
-from datasets import Dataset
+from datasets import Dataset, disable_progress_bars
 from tqdm.auto import trange
 
 from . import perturber_utils as pu
 from .emb_extractor import get_embs
-from .perturber_utils import TOKEN_DICTIONARY_FILE
-
-
-sns.set()
+from . import TOKEN_DICTIONARY_FILE
 
+disable_progress_bars()
 
 logger = logging.getLogger(__name__)
 

geneformer/in_silico_perturber_stats.py

@@ -38,9 +38,7 @@ from sklearn.mixture import GaussianMixture
 from tqdm.auto import tqdm, trange
 
 from .perturber_utils import flatten_list, validate_cell_states_to_model
-from .tokenizer import TOKEN_DICTIONARY_FILE
-
-GENE_NAME_ID_DICTIONARY_FILE = Path(__file__).parent / "gene_name_id_dict.pkl"
+from . import TOKEN_DICTIONARY_FILE, ENSEMBL_DICTIONARY_FILE
 
 logger = logging.getLogger(__name__)
 
@@ -673,7 +671,7 @@ class InSilicoPerturberStats:
         cell_states_to_model=None,
         pickle_suffix="_raw.pickle",
         token_dictionary_file=TOKEN_DICTIONARY_FILE,
-        gene_name_id_dictionary_file=GENE_NAME_ID_DICTIONARY_FILE,
+        gene_name_id_dictionary_file=ENSEMBL_DICTIONARY_FILE,
     ):
         """
         Initialize in silico perturber stats generator.

geneformer/perturber_utils.py

@@ -18,13 +18,9 @@ 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"
+from . import GENE_MEDIAN_FILE, TOKEN_DICTIONARY_FILE, ENSEMBL_DICTIONARY_FILE
 
 
-sns.set()
-
 logger = logging.getLogger(__name__)
 
 

geneformer/pretrainer.py

@@ -32,7 +32,7 @@ from transformers.training_args import ParallelMode
 from transformers.utils import is_tf_available, is_torch_available, logging, to_py_obj
 from transformers.utils.generic import _is_tensorflow, _is_torch
 
-from .tokenizer import TOKEN_DICTIONARY_FILE
+from . import TOKEN_DICTIONARY_FILE
 
 logger = logging.get_logger(__name__)
 EncodedInput = List[int]

geneformer/tokenizer.py

@@ -52,7 +52,7 @@ import loompy as lp  # noqa
 
 logger = logging.getLogger(__name__)
 
-from .perturber_utils import GENE_MEDIAN_FILE, TOKEN_DICTIONARY_FILE
+from . import GENE_MEDIAN_FILE, TOKEN_DICTIONARY_FILE
 
 
 def rank_genes(gene_vector, gene_tokens):