precommit formatting

.gitattributes

@@ -18,7 +18,7 @@
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text 
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text

examples/pretraining_new_model/pretrain_geneformer_w_deepspeed.py

@@ -138,7 +138,9 @@ training_args = {
     "per_device_train_batch_size": geneformer_batch_size,
     "num_train_epochs": epochs,
     "save_strategy": "steps",
-    "save_steps": np.floor(num_examples / geneformer_batch_size / 8),  # 8 saves per epoch
+    "save_steps": np.floor(
+        num_examples / geneformer_batch_size / 8
+    ),  # 8 saves per epoch
     "logging_steps": 1000,
     "output_dir": training_output_dir,
     "logging_dir": logging_dir,

geneformer/__init__.py

@@ -1,6 +1,7 @@
 # ruff: noqa: F401
-from pathlib import Path
 import warnings
+from pathlib import Path
+
 warnings.filterwarnings("ignore", message=".*The 'nopython' keyword.*")  # noqa # isort:skip
 
 GENE_MEDIAN_FILE = Path(__file__).parent / "gene_median_dictionary_gc95M.pkl"
@@ -30,4 +31,4 @@ from . import classifier  # noqa # isort:skip
 from .classifier import Classifier  # noqa # isort:skip
 
 from . import mtl_classifier  # noqa # isort:skip
-from .mtl_classifier import MTLClassifier  # noqa # isort:skip
+from .mtl_classifier import MTLClassifier  # noqa # isort:skip

geneformer/classifier.py

@@ -57,11 +57,14 @@ from tqdm.auto import tqdm, trange
 from transformers import Trainer
 from transformers.training_args import TrainingArguments
 
-from . import DataCollatorForCellClassification, DataCollatorForGeneClassification
+from . import (
+    TOKEN_DICTIONARY_FILE,
+    DataCollatorForCellClassification,
+    DataCollatorForGeneClassification,
+)
 from . import classifier_utils as cu
 from . import evaluation_utils as eu
 from . import perturber_utils as pu
-from . import TOKEN_DICTIONARY_FILE
 
 sns.set()
 
@@ -413,7 +416,7 @@ class Classifier:
                     "Column name 'labels' must be reserved for class IDs. Please rename column."
                 )
                 raise
-                
+
         if (attr_to_split is not None) and (attr_to_balance is None):
             logger.error(
                 "Splitting by attribute while balancing confounders requires both attr_to_split and attr_to_balance to be defined."
@@ -548,17 +551,19 @@ class Classifier:
         gene_balance : None, bool
             | Whether to automatically balance genes in training set.
             | Only available for binary gene classifications.
-            
+
         **Output**
 
         Returns trainer after fine-tuning with all data.
 
         """
 
-        if (gene_balance is True) and (len(self.gene_class_dict.values())!=2):
-            logger.error("Automatically balancing gene sets for training is only available for binary gene classifications.")
+        if (gene_balance is True) and (len(self.gene_class_dict.values()) != 2):
+            logger.error(
+                "Automatically balancing gene sets for training is only available for binary gene classifications."
+            )
             raise
-        
+
         ##### Load data and prepare output directory #####
         # load numerical id to class dictionary (id:class)
         with open(id_class_dict_file, "rb") as f:
@@ -679,11 +684,13 @@ class Classifier:
         if self.num_crossval_splits == 0:
             logger.error("num_crossval_splits must be 1 or 5 to validate.")
             raise
-            
-        if (gene_balance is True) and (len(self.gene_class_dict.values())!=2):
-            logger.error("Automatically balancing gene sets for training is only available for binary gene classifications.")
+
+        if (gene_balance is True) and (len(self.gene_class_dict.values()) != 2):
+            logger.error(
+                "Automatically balancing gene sets for training is only available for binary gene classifications."
+            )
             raise
-        
+
         # ensure number of genes in each class is > 5 if validating model
         if self.classifier == "gene":
             insuff_classes = [k for k, v in self.gene_class_dict.items() if len(v) < 5]
@@ -845,17 +852,18 @@ class Classifier:
                     self.nproc,
                     gene_balance,
                 )
-                
+
                 if save_gene_split_datasets is True:
                     for split_name in ["train", "valid"]:
                         labeled_dataset_output_path = (
-                    Path(output_dir) / f"{output_prefix}_{split_name}_gene_labeled_ksplit{iteration_num}"
-                ).with_suffix(".dataset")
+                            Path(output_dir)
+                            / f"{output_prefix}_{split_name}_gene_labeled_ksplit{iteration_num}"
+                        ).with_suffix(".dataset")
                         if split_name == "train":
                             train_data.save_to_disk(str(labeled_dataset_output_path))
                         elif split_name == "valid":
                             eval_data.save_to_disk(str(labeled_dataset_output_path))
-                            
+
                 if self.oos_test_size > 0:
                     test_data = cu.prep_gene_classifier_split(
                         data,
@@ -869,11 +877,14 @@ class Classifier:
                     )
                     if save_gene_split_datasets is True:
                         test_labeled_dataset_output_path = (
-                    Path(output_dir) / f"{output_prefix}_test_gene_labeled_ksplit{iteration_num}"
-                ).with_suffix(".dataset")
+                            Path(output_dir)
+                            / f"{output_prefix}_test_gene_labeled_ksplit{iteration_num}"
+                        ).with_suffix(".dataset")
                         test_data.save_to_disk(str(test_labeled_dataset_output_path))
                 if debug_gene_split_datasets is True:
-                    logger.error("Exiting after saving gene split datasets given debug_gene_split_datasets = True.")
+                    logger.error(
+                        "Exiting after saving gene split datasets given debug_gene_split_datasets = True."
+                    )
                     raise
                 if n_hyperopt_trials == 0:
                     trainer = self.train_classifier(
@@ -1023,7 +1034,13 @@ class Classifier:
         subprocess.call(f"mkdir {output_directory}", shell=True)
 
         ##### Load model and training args #####
-        model = pu.load_model(self.model_type, num_classes, model_directory, "train", quantize=self.quantize)
+        model = pu.load_model(
+            self.model_type,
+            num_classes,
+            model_directory,
+            "train",
+            quantize=self.quantize,
+        )
         def_training_args, def_freeze_layers = cu.get_default_train_args(
             model, self.classifier, train_data, output_directory
         )
@@ -1047,14 +1064,22 @@ class Classifier:
         ##### Fine-tune the model #####
         # define the data collator
         if self.classifier == "cell":
-            data_collator = DataCollatorForCellClassification(token_dictionary=self.token_dictionary)
+            data_collator = DataCollatorForCellClassification(
+                token_dictionary=self.token_dictionary
+            )
         elif self.classifier == "gene":
-            data_collator = DataCollatorForGeneClassification(token_dictionary=self.token_dictionary)
+            data_collator = DataCollatorForGeneClassification(
+                token_dictionary=self.token_dictionary
+            )
 
         # define function to initiate model
         def model_init():
             model = pu.load_model(
-                self.model_type, num_classes, model_directory, "train", quantize=self.quantize
+                self.model_type,
+                num_classes,
+                model_directory,
+                "train",
+                quantize=self.quantize,
             )
 
             if self.freeze_layers is not None:
@@ -1180,7 +1205,13 @@ class Classifier:
         subprocess.call(f"mkdir {output_directory}", shell=True)
 
         ##### Load model and training args #####
-        model = pu.load_model(self.model_type, num_classes, model_directory, "train", quantize=self.quantize)
+        model = pu.load_model(
+            self.model_type,
+            num_classes,
+            model_directory,
+            "train",
+            quantize=self.quantize,
+        )
 
         def_training_args, def_freeze_layers = cu.get_default_train_args(
             model, self.classifier, train_data, output_directory
@@ -1210,9 +1241,13 @@ class Classifier:
         ##### Fine-tune the model #####
         # define the data collator
         if self.classifier == "cell":
-            data_collator = DataCollatorForCellClassification(token_dictionary=self.token_dictionary)
+            data_collator = DataCollatorForCellClassification(
+                token_dictionary=self.token_dictionary
+            )
         elif self.classifier == "gene":
-            data_collator = DataCollatorForGeneClassification(token_dictionary=self.token_dictionary)
+            data_collator = DataCollatorForGeneClassification(
+                token_dictionary=self.token_dictionary
+            )
 
         # create the trainer
         trainer = Trainer(
@@ -1334,7 +1369,13 @@ class Classifier:
         test_data = pu.load_and_filter(None, self.nproc, test_data_file)
 
         # load previously fine-tuned model
-        model = pu.load_model(self.model_type, num_classes, model_directory, "eval", quantize=self.quantize)
+        model = pu.load_model(
+            self.model_type,
+            num_classes,
+            model_directory,
+            "eval",
+            quantize=self.quantize,
+        )
 
         # evaluate the model
         result = self.evaluate_model(

geneformer/classifier_utils.py

@@ -137,22 +137,53 @@ def label_gene_classes(example, class_id_dict, gene_class_dict):
 
 
 def prep_gene_classifier_train_eval_split(
-    data, targets, labels, train_index, eval_index, max_ncells, iteration_num, num_proc, balance=False
+    data,
+    targets,
+    labels,
+    train_index,
+    eval_index,
+    max_ncells,
+    iteration_num,
+    num_proc,
+    balance=False,
 ):
     # generate cross-validation splits
     train_data = prep_gene_classifier_split(
-        data, targets, labels, train_index, "train", max_ncells, iteration_num, num_proc, balance
+        data,
+        targets,
+        labels,
+        train_index,
+        "train",
+        max_ncells,
+        iteration_num,
+        num_proc,
+        balance,
     )
     eval_data = prep_gene_classifier_split(
-        data, targets, labels, eval_index, "eval", max_ncells, iteration_num, num_proc, balance
+        data,
+        targets,
+        labels,
+        eval_index,
+        "eval",
+        max_ncells,
+        iteration_num,
+        num_proc,
+        balance,
     )
     return train_data, eval_data
 
 
 def prep_gene_classifier_split(
-    data, targets, labels, index, subset_name, max_ncells, iteration_num, num_proc, balance=False
+    data,
+    targets,
+    labels,
+    index,
+    subset_name,
+    max_ncells,
+    iteration_num,
+    num_proc,
+    balance=False,
 ):
-
     # generate cross-validation splits
     targets = np.array(targets)
     labels = np.array(labels)
@@ -175,7 +206,9 @@ def prep_gene_classifier_split(
 
     # balance gene subsets if train
     if (subset_name == "train") and (balance is True):
-        subset_data, label_dict_subset = balance_gene_split(subset_data, label_dict_subset, num_proc)
+        subset_data, label_dict_subset = balance_gene_split(
+            subset_data, label_dict_subset, num_proc
+        )
 
     # subsample to max_ncells
     subset_data = downsample_and_shuffle(subset_data, max_ncells, None, None)
@@ -192,7 +225,9 @@ def prep_gene_classifier_split(
     return subset_data
 
 
-def prep_gene_classifier_all_data(data, targets, labels, max_ncells, num_proc, balance=False):
+def prep_gene_classifier_all_data(
+    data, targets, labels, max_ncells, num_proc, balance=False
+):
     targets = np.array(targets)
     labels = np.array(labels)
     label_dict_train = dict(zip(targets, labels))
@@ -211,8 +246,10 @@ def prep_gene_classifier_all_data(data, targets, labels, max_ncells, num_proc, b
     )
 
     if balance is True:
-        train_data, label_dict_train = balance_gene_split(train_data, label_dict_train, num_proc)
-    
+        train_data, label_dict_train = balance_gene_split(
+            train_data, label_dict_train, num_proc
+        )
+
     # subsample to max_ncells
     train_data = downsample_and_shuffle(train_data, max_ncells, None, None)
 
@@ -230,62 +267,97 @@ def prep_gene_classifier_all_data(data, targets, labels, max_ncells, num_proc, b
 
 def balance_gene_split(subset_data, label_dict_subset, num_proc):
     # count occurrence of genes in each label category
-    label0_counts, label1_counts = count_genes_for_balancing(subset_data, label_dict_subset, num_proc)
-    label_ratio_0to1 = label0_counts/label1_counts
-    
-    if 8/10 <= label_ratio_0to1 <= 10/8:
+    label0_counts, label1_counts = count_genes_for_balancing(
+        subset_data, label_dict_subset, num_proc
+    )
+    label_ratio_0to1 = label0_counts / label1_counts
+
+    if 8 / 10 <= label_ratio_0to1 <= 10 / 8:
         # gene sets already balanced
         logger.info(
             "Gene sets were already balanced within 0.8-1.25 fold and did not require balancing.\n"
         )
         return subset_data, label_dict_subset
     else:
-        label_ratio_0to1_orig = label_ratio_0to1+0
+        label_ratio_0to1_orig = label_ratio_0to1 + 0
         label_dict_subset_orig = label_dict_subset.copy()
         # balance gene sets
         max_ntrials = 25
         boost = 1
-        if label_ratio_0to1 > 10/8:
+        if label_ratio_0to1 > 10 / 8:
             # downsample label 0
             for i in range(max_ntrials):
                 label0 = 0
-                label0_genes = [k for k,v in label_dict_subset.items() if v == label0]
+                label0_genes = [k for k, v in label_dict_subset.items() if v == label0]
                 label0_ngenes = len(label0_genes)
-                label0_nremove = max(1,int(np.floor(label0_ngenes - label0_ngenes/(label_ratio_0to1*boost))))
+                label0_nremove = max(
+                    1,
+                    int(
+                        np.floor(
+                            label0_ngenes - label0_ngenes / (label_ratio_0to1 * boost)
+                        )
+                    ),
+                )
                 random.seed(i)
                 label0_remove_genes = random.sample(label0_genes, label0_nremove)
-                label_dict_subset_new = {k:v for k,v in label_dict_subset.items() if k not in label0_remove_genes}
-                label0_counts, label1_counts = count_genes_for_balancing(subset_data, label_dict_subset_new, num_proc)
-                label_ratio_0to1 = label0_counts/label1_counts
-                if 8/10 <= label_ratio_0to1 <= 10/8:
+                label_dict_subset_new = {
+                    k: v
+                    for k, v in label_dict_subset.items()
+                    if k not in label0_remove_genes
+                }
+                label0_counts, label1_counts = count_genes_for_balancing(
+                    subset_data, label_dict_subset_new, num_proc
+                )
+                label_ratio_0to1 = label0_counts / label1_counts
+                if 8 / 10 <= label_ratio_0to1 <= 10 / 8:
                     # if gene sets now balanced, return new filtered data and new label_dict_subset
-                    return filter_data_balanced_genes(subset_data, label_dict_subset_new, num_proc)
-                elif label_ratio_0to1 > 10/8:
-                    boost = boost*1.1
-                elif label_ratio_0to1 < 8/10:
-                    boost = boost*0.9
+                    return filter_data_balanced_genes(
+                        subset_data, label_dict_subset_new, num_proc
+                    )
+                elif label_ratio_0to1 > 10 / 8:
+                    boost = boost * 1.1
+                elif label_ratio_0to1 < 8 / 10:
+                    boost = boost * 0.9
         else:
             # downsample label 1
             for i in range(max_ntrials):
                 label1 = 1
-                label1_genes = [k for k,v in label_dict_subset.items() if v == label1]
+                label1_genes = [k for k, v in label_dict_subset.items() if v == label1]
                 label1_ngenes = len(label1_genes)
-                label1_nremove = max(1,int(np.floor(label1_ngenes - label1_ngenes/((1/label_ratio_0to1)*boost))))
+                label1_nremove = max(
+                    1,
+                    int(
+                        np.floor(
+                            label1_ngenes
+                            - label1_ngenes / ((1 / label_ratio_0to1) * boost)
+                        )
+                    ),
+                )
                 random.seed(i)
                 label1_remove_genes = random.sample(label1_genes, label1_nremove)
-                label_dict_subset_new = {k:v for k,v in label_dict_subset.items() if k not in label1_remove_genes}
-                label0_counts, label1_counts = count_genes_for_balancing(subset_data, label_dict_subset_new, num_proc)
-                label_ratio_0to1 = label0_counts/label1_counts
-                if 8/10 <= label_ratio_0to1 <= 10/8:
+                label_dict_subset_new = {
+                    k: v
+                    for k, v in label_dict_subset.items()
+                    if k not in label1_remove_genes
+                }
+                label0_counts, label1_counts = count_genes_for_balancing(
+                    subset_data, label_dict_subset_new, num_proc
+                )
+                label_ratio_0to1 = label0_counts / label1_counts
+                if 8 / 10 <= label_ratio_0to1 <= 10 / 8:
                     # if gene sets now balanced, return new filtered data and new label_dict_subset
-                    return filter_data_balanced_genes(subset_data, label_dict_subset_new, num_proc)
-                elif label_ratio_0to1 < 8/10:
-                    boost = boost*1.1
-                elif label_ratio_0to1 > 10/8:
-                    boost = boost*0.9
-                    
-        assert i+1 == max_ntrials
-        if (label_ratio_0to1 <= label_ratio_0to1_orig < 8/10) or (10/8 > label_ratio_0to1_orig >= label_ratio_0to1):
+                    return filter_data_balanced_genes(
+                        subset_data, label_dict_subset_new, num_proc
+                    )
+                elif label_ratio_0to1 < 8 / 10:
+                    boost = boost * 1.1
+                elif label_ratio_0to1 > 10 / 8:
+                    boost = boost * 0.9
+
+        assert i + 1 == max_ntrials
+        if (label_ratio_0to1 <= label_ratio_0to1_orig < 8 / 10) or (
+            10 / 8 > label_ratio_0to1_orig >= label_ratio_0to1
+        ):
             label_ratio_0to1 = label_ratio_0to1_orig
             label_dict_subset_new = label_dict_subset_orig
         logger.warning(
@@ -301,11 +373,11 @@ def count_genes_for_balancing(subset_data, label_dict_subset, num_proc):
         ]
         counter_labels = Counter(labels)
         # get count of labels 0 or 1, or if absent, return 0
-        example["labels_counts"] = [counter_labels.get(0,0),counter_labels.get(1,0)]
+        example["labels_counts"] = [counter_labels.get(0, 0), counter_labels.get(1, 0)]
         return example
-        
+
     subset_data = subset_data.map(count_targets, num_proc=num_proc)
- 
+
     label0_counts = sum([counts[0] for counts in subset_data["labels_counts"]])
     label1_counts = sum([counts[1] for counts in subset_data["labels_counts"]])
 

geneformer/collator_for_classification.py

@@ -3,17 +3,17 @@ 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
 from typing import Dict, List, Optional, Union
 
+import numpy as np
+import torch
 from transformers import (
+    BatchEncoding,
     DataCollatorForTokenClassification,
     SpecialTokensMixin,
-    BatchEncoding,
 )
 from transformers.utils import is_tf_available, is_torch_available, logging, to_py_obj
 from transformers.utils.generic import _is_tensorflow, _is_torch
@@ -29,6 +29,7 @@ LARGE_INTEGER = int(
 
 # precollator functions
 
+
 class ExplicitEnum(Enum):
     """
     Enum with more explicit error message for missing values.
@@ -41,6 +42,7 @@ class ExplicitEnum(Enum):
             % (value, cls.__name__, str(list(cls._value2member_map_.keys())))
         )
 
+
 class TruncationStrategy(ExplicitEnum):
     """
     Possible values for the ``truncation`` argument in :meth:`PreTrainedTokenizerBase.__call__`. Useful for
@@ -53,7 +55,6 @@ class TruncationStrategy(ExplicitEnum):
     DO_NOT_TRUNCATE = "do_not_truncate"
 
 
-
 class PaddingStrategy(ExplicitEnum):
     """
     Possible values for the ``padding`` argument in :meth:`PreTrainedTokenizerBase.__call__`. Useful for tab-completion
@@ -65,7 +66,6 @@ class PaddingStrategy(ExplicitEnum):
     DO_NOT_PAD = "do_not_pad"
 
 
-
 class TensorType(ExplicitEnum):
     """
     Possible values for the ``return_tensors`` argument in :meth:`PreTrainedTokenizerBase.__call__`. Useful for
@@ -77,7 +77,7 @@ class TensorType(ExplicitEnum):
     NUMPY = "np"
     JAX = "jax"
 
-    
+
 class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(mask_token="<mask>", pad_token="<pad>")
@@ -89,11 +89,17 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
         self.pad_token_id = self.token_dictionary.get("<pad>")
         self.all_special_ids = [
             self.token_dictionary.get("<mask>"),
-            self.token_dictionary.get("<pad>")
+            self.token_dictionary.get("<pad>"),
         ]
 
     def _get_padding_truncation_strategies(
-        self, padding=True, truncation=False, max_length=None, pad_to_multiple_of=None, verbose=True, **kwargs
+        self,
+        padding=True,
+        truncation=False,
+        max_length=None,
+        pad_to_multiple_of=None,
+        verbose=True,
+        **kwargs,
     ):
         """
         Find the correct padding/truncation strategy with backward compatibility for old arguments (truncation_strategy
@@ -106,7 +112,9 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
         # If you only set max_length, it activates truncation for max_length
         if max_length is not None and padding is False and truncation is False:
             if verbose:
-                if not self.deprecation_warnings.get("Truncation-not-explicitly-activated", False):
+                if not self.deprecation_warnings.get(
+                    "Truncation-not-explicitly-activated", False
+                ):
                     logger.warning(
                         "Truncation was not explicitly activated but `max_length` is provided a specific value, "
                         "please use `truncation=True` to explicitly truncate examples to max length. "
@@ -134,7 +142,9 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
                 padding_strategy = PaddingStrategy.MAX_LENGTH
         elif padding is not False:
             if padding is True:
-                padding_strategy = PaddingStrategy.LONGEST  # Default to pad to the longest sequence in the batch
+                padding_strategy = (
+                    PaddingStrategy.LONGEST
+                )  # Default to pad to the longest sequence in the batch
             elif not isinstance(padding, PaddingStrategy):
                 padding_strategy = PaddingStrategy(padding)
             elif isinstance(padding, PaddingStrategy):
@@ -174,7 +184,9 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
             if padding_strategy == PaddingStrategy.MAX_LENGTH:
                 if self.model_max_length > LARGE_INTEGER:
                     if verbose:
-                        if not self.deprecation_warnings.get("Asking-to-pad-to-max_length", False):
+                        if not self.deprecation_warnings.get(
+                            "Asking-to-pad-to-max_length", False
+                        ):
                             logger.warning(
                                 "Asking to pad to max_length but no maximum length is provided and the model has no predefined maximum length. "
                                 "Default to no padding."
@@ -187,18 +199,24 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
             if truncation_strategy != TruncationStrategy.DO_NOT_TRUNCATE:
                 if self.model_max_length > LARGE_INTEGER:
                     if verbose:
-                        if not self.deprecation_warnings.get("Asking-to-truncate-to-max_length", False):
+                        if not self.deprecation_warnings.get(
+                            "Asking-to-truncate-to-max_length", False
+                        ):
                             logger.warning(
                                 "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. "
                                 "Default to no truncation."
                             )
-                        self.deprecation_warnings["Asking-to-truncate-to-max_length"] = True
+                        self.deprecation_warnings[
+                            "Asking-to-truncate-to-max_length"
+                        ] = True
                     truncation_strategy = TruncationStrategy.DO_NOT_TRUNCATE
                 else:
                     max_length = self.model_max_length
 
         # Test if we have a padding token
-        if padding_strategy != PaddingStrategy.DO_NOT_PAD and (not self.pad_token or self.pad_token_id < 0):
+        if padding_strategy != PaddingStrategy.DO_NOT_PAD and (
+            not self.pad_token or self.pad_token_id < 0
+        ):
             raise ValueError(
                 "Asking to pad but the tokenizer does not have a padding token. "
                 "Please select a token to use as `pad_token` `(tokenizer.pad_token = tokenizer.eos_token e.g.)` "
@@ -229,7 +247,7 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
             Dict[str, List[EncodedInput]],
             List[Dict[str, EncodedInput]],
         ],
-        class_type, # options: "gene" or "cell"
+        class_type,  # options: "gene" or "cell"
         padding: Union[bool, str, PaddingStrategy] = True,
         max_length: Optional[int] = None,
         pad_to_multiple_of: Optional[int] = None,
@@ -292,8 +310,13 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
         """
         # If we have a list of dicts, let's convert it in a dict of lists
         # We do this to allow using this method as a collate_fn function in PyTorch Dataloader
-        if isinstance(encoded_inputs, (list, tuple)) and isinstance(encoded_inputs[0], (dict, BatchEncoding)):
-            encoded_inputs = {key: [example[key] for example in encoded_inputs] for key in encoded_inputs[0].keys()}
+        if isinstance(encoded_inputs, (list, tuple)) and isinstance(
+            encoded_inputs[0], (dict, BatchEncoding)
+        ):
+            encoded_inputs = {
+                key: [example[key] for example in encoded_inputs]
+                for key in encoded_inputs[0].keys()
+            }
 
         # The model's main input name, usually `input_ids`, has be passed for padding
         if self.model_input_names[0] not in encoded_inputs:
@@ -387,7 +410,7 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
     def _pad(
         self,
         encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding],
-        class_type, # options: "gene" or "cell"
+        class_type,  # options: "gene" or "cell"
         max_length: Optional[int] = None,
         padding_strategy: PaddingStrategy = PaddingStrategy.LONGEST,
         pad_to_multiple_of: Optional[int] = None,
@@ -423,46 +446,73 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
         if padding_strategy == PaddingStrategy.LONGEST:
             max_length = len(required_input)
 
-        if max_length is not None and pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
+        if (
+            max_length is not None
+            and pad_to_multiple_of is not None
+            and (max_length % pad_to_multiple_of != 0)
+        ):
             max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
 
-        needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len(required_input) != max_length
+        needs_to_be_padded = (
+            padding_strategy != PaddingStrategy.DO_NOT_PAD
+            and len(required_input) != max_length
+        )
 
         if needs_to_be_padded:
             difference = max_length - len(required_input)
             if self.padding_side == "right":
                 if return_attention_mask:
-                    encoded_inputs["attention_mask"] = [1] * len(required_input) + [0] * difference
+                    encoded_inputs["attention_mask"] = [1] * len(required_input) + [
+                        0
+                    ] * difference
                 if "token_type_ids" in encoded_inputs:
                     encoded_inputs["token_type_ids"] = (
-                        encoded_inputs["token_type_ids"] + [self.pad_token_type_id] * difference
+                        encoded_inputs["token_type_ids"]
+                        + [self.pad_token_type_id] * difference
                     )
                 if "special_tokens_mask" in encoded_inputs:
-                    encoded_inputs["special_tokens_mask"] = encoded_inputs["special_tokens_mask"] + [1] * difference
-                encoded_inputs[self.model_input_names[0]] = required_input + [self.pad_token_id] * difference
+                    encoded_inputs["special_tokens_mask"] = (
+                        encoded_inputs["special_tokens_mask"] + [1] * difference
+                    )
+                encoded_inputs[self.model_input_names[0]] = (
+                    required_input + [self.pad_token_id] * difference
+                )
                 if class_type == "gene":
-                    encoded_inputs["labels"] = encoded_inputs["labels"] + [-100] * difference
+                    encoded_inputs["labels"] = (
+                        encoded_inputs["labels"] + [-100] * difference
+                    )
             elif self.padding_side == "left":
                 if return_attention_mask:
-                    encoded_inputs["attention_mask"] = [0] * difference + [1] * len(required_input)
+                    encoded_inputs["attention_mask"] = [0] * difference + [1] * len(
+                        required_input
+                    )
                 if "token_type_ids" in encoded_inputs:
-                    encoded_inputs["token_type_ids"] = [self.pad_token_type_id] * difference + encoded_inputs[
-                        "token_type_ids"
-                    ]
+                    encoded_inputs["token_type_ids"] = [
+                        self.pad_token_type_id
+                    ] * difference + encoded_inputs["token_type_ids"]
                 if "special_tokens_mask" in encoded_inputs:
-                    encoded_inputs["special_tokens_mask"] = [1] * difference + encoded_inputs["special_tokens_mask"]
-                encoded_inputs[self.model_input_names[0]] = [self.pad_token_id] * difference + required_input
+                    encoded_inputs["special_tokens_mask"] = [
+                        1
+                    ] * difference + encoded_inputs["special_tokens_mask"]
+                encoded_inputs[self.model_input_names[0]] = [
+                    self.pad_token_id
+                ] * difference + required_input
                 if class_type == "gene":
-                    encoded_inputs["labels"] = [-100] * difference + encoded_inputs["labels"]
+                    encoded_inputs["labels"] = [-100] * difference + encoded_inputs[
+                        "labels"
+                    ]
             else:
                 raise ValueError("Invalid padding strategy:" + str(self.padding_side))
         elif return_attention_mask and "attention_mask" not in encoded_inputs:
             encoded_inputs["attention_mask"] = [1] * len(required_input)
-        
+
         return encoded_inputs
 
     def get_special_tokens_mask(
-        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None,
+        already_has_special_tokens: bool = False,
     ) -> List[int]:
         """
         Retrieves sequence ids from a token list that has no special tokens added. This method is called when adding
@@ -486,11 +536,15 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
 
         all_special_ids = self.all_special_ids  # cache the property
 
-        special_tokens_mask = [1 if token in all_special_ids else 0 for token in token_ids_0]
+        special_tokens_mask = [
+            1 if token in all_special_ids else 0 for token in token_ids_0
+        ]
 
         return special_tokens_mask
 
-    def convert_tokens_to_ids(self, tokens: Union[str, List[str]]) -> Union[int, List[int]]:
+    def convert_tokens_to_ids(
+        self, tokens: Union[str, List[str]]
+    ) -> Union[int, List[int]]:
         """
         Converts a token string (or a sequence of tokens) in a single integer id (or a sequence of ids), using the
         vocabulary.
@@ -514,14 +568,15 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
         if token is None:
             return None
 
-        return token_dictionary.get(token)
+        return self.token_dictionary.get(token)
 
     def __len__(self):
-        return len(token_dictionary)
+        return len(self.token_dictionary)
 
 
 # collator functions
 
+
 class DataCollatorForGeneClassification(DataCollatorForTokenClassification):
     """
     Data collator that will dynamically pad the inputs received, as well as the labels.
@@ -546,26 +601,34 @@ class DataCollatorForGeneClassification(DataCollatorForTokenClassification):
         label_pad_token_id (:obj:`int`, `optional`, defaults to -100):
             The id to use when padding the labels (-100 will be automatically ignore by PyTorch loss functions).
     """
-   
+
     class_type = "gene"
     padding: Union[bool, str, PaddingStrategy] = True
     max_length: Optional[int] = None
     pad_to_multiple_of: Optional[int] = None
     label_pad_token_id: int = -100
-    
+
     def __init__(self, *args, **kwargs) -> None:
         self.token_dictionary = kwargs.pop("token_dictionary")
         super().__init__(
-            tokenizer=PrecollatorForGeneAndCellClassification(token_dictionary=self.token_dictionary),
+            tokenizer=PrecollatorForGeneAndCellClassification(
+                token_dictionary=self.token_dictionary
+            ),
             padding=self.padding,
             max_length=self.max_length,
             pad_to_multiple_of=self.pad_to_multiple_of,
             label_pad_token_id=self.label_pad_token_id,
-            *args, **kwargs)
+            *args,
+            **kwargs,
+        )
 
     def _prepare_batch(self, features):
         label_name = "label" if "label" in features[0].keys() else "labels"
-        labels = [feature[label_name] for feature in features] if label_name in features[0].keys() else None
+        labels = (
+            [feature[label_name] for feature in features]
+            if label_name in features[0].keys()
+            else None
+        )
         batch = self.tokenizer.pad(
             features,
             class_type=self.class_type,
@@ -575,29 +638,31 @@ class DataCollatorForGeneClassification(DataCollatorForTokenClassification):
             return_tensors="pt",
         )
         return batch
-    
+
     def __call__(self, features):
         batch = self._prepare_batch(features)
 
         batch = {k: torch.tensor(v, dtype=torch.int64) for k, v in batch.items()}
         return batch
 
-    
-class DataCollatorForCellClassification(DataCollatorForGeneClassification):
 
+class DataCollatorForCellClassification(DataCollatorForGeneClassification):
     class_type = "cell"
 
     def _prepare_batch(self, features):
-        
         batch = super()._prepare_batch(features)
-        
+
         # Special handling for labels.
         # Ensure that tensor is created with the correct type
         # (it should be automatically the case, but let's make sure of it.)
         first = features[0]
         if "label" in first and first["label"] is not None:
-            label = first["label"].item() if isinstance(first["label"], torch.Tensor) else first["label"]
+            label = (
+                first["label"].item()
+                if isinstance(first["label"], torch.Tensor)
+                else first["label"]
+            )
             dtype = torch.long if isinstance(label, int) else torch.float
             batch["labels"] = torch.tensor([f["label"] for f in features], dtype=dtype)
-            
+
         return batch

geneformer/emb_extractor.py

@@ -24,8 +24,8 @@ import torch
 from tdigest import TDigest
 from tqdm.auto import trange
 
-from . import perturber_utils as pu
 from . import TOKEN_DICTIONARY_FILE
+from . import perturber_utils as pu
 
 logger = logging.getLogger(__name__)
 
@@ -45,7 +45,7 @@ def get_embs(
 ):
     model_input_size = pu.get_model_input_size(model)
     total_batch_length = len(filtered_input_data)
-    
+
     if summary_stat is None:
         embs_list = []
     elif summary_stat is not None:
@@ -73,17 +73,23 @@ def get_embs(
     if emb_mode == "cls":
         assert cls_present, "<cls> token missing in token dictionary"
         # Check to make sure that the first token of the filtered input data is cls token
-        gene_token_dict = {v:k for k,v in token_gene_dict.items()}
+        gene_token_dict = {v: k for k, v in token_gene_dict.items()}
         cls_token_id = gene_token_dict["<cls>"]
-        assert filtered_input_data["input_ids"][0][0] == cls_token_id, "First token is not <cls> token value"
+        assert (
+            filtered_input_data["input_ids"][0][0] == cls_token_id
+        ), "First token is not <cls> token value"
     elif emb_mode == "cell":
         if cls_present:
-            logger.warning("CLS token present in token dictionary, excluding from average.")    
+            logger.warning(
+                "CLS token present in token dictionary, excluding from average."
+            )
         if eos_present:
-            logger.warning("EOS token present in token dictionary, excluding from average.")
-            
+            logger.warning(
+                "EOS token present in token dictionary, excluding from average."
+            )
+
     overall_max_len = 0
-        
+
     for i in trange(0, total_batch_length, forward_batch_size, leave=(not silent)):
         max_range = min(i + forward_batch_size, total_batch_length)
 
@@ -108,7 +114,7 @@ def get_embs(
 
         if emb_mode == "cell":
             if cls_present:
-                non_cls_embs = embs_i[:, 1:, :] # Get all layers except the embs
+                non_cls_embs = embs_i[:, 1:, :]  # Get all layers except the embs
                 if eos_present:
                     mean_embs = pu.mean_nonpadding_embs(non_cls_embs, original_lens - 2)
                 else:
@@ -146,10 +152,10 @@ def get_embs(
                     del embs_h
                     del dict_h
         elif emb_mode == "cls":
-            cls_embs = embs_i[:,0,:].clone().detach() # CLS token layer
+            cls_embs = embs_i[:, 0, :].clone().detach()  # CLS token layer
             embs_list.append(cls_embs)
             del cls_embs
-            
+
         overall_max_len = max(overall_max_len, max_len)
         del outputs
         del minibatch
@@ -157,8 +163,7 @@ def get_embs(
         del embs_i
 
         torch.cuda.empty_cache()
-        
-        
+
     if summary_stat is None:
         if (emb_mode == "cell") or (emb_mode == "cls"):
             embs_stack = torch.cat(embs_list, dim=0)
@@ -204,6 +209,7 @@ def accumulate_tdigests(embs_tdigests, mean_embs, emb_dims):
         for j in range(emb_dims)
     ]
 
+
 def update_tdigest_dict(embs_tdigests_dict, gene, gene_embs, emb_dims):
     embs_tdigests_dict[gene] = accumulate_tdigests(
         embs_tdigests_dict[gene], gene_embs, emb_dims
@@ -294,11 +300,13 @@ def plot_umap(embs_df, emb_dims, label, output_file, kwargs_dict, seed=0):
 
     with plt.rc_context():
         ax = sc.pl.umap(adata, color=label, show=False, **default_kwargs_dict)
-        ax.legend(markerscale=2,
-                  frameon=False,
-                  loc="center left",
-                  bbox_to_anchor=(1, 0.5),
-                  ncol=(1 if len(cats) <= 14 else 2 if len(cats) <= 30 else 3))
+        ax.legend(
+            markerscale=2,
+            frameon=False,
+            loc="center left",
+            bbox_to_anchor=(1, 0.5),
+            ncol=(1 if len(cats) <= 14 else 2 if len(cats) <= 30 else 3),
+        )
         plt.show()
         plt.savefig(output_file, bbox_inches="tight")
 
@@ -394,7 +402,7 @@ class EmbExtractor:
         "emb_label": {None, list},
         "labels_to_plot": {None, list},
         "forward_batch_size": {int},
-        "token_dictionary_file" : {None, str},
+        "token_dictionary_file": {None, str},
         "nproc": {int},
         "summary_stat": {None, "mean", "median", "exact_mean", "exact_median"},
     }
@@ -631,13 +639,15 @@ class EmbExtractor:
             embs = embs.mean(dim=0)
             emb_dims = pu.get_model_emb_dims(model)
             embs_df = pd.DataFrame(
-                embs_df[0:emb_dims-1].mean(axis="rows"), columns=[self.exact_summary_stat]
+                embs_df[0 : emb_dims - 1].mean(axis="rows"),
+                columns=[self.exact_summary_stat],
             ).T
         elif self.exact_summary_stat == "exact_median":
             embs = torch.median(embs, dim=0)[0]
             emb_dims = pu.get_model_emb_dims(model)
             embs_df = pd.DataFrame(
-                embs_df[0:emb_dims-1].median(axis="rows"), columns=[self.exact_summary_stat]
+                embs_df[0 : emb_dims - 1].median(axis="rows"),
+                columns=[self.exact_summary_stat],
             ).T
 
         if cell_state is not None:
@@ -838,4 +848,4 @@ class EmbExtractor:
                 output_file = (
                     Path(output_directory) / output_prefix_label
                 ).with_suffix(".pdf")
-                plot_heatmap(embs, emb_dims, label, output_file, kwargs_dict)
+                plot_heatmap(embs, emb_dims, label, output_file, kwargs_dict)

geneformer/evaluation_utils.py

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

geneformer/in_silico_perturber.py

@@ -38,16 +38,15 @@ import logging
 import os
 import pickle
 from collections import defaultdict
-from multiprocess import set_start_method
-from typing import List
 
 import torch
 from datasets import Dataset, disable_progress_bars
+from multiprocess import set_start_method
 from tqdm.auto import trange
 
+from . import TOKEN_DICTIONARY_FILE
 from . import perturber_utils as pu
 from .emb_extractor import get_embs
-from . import TOKEN_DICTIONARY_FILE
 
 disable_progress_bars()
 
@@ -71,7 +70,7 @@ class InSilicoPerturber:
         "max_ncells": {None, int},
         "cell_inds_to_perturb": {"all", dict},
         "emb_layer": {-1, 0},
-        "token_dictionary_file" : {None, str},
+        "token_dictionary_file": {None, str},
         "forward_batch_size": {int},
         "nproc": {int},
     }
@@ -239,13 +238,14 @@ class InSilicoPerturber:
         self.cls_token_id = self.gene_token_dict.get("<cls>")
         self.eos_token_id = self.gene_token_dict.get("<eos>")
 
-
         # Identify if special token is present in the token dictionary
         if (self.cls_token_id is not None) and (self.eos_token_id is not None):
             self.special_token = True
         else:
             if "cls" in self.emb_mode:
-                logger.error(f"emb_mode set to {self.emb_mode} but <cls> or <eos> token not in token dictionary.")
+                logger.error(
+                    f"emb_mode set to {self.emb_mode} but <cls> or <eos> token not in token dictionary."
+                )
                 raise
             self.special_token = False
 
@@ -454,17 +454,21 @@ class InSilicoPerturber:
 
         # Ensure emb_mode is cls if first token of the filtered input data is cls token
         if self.special_token:
-            if (filtered_input_data["input_ids"][0][0] == self.cls_token_id) and ("cls" not in self.emb_mode):
+            if (filtered_input_data["input_ids"][0][0] == self.cls_token_id) and (
+                "cls" not in self.emb_mode
+            ):
                 logger.error(
-                            "Emb mode 'cls' or 'cls_and_gene' required when first token is <cls>."
-                        )
+                    "Emb mode 'cls' or 'cls_and_gene' required when first token is <cls>."
+                )
                 raise
-            if ("cls" in self.emb_mode):
-                if (filtered_input_data["input_ids"][0][0] != self.cls_token_id) or (filtered_input_data["input_ids"][0][-1] != self.eos_token_id):
+            if "cls" in self.emb_mode:
+                if (filtered_input_data["input_ids"][0][0] != self.cls_token_id) or (
+                    filtered_input_data["input_ids"][0][-1] != self.eos_token_id
+                ):
                     logger.error(
-                                "Emb mode 'cls' and 'cls_and_gene' require that first token is <cls> and last token is <eos>."
-                            )
-                    raise                    
+                        "Emb mode 'cls' and 'cls_and_gene' require that first token is <cls> and last token is <eos>."
+                    )
+                    raise
 
         filtered_input_data = self.apply_additional_filters(filtered_input_data)
 
@@ -530,7 +534,6 @@ class InSilicoPerturber:
         layer_to_quant: int,
         output_path_prefix: str,
     ):
-        
         def make_group_perturbation_batch(example):
             example_input_ids = example["input_ids"]
             example["tokens_to_perturb"] = self.tokens_to_perturb
@@ -549,7 +552,9 @@ class InSilicoPerturber:
             if self.perturb_type == "delete":
                 example = pu.delete_indices(example)
             elif self.perturb_type == "overexpress":
-                example = pu.overexpress_tokens(example, self.max_len, self.special_token)
+                example = pu.overexpress_tokens(
+                    example, self.max_len, self.special_token
+                )
                 example["n_overflow"] = pu.calc_n_overflow(
                     self.max_len,
                     example["length"],
@@ -570,7 +575,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"]
@@ -752,7 +757,6 @@ class InSilicoPerturber:
                 f"{output_path_prefix}_gene_embs_dict_{self.tokens_to_perturb}",
             )
 
-
     def isp_perturb_set_special(
         self,
         model,
@@ -760,7 +764,6 @@ class InSilicoPerturber:
         layer_to_quant: int,
         output_path_prefix: str,
     ):
-        
         def make_group_perturbation_batch(example):
             example_input_ids = example["input_ids"]
             example["tokens_to_perturb"] = self.tokens_to_perturb
@@ -779,7 +782,9 @@ class InSilicoPerturber:
             if self.perturb_type == "delete":
                 example = pu.delete_indices(example)
             elif self.perturb_type == "overexpress":
-                example = pu.overexpress_tokens(example, self.max_len, self.special_token)
+                example = pu.overexpress_tokens(
+                    example, self.max_len, self.special_token
+                )
                 example["n_overflow"] = pu.calc_n_overflow(
                     self.max_len,
                     example["length"],
@@ -808,7 +813,7 @@ class InSilicoPerturber:
             filtered_input_data = filtered_input_data.map(
                 pu.truncate_by_n_overflow_special, num_proc=self.nproc
             )
-        
+
         if self.emb_mode == "cls_and_gene":
             stored_gene_embs_dict = defaultdict(list)
 
@@ -847,28 +852,29 @@ class InSilicoPerturber:
                     summary_stat=None,
                     silent=True,
                 )
-                
-                # Calculate the cosine similarities 
+
+                # Calculate the cosine similarities
                 cls_cos_sims = pu.quant_cos_sims(
                     perturbation_cls_emb,
                     original_cls_emb,
                     self.cell_states_to_model,
                     self.state_embs_dict,
-                    emb_mode="cell")
-                    
+                    emb_mode="cell",
+                )
+
                 # Update perturbation dictionary
                 if self.cell_states_to_model is None:
                     cos_sims_dict = self.update_perturbation_dictionary(
                         cos_sims_dict,
                         cls_cos_sims,
-                        gene_list = None,
+                        gene_list=None,
                     )
                 else:
                     for state in cos_sims_dict.keys():
                         cos_sims_dict[state] = self.update_perturbation_dictionary(
                             cos_sims_dict[state],
                             cls_cos_sims[state],
-                            gene_list = None,
+                            gene_list=None,
                         )
 
             ##### CLS and Gene Embedding Mode #####
@@ -908,9 +914,13 @@ class InSilicoPerturber:
                 # remove special tokens and padding
                 original_emb = original_emb[:, 1:-1, :]
                 if self.perturb_type == "overexpress":
-                    perturbation_emb = full_perturbation_emb[:,1+len(self.tokens_to_perturb):-1,:]
+                    perturbation_emb = full_perturbation_emb[
+                        :, 1 + len(self.tokens_to_perturb) : -1, :
+                    ]
                 elif self.perturb_type == "delete":
-                    perturbation_emb = full_perturbation_emb[:,1:max(perturbation_batch["length"])-1,:]
+                    perturbation_emb = full_perturbation_emb[
+                        :, 1 : max(perturbation_batch["length"]) - 1, :
+                    ]
 
                 n_perturbation_genes = perturbation_emb.size()[1]
 
@@ -923,8 +933,8 @@ class InSilicoPerturber:
                 )
 
                 # get cls emb
-                original_cls_emb = full_original_emb[:,0,:]
-                perturbation_cls_emb = full_perturbation_emb[:,0,:]
+                original_cls_emb = full_original_emb[:, 0, :]
+                perturbation_cls_emb = full_perturbation_emb[:, 0, :]
 
                 cls_cos_sims = pu.quant_cos_sims(
                     perturbation_cls_emb,
@@ -939,7 +949,10 @@ class InSilicoPerturber:
 
                 gene_list = minibatch["input_ids"]
                 # need to truncate gene_list
-                genes_to_exclude = self.tokens_to_perturb + [self.cls_token_id, self.eos_token_id]
+                genes_to_exclude = self.tokens_to_perturb + [
+                    self.cls_token_id,
+                    self.eos_token_id,
+                ]
                 gene_list = [
                     [g for g in genes if g not in genes_to_exclude][
                         :n_perturbation_genes
@@ -968,20 +981,20 @@ class InSilicoPerturber:
                     cos_sims_dict = self.update_perturbation_dictionary(
                         cos_sims_dict,
                         cls_cos_sims,
-                        gene_list = None,
+                        gene_list=None,
                     )
                 else:
                     for state in cos_sims_dict.keys():
                         cos_sims_dict[state] = self.update_perturbation_dictionary(
                             cos_sims_dict[state],
                             cls_cos_sims[state],
-                            gene_list = None,
+                            gene_list=None,
                         )
                 del full_original_emb
                 del original_emb
                 del full_perturbation_emb
-                del perturbation_emb 
-                del gene_cos_sims    
+                del perturbation_emb
+                del gene_cos_sims
 
             del original_cls_emb
             del perturbation_cls_emb
@@ -1035,12 +1048,12 @@ class InSilicoPerturber:
                 summary_stat=None,
                 silent=True,
             )
-            
+
             if self.cell_states_to_model is not None:
                 original_cell_emb = pu.compute_nonpadded_cell_embedding(
                     full_original_emb, "mean_pool"
                 )
-                
+
             # gene_list is used to assign cos sims back to genes
             gene_list = example_cell["input_ids"][0][:]
             # need to remove the anchor gene
@@ -1049,15 +1062,13 @@ class InSilicoPerturber:
                     gene_list.remove(token)
             # index 0 is not overexpressed so remove
             if self.perturb_type == "overexpress":
-                gene_list = gene_list[
-                    num_inds_perturbed:
-                ]
+                gene_list = gene_list[num_inds_perturbed:]
             # remove perturbed index for gene list dict
             perturbed_gene_dict = {
                 gene: gene_list[:i] + gene_list[i + 1 :]
                 for i, gene in enumerate(gene_list)
             }
-            
+
             perturbation_batch, indices_to_perturb = pu.make_perturbation_batch(
                 example_cell,
                 self.perturb_type,
@@ -1068,11 +1079,19 @@ class InSilicoPerturber:
             )
 
             ispall_total_batch_length = len(perturbation_batch)
-            for i in trange(0, ispall_total_batch_length, self.forward_batch_size, leave=False):
-                ispall_max_range = min(i + self.forward_batch_size, ispall_total_batch_length)        
-                perturbation_minibatch = perturbation_batch.select([i for i in range(i, ispall_max_range)])
-                indices_to_perturb_mini = indices_to_perturb[i : ispall_max_range]
-                gene_list_mini = gene_list[i : ispall_max_range] # only perturbed genes from this minibatch
+            for i in trange(
+                0, ispall_total_batch_length, self.forward_batch_size, leave=False
+            ):
+                ispall_max_range = min(
+                    i + self.forward_batch_size, ispall_total_batch_length
+                )
+                perturbation_minibatch = perturbation_batch.select(
+                    [i for i in range(i, ispall_max_range)]
+                )
+                indices_to_perturb_mini = indices_to_perturb[i:ispall_max_range]
+                gene_list_mini = gene_list[
+                    i:ispall_max_range
+                ]  # only perturbed genes from this minibatch
 
                 full_perturbation_emb = get_embs(
                     model,
@@ -1085,18 +1104,20 @@ class InSilicoPerturber:
                     summary_stat=None,
                     silent=True,
                 )
-    
+
                 del perturbation_minibatch
-    
+
                 # need to remove overexpressed gene to quantify cosine shifts
                 if self.perturb_type == "overexpress":
                     perturbation_emb = full_perturbation_emb[:, num_inds_perturbed:, :]
-    
+
                 elif self.perturb_type == "delete":
                     perturbation_emb = full_perturbation_emb
-    
-    
-                if self.cell_states_to_model is None or self.emb_mode == "cell_and_gene":
+
+                if (
+                    self.cell_states_to_model is None
+                    or self.emb_mode == "cell_and_gene"
+                ):
                     original_emb_minibatch = pu.make_comparison_batch(
                         full_original_emb, indices_to_perturb_mini, perturb_group=False
                     )
@@ -1108,12 +1129,12 @@ class InSilicoPerturber:
                         emb_mode="gene",
                     )
                     del original_emb_minibatch
-    
+
                 if self.cell_states_to_model is not None:
                     perturbation_cell_emb = pu.compute_nonpadded_cell_embedding(
                         full_perturbation_emb, "mean_pool"
                     )
-    
+
                     cell_cos_sims = pu.quant_cos_sims(
                         perturbation_cell_emb,
                         original_cell_emb,
@@ -1122,9 +1143,8 @@ class InSilicoPerturber:
                         emb_mode="cell",
                     )
                     del perturbation_cell_emb
-    
-                if self.emb_mode == "cell_and_gene":
 
+                if self.emb_mode == "cell_and_gene":
                     for perturbation_i, perturbed_gene in enumerate(gene_list_mini):
                         for gene_j, affected_gene in enumerate(
                             perturbed_gene_dict[perturbed_gene]
@@ -1137,9 +1157,9 @@ class InSilicoPerturber:
                                 stored_gene_embs_dict[
                                     (perturbed_gene, affected_gene)
                                 ] = gene_cos_sims[perturbation_i, gene_j].item()
-    
+
                 del full_perturbation_emb
-    
+
                 if self.cell_states_to_model is None:
                     cos_sims_data = torch.mean(gene_cos_sims, dim=1)
                     cos_sims_dict = self.update_perturbation_dictionary(
@@ -1155,9 +1175,9 @@ class InSilicoPerturber:
                             cos_sims_data[state],
                             gene_list_mini,
                         )
-                
+
                 # save dict to disk every self.clear_mem_ncells/10 (default 100) simulated cells
-                if i % self.clear_mem_ncells/10 == 0:
+                if i % self.clear_mem_ncells / 10 == 0:
                     pu.write_perturbation_dictionary(
                         cos_sims_dict,
                         f"{output_path_prefix}_dict_cell_embs_{h}batch{pickle_batch}",
@@ -1167,7 +1187,7 @@ class InSilicoPerturber:
                             stored_gene_embs_dict,
                             f"{output_path_prefix}_dict_gene_embs_{h}batch{pickle_batch}",
                         )
-    
+
                 # reset and clear memory every self.clear_mem_ncells (default 1000) simulated cells or at the end of the example cell
                 if i % self.clear_mem_ncells == 0:
                     pickle_batch += 1
@@ -1176,18 +1196,21 @@ class InSilicoPerturber:
                     else:
                         cos_sims_dict = {
                             state: defaultdict(list)
-                            for state in pu.get_possible_states(self.cell_states_to_model)
+                            for state in pu.get_possible_states(
+                                self.cell_states_to_model
+                            )
                         }
-    
+
                     if self.emb_mode == "cell_and_gene":
                         stored_gene_embs_dict = defaultdict(list)
-    
+
                     torch.cuda.empty_cache()
-        
+
             pu.write_perturbation_dictionary(
-                cos_sims_dict, f"{output_path_prefix}_dict_cell_embs_{h}batch{pickle_batch}"
+                cos_sims_dict,
+                f"{output_path_prefix}_dict_cell_embs_{h}batch{pickle_batch}",
             )
-    
+
             if self.emb_mode == "cell_and_gene":
                 pu.write_perturbation_dictionary(
                     stored_gene_embs_dict,
@@ -1212,7 +1235,7 @@ class InSilicoPerturber:
             if self.cell_states_to_model is not None:
                 del original_cell_emb
             torch.cuda.empty_cache()
-    
+
     def isp_perturb_all_special(
         self,
         model,
@@ -1248,7 +1271,7 @@ class InSilicoPerturber:
                     self.token_gene_dict,
                     summary_stat=None,
                     silent=True,
-                )                
+                )
             elif self.emb_mode == "cls_and_gene":
                 full_original_emb = get_embs(
                     model,
@@ -1261,8 +1284,8 @@ class InSilicoPerturber:
                     summary_stat=None,
                     silent=True,
                 )
-                original_cls_emb = full_original_emb[:,0,:].clone().detach()
-            
+                original_cls_emb = full_original_emb[:, 0, :].clone().detach()
+
             # gene_list is used to assign cos sims back to genes
             gene_list = example_cell["input_ids"][0][:]
 
@@ -1275,9 +1298,7 @@ class InSilicoPerturber:
                     gene_list.remove(token)
             # index 0 is not overexpressed so remove
             if self.perturb_type == "overexpress":
-                gene_list = gene_list[
-                    num_inds_perturbed:
-                ]
+                gene_list = gene_list[num_inds_perturbed:]
             # remove perturbed index for gene list dict
             perturbed_gene_dict = {
                 gene: gene_list[:i] + gene_list[i + 1 :]
@@ -1294,12 +1315,20 @@ class InSilicoPerturber:
             )
 
             ispall_total_batch_length = len(perturbation_batch)
-            for i in trange(0, ispall_total_batch_length, self.forward_batch_size, leave=False):
-                ispall_max_range = min(i + self.forward_batch_size, ispall_total_batch_length)        
-                perturbation_minibatch = perturbation_batch.select([i for i in range(i, ispall_max_range)])
-                indices_to_perturb_mini = indices_to_perturb[i : ispall_max_range]
-                gene_list_mini = gene_list[i : ispall_max_range] # only perturbed genes from this minibatch 
-                
+            for i in trange(
+                0, ispall_total_batch_length, self.forward_batch_size, leave=False
+            ):
+                ispall_max_range = min(
+                    i + self.forward_batch_size, ispall_total_batch_length
+                )
+                perturbation_minibatch = perturbation_batch.select(
+                    [i for i in range(i, ispall_max_range)]
+                )
+                indices_to_perturb_mini = indices_to_perturb[i:ispall_max_range]
+                gene_list_mini = gene_list[
+                    i:ispall_max_range
+                ]  # only perturbed genes from this minibatch
+
                 ##### CLS Embedding Mode #####
                 if self.emb_mode == "cls":
                     # Extract cls embeddings from perturbed cells
@@ -1314,7 +1343,7 @@ class InSilicoPerturber:
                         summary_stat=None,
                         silent=True,
                     )
-    
+
                     # Calculate cosine similarities
                     cls_cos_sims = pu.quant_cos_sims(
                         perturbation_cls_emb,
@@ -1322,8 +1351,8 @@ class InSilicoPerturber:
                         self.cell_states_to_model,
                         self.state_embs_dict,
                         emb_mode="cell",
-                    )           
-                    
+                    )
+
                     if self.cell_states_to_model is None:
                         cos_sims_dict = self.update_perturbation_dictionary(
                             cos_sims_dict,
@@ -1331,20 +1360,19 @@ class InSilicoPerturber:
                             gene_list_mini,
                         )
                     else:
-    
                         for state in cos_sims_dict.keys():
                             cos_sims_dict[state] = self.update_perturbation_dictionary(
                                 cos_sims_dict[state],
                                 cls_cos_sims[state],
                                 gene_list_mini,
                             )
-                    
+
                     del perturbation_minibatch
                     del perturbation_cls_emb
                     del cls_cos_sims
-    
+
                 ##### CLS and Gene Embedding Mode #####
-                elif self.emb_mode == "cls_and_gene":                
+                elif self.emb_mode == "cls_and_gene":
                     full_perturbation_emb = get_embs(
                         model,
                         perturbation_minibatch,
@@ -1356,18 +1384,26 @@ class InSilicoPerturber:
                         summary_stat=None,
                         silent=True,
                     )
-                    
+
                     # need to remove overexpressed gene and cls/eos to quantify cosine shifts
                     if self.perturb_type == "overexpress":
-                        perturbation_emb = full_perturbation_emb[:, 1+num_inds_perturbed:-1, :].clone().detach()
+                        perturbation_emb = (
+                            full_perturbation_emb[:, 1 + num_inds_perturbed : -1, :]
+                            .clone()
+                            .detach()
+                        )
                     elif self.perturb_type == "delete":
-                        perturbation_emb = full_perturbation_emb[:, 1:-1, :].clone().detach()
-    
+                        perturbation_emb = (
+                            full_perturbation_emb[:, 1:-1, :].clone().detach()
+                        )
+
                     original_emb_minibatch = pu.make_comparison_batch(
                         full_original_emb, indices_to_perturb_mini, perturb_group=False
                     )
-    
-                    original_emb_minibatch = original_emb_minibatch[:, 1:-1, :].clone().detach()
+
+                    original_emb_minibatch = (
+                        original_emb_minibatch[:, 1:-1, :].clone().detach()
+                    )
                     gene_cos_sims = pu.quant_cos_sims(
                         perturbation_emb,
                         original_emb_minibatch,
@@ -1375,7 +1411,7 @@ class InSilicoPerturber:
                         self.state_embs_dict,
                         emb_mode="gene",
                     )
-                       
+
                     for perturbation_i, perturbed_gene in enumerate(gene_list_mini):
                         for gene_j, affected_gene in enumerate(
                             perturbed_gene_dict[perturbed_gene]
@@ -1388,10 +1424,12 @@ class InSilicoPerturber:
                                 stored_gene_embs_dict[
                                     (perturbed_gene, affected_gene)
                                 ] = gene_cos_sims[perturbation_i, gene_j].item()
-                    
+
                     # get cls emb
-                    perturbation_cls_emb = full_perturbation_emb[:,0,:].clone().detach()
-    
+                    perturbation_cls_emb = (
+                        full_perturbation_emb[:, 0, :].clone().detach()
+                    )
+
                     cls_cos_sims = pu.quant_cos_sims(
                         perturbation_cls_emb,
                         original_cls_emb,
@@ -1399,7 +1437,7 @@ class InSilicoPerturber:
                         self.state_embs_dict,
                         emb_mode="cell",
                     )
-    
+
                     if self.cell_states_to_model is None:
                         cos_sims_dict = self.update_perturbation_dictionary(
                             cos_sims_dict,
@@ -1413,7 +1451,7 @@ class InSilicoPerturber:
                                 cls_cos_sims[state],
                                 gene_list_mini,
                             )
-    
+
                     del perturbation_minibatch
                     del original_emb_minibatch
                     del full_perturbation_emb
@@ -1421,9 +1459,9 @@ class InSilicoPerturber:
                     del perturbation_cls_emb
                     del cls_cos_sims
                     del gene_cos_sims
-    
+
                 # save dict to disk every self.clear_mem_ncells/10 (default 100) simulated cells
-                if i % max(1,self.clear_mem_ncells/10) == 0:
+                if i % max(1, self.clear_mem_ncells / 10) == 0:
                     pu.write_perturbation_dictionary(
                         cos_sims_dict,
                         f"{output_path_prefix}_dict_cell_embs_{h}batch{pickle_batch}",
@@ -1433,7 +1471,7 @@ class InSilicoPerturber:
                             stored_gene_embs_dict,
                             f"{output_path_prefix}_dict_gene_embs_{h}batch{pickle_batch}",
                         )
-    
+
                 # reset and clear memory every self.clear_mem_ncells (default 1000) simulated cells or at the end of the example cell
                 if i % self.clear_mem_ncells == 0:
                     pickle_batch += 1
@@ -1442,18 +1480,21 @@ class InSilicoPerturber:
                     else:
                         cos_sims_dict = {
                             state: defaultdict(list)
-                            for state in pu.get_possible_states(self.cell_states_to_model)
+                            for state in pu.get_possible_states(
+                                self.cell_states_to_model
+                            )
                         }
-    
+
                     if self.emb_mode == "cls_and_gene":
                         stored_gene_embs_dict = defaultdict(list)
-    
+
                     torch.cuda.empty_cache()
-                
+
             pu.write_perturbation_dictionary(
-                cos_sims_dict, f"{output_path_prefix}_dict_cell_embs_{h}batch{pickle_batch}"
+                cos_sims_dict,
+                f"{output_path_prefix}_dict_cell_embs_{h}batch{pickle_batch}",
             )
-    
+
             if self.emb_mode == "cls_and_gene":
                 pu.write_perturbation_dictionary(
                     stored_gene_embs_dict,
@@ -1470,8 +1511,8 @@ class InSilicoPerturber:
                 }
 
             if self.emb_mode == "cls_and_gene":
-                stored_gene_embs_dict = defaultdict(list)            
-            
+                stored_gene_embs_dict = defaultdict(list)
+
             # clear memory between cells
             del perturbation_batch
             del original_cls_emb
@@ -1479,7 +1520,6 @@ class InSilicoPerturber:
                 del full_original_emb
             torch.cuda.empty_cache()
 
-    
     def update_perturbation_dictionary(
         self,
         cos_sims_dict: defaultdict,
@@ -1514,4 +1554,4 @@ class InSilicoPerturber:
             for i, cos in enumerate(cos_sims_data.tolist()):
                 cos_sims_dict[(gene_list[i], "cell_emb")].append(cos)
 
-        return cos_sims_dict
+        return cos_sims_dict

geneformer/in_silico_perturber_stats.py

@@ -37,8 +37,8 @@ from scipy.stats import ranksums
 from sklearn.mixture import GaussianMixture
 from tqdm.auto import tqdm, trange
 
+from . import ENSEMBL_DICTIONARY_FILE, TOKEN_DICTIONARY_FILE
 from .perturber_utils import flatten_list, validate_cell_states_to_model
-from . import TOKEN_DICTIONARY_FILE, ENSEMBL_DICTIONARY_FILE
 
 logger = logging.getLogger(__name__)
 
@@ -194,23 +194,29 @@ def get_impact_component(test_value, gaussian_mixture_model):
 def isp_aggregate_grouped_perturb(cos_sims_df, dict_list, genes_perturbed):
     names = ["Cosine_sim", "Gene"]
     cos_sims_full_dfs = []
-    if isinstance(genes_perturbed,list):
-        if len(genes_perturbed)>1:
-            gene_ids_df = cos_sims_df.loc[np.isin([set(idx) for idx in cos_sims_df["Ensembl_ID"]], set(genes_perturbed)), :]
+    if isinstance(genes_perturbed, list):
+        if len(genes_perturbed) > 1:
+            gene_ids_df = cos_sims_df.loc[
+                np.isin(
+                    [set(idx) for idx in cos_sims_df["Ensembl_ID"]],
+                    set(genes_perturbed),
+                ),
+                :,
+            ]
         else:
-            gene_ids_df = cos_sims_df.loc[np.isin(cos_sims_df["Ensembl_ID"], genes_perturbed), :]
+            gene_ids_df = cos_sims_df.loc[
+                np.isin(cos_sims_df["Ensembl_ID"], genes_perturbed), :
+            ]
     else:
         logger.error(
-                        "aggregate_data is for perturbation of single gene or single group of genes. genes_to_perturb should be formatted as list."
-                    )
-        raise        
+            "aggregate_data is for perturbation of single gene or single group of genes. genes_to_perturb should be formatted as list."
+        )
+        raise
 
     if gene_ids_df.empty:
-        logger.error(
-                        "genes_to_perturb not found in data."
-                    )
+        logger.error("genes_to_perturb not found in data.")
         raise
-        
+
     tokens = gene_ids_df["Gene"]
     symbols = gene_ids_df["Gene_name"]
 
@@ -223,7 +229,7 @@ def isp_aggregate_grouped_perturb(cos_sims_df, dict_list, genes_perturbed):
         df["Cosine_sim"] = cos_shift_data
         df["Gene"] = symbol
         cos_sims_full_dfs.append(df)
-    
+
     return pd.concat(cos_sims_full_dfs)
 
 
@@ -1018,7 +1024,7 @@ class InSilicoPerturberStats:
             },
             index=[i for i in range(len(gene_list))],
         )
-        
+
         if self.mode == "goal_state_shift":
             cos_sims_df = isp_stats_to_goal_state(
                 cos_sims_df_initial,
@@ -1045,12 +1051,16 @@ class InSilicoPerturberStats:
                 cos_sims_df_initial, dict_list, self.combos, self.anchor_token
             )
 
-        elif self.mode == "aggregate_data":        
-            cos_sims_df = isp_aggregate_grouped_perturb(cos_sims_df_initial, dict_list, self.genes_perturbed)
+        elif self.mode == "aggregate_data":
+            cos_sims_df = isp_aggregate_grouped_perturb(
+                cos_sims_df_initial, dict_list, self.genes_perturbed
+            )
 
         elif self.mode == "aggregate_gene_shifts":
             if (self.genes_perturbed == "all") and (self.combos == 0):
-                tuple_types = [True if isinstance(genes, tuple) else False for genes in gene_list]
+                tuple_types = [
+                    True if isinstance(genes, tuple) else False for genes in gene_list
+                ]
                 if all(tuple_types):
                     token_dtype = "tuple"
                 elif not any(tuple_types):
@@ -1059,13 +1069,13 @@ class InSilicoPerturberStats:
                     token_dtype = "mix"
             else:
                 token_dtype = "mix"
-            
+
             cos_sims_df = isp_aggregate_gene_shifts(
                 cos_sims_df_initial,
                 dict_list,
                 self.gene_token_id_dict,
                 self.gene_id_name_dict,
-                token_dtype
+                token_dtype,
             )
 
         # save perturbation stats to output_path

geneformer/mtl/collators.py

@@ -1,4 +1,4 @@
-#imports
+# imports
 import torch
 
 from ..collator_for_classification import DataCollatorForGeneClassification
@@ -7,6 +7,7 @@ from ..collator_for_classification import DataCollatorForGeneClassification
 Geneformer collator for multi-task cell classification.
 """
 
+
 class DataCollatorForMultitaskCellClassification(DataCollatorForGeneClassification):
     class_type = "cell"
 
@@ -47,7 +48,10 @@ class DataCollatorForMultitaskCellClassification(DataCollatorForGeneClassificati
             batch["labels"] = labels
         else:
             # If no labels are present, create empty labels for all tasks
-            batch["labels"] = {task: torch.tensor([], dtype=torch.long) for task in features[0]["input_ids"].keys()}
+            batch["labels"] = {
+                task: torch.tensor([], dtype=torch.long)
+                for task in features[0]["input_ids"].keys()
+            }
 
         return batch
 
@@ -59,8 +63,11 @@ class DataCollatorForMultitaskCellClassification(DataCollatorForGeneClassificati
                 batch[k] = v.clone().detach()
             elif isinstance(v, dict):
                 # Assuming nested structure needs conversion
-                batch[k] = {task: torch.tensor(labels, dtype=torch.int64) for task, labels in v.items()}
+                batch[k] = {
+                    task: torch.tensor(labels, dtype=torch.int64)
+                    for task, labels in v.items()
+                }
             else:
                 batch[k] = torch.tensor(v, dtype=torch.int64)
 
-        return batch
+        return batch

geneformer/mtl/data.py

@@ -1,6 +1,8 @@
-from .imports import *
 import os
+
 from .collators import DataCollatorForMultitaskCellClassification
+from .imports import *
+
 
 def load_and_preprocess_data(dataset_path, config, is_test=False, dataset_type=""):
     try:
@@ -14,7 +16,9 @@ def load_and_preprocess_data(dataset_path, config, is_test=False, dataset_type="
             available_columns = set(dataset.column_names)
             for column in task_to_column.values():
                 if column not in available_columns:
-                    raise KeyError(f"Column {column} not found in the dataset. Available columns: {list(available_columns)}")
+                    raise KeyError(
+                        f"Column {column} not found in the dataset. Available columns: {list(available_columns)}"
+                    )
 
         label_mappings = {}
         task_label_mappings = {}
@@ -25,13 +29,17 @@ def load_and_preprocess_data(dataset_path, config, is_test=False, dataset_type="
         if not is_test:
             for task, column in task_to_column.items():
                 unique_values = sorted(set(dataset[column]))  # Ensure consistency
-                label_mappings[column] = {label: idx for idx, label in enumerate(unique_values)}
+                label_mappings[column] = {
+                    label: idx for idx, label in enumerate(unique_values)
+                }
                 task_label_mappings[task] = label_mappings[column]
                 num_labels_list.append(len(unique_values))
 
             # Print the mappings for each task with dataset type prefix
             for task, mapping in task_label_mappings.items():
-                print(f"{dataset_type.capitalize()} mapping for {task}: {mapping}")  # sanity check, for train/validation splits
+                print(
+                    f"{dataset_type.capitalize()} mapping for {task}: {mapping}"
+                )  # sanity check, for train/validation splits
 
             # Save the task label mappings as a pickle file
             with open(f"{config['results_dir']}/task_label_mappings.pkl", "wb") as f:
@@ -40,24 +48,26 @@ def load_and_preprocess_data(dataset_path, config, is_test=False, dataset_type="
             # Load task label mappings from pickle file for test data
             with open(f"{config['results_dir']}/task_label_mappings.pkl", "rb") as f:
                 task_label_mappings = pickle.load(f)
-            
+
             # Infer num_labels_list from task_label_mappings
             for task, mapping in task_label_mappings.items():
                 num_labels_list.append(len(mapping))
 
         # Store unique cell IDs in a separate dictionary
         for idx, record in enumerate(dataset):
-            cell_id = record.get('unique_cell_id', idx)
+            cell_id = record.get("unique_cell_id", idx)
             cell_id_mapping[idx] = cell_id
 
         # Transform records to the desired format
         transformed_dataset = []
         for idx, record in enumerate(dataset):
             transformed_record = {}
-            transformed_record['input_ids'] = torch.tensor(record['input_ids'], dtype=torch.long)
-            
+            transformed_record["input_ids"] = torch.tensor(
+                record["input_ids"], dtype=torch.long
+            )
+
             # Use index-based cell ID for internal tracking
-            transformed_record['cell_id'] = idx
+            transformed_record["cell_id"] = idx
 
             if not is_test:
                 # Prepare labels
@@ -66,11 +76,11 @@ def load_and_preprocess_data(dataset_path, config, is_test=False, dataset_type="
                     label_value = record[column]
                     label_index = task_label_mappings[task][label_value]
                     label_dict[task] = label_index
-                transformed_record['label'] = label_dict
+                transformed_record["label"] = label_dict
             else:
                 # Create dummy labels for test data
                 label_dict = {task: -1 for task in config["task_names"]}
-                transformed_record['label'] = label_dict
+                transformed_record["label"] = label_dict
 
             transformed_dataset.append(transformed_record)
 
@@ -81,36 +91,60 @@ def load_and_preprocess_data(dataset_path, config, is_test=False, dataset_type="
         print(f"An error occurred while loading or preprocessing data: {e}")
         return None, None, None
 
+
 def preload_and_process_data(config):
     # Load and preprocess data once
-    train_dataset, train_cell_id_mapping, num_labels_list = load_and_preprocess_data(config["train_path"], config, dataset_type="train")
-    val_dataset, val_cell_id_mapping, _ = load_and_preprocess_data(config["val_path"], config, dataset_type="validation")
-    return train_dataset, train_cell_id_mapping, val_dataset, val_cell_id_mapping, num_labels_list
+    train_dataset, train_cell_id_mapping, num_labels_list = load_and_preprocess_data(
+        config["train_path"], config, dataset_type="train"
+    )
+    val_dataset, val_cell_id_mapping, _ = load_and_preprocess_data(
+        config["val_path"], config, dataset_type="validation"
+    )
+    return (
+        train_dataset,
+        train_cell_id_mapping,
+        val_dataset,
+        val_cell_id_mapping,
+        num_labels_list,
+    )
+
 
 def get_data_loader(preprocessed_dataset, batch_size):
-    nproc = os.cpu_count() ### I/O operations
-    
+    nproc = os.cpu_count()  ### I/O operations
+
     data_collator = DataCollatorForMultitaskCellClassification()
-    
-    loader = DataLoader(preprocessed_dataset, batch_size=batch_size, shuffle=True,
-                        collate_fn=data_collator, num_workers=nproc, pin_memory=True)
+
+    loader = DataLoader(
+        preprocessed_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        collate_fn=data_collator,
+        num_workers=nproc,
+        pin_memory=True,
+    )
     return loader
+
+
 def preload_data(config):
     # Preprocessing the data before the Optuna trials start
     train_loader = get_data_loader("train", config)
     val_loader = get_data_loader("val", config)
     return train_loader, val_loader
 
+
 def load_and_preprocess_test_data(config):
     """
     Load and preprocess test data, treating it as unlabeled.
     """
     return load_and_preprocess_data(config["test_path"], config, is_test=True)
 
+
 def prepare_test_loader(config):
     """
     Prepare DataLoader for the test dataset.
     """
-    test_dataset, cell_id_mapping, num_labels_list = load_and_preprocess_test_data(config)
-    test_loader = get_data_loader(test_dataset, config['batch_size'])
-    return test_loader, cell_id_mapping, num_labels_list
+    test_dataset, cell_id_mapping, num_labels_list = load_and_preprocess_test_data(
+        config
+    )
+    test_loader = get_data_loader(test_dataset, config["batch_size"])
+    return test_loader, cell_id_mapping, num_labels_list

geneformer/mtl/eval_utils.py

@@ -1,29 +1,33 @@
-from .imports import *
 import pandas as pd
+
 from .data import prepare_test_loader
+from .imports import *
 from .model import GeneformerMultiTask
 
+
 def evaluate_test_dataset(model, device, test_loader, cell_id_mapping, config):
     task_pred_labels = {task_name: [] for task_name in config["task_names"]}
     task_pred_probs = {task_name: [] for task_name in config["task_names"]}
     cell_ids = []
 
-    # Load task label mappings from pickle file
-    with open(f"{config['results_dir']}/task_label_mappings.pkl", "rb") as f:
-        task_label_mappings = pickle.load(f)
+    # # Load task label mappings from pickle file
+    # with open(f"{config['results_dir']}/task_label_mappings.pkl", "rb") as f:
+    #     task_label_mappings = pickle.load(f)
 
     model.eval()
     with torch.no_grad():
         for batch in test_loader:
-            input_ids = batch['input_ids'].to(device)
-            attention_mask = batch['attention_mask'].to(device)
+            input_ids = batch["input_ids"].to(device)
+            attention_mask = batch["attention_mask"].to(device)
             _, logits, _ = model(input_ids, attention_mask)
-            for sample_idx in range(len(batch['input_ids'])):
-                cell_id = cell_id_mapping[batch['cell_id'][sample_idx].item()]
+            for sample_idx in range(len(batch["input_ids"])):
+                cell_id = cell_id_mapping[batch["cell_id"][sample_idx].item()]
                 cell_ids.append(cell_id)
                 for i, task_name in enumerate(config["task_names"]):
                     pred_label = torch.argmax(logits[i][sample_idx], dim=-1).item()
-                    pred_prob = torch.softmax(logits[i][sample_idx], dim=-1).cpu().numpy()
+                    pred_prob = (
+                        torch.softmax(logits[i][sample_idx], dim=-1).cpu().numpy()
+                    )
                     task_pred_labels[task_name].append(pred_label)
                     task_pred_probs[task_name].append(pred_prob)
 
@@ -31,19 +35,22 @@ def evaluate_test_dataset(model, device, test_loader, cell_id_mapping, config):
     test_results_dir = config["results_dir"]
     os.makedirs(test_results_dir, exist_ok=True)
     test_preds_file = os.path.join(test_results_dir, "test_preds.csv")
-    
+
     rows = []
     for sample_idx in range(len(cell_ids)):
-        row = {'Cell ID': cell_ids[sample_idx]}
+        row = {"Cell ID": cell_ids[sample_idx]}
         for task_name in config["task_names"]:
-            row[f'{task_name} Prediction'] = task_pred_labels[task_name][sample_idx]
-            row[f'{task_name} Probabilities'] = ','.join(map(str, task_pred_probs[task_name][sample_idx]))
+            row[f"{task_name} Prediction"] = task_pred_labels[task_name][sample_idx]
+            row[f"{task_name} Probabilities"] = ",".join(
+                map(str, task_pred_probs[task_name][sample_idx])
+            )
         rows.append(row)
-    
+
     df = pd.DataFrame(rows)
     df.to_csv(test_preds_file, index=False)
     print(f"Test predictions saved to {test_preds_file}")
 
+
 def load_and_evaluate_test_model(config):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     test_loader, cell_id_mapping, num_labels_list = prepare_test_loader(config)
@@ -51,7 +58,7 @@ def load_and_evaluate_test_model(config):
     hyperparams_path = os.path.join(model_directory, "hyperparameters.json")
 
     # Load the saved best hyperparameters
-    with open(hyperparams_path, 'r') as f:
+    with open(hyperparams_path, "r") as f:
         best_hyperparams = json.load(f)
 
     # Extract the task weights if present, otherwise set to None
@@ -72,7 +79,7 @@ def load_and_evaluate_test_model(config):
         num_labels_list,
         dropout_rate=best_hyperparams["dropout_rate"],
         use_task_weights=config["use_task_weights"],
-        task_weights=normalized_task_weights
+        task_weights=normalized_task_weights,
     )
     best_model.load_state_dict(torch.load(best_model_path))
     best_model.to(device)

geneformer/mtl/imports.py

@@ -1,46 +1,43 @@
-import numpy as np
+import functools
+import gc
+import json
+import os
 import pickle
-import torch
-import torch.nn as nn
-import torch.optim as optim
-import torch.nn.functional as F
-from torch.utils.data import DataLoader
-
-from itertools import chain
+import sys
 import warnings
 from enum import Enum
+from itertools import chain
 from typing import Dict, List, Optional, Union
-import sys
-import os
-import json
-import gc
-import functools
-import pandas as pd
-
-from sklearn.metrics import f1_score, accuracy_score, roc_auc_score, roc_curve
-from sklearn.preprocessing import LabelEncoder
-from sklearn.model_selection import train_test_split
 
+import numpy as np
 import optuna
-
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from datasets import load_from_disk
+from sklearn.metrics import accuracy_score, f1_score, roc_auc_score, roc_curve
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import LabelEncoder
+from torch.utils.data import DataLoader
 from transformers import (
+    AdamW,
+    BatchEncoding,
     BertConfig,
     BertModel,
-    AdamW,
-    get_linear_schedule_with_warmup,
-    get_cosine_schedule_with_warmup,
     DataCollatorForTokenClassification,
     SpecialTokensMixin,
-    BatchEncoding,
+    get_cosine_schedule_with_warmup,
+    get_linear_schedule_with_warmup,
     get_scheduler,
 )
 from transformers.utils import logging, to_py_obj
 
-from datasets import load_from_disk
+from .collators import DataCollatorForMultitaskCellClassification
 
 # local modules
-from .data import preload_and_process_data, get_data_loader
+from .data import get_data_loader, preload_and_process_data
 from .model import GeneformerMultiTask
-from .utils import save_model
 from .optuna_utils import create_optuna_study
-from .collators import DataCollatorForMultitaskCellClassification
+from .utils import save_model

geneformer/mtl/model.py

@@ -1,13 +1,17 @@
-from transformers import BertModel, BertConfig
 import torch
 import torch.nn as nn
+from transformers import BertConfig, BertModel
+
 
 class AttentionPool(nn.Module):
     """Attention-based pooling layer."""
+
     def __init__(self, hidden_size):
         super(AttentionPool, self).__init__()
         self.attention_weights = nn.Parameter(torch.randn(hidden_size, 1))
-        nn.init.xavier_uniform_(self.attention_weights)  # https://pytorch.org/docs/stable/nn.init.html
+        nn.init.xavier_uniform_(
+            self.attention_weights
+        )  # https://pytorch.org/docs/stable/nn.init.html
 
     def forward(self, hidden_states):
         attention_scores = torch.matmul(hidden_states, self.attention_weights)
@@ -15,8 +19,18 @@ class AttentionPool(nn.Module):
         pooled_output = torch.sum(hidden_states * attention_scores, dim=1)
         return pooled_output
 
+
 class GeneformerMultiTask(nn.Module):
-    def __init__(self, pretrained_path, num_labels_list, dropout_rate=0.1, use_task_weights=False, task_weights=None, max_layers_to_freeze=0, use_attention_pooling=False):
+    def __init__(
+        self,
+        pretrained_path,
+        num_labels_list,
+        dropout_rate=0.1,
+        use_task_weights=False,
+        task_weights=None,
+        max_layers_to_freeze=0,
+        use_attention_pooling=False,
+    ):
         super(GeneformerMultiTask, self).__init__()
         self.config = BertConfig.from_pretrained(pretrained_path)
         self.bert = BertModel(self.config)
@@ -25,20 +39,31 @@ class GeneformerMultiTask(nn.Module):
         self.dropout = nn.Dropout(dropout_rate)
         self.use_attention_pooling = use_attention_pooling
 
-        if use_task_weights and (task_weights is None or len(task_weights) != len(num_labels_list)):
-            raise ValueError("Task weights must be defined and match the number of tasks when 'use_task_weights' is True.")
-        self.task_weights = task_weights if use_task_weights else [1.0] * len(num_labels_list)
+        if use_task_weights and (
+            task_weights is None or len(task_weights) != len(num_labels_list)
+        ):
+            raise ValueError(
+                "Task weights must be defined and match the number of tasks when 'use_task_weights' is True."
+            )
+        self.task_weights = (
+            task_weights if use_task_weights else [1.0] * len(num_labels_list)
+        )
 
         # Freeze the specified initial layers
         for layer in self.bert.encoder.layer[:max_layers_to_freeze]:
             for param in layer.parameters():
                 param.requires_grad = False
 
-        self.attention_pool = AttentionPool(self.config.hidden_size) if use_attention_pooling else None
+        self.attention_pool = (
+            AttentionPool(self.config.hidden_size) if use_attention_pooling else None
+        )
 
-        self.classification_heads = nn.ModuleList([
-            nn.Linear(self.config.hidden_size, num_labels) for num_labels in num_labels_list
-        ])
+        self.classification_heads = nn.ModuleList(
+            [
+                nn.Linear(self.config.hidden_size, num_labels)
+                for num_labels in num_labels_list
+            ]
+        )
         # initialization of the classification heads: https://pytorch.org/docs/stable/nn.init.html
         for head in self.classification_heads:
             nn.init.xavier_uniform_(head.weight)
@@ -53,7 +78,11 @@ class GeneformerMultiTask(nn.Module):
         sequence_output = outputs.last_hidden_state
 
         try:
-            pooled_output = self.attention_pool(sequence_output) if self.use_attention_pooling else sequence_output[:, 0, :]
+            pooled_output = (
+                self.attention_pool(sequence_output)
+                if self.use_attention_pooling
+                else sequence_output[:, 0, :]
+            )
             pooled_output = self.dropout(pooled_output)
         except Exception as e:
             raise RuntimeError(f"Error during pooling and dropout: {e}")
@@ -62,23 +91,31 @@ class GeneformerMultiTask(nn.Module):
         logits = []
         losses = []
 
-        for task_id, (head, num_labels) in enumerate(zip(self.classification_heads, self.num_labels_list)):
+        for task_id, (head, num_labels) in enumerate(
+            zip(self.classification_heads, self.num_labels_list)
+        ):
             try:
                 task_logits = head(pooled_output)
             except Exception as e:
-                raise RuntimeError(f"Error during forward pass of classification head {task_id}: {e}")
+                raise RuntimeError(
+                    f"Error during forward pass of classification head {task_id}: {e}"
+                )
 
             logits.append(task_logits)
 
             if labels is not None:
                 try:
                     loss_fct = nn.CrossEntropyLoss()
-                    task_loss = loss_fct(task_logits.view(-1, num_labels), labels[task_id].view(-1))
+                    task_loss = loss_fct(
+                        task_logits.view(-1, num_labels), labels[task_id].view(-1)
+                    )
                     if self.use_task_weights:
                         task_loss *= self.task_weights[task_id]
                     total_loss += task_loss
                     losses.append(task_loss.item())
                 except Exception as e:
-                    raise RuntimeError(f"Error during loss computation for task {task_id}: {e}")
+                    raise RuntimeError(
+                        f"Error during loss computation for task {task_id}: {e}"
+                    )
 
         return total_loss, logits, losses if labels is not None else logits

geneformer/mtl/optuna_utils.py

@@ -1,21 +1,27 @@
 import optuna
 from optuna.integration import TensorBoardCallback
 
+
 def save_trial_callback(study, trial, trials_result_path):
     with open(trials_result_path, "a") as f:
-        f.write(f"Trial {trial.number}: Value (F1 Macro): {trial.value}, Params: {trial.params}\n")
+        f.write(
+            f"Trial {trial.number}: Value (F1 Macro): {trial.value}, Params: {trial.params}\n"
+        )
+
 
 def create_optuna_study(objective, n_trials, trials_result_path, tensorboard_log_dir):
     study = optuna.create_study(direction="maximize")
-    
+
     # init TensorBoard callback
-    tensorboard_callback = TensorBoardCallback(dirname=tensorboard_log_dir, metric_name="F1 Macro")
-    
+    tensorboard_callback = TensorBoardCallback(
+        dirname=tensorboard_log_dir, metric_name="F1 Macro"
+    )
+
     # callback and TensorBoard callback
     callbacks = [
         lambda study, trial: save_trial_callback(study, trial, trials_result_path),
-        tensorboard_callback
+        tensorboard_callback,
     ]
-    
+
     study.optimize(objective, n_trials=n_trials, callbacks=callbacks)
-    return study
+    return study

geneformer/mtl/train.py

@@ -1,14 +1,15 @@
-from .imports import *
-from .data import preload_and_process_data, get_data_loader
-from .model import GeneformerMultiTask
-from .utils import calculate_task_specific_metrics
-from torch.utils.tensorboard import SummaryWriter
-import pandas as pd
 import os
-from tqdm import tqdm
 import random
+
 import numpy as np
+import pandas as pd
 import torch
+from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+from .imports import *
+from .model import GeneformerMultiTask
+from .utils import calculate_task_specific_metrics
 
 
 def set_seed(seed):
@@ -19,13 +20,18 @@ def set_seed(seed):
     torch.backends.cudnn.deterministic = True
     torch.backends.cudnn.benchmark = False
 
+
 def initialize_wandb(config):
     if config.get("use_wandb", False):
         import wandb
+
         wandb.init(project=config["wandb_project"], config=config)
         print("Weights & Biases (wandb) initialized and will be used for logging.")
     else:
-        print("Weights & Biases (wandb) is not enabled. Logging will use other methods.")
+        print(
+            "Weights & Biases (wandb) is not enabled. Logging will use other methods."
+        )
+
 
 def create_model(config, num_labels_list, device):
     model = GeneformerMultiTask(
@@ -35,31 +41,48 @@ def create_model(config, num_labels_list, device):
         use_task_weights=config["use_task_weights"],
         task_weights=config["task_weights"],
         max_layers_to_freeze=config["max_layers_to_freeze"],
-        use_attention_pooling=config["use_attention_pooling"]
+        use_attention_pooling=config["use_attention_pooling"],
     )
     if config["use_data_parallel"]:
         model = nn.DataParallel(model)
     return model.to(device)
 
+
 def setup_optimizer_and_scheduler(model, config, total_steps):
-    optimizer = AdamW(model.parameters(), lr=config["learning_rate"], weight_decay=config["weight_decay"])
+    optimizer = AdamW(
+        model.parameters(),
+        lr=config["learning_rate"],
+        weight_decay=config["weight_decay"],
+    )
     warmup_steps = int(config["warmup_ratio"] * total_steps)
 
     if config["lr_scheduler_type"] == "linear":
-        scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps)
+        scheduler = get_linear_schedule_with_warmup(
+            optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps
+        )
     elif config["lr_scheduler_type"] == "cosine":
-        scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps, num_cycles=0.5)
+        scheduler = get_cosine_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=warmup_steps,
+            num_training_steps=total_steps,
+            num_cycles=0.5,
+        )
 
     return optimizer, scheduler
 
-def train_epoch(model, train_loader, optimizer, scheduler, device, config, writer, epoch):
+
+def train_epoch(
+    model, train_loader, optimizer, scheduler, device, config, writer, epoch
+):
     model.train()
     progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{config['epochs']}")
     for batch_idx, batch in enumerate(progress_bar):
         optimizer.zero_grad()
-        input_ids = batch['input_ids'].to(device)
-        attention_mask = batch['attention_mask'].to(device)
-        labels = [batch['labels'][task_name].to(device) for task_name in config["task_names"]]
+        input_ids = batch["input_ids"].to(device)
+        attention_mask = batch["attention_mask"].to(device)
+        labels = [
+            batch["labels"][task_name].to(device) for task_name in config["task_names"]
+        ]
 
         loss, _, _ = model(input_ids, attention_mask, labels)
         loss.backward()
@@ -70,15 +93,20 @@ def train_epoch(model, train_loader, optimizer, scheduler, device, config, write
         optimizer.step()
         scheduler.step()
 
-        writer.add_scalar('Training Loss', loss.item(), epoch * len(train_loader) + batch_idx)
+        writer.add_scalar(
+            "Training Loss", loss.item(), epoch * len(train_loader) + batch_idx
+        )
         if config.get("use_wandb", False):
-            wandb.log({'Training Loss': loss.item()})
+            import wandb
+
+            wandb.log({"Training Loss": loss.item()})
 
         # Update progress bar
-        progress_bar.set_postfix({'loss': f"{loss.item():.4f}"})
+        progress_bar.set_postfix({"loss": f"{loss.item():.4f}"})
 
     return loss.item()  # Return the last batch loss
 
+
 def validate_model(model, val_loader, device, config):
     model.eval()
     val_loss = 0.0
@@ -88,17 +116,22 @@ def validate_model(model, val_loader, device, config):
 
     with torch.no_grad():
         for batch in val_loader:
-            input_ids = batch['input_ids'].to(device)
-            attention_mask = batch['attention_mask'].to(device)
-            labels = [batch['labels'][task_name].to(device) for task_name in config["task_names"]]
+            input_ids = batch["input_ids"].to(device)
+            attention_mask = batch["attention_mask"].to(device)
+            labels = [
+                batch["labels"][task_name].to(device)
+                for task_name in config["task_names"]
+            ]
             loss, logits, _ = model(input_ids, attention_mask, labels)
             val_loss += loss.item()
 
-            for sample_idx in range(len(batch['input_ids'])):
+            for sample_idx in range(len(batch["input_ids"])):
                 for i, task_name in enumerate(config["task_names"]):
-                    true_label = batch['labels'][task_name][sample_idx].item()
+                    true_label = batch["labels"][task_name][sample_idx].item()
                     pred_label = torch.argmax(logits[i][sample_idx], dim=-1).item()
-                    pred_prob = torch.softmax(logits[i][sample_idx], dim=-1).cpu().numpy()
+                    pred_prob = (
+                        torch.softmax(logits[i][sample_idx], dim=-1).cpu().numpy()
+                    )
                     task_true_labels[task_name].append(true_label)
                     task_pred_labels[task_name].append(pred_label)
                     task_pred_probs[task_name].append(pred_prob)
@@ -106,44 +139,70 @@ def validate_model(model, val_loader, device, config):
     val_loss /= len(val_loader)
     return val_loss, task_true_labels, task_pred_labels, task_pred_probs
 
+
 def log_metrics(task_metrics, val_loss, config, writer, epochs):
     for task_name, metrics in task_metrics.items():
-        print(f"{task_name} - Validation F1 Macro: {metrics['f1']:.4f}, Validation Accuracy: {metrics['accuracy']:.4f}")
+        print(
+            f"{task_name} - Validation F1 Macro: {metrics['f1']:.4f}, Validation Accuracy: {metrics['accuracy']:.4f}"
+        )
         if config.get("use_wandb", False):
             import wandb
-            wandb.log({
-                f'{task_name} Validation F1 Macro': metrics['f1'],
-                f'{task_name} Validation Accuracy': metrics['accuracy']
-            })
 
-    writer.add_scalar('Validation Loss', val_loss, epochs)
-    for task_name, metrics in task_metrics.items():
-        writer.add_scalar(f'{task_name} - Validation F1 Macro', metrics['f1'], epochs)
-        writer.add_scalar(f'{task_name} - Validation Accuracy', metrics['accuracy'], epochs)
+            wandb.log(
+                {
+                    f"{task_name} Validation F1 Macro": metrics["f1"],
+                    f"{task_name} Validation Accuracy": metrics["accuracy"],
+                }
+            )
 
-def save_validation_predictions(val_cell_id_mapping, task_true_labels, task_pred_labels, task_pred_probs, config, trial_number=None):
+    writer.add_scalar("Validation Loss", val_loss, epochs)
+    for task_name, metrics in task_metrics.items():
+        writer.add_scalar(f"{task_name} - Validation F1 Macro", metrics["f1"], epochs)
+        writer.add_scalar(
+            f"{task_name} - Validation Accuracy", metrics["accuracy"], epochs
+        )
+
+
+def save_validation_predictions(
+    val_cell_id_mapping,
+    task_true_labels,
+    task_pred_labels,
+    task_pred_probs,
+    config,
+    trial_number=None,
+):
     if trial_number is not None:
         trial_results_dir = os.path.join(config["results_dir"], f"trial_{trial_number}")
         os.makedirs(trial_results_dir, exist_ok=True)
         val_preds_file = os.path.join(trial_results_dir, "val_preds.csv")
     else:
         val_preds_file = os.path.join(config["results_dir"], "manual_run_val_preds.csv")
-    
+
     rows = []
     for sample_idx in range(len(val_cell_id_mapping)):
-        row = {'Cell ID': val_cell_id_mapping[sample_idx]}
+        row = {"Cell ID": val_cell_id_mapping[sample_idx]}
         for task_name in config["task_names"]:
-            row[f'{task_name} True'] = task_true_labels[task_name][sample_idx]
-            row[f'{task_name} Pred'] = task_pred_labels[task_name][sample_idx]
-            row[f'{task_name} Probabilities'] = ','.join(map(str, task_pred_probs[task_name][sample_idx]))
+            row[f"{task_name} True"] = task_true_labels[task_name][sample_idx]
+            row[f"{task_name} Pred"] = task_pred_labels[task_name][sample_idx]
+            row[f"{task_name} Probabilities"] = ",".join(
+                map(str, task_pred_probs[task_name][sample_idx])
+            )
         rows.append(row)
-    
+
     df = pd.DataFrame(rows)
     df.to_csv(val_preds_file, index=False)
     print(f"Validation predictions saved to {val_preds_file}")
 
 
-def train_model(config, device, train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list):
+def train_model(
+    config,
+    device,
+    train_loader,
+    val_loader,
+    train_cell_id_mapping,
+    val_cell_id_mapping,
+    num_labels_list,
+):
     set_seed(config["seed"])
     initialize_wandb(config)
 
@@ -156,46 +215,97 @@ def train_model(config, device, train_loader, val_loader, train_cell_id_mapping,
 
     epoch_progress = tqdm(range(config["epochs"]), desc="Training Progress")
     for epoch in epoch_progress:
-        last_loss = train_epoch(model, train_loader, optimizer, scheduler, device, config, writer, epoch)
-        epoch_progress.set_postfix({'last_loss': f"{last_loss:.4f}"})
+        last_loss = train_epoch(
+            model, train_loader, optimizer, scheduler, device, config, writer, epoch
+        )
+        epoch_progress.set_postfix({"last_loss": f"{last_loss:.4f}"})
 
-    val_loss, task_true_labels, task_pred_labels, task_pred_probs = validate_model(model, val_loader, device, config)
+    val_loss, task_true_labels, task_pred_labels, task_pred_probs = validate_model(
+        model, val_loader, device, config
+    )
     task_metrics = calculate_task_specific_metrics(task_true_labels, task_pred_labels)
-    
+
     log_metrics(task_metrics, val_loss, config, writer, config["epochs"])
     writer.close()
 
-    save_validation_predictions(val_cell_id_mapping, task_true_labels, task_pred_labels, task_pred_probs, config)
+    save_validation_predictions(
+        val_cell_id_mapping, task_true_labels, task_pred_labels, task_pred_probs, config
+    )
 
     if config.get("use_wandb", False):
         import wandb
+
         wandb.finish()
 
     print(f"\nFinal Validation Loss: {val_loss:.4f}")
     return val_loss, model  # Return both the validation loss and the trained model
 
-def objective(trial, train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list, config, device):
+
+def objective(
+    trial,
+    train_loader,
+    val_loader,
+    train_cell_id_mapping,
+    val_cell_id_mapping,
+    num_labels_list,
+    config,
+    device,
+):
     set_seed(config["seed"])  # Set the seed before each trial
     initialize_wandb(config)
 
     # Hyperparameters
-    config["learning_rate"] = trial.suggest_float("learning_rate", config["hyperparameters"]["learning_rate"]["low"], config["hyperparameters"]["learning_rate"]["high"], log=config["hyperparameters"]["learning_rate"]["log"])
-    config["warmup_ratio"] = trial.suggest_float("warmup_ratio", config["hyperparameters"]["warmup_ratio"]["low"], config["hyperparameters"]["warmup_ratio"]["high"])
-    config["weight_decay"] = trial.suggest_float("weight_decay", config["hyperparameters"]["weight_decay"]["low"], config["hyperparameters"]["weight_decay"]["high"])
-    config["dropout_rate"] = trial.suggest_float("dropout_rate", config["hyperparameters"]["dropout_rate"]["low"], config["hyperparameters"]["dropout_rate"]["high"])
-    config["lr_scheduler_type"] = trial.suggest_categorical("lr_scheduler_type", config["hyperparameters"]["lr_scheduler_type"]["choices"])
-    config["use_attention_pooling"] = trial.suggest_categorical("use_attention_pooling", [True, False])
+    config["learning_rate"] = trial.suggest_float(
+        "learning_rate",
+        config["hyperparameters"]["learning_rate"]["low"],
+        config["hyperparameters"]["learning_rate"]["high"],
+        log=config["hyperparameters"]["learning_rate"]["log"],
+    )
+    config["warmup_ratio"] = trial.suggest_float(
+        "warmup_ratio",
+        config["hyperparameters"]["warmup_ratio"]["low"],
+        config["hyperparameters"]["warmup_ratio"]["high"],
+    )
+    config["weight_decay"] = trial.suggest_float(
+        "weight_decay",
+        config["hyperparameters"]["weight_decay"]["low"],
+        config["hyperparameters"]["weight_decay"]["high"],
+    )
+    config["dropout_rate"] = trial.suggest_float(
+        "dropout_rate",
+        config["hyperparameters"]["dropout_rate"]["low"],
+        config["hyperparameters"]["dropout_rate"]["high"],
+    )
+    config["lr_scheduler_type"] = trial.suggest_categorical(
+        "lr_scheduler_type", config["hyperparameters"]["lr_scheduler_type"]["choices"]
+    )
+    config["use_attention_pooling"] = trial.suggest_categorical(
+        "use_attention_pooling", [True, False]
+    )
 
     if config["use_task_weights"]:
-        config["task_weights"] = [trial.suggest_float(f"task_weight_{i}", config["hyperparameters"]["task_weights"]["low"], config["hyperparameters"]["task_weights"]["high"]) for i in range(len(num_labels_list))]
+        config["task_weights"] = [
+            trial.suggest_float(
+                f"task_weight_{i}",
+                config["hyperparameters"]["task_weights"]["low"],
+                config["hyperparameters"]["task_weights"]["high"],
+            )
+            for i in range(len(num_labels_list))
+        ]
         weight_sum = sum(config["task_weights"])
-        config["task_weights"] = [weight / weight_sum for weight in config["task_weights"]]
+        config["task_weights"] = [
+            weight / weight_sum for weight in config["task_weights"]
+        ]
     else:
         config["task_weights"] = None
 
     # Fix for max_layers_to_freeze
     if isinstance(config["max_layers_to_freeze"], dict):
-        config["max_layers_to_freeze"] = trial.suggest_int("max_layers_to_freeze", config["max_layers_to_freeze"]["min"], config["max_layers_to_freeze"]["max"])
+        config["max_layers_to_freeze"] = trial.suggest_int(
+            "max_layers_to_freeze",
+            config["max_layers_to_freeze"]["min"],
+            config["max_layers_to_freeze"]["max"],
+        )
     elif isinstance(config["max_layers_to_freeze"], int):
         # If it's already an int, we don't need to suggest it
         pass
@@ -210,15 +320,26 @@ def objective(trial, train_loader, val_loader, train_cell_id_mapping, val_cell_i
     writer = SummaryWriter(log_dir=log_dir)
 
     for epoch in range(config["epochs"]):
-        train_epoch(model, train_loader, optimizer, scheduler, device, config, writer, epoch)
+        train_epoch(
+            model, train_loader, optimizer, scheduler, device, config, writer, epoch
+        )
 
-    val_loss, task_true_labels, task_pred_labels, task_pred_probs = validate_model(model, val_loader, device, config)
+    val_loss, task_true_labels, task_pred_labels, task_pred_probs = validate_model(
+        model, val_loader, device, config
+    )
     task_metrics = calculate_task_specific_metrics(task_true_labels, task_pred_labels)
-    
+
     log_metrics(task_metrics, val_loss, config, writer, config["epochs"])
     writer.close()
 
-    save_validation_predictions(val_cell_id_mapping, task_true_labels, task_pred_labels, task_pred_probs, config, trial.number)
+    save_validation_predictions(
+        val_cell_id_mapping,
+        task_true_labels,
+        task_pred_labels,
+        task_pred_probs,
+        config,
+        trial.number,
+    )
 
     trial.set_user_attr("model_state_dict", model.state_dict())
     trial.set_user_attr("task_weights", config["task_weights"])
@@ -230,13 +351,35 @@ def objective(trial, train_loader, val_loader, train_cell_id_mapping, val_cell_i
 
     if config.get("use_wandb", False):
         import wandb
-        wandb.log({
-            "trial_number": trial.number,
-            "val_loss": val_loss,
-            **{f"{task_name}_f1": metrics['f1'] for task_name, metrics in task_metrics.items()},
-            **{f"{task_name}_accuracy": metrics['accuracy'] for task_name, metrics in task_metrics.items()},
-            **{k: v for k, v in config.items() if k in ["learning_rate", "warmup_ratio", "weight_decay", "dropout_rate", "lr_scheduler_type", "use_attention_pooling", "max_layers_to_freeze"]}
-        })
+
+        wandb.log(
+            {
+                "trial_number": trial.number,
+                "val_loss": val_loss,
+                **{
+                    f"{task_name}_f1": metrics["f1"]
+                    for task_name, metrics in task_metrics.items()
+                },
+                **{
+                    f"{task_name}_accuracy": metrics["accuracy"]
+                    for task_name, metrics in task_metrics.items()
+                },
+                **{
+                    k: v
+                    for k, v in config.items()
+                    if k
+                    in [
+                        "learning_rate",
+                        "warmup_ratio",
+                        "weight_decay",
+                        "dropout_rate",
+                        "lr_scheduler_type",
+                        "use_attention_pooling",
+                        "max_layers_to_freeze",
+                    ]
+                },
+            }
+        )
         wandb.finish()
 
-    return val_loss
+    return val_loss

geneformer/mtl/train_utils.py

@@ -1,9 +1,11 @@
+import random
+
+from .data import get_data_loader, preload_and_process_data
 from .imports import *
-from .data import preload_and_process_data, get_data_loader
-from .train import objective, train_model
 from .model import GeneformerMultiTask
+from .train import objective, train_model
 from .utils import save_model
-import random
+
 
 def set_seed(seed):
     random.seed(seed)
@@ -12,15 +14,22 @@ def set_seed(seed):
     torch.cuda.manual_seed_all(seed)
     torch.backends.cudnn.deterministic = True
     torch.backends.cudnn.benchmark = False
-    
+
+
 def run_manual_tuning(config):
     # Set seed for reproducibility
     set_seed(config["seed"])
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    train_dataset, train_cell_id_mapping, val_dataset, val_cell_id_mapping, num_labels_list = preload_and_process_data(config)
-    train_loader = get_data_loader(train_dataset, config['batch_size'])
-    val_loader = get_data_loader(val_dataset, config['batch_size'])
+    (
+        train_dataset,
+        train_cell_id_mapping,
+        val_dataset,
+        val_cell_id_mapping,
+        num_labels_list,
+    ) = preload_and_process_data(config)
+    train_loader = get_data_loader(train_dataset, config["batch_size"])
+    val_loader = get_data_loader(val_dataset, config["batch_size"])
 
     # Print the manual hyperparameters being used
     print("\nManual hyperparameters being used:")
@@ -33,12 +42,22 @@ def run_manual_tuning(config):
         config[key] = value
 
     # Train the model
-    val_loss, trained_model = train_model(config, device, train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list)
+    val_loss, trained_model = train_model(
+        config,
+        device,
+        train_loader,
+        val_loader,
+        train_cell_id_mapping,
+        val_cell_id_mapping,
+        num_labels_list,
+    )
 
     print(f"\nValidation loss with manual hyperparameters: {val_loss}")
 
     # Save the trained model
-    model_save_directory = os.path.join(config["model_save_path"], "GeneformerMultiTask")
+    model_save_directory = os.path.join(
+        config["model_save_path"], "GeneformerMultiTask"
+    )
     save_model(trained_model, model_save_directory)
 
     # Save the hyperparameters
@@ -48,26 +67,41 @@ def run_manual_tuning(config):
         "use_task_weights": config["use_task_weights"],
         "task_weights": config["task_weights"],
         "max_layers_to_freeze": config["max_layers_to_freeze"],
-        "use_attention_pooling": config["use_attention_pooling"]
+        "use_attention_pooling": config["use_attention_pooling"],
     }
     hyperparams_path = os.path.join(model_save_directory, "hyperparameters.json")
-    with open(hyperparams_path, 'w') as f:
+    with open(hyperparams_path, "w") as f:
         json.dump(hyperparams_to_save, f)
     print(f"Manual hyperparameters saved to {hyperparams_path}")
 
     return val_loss
 
+
 def run_optuna_study(config):
     # Set seed for reproducibility
     set_seed(config["seed"])
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    train_dataset, train_cell_id_mapping, val_dataset, val_cell_id_mapping, num_labels_list = preload_and_process_data(config)
-    train_loader = get_data_loader(train_dataset, config['batch_size'])
-    val_loader = get_data_loader(val_dataset, config['batch_size'])
+    (
+        train_dataset,
+        train_cell_id_mapping,
+        val_dataset,
+        val_cell_id_mapping,
+        num_labels_list,
+    ) = preload_and_process_data(config)
+    train_loader = get_data_loader(train_dataset, config["batch_size"])
+    val_loader = get_data_loader(val_dataset, config["batch_size"])
 
     if config["use_manual_hyperparameters"]:
-        train_model(config, device, train_loader, val_loader, train_cell_id_mapping, val_cell_id_mapping, num_labels_list)
+        train_model(
+            config,
+            device,
+            train_loader,
+            val_loader,
+            train_cell_id_mapping,
+            val_cell_id_mapping,
+            num_labels_list,
+        )
     else:
         objective_with_config_and_data = functools.partial(
             objective,
@@ -77,20 +111,17 @@ def run_optuna_study(config):
             val_cell_id_mapping=val_cell_id_mapping,
             num_labels_list=num_labels_list,
             config=config,
-            device=device
+            device=device,
         )
 
         study = optuna.create_study(
-            direction='minimize',  # Minimize validation loss
+            direction="minimize",  # Minimize validation loss
             study_name=config["study_name"],
-            #storage=config["storage"],
-            load_if_exists=True
+            # storage=config["storage"],
+            load_if_exists=True,
         )
 
-        study.optimize(
-            objective_with_config_and_data,
-            n_trials=config["n_trials"]
-        )
+        study.optimize(objective_with_config_and_data, n_trials=config["n_trials"])
 
         # After finding the best trial
         best_params = study.best_trial.params
@@ -103,24 +134,28 @@ def run_optuna_study(config):
             num_labels_list,
             dropout_rate=best_params["dropout_rate"],
             use_task_weights=config["use_task_weights"],
-            task_weights=best_task_weights
+            task_weights=best_task_weights,
         )
-        
+
         # Get the best model state dictionary
         best_model_state_dict = study.best_trial.user_attrs["model_state_dict"]
 
         # Remove the "module." prefix from the state dictionary keys if present
-        best_model_state_dict = {k.replace("module.", ""): v for k, v in best_model_state_dict.items()}
+        best_model_state_dict = {
+            k.replace("module.", ""): v for k, v in best_model_state_dict.items()
+        }
 
         # Load the modified state dictionary into the model, skipping unexpected keys
         best_model.load_state_dict(best_model_state_dict, strict=False)
 
-        model_save_directory = os.path.join(config["model_save_path"], "GeneformerMultiTask")
+        model_save_directory = os.path.join(
+            config["model_save_path"], "GeneformerMultiTask"
+        )
         save_model(best_model, model_save_directory)
 
         # Additionally, save the best hyperparameters and task weights
         hyperparams_path = os.path.join(model_save_directory, "hyperparameters.json")
-        
-        with open(hyperparams_path, 'w') as f:
+
+        with open(hyperparams_path, "w") as f:
             json.dump({**best_params, "task_weights": best_task_weights}, f)
-        print(f"Best hyperparameters and task weights saved to {hyperparams_path}")
+        print(f"Best hyperparameters and task weights saved to {hyperparams_path}")

geneformer/mtl/utils.py

@@ -1,44 +1,55 @@
-from .imports import *
-from sklearn.metrics import f1_score, accuracy_score
-from sklearn.preprocessing import LabelEncoder
-from transformers import BertModel, BertConfig, AutoConfig
 import os
 import shutil
 
+from sklearn.metrics import accuracy_score, f1_score
+from sklearn.preprocessing import LabelEncoder
+from transformers import AutoConfig, BertConfig, BertModel
+
+from .imports import *
+
+
 def save_model(model, model_save_directory):
     if not os.path.exists(model_save_directory):
         os.makedirs(model_save_directory)
-    
+
     # Get the state dict
     if isinstance(model, nn.DataParallel):
-        model_state_dict = model.module.state_dict()  # Use model.module to access the underlying model
+        model_state_dict = (
+            model.module.state_dict()
+        )  # Use model.module to access the underlying model
     else:
         model_state_dict = model.state_dict()
-    
+
     # Remove the "module." prefix from the keys if present
-    model_state_dict = {k.replace("module.", ""): v for k, v in model_state_dict.items()}
-    
+    model_state_dict = {
+        k.replace("module.", ""): v for k, v in model_state_dict.items()
+    }
+
     model_save_path = os.path.join(model_save_directory, "pytorch_model.bin")
     torch.save(model_state_dict, model_save_path)
-    
+
     # Save the model configuration
     if isinstance(model, nn.DataParallel):
-        model.module.config.to_json_file(os.path.join(model_save_directory, "config.json"))
+        model.module.config.to_json_file(
+            os.path.join(model_save_directory, "config.json")
+        )
     else:
         model.config.to_json_file(os.path.join(model_save_directory, "config.json"))
-    
+
     print(f"Model and configuration saved to {model_save_directory}")
 
+
 def calculate_task_specific_metrics(task_true_labels, task_pred_labels):
     task_metrics = {}
     for task_name in task_true_labels.keys():
         true_labels = task_true_labels[task_name]
         pred_labels = task_pred_labels[task_name]
-        f1 = f1_score(true_labels, pred_labels, average='macro')
+        f1 = f1_score(true_labels, pred_labels, average="macro")
         accuracy = accuracy_score(true_labels, pred_labels)
-        task_metrics[task_name] = {'f1': f1, 'accuracy': accuracy}
+        task_metrics[task_name] = {"f1": f1, "accuracy": accuracy}
     return task_metrics
 
+
 def calculate_combined_f1(combined_labels, combined_preds):
     # Initialize the LabelEncoder
     le = LabelEncoder()
@@ -57,10 +68,11 @@ def calculate_combined_f1(combined_labels, combined_preds):
     accuracy = accuracy_score(encoded_true_labels, encoded_pred_labels)
 
     # Calculate F1 Macro score
-    f1 = f1_score(encoded_true_labels, encoded_pred_labels, average='macro')
+    f1 = f1_score(encoded_true_labels, encoded_pred_labels, average="macro")
 
     return f1, accuracy
 
+
 def save_model_without_heads(original_model_save_directory):
     # Create a new directory for the model without heads
     new_model_save_directory = original_model_save_directory + "_No_Heads"
@@ -68,25 +80,36 @@ def save_model_without_heads(original_model_save_directory):
         os.makedirs(new_model_save_directory)
 
     # Load the model state dictionary
-    model_state_dict = torch.load(os.path.join(original_model_save_directory, "pytorch_model.bin"))
+    model_state_dict = torch.load(
+        os.path.join(original_model_save_directory, "pytorch_model.bin")
+    )
 
     # Initialize a new BERT model without the classification heads
-    config = BertConfig.from_pretrained(os.path.join(original_model_save_directory, "config.json"))
+    config = BertConfig.from_pretrained(
+        os.path.join(original_model_save_directory, "config.json")
+    )
     model_without_heads = BertModel(config)
-    
+
     # Filter the state dict to exclude classification heads
-    model_without_heads_state_dict = {k: v for k, v in model_state_dict.items() if not k.startswith("classification_heads")}
-    
+    model_without_heads_state_dict = {
+        k: v
+        for k, v in model_state_dict.items()
+        if not k.startswith("classification_heads")
+    }
+
     # Load the filtered state dict into the model
     model_without_heads.load_state_dict(model_without_heads_state_dict, strict=False)
-    
+
     # Save the model without heads
     model_save_path = os.path.join(new_model_save_directory, "pytorch_model.bin")
     torch.save(model_without_heads.state_dict(), model_save_path)
-    
+
     # Copy the configuration file
-    shutil.copy(os.path.join(original_model_save_directory, "config.json"), new_model_save_directory)
-    
+    shutil.copy(
+        os.path.join(original_model_save_directory, "config.json"),
+        new_model_save_directory,
+    )
+
     print(f"Model without classification heads saved to {new_model_save_directory}")
 
 
@@ -100,7 +123,7 @@ def get_layer_freeze_range(pretrained_path):
     """
     if pretrained_path:
         config = AutoConfig.from_pretrained(pretrained_path)
-        total_layers = config.num_hidden_layers  
-        return {"min": 0, "max": total_layers - 1} 
+        total_layers = config.num_hidden_layers
+        return {"min": 0, "max": total_layers - 1}
     else:
-        return {"min": 0, "max": 0} 
+        return {"min": 0, "max": 0}

geneformer/mtl_classifier.py

@@ -28,9 +28,8 @@ Geneformer multi-task cell classifier.
 
 import logging
 import os
-from .mtl import train_utils
-from .mtl import utils
-from .mtl import eval_utils
+
+from .mtl import eval_utils, train_utils, utils
 
 logger = logging.getLogger(__name__)
 
@@ -90,9 +89,8 @@ class MTLClassifier:
         wandb_project=None,
         gradient_clipping=False,
         max_grad_norm=None,
-        seed=42  # Default seed value
+        seed=42,  # Default seed value
     ):
-
         """
         Initialize Geneformer multi-task classifier.
         **Parameters:**
@@ -165,11 +163,11 @@ class MTLClassifier:
         self.batch_size = batch_size
         self.n_trials = n_trials
         self.study_name = study_name
-        
+
         if max_layers_to_freeze is None:
             # Dynamically determine the range of layers to freeze
             layer_freeze_range = utils.get_layer_freeze_range(pretrained_path)
-            self.max_layers_to_freeze = {"min": 1, "max": layer_freeze_range['max']}
+            self.max_layers_to_freeze = {"min": 1, "max": layer_freeze_range["max"]}
         else:
             self.max_layers_to_freeze = max_layers_to_freeze
 
@@ -178,48 +176,37 @@ class MTLClassifier:
         self.use_data_parallel = use_data_parallel
         self.use_attention_pooling = use_attention_pooling
         self.use_task_weights = use_task_weights
-        self.hyperparameters = hyperparameters if hyperparameters is not None else {
-            "learning_rate": {
-                "type": "float",
-                "low": 1e-5,
-                "high": 1e-3,
-                "log": True
-            },
-            "warmup_ratio": {
-                "type": "float",
-                "low": 0.005,
-                "high": 0.01
-            },
-            "weight_decay": {
-                "type": "float",
-                "low": 0.01,
-                "high": 0.1
-            },
-            "dropout_rate": {
-                "type": "float",
-                "low": 0.0,
-                "high": 0.7
-            },
-            "lr_scheduler_type": {
-                "type": "categorical",
-                "choices": ["cosine"]
-            },
-            "task_weights": {
-                "type": "float",
-                "low": 0.1,
-                "high": 2.0
+        self.hyperparameters = (
+            hyperparameters
+            if hyperparameters is not None
+            else {
+                "learning_rate": {
+                    "type": "float",
+                    "low": 1e-5,
+                    "high": 1e-3,
+                    "log": True,
+                },
+                "warmup_ratio": {"type": "float", "low": 0.005, "high": 0.01},
+                "weight_decay": {"type": "float", "low": 0.01, "high": 0.1},
+                "dropout_rate": {"type": "float", "low": 0.0, "high": 0.7},
+                "lr_scheduler_type": {"type": "categorical", "choices": ["cosine"]},
+                "task_weights": {"type": "float", "low": 0.1, "high": 2.0},
             }
-        }
-        self.manual_hyperparameters = manual_hyperparameters if manual_hyperparameters is not None else {
-            "learning_rate": 0.001,
-            "warmup_ratio": 0.01,
-            "weight_decay": 0.1,
-            "dropout_rate": 0.1,
-            "lr_scheduler_type": "cosine",
-            "use_attention_pooling": False,
-            "task_weights": [1, 1],
-            "max_layers_to_freeze": 2
-        }
+        )
+        self.manual_hyperparameters = (
+            manual_hyperparameters
+            if manual_hyperparameters is not None
+            else {
+                "learning_rate": 0.001,
+                "warmup_ratio": 0.01,
+                "weight_decay": 0.1,
+                "dropout_rate": 0.1,
+                "lr_scheduler_type": "cosine",
+                "use_attention_pooling": False,
+                "task_weights": [1, 1],
+                "max_layers_to_freeze": 2,
+            }
+        )
         self.use_manual_hyperparameters = use_manual_hyperparameters
         self.use_wandb = use_wandb
         self.wandb_project = wandb_project
@@ -236,13 +223,19 @@ class MTLClassifier:
 
         # set up output directories
         if self.results_dir is not None:
-            self.trials_results_path = f"{self.results_dir}/results.txt".replace("//","/")
-        
+            self.trials_results_path = f"{self.results_dir}/results.txt".replace(
+                "//", "/"
+            )
+
         for output_dir in [self.model_save_path, self.results_dir]:
             if not os.path.exists(output_dir):
                 os.makedirs(output_dir)
 
-        self.config = {key: value for key, value in self.__dict__.items() if key in self.valid_option_dict}
+        self.config = {
+            key: value
+            for key, value in self.__dict__.items()
+            if key in self.valid_option_dict
+        }
 
     def validate_options(self):
         # confirm arguments are within valid options and compatible with each other
@@ -264,19 +257,35 @@ class MTLClassifier:
                 f"Invalid option for {attr_name}. "
                 f"Valid options for {attr_name}: {valid_options}"
             )
-            raise ValueError(f"Invalid option for {attr_name}. Valid options for {attr_name}: {valid_options}")
+            raise ValueError(
+                f"Invalid option for {attr_name}. Valid options for {attr_name}: {valid_options}"
+            )
 
     def run_manual_tuning(self):
         """
         Manual hyperparameter tuning and multi-task fine-tuning of pretrained model.
         """
-        required_variable_names = ["train_path", "val_path", "pretrained_path", "model_save_path", "results_dir"]
-        required_variables = [self.train_path, self.val_path, self.pretrained_path, self.model_save_path, self.results_dir]
+        required_variable_names = [
+            "train_path",
+            "val_path",
+            "pretrained_path",
+            "model_save_path",
+            "results_dir",
+        ]
+        required_variables = [
+            self.train_path,
+            self.val_path,
+            self.pretrained_path,
+            self.model_save_path,
+            self.results_dir,
+        ]
         req_var_dict = dict(zip(required_variable_names, required_variables))
         self.validate_additional_options(req_var_dict)
 
         if not self.use_manual_hyperparameters:
-            raise ValueError("Manual hyperparameters are not enabled. Set use_manual_hyperparameters to True.")
+            raise ValueError(
+                "Manual hyperparameters are not enabled. Set use_manual_hyperparameters to True."
+            )
 
         # Ensure manual_hyperparameters are set in the config
         self.config["manual_hyperparameters"] = self.manual_hyperparameters
@@ -302,8 +311,20 @@ class MTLClassifier:
         Hyperparameter optimization and/or multi-task fine-tuning of pretrained model.
         """
 
-        required_variable_names = ["train_path", "val_path", "pretrained_path", "model_save_path", "results_dir"]
-        required_variables = [self.train_path, self.val_path, self.pretrained_path, self.model_save_path, self.results_dir]
+        required_variable_names = [
+            "train_path",
+            "val_path",
+            "pretrained_path",
+            "model_save_path",
+            "results_dir",
+        ]
+        required_variables = [
+            self.train_path,
+            self.val_path,
+            self.pretrained_path,
+            self.model_save_path,
+            self.results_dir,
+        ]
         req_var_dict = dict(zip(required_variable_names, required_variables))
         self.validate_additional_options(req_var_dict)
 
@@ -322,7 +343,7 @@ class MTLClassifier:
         self.validate_additional_options(req_var_dict)
 
         eval_utils.load_and_evaluate_test_model(self.config)
-        
+
     def save_model_without_heads(
         self,
     ):
@@ -335,4 +356,6 @@ class MTLClassifier:
         req_var_dict = dict(zip(required_variable_names, required_variables))
         self.validate_additional_options(req_var_dict)
 
-        utils.save_model_without_heads(os.path.join(self.model_save_path, "GeneformerMultiTask"))
+        utils.save_model_without_heads(
+            os.path.join(self.model_save_path, "GeneformerMultiTask")
+        )

geneformer/perturber_utils.py

@@ -1,18 +1,15 @@
 import itertools as it
 import logging
 import pickle
-import re
 from collections import defaultdict
-from typing import List
 from pathlib import Path
-
+from typing import List
 
 import numpy as np
 import pandas as pd
-import seaborn as sns
 import torch
 from datasets import Dataset, load_from_disk
-from peft import LoraConfig, get_peft_model 
+from peft import LoraConfig, get_peft_model
 from transformers import (
     BertForMaskedLM,
     BertForSequenceClassification,
@@ -119,7 +116,7 @@ def load_model(model_type, num_classes, model_directory, mode, quantize=False):
     if model_type == "MTLCellClassifier-Quantized":
         model_type = "MTLCellClassifier"
         quantize = True
-    
+
     if mode == "eval":
         output_hidden_states = True
     elif mode == "train":
@@ -131,7 +128,7 @@ def load_model(model_type, num_classes, model_directory, mode, quantize=False):
                 "peft_config": None,
                 "bnb_config": BitsAndBytesConfig(
                     load_in_8bit=True,
-                )
+                ),
             }
         else:
             quantize = {
@@ -141,13 +138,13 @@ def load_model(model_type, num_classes, model_directory, mode, quantize=False):
                     r=64,
                     bias="none",
                     task_type="TokenClassification",
-                  ),
+                ),
                 "bnb_config": BitsAndBytesConfig(
                     load_in_4bit=True,
                     bnb_4bit_use_double_quant=True,
                     bnb_4bit_quant_type="nf4",
-                    bnb_4bit_compute_dtype=torch.bfloat16
-                  )
+                    bnb_4bit_compute_dtype=torch.bfloat16,
+                ),
             }
     elif quantize is False:
         quantize = {"bnb_config": None}
@@ -186,7 +183,11 @@ def load_model(model_type, num_classes, model_directory, mode, quantize=False):
     # if eval mode, put the model in eval mode for fwd pass
     if mode == "eval":
         model.eval()
-    if (quantize is False) or (quantize == {'bnb_config': None}) or (model_type == "MTLCellClassifier"):
+    if (
+        (quantize is False)
+        or (quantize == {"bnb_config": None})
+        or (model_type == "MTLCellClassifier")
+    ):
         model = model.to("cuda")
     else:
         model.enable_input_require_grads()
@@ -279,18 +280,20 @@ def overexpress_indices(example):
     example["length"] = len(example["input_ids"])
     return example
 
+
 # if CLS token present, move to 1st rather than 0th position
 def overexpress_indices_special(example):
     indices = example["perturb_index"]
     if any(isinstance(el, list) for el in indices):
         indices = flatten_list(indices)
-    insert_pos = 1 # Insert starting after CLS token
+    insert_pos = 1  # Insert starting after CLS token
     for index in sorted(indices, reverse=False):
         example["input_ids"].insert(insert_pos, example["input_ids"].pop(index))
         insert_pos += 1
     example["length"] = len(example["input_ids"])
     return example
 
+
 # for genes_to_perturb = list of genes to overexpress that are not necessarily expressed in cell
 def overexpress_tokens(example, max_len, special_token):
     # -100 indicates tokens to overexpress are not present in rank value encoding
@@ -310,7 +313,9 @@ def overexpress_tokens(example, max_len, special_token):
     # truncate to max input size, must also truncate original emb to be comparable
     if len(example["input_ids"]) > max_len:
         if special_token:
-            example["input_ids"] = example["input_ids"][0:max_len-1]+[example["input_ids"][-1]]
+            example["input_ids"] = example["input_ids"][0 : max_len - 1] + [
+                example["input_ids"][-1]
+            ]
         else:
             example["input_ids"] = example["input_ids"][0:max_len]
     example["length"] = len(example["input_ids"])
@@ -329,10 +334,13 @@ def truncate_by_n_overflow(example):
     example["length"] = len(example["input_ids"])
     return example
 
+
 def truncate_by_n_overflow_special(example):
     if example["n_overflow"] > 0:
         new_max_len = example["length"] - example["n_overflow"]
-        example["input_ids"] = example["input_ids"][0:new_max_len-1]+[example["input_ids"][-1]]
+        example["input_ids"] = example["input_ids"][0 : new_max_len - 1] + [
+            example["input_ids"][-1]
+        ]
         example["length"] = len(example["input_ids"])
     return example
 
@@ -477,19 +485,24 @@ def make_perturbation_batch_special(
             range_start = 1
         elif perturb_type in ["delete", "inhibit"]:
             range_start = 0
-        range_start += 1 # Starting after the CLS token
+        range_start += 1  # Starting after the CLS token
         indices_to_perturb = [
-            [i] for i in range(range_start, example_cell["length"][0]-1) # And excluding the EOS token
+            [i]
+            for i in range(
+                range_start, example_cell["length"][0] - 1
+            )  # And excluding the EOS token
         ]
 
     # elif combo_lvl > 0 and anchor_token is None:
     ## to implement
-    elif combo_lvl > 0 and (anchor_token is not None): 
+    elif combo_lvl > 0 and (anchor_token is not None):
         example_input_ids = example_cell["input_ids"][0]
         anchor_index = example_input_ids.index(anchor_token[0])
         indices_to_perturb = [
             sorted([anchor_index, i]) if i != anchor_index else None
-            for i in range(1, example_cell["length"][0]-1) # Exclude CLS and EOS tokens
+            for i in range(
+                1, example_cell["length"][0] - 1
+            )  # Exclude CLS and EOS tokens
         ]
         indices_to_perturb = [item for item in indices_to_perturb if item is not None]
     else:
@@ -508,7 +521,9 @@ def make_perturbation_batch_special(
                     list(x) for x in it.combinations(indices_to_perturb, combo_lvl + 1)
                 ]
         else:
-            all_indices = [[i] for i in range(1, example_cell["length"][0]-1)] # Exclude CLS and EOS tokens
+            all_indices = [
+                [i] for i in range(1, example_cell["length"][0] - 1)
+            ]  # Exclude CLS and EOS tokens
             all_indices = [
                 index for index in all_indices if index not in indices_to_perturb
             ]
@@ -535,7 +550,7 @@ def make_perturbation_batch_special(
         )
     elif perturb_type == "overexpress":
         perturbation_dataset = perturbation_dataset.map(
-                overexpress_indices_special, num_proc=num_proc_i
+            overexpress_indices_special, num_proc=num_proc_i
         )
 
     perturbation_dataset = perturbation_dataset.map(measure_length, num_proc=num_proc_i)
@@ -743,7 +758,7 @@ def quant_cos_sims(
     # against original cell
     if cell_states_to_model is None or emb_mode == "gene":
         cos_sims = cos(perturbation_emb, original_emb).to("cuda")
-        
+
     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))]))
@@ -867,27 +882,28 @@ def validate_cell_states_to_model(cell_states_to_model):
             )
             raise
 
+
 class GeneIdHandler:
     def __init__(self, raise_errors=False):
         def invert_dict(dict_obj):
-            return {v:k for k,v in dict_obj.items()}
-        
+            return {v: k for k, v in dict_obj.items()}
+
         self.raise_errors = raise_errors
-        
-        with open(TOKEN_DICTIONARY_FILE, 'rb') as f:
+
+        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:
+        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)
@@ -899,15 +915,15 @@ class GeneIdHandler:
             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))
+        return self.ens_to_token(self.symbol_to_ens(symbol))

geneformer/tokenizer.py

@@ -22,28 +22,28 @@ Geneformer tokenizer.
 
 from __future__ import annotations
 
-import os
 import logging
+import os
 import pickle
 import warnings
+from collections import Counter
 from pathlib import Path
 from typing import Literal
-from tqdm import tqdm
-from collections import Counter
 
-import numpy as np
-import scanpy as sc
 import loompy as lp
+import numpy as np
 import pandas as pd
+import scanpy as sc
 import scipy.sparse as sp
 from datasets import Dataset
+from tqdm import tqdm
 
 warnings.filterwarnings("ignore", message=".*The 'nopython' keyword.*")  # noqa
 import loompy as lp  # noqa
 
 logger = logging.getLogger(__name__)
 
-from . import GENE_MEDIAN_FILE, TOKEN_DICTIONARY_FILE, ENSEMBL_MAPPING_FILE
+from . import ENSEMBL_MAPPING_FILE, GENE_MEDIAN_FILE, TOKEN_DICTIONARY_FILE
 
 
 def rank_genes(gene_vector, gene_tokens):
@@ -65,20 +65,26 @@ def tokenize_cell(gene_vector, gene_tokens):
     # rank by median-scaled gene values
     return rank_genes(gene_vector[nonzero_mask], gene_tokens[nonzero_mask])
 
-def sum_ensembl_ids(data_directory,
-                    collapse_gene_ids,
-                    gene_mapping_dict,
-                    gene_token_dict,
-                    file_format = "loom",
-                    chunk_size = 512):
 
+def sum_ensembl_ids(
+    data_directory,
+    collapse_gene_ids,
+    gene_mapping_dict,
+    gene_token_dict,
+    file_format="loom",
+    chunk_size=512,
+):
     if file_format == "loom":
-        """ 
+        """
         Map Ensembl IDs from gene mapping dictionary. If duplicate Ensembl IDs are found, sum counts together.
         """
         with lp.connect(data_directory) as data:
-            assert "ensembl_id" in data.ra.keys(), "'ensembl_id' column missing from data.ra.keys()"
-            gene_ids_in_dict = [gene for gene in data.ra.ensembl_id if gene in gene_token_dict.keys()]
+            assert (
+                "ensembl_id" in data.ra.keys()
+            ), "'ensembl_id' column missing from data.ra.keys()"
+            gene_ids_in_dict = [
+                gene for gene in data.ra.ensembl_id if gene in gene_token_dict.keys()
+            ]
             if len(gene_ids_in_dict) == len(set(gene_ids_in_dict)):
                 token_genes_unique = True
             else:
@@ -89,45 +95,80 @@ def sum_ensembl_ids(data_directory,
                 else:
                     raise ValueError("Error: data Ensembl IDs non-unique.")
 
-            gene_ids_collapsed = [gene_mapping_dict.get(gene_id.upper()) for gene_id in data.ra.ensembl_id]
-            gene_ids_collapsed_in_dict = [gene for gene in gene_ids_collapsed if gene in gene_token_dict.keys()]
+            gene_ids_collapsed = [
+                gene_mapping_dict.get(gene_id.upper()) for gene_id in data.ra.ensembl_id
+            ]
+            gene_ids_collapsed_in_dict = [
+                gene for gene in gene_ids_collapsed if gene in gene_token_dict.keys()
+            ]
 
-            if (len(set(gene_ids_collapsed_in_dict)) == len(set(gene_ids_in_dict))) and token_genes_unique:
+            if (
+                len(set(gene_ids_collapsed_in_dict)) == len(set(gene_ids_in_dict))
+            ) and token_genes_unique:
                 return data_directory
             else:
-                dedup_filename = data_directory.with_name(data_directory.stem + "__dedup.loom")
+                dedup_filename = data_directory.with_name(
+                    data_directory.stem + "__dedup.loom"
+                )
                 data.ra["gene_ids_collapsed"] = gene_ids_collapsed
-                dup_genes = [idx for idx, count in Counter(data.ra["gene_ids_collapsed"]).items() if count > 1]
+                dup_genes = [
+                    idx
+                    for idx, count in Counter(data.ra["gene_ids_collapsed"]).items()
+                    if count > 1
+                ]
                 num_chunks = int(np.ceil(data.shape[1] / chunk_size))
                 first_chunk = True
-                for _, _, view in tqdm(data.scan(axis = 1, batch_size = chunk_size), total = num_chunks):
+                for _, _, view in tqdm(
+                    data.scan(axis=1, batch_size=chunk_size), total=num_chunks
+                ):
+
                     def process_chunk(view, duplic_genes):
-                        data_count_view = pd.DataFrame(view, index=data.ra["gene_ids_collapsed"])
-                        unique_data_df = data_count_view.loc[~data_count_view.index.isin(duplic_genes)]
-                        dup_data_df = data_count_view.loc[data_count_view.index.isin([i for i in duplic_genes if "None" not in i])]
+                        data_count_view = pd.DataFrame(
+                            view, index=data.ra["gene_ids_collapsed"]
+                        )
+                        unique_data_df = data_count_view.loc[
+                            ~data_count_view.index.isin(duplic_genes)
+                        ]
+                        dup_data_df = data_count_view.loc[
+                            data_count_view.index.isin(
+                                [i for i in duplic_genes if "None" not in i]
+                            )
+                        ]
                         summed_data = dup_data_df.groupby(dup_data_df.index).sum()
                         if not summed_data.index.is_unique:
-                            raise ValueError("Error: Ensembl IDs in summed data frame non-unique.")
-                        data_count_view = pd.concat([unique_data_df, summed_data], axis=0)
+                            raise ValueError(
+                                "Error: Ensembl IDs in summed data frame non-unique."
+                            )
+                        data_count_view = pd.concat(
+                            [unique_data_df, summed_data], axis=0
+                        )
                         if not data_count_view.index.is_unique:
-                            raise ValueError("Error: Ensembl IDs in final data frame non-unique.")
+                            raise ValueError(
+                                "Error: Ensembl IDs in final data frame non-unique."
+                            )
                         return data_count_view
+
                     processed_chunk = process_chunk(view[:, :], dup_genes)
                     processed_array = processed_chunk.to_numpy()
                     new_row_attrs = {"ensembl_id": processed_chunk.index.to_numpy()}
 
                     if "n_counts" not in view.ca.keys():
-                        total_count_view = np.sum(view[:,:], axis=0).astype(int)
+                        total_count_view = np.sum(view[:, :], axis=0).astype(int)
                         view.ca["n_counts"] = total_count_view
 
-                    if first_chunk: # Create the Loom file with the first chunk
-                        lp.create(f"{dedup_filename}", processed_array, row_attrs=new_row_attrs, col_attrs=view.ca)
+                    if first_chunk:  # Create the Loom file with the first chunk
+                        lp.create(
+                            f"{dedup_filename}",
+                            processed_array,
+                            row_attrs=new_row_attrs,
+                            col_attrs=view.ca,
+                        )
                         first_chunk = False
-                    else: # Append subsequent chunks
-                        with lp.connect(dedup_filename, mode='r+') as dsout:
+                    else:  # Append subsequent chunks
+                        with lp.connect(dedup_filename, mode="r+") as dsout:
                             dsout.add_columns(processed_array, col_attrs=view.ca)
                 return dedup_filename
-                
+
     elif file_format == "h5ad":
         """
         Map Ensembl IDs from gene mapping dictionary. If duplicate Ensembl IDs are found, sum counts together.
@@ -136,8 +177,12 @@ def sum_ensembl_ids(data_directory,
 
         data = sc.read_h5ad(str(data_directory))
 
-        assert "ensembl_id" in data.var.columns, "'ensembl_id' column missing from data.var"
-        gene_ids_in_dict = [gene for gene in data.var.ensembl_id if gene in gene_token_dict.keys()]
+        assert (
+            "ensembl_id" in data.var.columns
+        ), "'ensembl_id' column missing from data.var"
+        gene_ids_in_dict = [
+            gene for gene in data.var.ensembl_id if gene in gene_token_dict.keys()
+        ]
         if len(gene_ids_in_dict) == len(set(gene_ids_in_dict)):
             token_genes_unique = True
         else:
@@ -148,22 +193,29 @@ def sum_ensembl_ids(data_directory,
             else:
                 raise ValueError("Error: data Ensembl IDs non-unique.")
 
-        gene_ids_collapsed = [gene_mapping_dict.get(gene_id.upper()) for gene_id in data.var.ensembl_id]
-        gene_ids_collapsed_in_dict = [gene for gene in gene_ids_collapsed if gene in gene_token_dict.keys()]
-        if (len(set(gene_ids_collapsed_in_dict)) == len(set(gene_ids_in_dict))) and token_genes_unique:
+        gene_ids_collapsed = [
+            gene_mapping_dict.get(gene_id.upper()) for gene_id in data.var.ensembl_id
+        ]
+        gene_ids_collapsed_in_dict = [
+            gene for gene in gene_ids_collapsed if gene in gene_token_dict.keys()
+        ]
+        if (
+            len(set(gene_ids_collapsed_in_dict)) == len(set(gene_ids_in_dict))
+        ) and token_genes_unique:
             return data
-        
+
         else:
             data.var["gene_ids_collapsed"] = gene_ids_collapsed
             data.var_names = gene_ids_collapsed
             data = data[:, ~data.var.index.isna()]
-            dup_genes = [idx for idx, count in Counter(data.var_names).items() if count > 1]
+            dup_genes = [
+                idx for idx, count in Counter(data.var_names).items() if count > 1
+            ]
 
             num_chunks = int(np.ceil(data.shape[0] / chunk_size))
 
             processed_genes = []
             for i in tqdm(range(num_chunks)):
-                
                 start_idx = i * chunk_size
                 end_idx = min((i + 1) * chunk_size, data.shape[0])
                 data_chunk = data[start_idx:end_idx, :]
@@ -171,29 +223,32 @@ def sum_ensembl_ids(data_directory,
                 processed_chunks = []
                 for dup_gene in dup_genes:
                     data_dup_gene = data_chunk[:, data_chunk.var_names == dup_gene]
-                    df = pd.DataFrame.sparse.from_spmatrix(data_dup_gene.X,
-                                                           index=data_dup_gene.obs_names,
-                                                           columns=data_dup_gene.var_names)
+                    df = pd.DataFrame.sparse.from_spmatrix(
+                        data_dup_gene.X,
+                        index=data_dup_gene.obs_names,
+                        columns=data_dup_gene.var_names,
+                    )
                     df_sum = pd.DataFrame(df.sum(axis=1))
-                    df_sum.columns = [dup_gene] 
+                    df_sum.columns = [dup_gene]
                     df_sum.index = data_dup_gene.obs.index
                     processed_chunks.append(df_sum)
 
-                processed_chunks = pd.concat(processed_chunks, axis=1) 
+                processed_chunks = pd.concat(processed_chunks, axis=1)
                 processed_genes.append(processed_chunks)
-            processed_genes = pd.concat(processed_genes, axis = 0)
-            var_df = pd.DataFrame({"gene_ids_collapsed" : processed_genes.columns})
+            processed_genes = pd.concat(processed_genes, axis=0)
+            var_df = pd.DataFrame({"gene_ids_collapsed": processed_genes.columns})
             var_df.index = processed_genes.columns
-            processed_genes = sc.AnnData(X = processed_genes,
-                                        obs = data.obs,
-                                        var = var_df)
+            processed_genes = sc.AnnData(X=processed_genes, obs=data.obs, var=var_df)
 
-            data_dedup = data[:, ~data.var.index.isin(dup_genes)] # Deduplicated data
-            data_dedup = sc.concat([data_dedup, processed_genes], axis = 1)
+            data_dedup = data[:, ~data.var.index.isin(dup_genes)]  # Deduplicated data
+            data_dedup = sc.concat([data_dedup, processed_genes], axis=1)
             data_dedup.obs = data.obs
-            data_dedup.var = data_dedup.var.rename(columns = {"gene_ids_collapsed" : "ensembl_id"})
+            data_dedup.var = data_dedup.var.rename(
+                columns={"gene_ids_collapsed": "ensembl_id"}
+            )
             return data_dedup
 
+
 class TranscriptomeTokenizer:
     def __init__(
         self,
@@ -258,10 +313,12 @@ class TranscriptomeTokenizer:
 
         # check for special token in gene_token_dict
         if self.special_token:
-            if ("<cls>" not in self.gene_token_dict.keys()) and ("<eos>" not in self.gene_token_dict.keys()):
+            if ("<cls>" not in self.gene_token_dict.keys()) and (
+                "<eos>" not in self.gene_token_dict.keys()
+            ):
                 logger.error(
-                            "<cls> and <eos> required in gene_token_dict when special_token = True."
-                        )
+                    "<cls> and <eos> required in gene_token_dict when special_token = True."
+                )
                 raise
 
         # if collapsing duplicate gene IDs
@@ -272,14 +329,16 @@ class TranscriptomeTokenizer:
             with open(gene_mapping_file, "rb") as f:
                 self.gene_mapping_dict = pickle.load(f)
         else:
-            self.gene_mapping_dict = {k:k for k,_ in self.gene_token_dict.items()}
+            self.gene_mapping_dict = {k: k for k, _ in self.gene_token_dict.items()}
 
         # gene keys for full vocabulary
         self.gene_keys = list(self.gene_token_dict.keys())
 
         #  Filter gene mapping dict for items that exist in gene_token_dict
         gene_keys_set = set(self.gene_token_dict.keys())
-        self.gene_mapping_dict = {k: v for k, v in self.gene_mapping_dict.items() if v in gene_keys_set}
+        self.gene_mapping_dict = {
+            k: v for k, v in self.gene_mapping_dict.items() if v in gene_keys_set
+        }
 
         # protein-coding and miRNA gene list dictionary for selecting .loom rows for tokenization
         self.genelist_dict = dict(zip(self.gene_keys, [True] * len(self.gene_keys)))
@@ -355,7 +414,14 @@ class TranscriptomeTokenizer:
         return tokenized_cells, cell_metadata
 
     def tokenize_anndata(self, adata_file_path, target_sum=10_000):
-        adata = sum_ensembl_ids(adata_file_path, self.collapse_gene_ids, self.gene_mapping_dict, self.gene_token_dict, file_format = "h5ad", chunk_size = self.chunk_size)
+        adata = sum_ensembl_ids(
+            adata_file_path,
+            self.collapse_gene_ids,
+            self.gene_mapping_dict,
+            self.gene_token_dict,
+            file_format="h5ad",
+            chunk_size=self.chunk_size,
+        )
 
         if self.custom_attr_name_dict is not None:
             file_cell_metadata = {
@@ -397,7 +463,7 @@ class TranscriptomeTokenizer:
             idx = filter_pass_loc[i : i + self.chunk_size]
 
             n_counts = adata[idx].obs["n_counts"].values[:, None]
-            X_view0 = adata[idx,:].X
+            X_view0 = adata[idx, :].X
             X_view = X_view0[:, coding_miRNA_loc]
             X_norm = X_view / n_counts * target_sum / norm_factor_vector
             X_norm = sp.csr_matrix(X_norm)
@@ -423,7 +489,14 @@ class TranscriptomeTokenizer:
             }
 
         dedup_filename = loom_file_path.with_name(loom_file_path.stem + "__dedup.loom")
-        loom_file_path = sum_ensembl_ids(loom_file_path, self.collapse_gene_ids, self.gene_mapping_dict, self.gene_token_dict, file_format = "loom", chunk_size = self.chunk_size)
+        loom_file_path = sum_ensembl_ids(
+            loom_file_path,
+            self.collapse_gene_ids,
+            self.gene_mapping_dict,
+            self.gene_token_dict,
+            file_format="loom",
+            chunk_size=self.chunk_size,
+        )
 
         with lp.connect(str(loom_file_path)) as data:
             # define coordinates of detected protein-coding or miRNA genes and vector of their normalization factors
@@ -544,4 +617,4 @@ class TranscriptomeTokenizer:
         output_dataset_truncated = output_dataset.map(
             format_cell_features, num_proc=self.nproc
         )
-        return output_dataset_truncated
+        return output_dataset_truncated