genetransformer

.gitattributes

@@ -14,11 +14,10 @@
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.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
@@ -26,4 +25,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+geneformer/gene_name_id_dict.pkl filter=lfs diff=lfs merge=lfs -text
 model.safetensors filter=lfs diff=lfs merge=lfs -text

MANIFEST.in

@@ -1,4 +1,3 @@
-include geneformer/gene_median_dictionary_gc95M.pkl
-include geneformer/gene_name_id_dict_gc95M.pkl
-include geneformer/ensembl_mapping_dict_gc95M.pkl
-include geneformer/token_dictionary_gc95M.pkl
+include geneformer/gene_median_dictionary.pkl
+include geneformer/token_dictionary.pkl
+include geneformer/gene_name_id_dict.pkl

README.md

@@ -3,38 +3,23 @@ datasets: ctheodoris/Genecorpus-30M
 license: apache-2.0
 ---
 # Geneformer
-Geneformer is a foundational transformer model pretrained on a large-scale corpus of single cell transcriptomes to enable context-aware predictions in settings with limited data in network biology.
+Geneformer is a foundation transformer model pretrained on a large-scale corpus of ~30 million single cell transcriptomes to enable context-aware predictions in settings with limited data in network biology.
 
-- See [our manuscript](https://rdcu.be/ddrx0) for details of the original model trained on ~30 million transcriptomes in June 2021 and the initial report of our in silico perturbation and cell and gene classification strategies.
-- See [our manuscript](https://www.biorxiv.org/content/10.1101/2024.08.16.608180v1.full.pdf) for details of the expanded model trained on ~95 million transcriptomes in April 2024 and our continual learning, multitask learning, and quantization strategies.
+- See [our manuscript](https://rdcu.be/ddrx0) for details.
 - See [geneformer.readthedocs.io](https://geneformer.readthedocs.io) for documentation.
 
 # Model Description
-Geneformer is a foundational transformer model pretrained on a large-scale corpus of single cell transcriptomes representing a broad range of human tissues. Geneformer was originally pretrained in June 2021 on [Genecorpus-30M](https://huggingface.co/datasets/ctheodoris/Genecorpus-30M), a corpus comprised of ~30 million single cell transcriptomes. We excluded cells with high mutational burdens (e.g. malignant cells and immortalized cell lines) that could lead to substantial network rewiring without companion genome sequencing to facilitate interpretation. Then, in April 2024, Geneformer was pretrained on ~95 million non-cancer transcriptomes, followed by continual learning on ~14 million cancer transcriptomes to yield a cancer domain-tuned model. 
+Geneformer is a foundation transformer model pretrained on [Genecorpus-30M](https://huggingface.co/datasets/ctheodoris/Genecorpus-30M), a pretraining corpus comprised of ~30 million single cell transcriptomes from a broad range of human tissues. We excluded cells with high mutational burdens (e.g. malignant cells and immortalized cell lines) that could lead to substantial network rewiring without companion genome sequencing to facilitate interpretation. Each single cell’s transcriptome is presented to the model as a rank value encoding where genes are ranked by their expression in that cell normalized by their expression across the entire Genecorpus-30M. The rank value encoding provides a nonparametric representation of that cell’s transcriptome and takes advantage of the many observations of each gene’s expression across Genecorpus-30M to prioritize genes that distinguish cell state. Specifically, this method will deprioritize ubiquitously highly-expressed housekeeping genes by normalizing them to a lower rank. Conversely, genes such as transcription factors that may be lowly expressed when they are expressed but highly distinguish cell state will move to a higher rank within the encoding. Furthermore, this rank-based approach may be more robust against technical artifacts that may systematically bias the absolute transcript counts value while the overall relative ranking of genes within each cell remains more stable.
 
-Each single cell’s transcriptome is presented to the model as a rank value encoding where genes are ranked by their expression in that cell scaled by their expression across the entire Genecorpus-30M. The rank value encoding provides a nonparametric representation of that cell’s transcriptome and takes advantage of the many observations of each gene’s expression across the pretraining corpus to prioritize genes that distinguish cell state. Specifically, this method will deprioritize ubiquitously highly-expressed housekeeping genes by scaling them to a lower rank. Conversely, genes such as transcription factors that may be lowly expressed when they are expressed but highly distinguish cell state will move to a higher rank within the encoding. Furthermore, this rank-based approach may be more robust against technical artifacts that may systematically bias the absolute transcript counts value while the overall relative ranking of genes within each cell remains more stable.
-
-The rank value encoding of each single cell’s transcriptome then proceeds through N layers of transformer encoder units, where N varies dependent on the model size. Pretraining was accomplished using a masked learning objective where 15% of the genes within each transcriptome were masked and the model was trained to predict which gene should be within each masked position in that specific cell state using the context of the remaining unmasked genes. A major strength of this approach is that it is entirely self-supervised and can be accomplished on completely unlabeled data, which allows the inclusion of large amounts of training data without being restricted to samples with accompanying labels.
+The rank value encoding of each single cell��s transcriptome then proceeds through six transformer encoder units. Pretraining was accomplished using a masked learning objective where 15% of the genes within each transcriptome were masked and the model was trained to predict which gene should be within each masked position in that specific cell state using the context of the remaining unmasked genes. A major strength of this approach is that it is entirely self-supervised and can be accomplished on completely unlabeled data, which allows the inclusion of large amounts of training data without being restricted to samples with accompanying labels.
 
 We detail applications and results in [our manuscript](https://rdcu.be/ddrx0).
 
-During pretraining, Geneformer gained a fundamental understanding of network dynamics, encoding network hierarchy in the model’s attention weights in a completely self-supervised manner. With both zero-shot learning and fine-tuning with limited task-specific data, Geneformer consistently boosted predictive accuracy in a diverse panel of downstream tasks relevant to chromatin and network dynamics. In silico perturbation with zero-shot learning identified a novel transcription factor in cardiomyocytes that we experimentally validated to be critical to their ability to generate contractile force. In silico treatment with limited patient data revealed candidate therapeutic targets for cardiomyopathy that we experimentally validated to significantly improve the ability of cardiomyocytes to generate contractile force in an induced pluripotent stem cell (iPSC) model of the disease. Overall, Geneformer represents a foundational deep learning model pretrained on a large-scale corpus human single cell transcriptomes to gain a fundamental understanding of gene network dynamics that can now be democratized to a vast array of downstream tasks to accelerate discovery of key network regulators and candidate therapeutic targets.
-
-The repository includes the following pretrained models:
-
-L=layers\
-M=millions of cells used for pretraining\
-i=input size\
-(pretraining date)
+During pretraining, Geneformer gained a fundamental understanding of network dynamics, encoding network hierarchy in the model’s attention weights in a completely self-supervised manner. With both zero-shot learning and fine-tuning with limited task-specific data, Geneformer consistently boosted predictive accuracy in a diverse panel of downstream tasks relevant to chromatin and network dynamics. In silico perturbation with zero-shot learning identified a novel transcription factor in cardiomyocytes that we experimentally validated to be critical to their ability to generate contractile force. In silico treatment with limited patient data revealed candidate therapeutic targets for cardiomyopathy that we experimentally validated to significantly improve the ability of cardiomyocytes to generate contractile force in an iPSC model of the disease. Overall, Geneformer represents a foundational deep learning model pretrained on ~30 million human single cell transcriptomes to gain a fundamental understanding of gene network dynamics that can now be democratized to a vast array of downstream tasks to accelerate discovery of key network regulators and candidate therapeutic targets.
 
-- GF-6L-30M-i2048 (June 2021)
-- GF-12L-30M-i2048 (June 2021)
-- GF-12L-95M-i4096 (April 2024)
-- GF-20L-95M-i4096 (April 2024)
+In [our manuscript](https://rdcu.be/ddrx0), we report results for the 6 layer Geneformer model pretrained on Genecorpus-30M. We additionally provide within this repository a 12 layer Geneformer model, scaled up with retained width:depth aspect ratio, also pretrained on Genecorpus-30M.
 
-The current default model in the main directory of the repository is GF-12L-95M-i4096.
-
-The repository also contains fined tuned models in the fine_tuned_models directory and the cancer-tuned model following continual learning on ~14 million cancer cells, GF-12L-95M-i4096_CLcancer.
+Both the 6 and 12 layer Geneformer models were pretrained in June 2021.
 
 # Application
 The pretrained Geneformer model can be used directly for zero-shot learning, for example for in silico perturbation analysis, or by fine-tuning towards the relevant downstream task, such as gene or cell state classification.
@@ -64,7 +49,7 @@ Example applications demonstrated in [our manuscript](https://rdcu.be/ddrx0) inc
 - in silico perturbation to determine transcription factor cooperativity
 
 # Installation
-In addition to the pretrained model, contained herein are functions for tokenizing and collating data specific to single cell transcriptomics, pretraining the model, fine-tuning the model, extracting and plotting cell embeddings, and performing in silico pertrubation with either the pretrained or fine-tuned models. To install (~20s):
+In addition to the pretrained model, contained herein are functions for tokenizing and collating data specific to single cell transcriptomics, pretraining the model, fine-tuning the model, extracting and plotting cell embeddings, and performing in silico pertrubation with either the pretrained or fine-tuned models. To install:
 
 ```bash
 # Make sure you have git-lfs installed (https://git-lfs.com)
@@ -85,7 +70,3 @@ For usage, see [examples](https://huggingface.co/ctheodoris/Geneformer/tree/main
 Please note that the fine-tuning examples are meant to be generally applicable and the input datasets and labels will vary dependent on the downstream task. Example input files for a few of the downstream tasks demonstrated in the manuscript are located within the [example_input_files directory](https://huggingface.co/datasets/ctheodoris/Genecorpus-30M/tree/main/example_input_files) in the dataset repository, but these only represent a few example fine-tuning applications.
 
 Please note that GPU resources are required for efficient usage of Geneformer. Additionally, we strongly recommend tuning hyperparameters for each downstream fine-tuning application as this can significantly boost predictive potential in the downstream task (e.g. max learning rate, learning schedule, number of layers to freeze, etc.).
-
-# Citations
-- C V Theodoris#, L Xiao, A Chopra, M D Chaffin, Z R Al Sayed, M C Hill, H Mantineo, E Brydon, Z Zeng, X S Liu, P T Ellinor#. Transfer learning enables predictions in network biology. _**Nature**_, 31 May 2023. (#co-corresponding authors)
-- H Chen*, M S Venkatesh*, J Gomez Ortega, S V Mahesh, T Nandi, R Madduri, K Pelka†, C V Theodoris†#. Quantized multi-task learning for context-specific representations of gene network dynamics. _**bioRxiv**_, 19 Aug 2024. (*co-first authors, †co-senior authors, #corresponding author)

config.json

@@ -3,22 +3,21 @@
     "BertForMaskedLM"
   ],
   "attention_probs_dropout_prob": 0.02,
-  "classifier_dropout": null,
+  "gradient_checkpointing": false,
   "hidden_act": "relu",
   "hidden_dropout_prob": 0.02,
-  "hidden_size": 512,
+  "hidden_size": 256,
   "initializer_range": 0.02,
-  "intermediate_size": 1024,
+  "intermediate_size": 512,
   "layer_norm_eps": 1e-12,
-  "max_position_embeddings": 4096,
+  "max_position_embeddings": 2048,
   "model_type": "bert",
-  "num_attention_heads": 8,
-  "num_hidden_layers": 12,
+  "num_attention_heads": 4,
+  "num_hidden_layers": 6,
   "pad_token_id": 0,
   "position_embedding_type": "absolute",
-  "torch_dtype": "float32",
-  "transformers_version": "4.37.1",
+  "transformers_version": "4.6.0",
   "type_vocab_size": 2,
   "use_cache": true,
-  "vocab_size": 20275
+  "vocab_size": 25426
 }

docs/source/about.rst

@@ -4,13 +4,11 @@ About
 Model Description
 -----------------
 
-**Geneformer** is a context-aware, attention-based deep learning model pretrained on a large-scale corpus of single-cell transcriptomes to enable context-specific predictions in settings with limited data in network biology. During pretraining, Geneformer gained a fundamental understanding of network dynamics, encoding network hierarchy in the attention weights of the model in a completely self-supervised manner. With both zero-shot learning and fine-tuning with limited task-specific data, Geneformer consistently boosted predictive accuracy in a diverse panel of downstream tasks relevant to chromatin and network dynamics. In silico perturbation with zero-shot learning identified a novel transcription factor in cardiomyocytes that we experimentally validated to be critical to their ability to generate contractile force. In silico treatment with limited patient data revealed candidate therapeutic targets for cardiomyopathy that we experimentally validated to significantly improve the ability of cardiomyocytes to generate contractile force in an iPSC model of the disease. Overall, Geneformer represents a foundational deep learning model pretrained on a large-scale corpus of human single cell transcriptomes to gain a fundamental understanding of gene network dynamics that can now be democratized to a vast array of downstream tasks to accelerate discovery of key network regulators and candidate therapeutic targets.
+**Geneformer** is a context-aware, attention-based deep learning model pretrained on a large-scale corpus of ~30 million single-cell transcriptomes to enable context-specific predictions in settings with limited data in network biology. During pretraining, Geneformer gained a fundamental understanding of network dynamics, encoding network hierarchy in the attention weights of the model in a completely self-supervised manner. With both zero-shot learning and fine-tuning with limited task-specific data, Geneformer consistently boosted predictive accuracy in a diverse panel of downstream tasks relevant to chromatin and network dynamics. In silico perturbation with zero-shot learning identified a novel transcription factor in cardiomyocytes that we experimentally validated to be critical to their ability to generate contractile force. In silico treatment with limited patient data revealed candidate therapeutic targets for cardiomyopathy that we experimentally validated to significantly improve the ability of cardiomyocytes to generate contractile force in an iPSC model of the disease. Overall, Geneformer represents a foundational deep learning model pretrained on ~30 million human single cell transcriptomes to gain a fundamental understanding of gene network dynamics that can now be democratized to a vast array of downstream tasks to accelerate discovery of key network regulators and candidate therapeutic targets.
 
-In `our manuscript <https://rdcu.be/ddrx0>`_, we report results for the original 6 layer Geneformer model pretrained on Genecorpus-30M. We additionally provide within the repository a 12 layer Geneformer model, scaled up with retained width:depth aspect ratio, also pretrained on Genecorpus-30M.
+In `our manuscript <https://rdcu.be/ddrx0>`_, we report results for the 6 layer Geneformer model pretrained on Genecorpus-30M. We additionally provide within the repository a 12 layer Geneformer model, scaled up with retained width:depth aspect ratio, also pretrained on Genecorpus-30M.
 
-Both the `6 <https://huggingface.co/ctheodoris/Geneformer/blob/main/gf-6L-30M-i2048/model.safetensors>`_ and `12 <https://huggingface.co/ctheodoris/Geneformer/blob/main/gf-12L-30M-i2048/pytorch_model.bin>`_ layer Geneformer models were pretrained in June 2021.
-
-Also see `our 2024 manuscript <https://www.biorxiv.org/content/10.1101/2024.08.16.608180v1.full.pdf>`_, for details of the `expanded model <https://huggingface.co/ctheodoris/Geneformer/blob/main/model.safetensors>`_ trained on ~95 million transcriptomes in April 2024 and our continual learning, multitask learning, and quantization strategies.
+Both the `6 <https://huggingface.co/ctheodoris/Geneformer/blob/main/pytorch_model.bin>`_ and `12 <https://huggingface.co/ctheodoris/Geneformer/blob/main/geneformer-12L-30M/pytorch_model.bin>`_ layer Geneformer models were pretrained in June 2021.
 
 Application
 -----------
@@ -41,9 +39,7 @@ Example applications demonstrated in `our manuscript <https://rdcu.be/ddrx0>`_ i
 | - in silico perturbation to determine transcription factor targets
 | - in silico perturbation to determine transcription factor cooperativity
 
-Citations
----------
+Citation
+--------
 
 | C V Theodoris #, L Xiao, A Chopra, M D Chaffin, Z R Al Sayed, M C Hill, H Mantineo, E Brydon, Z Zeng, X S Liu, P T Ellinor #. `Transfer learning enables predictions in network biology. <https://rdcu.be/ddrx0>`_ *Nature*, 31 May 2023. (# co-corresponding authors)
-
-| H Chen \*, M S Venkatesh \*, J Gomez Ortega, S V Mahesh, T Nandi, R Madduri, K Pelka †, C V Theodoris † #. `Quantized multi-task learning for context-specific representations of gene network dynamics. <https://www.biorxiv.org/content/10.1101/2024.08.16.608180v1.full.pdf>`_ *bioRxiv*, 19 Aug 2024. (\* co-first authors, † co-senior authors, # corresponding author)

docs/source/api.rst

@@ -17,14 +17,6 @@ Classifier
 
    geneformer.classifier
 
-Multitask Classifier
---------------------
-
-.. toctree::
-   :maxdepth: 1
-
-   geneformer.mtl_classifier
-
 Embedding Extractor
 -------------------
 

docs/source/geneformer.mtl_classifier.rst

@@ -1,11 +0,0 @@
-geneformer.mtl\_classifier
-==========================
-
-.. automodule:: geneformer.mtl_classifier
-   :members:
-   :undoc-members:
-   :show-inheritance:
-   :exclude-members:
-        valid_option_dict,
-        validate_options,
-        validate_additional_options

docs/source/geneformer.tokenizer.rst

@@ -11,5 +11,4 @@ geneformer.tokenizer
         tokenize_files,
         tokenize_loom,
         rank_genes,
-        tokenize_cell,
-        sum_ensembl_ids
+        tokenize_cell

docs/source/index.rst

@@ -1,7 +1,7 @@
 Geneformer
 ==========
 
-Geneformer is a foundation transformer model pretrained on a large-scale corpus of single cell transcriptomes to enable context-aware predictions in network biology.
+Geneformer is a foundation transformer model pretrained on a large-scale corpus of ~30 million single cell transcriptomes to enable context-aware predictions in network biology.
 
 See `our manuscript <https://rdcu.be/ddrx0>`_ for details.
 

examples/multitask_cell_classification.ipynb

@@ -1,420 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "id": "866f100c-e11a-4e7b-a37c-831775d845a7",
-   "metadata": {},
-   "source": [
-    "# Geneformer Multi-Task Cell Classifier Tutorial\n",
-    "\n",
-    "This tutorial demonstrates how to use the Geneformer Multi-Task Cell Classifier and optimizatize hyperparameter for fine-tuning"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "311ba456-b44d-40c7-941d-3fc03bcda85a",
-   "metadata": {},
-   "source": [
-    "## 1. Installation and Imports\n",
-    "\n",
-    "First import the necessary modules."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "id": "cd9defdc-0524-4c3b-a741-27117ed3a5be",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from geneformer import MTLClassifier"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "790e9c3c-f6d9-44b3-b9a5-05725760f4fd",
-   "metadata": {},
-   "source": [
-    "## 2. Set up Paths and Parameters\n",
-    "\n",
-    "Now, let's set up the necessary paths and parameters for our classifier. We'll also define our task columns, which are specific columns from our dataset that represent the classification tasks we want to train the model on."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "04a04197-8e45-47f8-a86f-202209ea10ae",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Define paths\n",
-    "pretrained_path = \"/path/to/pretrained/Geneformer/model\" \n",
-    "# input data is tokenized rank value encodings generated by Geneformer tokenizer (see tokenizing_scRNAseq_data.ipynb)\n",
-    "train_path = \"/path/to/train/data.dataset\"\n",
-    "val_path = \"/path/to/val/data.dataset\"\n",
-    "test_path = \"/path/to/test/data.dataset\"\n",
-    "results_dir = \"/path/to/results/directory\"\n",
-    "model_save_path = \"/path/to/model/save/path\"\n",
-    "tensorboard_log_dir = \"/path/to/tensorboard/log/dir\"\n",
-    "\n",
-    "# Define tasks and hyperparameters\n",
-    "# task_columns should be a list of column names from your dataset\n",
-    "# Each column represents a specific classification task (e.g. cell type, disease state)\n",
-    "task_columns = [\"cell_type\", \"disease_state\"]  # Example task columns\n",
-    "\n",
-    "hyperparameters = {\n",
-    "    \"learning_rate\": {\"type\": \"float\", \"low\": 1e-5, \"high\": 1e-3, \"log\": True},\n",
-    "    \"warmup_ratio\": {\"type\": \"float\", \"low\": 0.005, \"high\": 0.01},\n",
-    "    \"weight_decay\": {\"type\": \"float\", \"low\": 0.01, \"high\": 0.1},\n",
-    "    \"dropout_rate\": {\"type\": \"float\", \"low\": 0.0, \"high\": 0.7},\n",
-    "    \"lr_scheduler_type\": {\"type\": \"categorical\", \"choices\": [\"cosine\"]},\n",
-    "    \"task_weights\": {\"type\": \"float\", \"low\": 0.1, \"high\": 2.0}\n",
-    "}"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "31857690-a739-435a-aefd-f171fafc1b78",
-   "metadata": {},
-   "source": [
-    "In the code above, we've defined `task_columns` as `[\"cell_type\", \"disease_state\"]`. This means our model will be trained to classify cells based on two tasks:\n",
-    "1. Identifying the cell type\n",
-    "2. Determining the disease state\n",
-    "3. Note: \"unique_cell_id\" is a required column in the dataset for logging and inference purposes\n",
-    "\n",
-    "These column names should correspond to actual columns in your dataset. Each column should contain the labels for that specific classification task.\n",
-    "\n",
-    "For example, your dataset might look something like this:\n",
-    "\n",
-    "    | unique_cell_id | input_ids | ... | cell_type | disease_state |\n",
-    "    |----------------|-----------|-----|-----------|---------------|\n",
-    "    | cell1          | ...       | ... | neuron    | healthy       |\n",
-    "    | cell2          | ...       | ... | astrocyte | diseased      |\n",
-    "    | ...            | ...       | ... | ...       | ...           |\n",
-    "The model will learn to predict classes within 'cell_type' and 'disease_state' "
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "b9e3050a-6162-4c01-b6fd-8784bf4ab1e4",
-   "metadata": {},
-   "source": [
-    "## 3. Initialize the MTLClassifier\n",
-    "\n",
-    "Now, let's create an instance of the MTLClassifier with our defined parameters and task columns."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "e27caac9-670c-409d-9313-50201c665cb9",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "mc = MTLClassifier(\n",
-    "    task_columns=task_columns,  # Our defined classification tasks\n",
-    "    study_name=\"MTLClassifier_example\",\n",
-    "    pretrained_path=pretrained_path,\n",
-    "    train_path=train_path,\n",
-    "    val_path=val_path,\n",
-    "    test_path=test_path,\n",
-    "    model_save_path=model_save_path,\n",
-    "    results_dir=results_dir,\n",
-    "    tensorboard_log_dir=tensorboard_log_dir,\n",
-    "    hyperparameters=hyperparameters,\n",
-    "    n_trials=15,  # Number of trials for hyperparameter optimization (at least 50 suggested)\n",
-    "    epochs=1,    # Number of training epochs (1 suggested to prevent overfitting)\n",
-    "    batch_size=8, # Adjust based on available GPU memory\n",
-    "    seed=42\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "0d729444-e3ad-4584-9659-0c464ac97462",
-   "metadata": {},
-   "source": [
-    "## 4. Run Hyperparameter Optimization\n",
-    "\n",
-    "Now, let's run the Optuna study to optimize our hyperparameters for both classification tasks."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "9298aa3e-6a52-4aa8-b9ff-b63d97beac93",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "mc.run_optuna_study()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "af23075d-d07b-43d3-bc5d-4df4d5d7199b",
-   "metadata": {},
-   "source": [
-    "## 5. Evaluate the Model on Test Data\n",
-    "\n",
-    "After optimization, we can evaluate our model on the test dataset. This will provide performance metrics for both classification tasks. CSV containing following keys will be generated in specified results directiory \"Cell ID, task(1...n) True,task(1.,.n) Pred,task(1...n) Probabilities\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "461bf8d3-b964-4ff4-994f-9f3d313d4614",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "mc.load_and_evaluate_test_model()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "31cfeb2d-6673-4b02-a79c-2533cc5e4d28",
-   "metadata": {},
-   "source": [
-    "## 6. (Optional) Manual Hyperparameter Tuning\n",
-    "\n",
-    "If you prefer to set hyperparameters manually, you can use the following approach:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "8ee6b99f-42e9-4abf-a292-aa9047735e0e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "manual_hyperparameters = {\n",
-    "    \"learning_rate\": 0.001,\n",
-    "    \"warmup_ratio\": 0.01,\n",
-    "    \"weight_decay\": 0.1,\n",
-    "    \"dropout_rate\": 0.1,\n",
-    "    \"lr_scheduler_type\": \"cosine\",\n",
-    "    \"task_weights\": [1, 1],  # Weights for each task (cell_type, disease_state)\n",
-    "    \"max_layers_to_freeze\": 2\n",
-    "}\n",
-    "\n",
-    "mc_manual = MTLClassifier(\n",
-    "    task_columns=task_columns,\n",
-    "    study_name=\"mtl_manual\",\n",
-    "    pretrained_path=pretrained_path,\n",
-    "    train_path=train_path,\n",
-    "    val_path=val_path,\n",
-    "    test_path=test_path,\n",
-    "    model_save_path=model_save_path,\n",
-    "    results_dir=results_dir,\n",
-    "    tensorboard_log_dir=tensorboard_log_dir,\n",
-    "    manual_hyperparameters=manual_hyperparameters,\n",
-    "    use_manual_hyperparameters=True,\n",
-    "    epochs=10,\n",
-    "    batch_size=32,\n",
-    "    seed=42\n",
-    ")\n",
-    "\n",
-    "mc_manual.run_manual_tuning()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "dbaac008-fc00-4b71-8e78-89b2d922d9d8",
-   "metadata": {},
-   "source": [
-    "# Geneformer In Silico Perturber Tutorial (MTL Quantized)\n",
-    "This demonstrates how to use the Geneformer In Silico Perturber with a Multi-Task Learning (MTL) model in a quantized configuration to optimize runtime and memory."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "2e15ad57-736c-48f0-be87-39cf5015bc5c",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from geneformer import InSilicoPerturber, EmbExtractor, InSilicoPerturberStats"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "43c18140-151e-4d44-95b4-a9b3a47172cf",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Define paths\n",
-    "model_directory = \"/path/to/model/save/path\"\n",
-    "input_data_file = \"/path/to/input/data.dataset\"\n",
-    "output_directory = \"/path/to/output/directory\"\n",
-    "output_prefix = \"mtl_quantized_perturbation\"\n",
-    "\n",
-    "# Define parameters\n",
-    "perturb_type = \"delete\"  # or \"overexpress\"\n",
-    "\n",
-    "# Define cell states to model\n",
-    "cell_states_to_model = {\n",
-    "    \"state_key\": \"disease_state\", \n",
-    "    \"start_state\": \"disease\", \n",
-    "    \"goal_state\": \"control\"\n",
-    "}\n",
-    "\n",
-    "# Define filter data\n",
-    "filter_data_dict = {\n",
-    "    \"cell_type\": [\"Fibroblast\"]\n",
-    "}"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "3010d0bf-b23c-45c1-ac12-8c472dc8b7a1",
-   "metadata": {},
-   "source": [
-    "## 3. Extract State Embeddings\n",
-    "\n",
-    "Before we initialize the InSilicoPerturber, we need to extract the state embeddings using the EmbExtractor."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "215f0a90-8041-417d-a5d3-b2483626c3b2",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Initialize EmbExtractor\n",
-    "embex = EmbExtractor(\n",
-    "    filter_data_dict=filter_data_dict,\n",
-    "    max_ncells=1000, # Number of cells to extract embeddings for\n",
-    "    emb_layer=0,  # Use the second to last layer\n",
-    "    emb_mode = \"cls\",\n",
-    "    summary_stat=\"exact_mean\",\n",
-    "    forward_batch_size=8, # Adjust based on available GPU memory\n",
-    "    nproc=4\n",
-    ")\n",
-    "\n",
-    "# Extract state embeddings\n",
-    "state_embs_dict = embex.get_state_embs(\n",
-    "    cell_states_to_model,\n",
-    "    model_directory=model_directory,\n",
-    "    input_data_file=input_data_file,\n",
-    "    output_directory=output_directory,\n",
-    "    output_prefix=output_prefix\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "23f14e36-4529-4fb2-8af9-7f4875cf81e3",
-   "metadata": {},
-   "source": [
-    "## 4. Initialize the InSilicoPerturber\n",
-    "\n",
-    "Now that we have our state embeddings, let's create an instance of the InSilicoPerturber with MTL and quantized configurations."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "09f985a1-91bc-4e8d-8001-a3663531b570",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Initialize InSilicoPerturber\n",
-    "isp = InSilicoPerturber(\n",
-    "    perturb_type=perturb_type,\n",
-    "    genes_to_perturb=\"all\",  # Perturb all genes\n",
-    "    model_type=\"MTLCellClassifier-Quantized\",  # Use quantized MTL model\n",
-    "    emb_mode=\"cls\",  # Use CLS token embedding\n",
-    "    cell_states_to_model=cell_states_to_model,\n",
-    "    state_embs_dict=state_embs_dict,\n",
-    "    max_ncells=1000,  # Number of cells to perturb (larger number increases power)\n",
-    "    emb_layer=0,  \n",
-    "    forward_batch_size=8, # Adjust based on available GPU memory\n",
-    "    nproc=1\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "cfcc2c1e-fd7f-4a36-99fc-ac7f43e5be6b",
-   "metadata": {},
-   "source": [
-    "## 5. Run In Silico Perturbation\n",
-    "\n",
-    "Run the in silico perturbation on the dataset."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "cf030c09-8ae4-45a7-aaf7-3fc2af4fe296",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Run perturbation and output intermediate files\n",
-    "isp.perturb_data(\n",
-    "    model_directory=model_directory,\n",
-    "    input_data_file=input_data_file,\n",
-    "    output_directory=output_directory,\n",
-    "    output_prefix=output_prefix\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "bb8ec074-6f2f-422b-a973-37ed32a15c38",
-   "metadata": {},
-   "source": [
-    "## 6. Process Results with InSilicoPerturberStats\n",
-    "\n",
-    "After running the perturbation, we'll use InSilicoPerturberStats to process the intermediate files and generate the final statistics."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "0a748043-43fc-47ad-ace5-f0ae3dd34674",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Initialize InSilicoPerturberStats\n",
-    "ispstats = InSilicoPerturberStats(\n",
-    "    mode=\"goal_state_shift\",\n",
-    "    genes_perturbed=\"all\",\n",
-    "    combos=0,\n",
-    "    anchor_gene=None,\n",
-    "    cell_states_to_model=cell_states_to_model\n",
-    ")\n",
-    "\n",
-    "# Process stats and output final .csv\n",
-    "ispstats.get_stats(\n",
-    "    input_data_file,\n",
-    "    None,\n",
-    "    output_directory,\n",
-    "    output_prefix\n",
-    ")"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.5"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

examples/pretraining_new_model/pretrain_geneformer_w_deepspeed.py

@@ -138,9 +138,7 @@ 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,

examples/tokenizing_scRNAseq_data.ipynb

@@ -25,11 +25,7 @@
     "\n",
     "#### Additionally, if the original .loom file contains a cell column attribute called \"filter_pass\", this column will be used as a binary indicator of whether to include these cells in the tokenized data. All cells with \"1\" in this attribute will be tokenized, whereas the others will be excluded. One may use this column to indicate QC filtering or other criteria for selection for inclusion in the final tokenized dataset.\n",
     "\n",
-    "#### If one's data is in other formats besides .loom or .h5ad, one can use the relevant tools (such as Anndata tools) to convert the file to a .loom or .h5ad format prior to running the transcriptome tokenizer.\n",
-    "\n",
-    "#### OF NOTE: PLEASE ENSURE THE CORRECT TOKEN DICTIONARY AND GENE MEDIAN FILE IS USED FOR THE CORRECT MODEL.\n",
-    "\n",
-    "#### The 95M model series also require the special_token argument to be set to True and model_input_size to be 4096."
+    "#### If one's data is in other formats besides .loom or .h5ad, one can use the relevant tools (such as Anndata tools) to convert the file to a .loom or .h5ad format prior to running the transcriptome tokenizer."
    ]
   },
   {

fine_tuned_models/gf-12L-95M-i4096_MTLCellClassifier_CELLxGENE_240522/config.json

@@ -1,24 +0,0 @@
-{
-  "architectures": [
-    "BertForMaskedLM"
-  ],
-  "attention_probs_dropout_prob": 0.02,
-  "classifier_dropout": null,
-  "hidden_act": "relu",
-  "hidden_dropout_prob": 0.02,
-  "hidden_size": 512,
-  "initializer_range": 0.02,
-  "intermediate_size": 1024,
-  "layer_norm_eps": 1e-12,
-  "max_position_embeddings": 4096,
-  "model_type": "bert",
-  "num_attention_heads": 8,
-  "num_hidden_layers": 12,
-  "pad_token_id": 0,
-  "position_embedding_type": "absolute",
-  "torch_dtype": "float32",
-  "transformers_version": "4.37.2",
-  "type_vocab_size": 2,
-  "use_cache": true,
-  "vocab_size": 20275
-}

fine_tuned_models/gf-12L-95M-i4096_MTLCellClassifier_CELLxGENE_240522/pytorch_model.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:07b28d8c7bb789d59755c42d32f6182cc04d2cf34aafaa6397aa50e4fdf1a9b4
-size 152363342

geneformer/__init__.py

@@ -1,14 +1,4 @@
 # ruff: noqa: F401
-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"
-TOKEN_DICTIONARY_FILE = Path(__file__).parent / "token_dictionary_gc95M.pkl"
-ENSEMBL_DICTIONARY_FILE = Path(__file__).parent / "gene_name_id_dict_gc95M.pkl"
-ENSEMBL_MAPPING_FILE = Path(__file__).parent / "ensembl_mapping_dict_gc95M.pkl"
-
 from . import (
     collator_for_classification,
     emb_extractor,
@@ -21,7 +11,7 @@ from .collator_for_classification import (
     DataCollatorForCellClassification,
     DataCollatorForGeneClassification,
 )
-from .emb_extractor import EmbExtractor, get_embs
+from .emb_extractor import EmbExtractor
 from .in_silico_perturber import InSilicoPerturber
 from .in_silico_perturber_stats import InSilicoPerturberStats
 from .pretrainer import GeneformerPretrainer
@@ -29,6 +19,3 @@ from .tokenizer import TranscriptomeTokenizer
 
 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

geneformer/classifier.py

@@ -53,18 +53,16 @@ from pathlib import Path
 import numpy as np
 import pandas as pd
 import seaborn as sns
+from sklearn.model_selection import StratifiedKFold
 from tqdm.auto import tqdm, trange
 from transformers import Trainer
 from transformers.training_args import TrainingArguments
 
-from . import (
-    TOKEN_DICTIONARY_FILE,
-    DataCollatorForCellClassification,
-    DataCollatorForGeneClassification,
-)
+from . import DataCollatorForCellClassification, DataCollatorForGeneClassification
 from . import classifier_utils as cu
 from . import evaluation_utils as eu
 from . import perturber_utils as pu
+from .tokenizer import TOKEN_DICTIONARY_FILE
 
 sns.set()
 
@@ -75,7 +73,6 @@ logger = logging.getLogger(__name__)
 class Classifier:
     valid_option_dict = {
         "classifier": {"cell", "gene"},
-        "quantize": {bool, dict},
         "cell_state_dict": {None, dict},
         "gene_class_dict": {None, dict},
         "filter_data": {None, dict},
@@ -89,7 +86,6 @@ class Classifier:
         "no_eval": {bool},
         "stratify_splits_col": {None, str},
         "forward_batch_size": {int},
-        "token_dictionary_file": {None, str},
         "nproc": {int},
         "ngpu": {int},
     }
@@ -97,7 +93,6 @@ class Classifier:
     def __init__(
         self,
         classifier=None,
-        quantize=False,
         cell_state_dict=None,
         gene_class_dict=None,
         filter_data=None,
@@ -112,7 +107,6 @@ class Classifier:
         stratify_splits_col=None,
         no_eval=False,
         forward_batch_size=100,
-        token_dictionary_file=None,
         nproc=4,
         ngpu=1,
     ):
@@ -123,13 +117,6 @@ class Classifier:
 
         classifier : {"cell", "gene"}
             | Whether to fine-tune a cell state or gene classifier.
-        quantize : bool, dict
-            | Whether to fine-tune a quantized model.
-            | If True and no config provided, will use default.
-            | Will use custom config if provided.
-            | Configs should be provided as dictionary of BitsAndBytesConfig (transformers) and LoraConfig (peft).
-            | For example: {"bnb_config": BitsAndBytesConfig(...),
-            |               "peft_config": LoraConfig(...)}
         cell_state_dict : None, dict
             | Cell states to fine-tune model to distinguish.
             | Two-item dictionary with keys: state_key and states
@@ -188,9 +175,6 @@ class Classifier:
             | Otherwise, will perform eval during training.
         forward_batch_size : int
             | Batch size for forward pass (for evaluation, not training).
-        token_dictionary_file : None, str
-            | Default is to use token dictionary file from Geneformer
-            | Otherwise, will load custom gene token dictionary.
         nproc : int
             | Number of CPU processes to use.
         ngpu : int
@@ -199,11 +183,6 @@ class Classifier:
         """
 
         self.classifier = classifier
-        if self.classifier == "cell":
-            self.model_type = "CellClassifier"
-        elif self.classifier == "gene":
-            self.model_type = "GeneClassifier"
-        self.quantize = quantize
         self.cell_state_dict = cell_state_dict
         self.gene_class_dict = gene_class_dict
         self.filter_data = filter_data
@@ -222,7 +201,6 @@ class Classifier:
         self.stratify_splits_col = stratify_splits_col
         self.no_eval = no_eval
         self.forward_batch_size = forward_batch_size
-        self.token_dictionary_file = token_dictionary_file
         self.nproc = nproc
         self.ngpu = ngpu
 
@@ -244,9 +222,7 @@ class Classifier:
                 ] = self.cell_state_dict["states"]
 
         # load token dictionary (Ensembl IDs:token)
-        if self.token_dictionary_file is None:
-            self.token_dictionary_file = TOKEN_DICTIONARY_FILE
-        with open(self.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 = {v: k for k, v in self.gene_token_dict.items()}
@@ -269,7 +245,7 @@ class Classifier:
                     f"Genes to classify {missing_genes} are not in token dictionary."
                 )
             self.gene_class_dict = {
-                k: list(set([self.gene_token_dict.get(gene) for gene in v]))
+                k: set([self.gene_token_dict.get(gene) for gene in v])
                 for k, v in self.gene_class_dict.items()
             }
             empty_classes = []
@@ -291,7 +267,7 @@ class Classifier:
                     continue
             valid_type = False
             for option in valid_options:
-                if (option in [int, float, list, dict, bool, str]) and isinstance(
+                if (option in [int, float, list, dict, bool]) and isinstance(
                     attr_value, option
                 ):
                     valid_type = True
@@ -417,15 +393,6 @@ class Classifier:
                 )
                 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."
-            )
-            raise
-
-        if not isinstance(attr_to_balance, list):
-            attr_to_balance = [attr_to_balance]
-
         if self.classifier == "cell":
             # remove cell states representing < rare_threshold of cells
             data = cu.remove_rare(
@@ -467,8 +434,8 @@ class Classifier:
             test_data_output_path = (
                 Path(output_directory) / f"{output_prefix}_labeled_test"
             ).with_suffix(".dataset")
-            data_dict["train"].save_to_disk(str(train_data_output_path))
-            data_dict["test"].save_to_disk(str(test_data_output_path))
+            data_dict["train"].save_to_disk(train_data_output_path)
+            data_dict["test"].save_to_disk(test_data_output_path)
         elif (test_size is not None) and (self.classifier == "cell"):
             if 1 > test_size > 0:
                 if attr_to_split is None:
@@ -483,8 +450,8 @@ class Classifier:
                     test_data_output_path = (
                         Path(output_directory) / f"{output_prefix}_labeled_test"
                     ).with_suffix(".dataset")
-                    data_dict["train"].save_to_disk(str(train_data_output_path))
-                    data_dict["test"].save_to_disk(str(test_data_output_path))
+                    data_dict["train"].save_to_disk(train_data_output_path)
+                    data_dict["test"].save_to_disk(test_data_output_path)
                 else:
                     data_dict, balance_df = cu.balance_attr_splits(
                         data,
@@ -505,19 +472,19 @@ class Classifier:
                     test_data_output_path = (
                         Path(output_directory) / f"{output_prefix}_labeled_test"
                     ).with_suffix(".dataset")
-                    data_dict["train"].save_to_disk(str(train_data_output_path))
-                    data_dict["test"].save_to_disk(str(test_data_output_path))
+                    data_dict["train"].save_to_disk(train_data_output_path)
+                    data_dict["test"].save_to_disk(test_data_output_path)
             else:
                 data_output_path = (
                     Path(output_directory) / f"{output_prefix}_labeled"
                 ).with_suffix(".dataset")
-                data.save_to_disk(str(data_output_path))
+                data.save_to_disk(data_output_path)
                 print(data_output_path)
         else:
             data_output_path = (
                 Path(output_directory) / f"{output_prefix}_labeled"
             ).with_suffix(".dataset")
-            data.save_to_disk(str(data_output_path))
+            data.save_to_disk(data_output_path)
 
     def train_all_data(
         self,
@@ -527,7 +494,6 @@ class Classifier:
         output_directory,
         output_prefix,
         save_eval_output=True,
-        gene_balance=False,
     ):
         """
         Train cell state or gene classifier using all data.
@@ -548,9 +514,6 @@ class Classifier:
         save_eval_output : bool
             | Whether to save cross-fold eval output
             | Saves as pickle file of dictionary of eval metrics
-        gene_balance : None, bool
-            | Whether to automatically balance genes in training set.
-            | Only available for binary gene classifications.
 
         **Output**
 
@@ -558,12 +521,6 @@ class Classifier:
 
         """
 
-        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:
@@ -595,7 +552,7 @@ class Classifier:
             )
             assert len(targets) == len(labels)
             data = cu.prep_gene_classifier_all_data(
-                data, targets, labels, self.max_ncells, self.nproc, gene_balance
+                data, targets, labels, self.max_ncells, self.nproc
             )
 
         trainer = self.train_classifier(
@@ -614,15 +571,12 @@ class Classifier:
         split_id_dict=None,
         attr_to_split=None,
         attr_to_balance=None,
-        gene_balance=False,
         max_trials=100,
         pval_threshold=0.1,
         save_eval_output=True,
         predict_eval=True,
         predict_trainer=False,
         n_hyperopt_trials=0,
-        save_gene_split_datasets=True,
-        debug_gene_split_datasets=False,
     ):
         """
         (Cross-)validate cell state or gene classifier.
@@ -657,9 +611,6 @@ class Classifier:
         attr_to_balance : None, list
             | List of attribute keys on which to balance data while splitting on attr_to_split
             | e.g. ["age", "sex"] for balancing these characteristics while splitting by patient
-        gene_balance : None, bool
-            | Whether to automatically balance genes in training set.
-            | Only available for binary gene classifications.
         max_trials : None, int
             | Maximum number of trials of random splitting to try to achieve balanced other attribute
             | If no split is found without significant (p < pval_threshold) differences in other attributes, will select best
@@ -678,19 +629,12 @@ class Classifier:
         n_hyperopt_trials : int
             | Number of trials to run for hyperparameter optimization
             | If 0, will not optimize hyperparameters
-        save_gene_split_datasets : bool
-            | Whether or not to save train, valid, and test gene-labeled datasets
         """
+
         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."
-            )
-            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]
@@ -771,7 +715,7 @@ class Classifier:
                 else:
                     # 5-fold cross-validate
                     num_cells = len(data)
-                    fifth_cells = int(np.floor(num_cells * 0.2))
+                    fifth_cells = num_cells * 0.2
                     num_eval = min((self.eval_size * num_cells), fifth_cells)
                     start = i * fifth_cells
                     end = start + num_eval
@@ -828,20 +772,17 @@ class Classifier:
                 ]
             )
             assert len(targets) == len(labels)
-            n_splits = int(1 / (1 - self.train_size))
-            skf = cu.StratifiedKFold3(n_splits=n_splits, random_state=0, shuffle=True)
+            n_splits = int(1 / self.eval_size)
+            skf = StratifiedKFold(n_splits=n_splits, random_state=0, shuffle=True)
             # (Cross-)validate
-            test_ratio = self.oos_test_size / (self.eval_size + self.oos_test_size)
-            for train_index, eval_index, test_index in tqdm(
-                skf.split(targets, labels, test_ratio)
-            ):
+            for train_index, eval_index in tqdm(skf.split(targets, labels)):
                 print(
                     f"****** Validation split: {iteration_num}/{self.num_crossval_splits} ******\n"
                 )
                 ksplit_output_dir = os.path.join(output_dir, f"ksplit{iteration_num}")
                 # filter data for examples containing classes for this split
                 # subsample to max_ncells and relabel data in column "labels"
-                train_data, eval_data = cu.prep_gene_classifier_train_eval_split(
+                train_data, eval_data = cu.prep_gene_classifier_split(
                     data,
                     targets,
                     labels,
@@ -850,42 +791,8 @@ class Classifier:
                     self.max_ncells,
                     iteration_num,
                     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")
-                        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,
-                        targets,
-                        labels,
-                        test_index,
-                        "test",
-                        self.max_ncells,
-                        iteration_num,
-                        self.nproc,
-                    )
-                    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")
-                        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."
-                    )
-                    raise
                 if n_hyperopt_trials == 0:
                     trainer = self.train_classifier(
                         model_directory,
@@ -895,15 +802,6 @@ class Classifier:
                         ksplit_output_dir,
                         predict_trainer,
                     )
-                    result = self.evaluate_model(
-                        trainer.model,
-                        num_classes,
-                        id_class_dict,
-                        eval_data,
-                        predict_eval,
-                        ksplit_output_dir,
-                        output_prefix,
-                    )
                 else:
                     trainer = self.hyperopt_classifier(
                         model_directory,
@@ -913,27 +811,20 @@ class Classifier:
                         ksplit_output_dir,
                         n_trials=n_hyperopt_trials,
                     )
-
-                    model = cu.load_best_model(
-                        ksplit_output_dir, self.model_type, num_classes
-                    )
-
-                    if self.oos_test_size > 0:
-                        result = self.evaluate_model(
-                            model,
-                            num_classes,
-                            id_class_dict,
-                            test_data,
-                            predict_eval,
-                            ksplit_output_dir,
-                            output_prefix,
-                        )
+                    if iteration_num == self.num_crossval_splits:
+                        return
                     else:
-                        if iteration_num == self.num_crossval_splits:
-                            return
-                        else:
-                            iteration_num = iteration_num + 1
-                            continue
+                        iteration_num = iteration_num + 1
+                        continue
+                result = self.evaluate_model(
+                    trainer.model,
+                    num_classes,
+                    id_class_dict,
+                    eval_data,
+                    predict_eval,
+                    ksplit_output_dir,
+                    output_prefix,
+                )
                 results += [result]
                 all_conf_mat = all_conf_mat + result["conf_mat"]
                 # break after 1 or 5 splits, each with train/eval proportions dictated by eval_size
@@ -1034,13 +925,12 @@ 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,
-        )
+        if self.classifier == "cell":
+            model_type = "CellClassifier"
+        elif self.classifier == "gene":
+            model_type = "GeneClassifier"
+
+        model = pu.load_model(model_type, num_classes, model_directory, "train")
         def_training_args, def_freeze_layers = cu.get_default_train_args(
             model, self.classifier, train_data, output_directory
         )
@@ -1056,31 +946,18 @@ class Classifier:
         if eval_data is None:
             def_training_args["evaluation_strategy"] = "no"
             def_training_args["load_best_model_at_end"] = False
-        def_training_args.update(
-            {"save_strategy": "epoch", "save_total_limit": 1}
-        )  # only save last model for each run
         training_args_init = TrainingArguments(**def_training_args)
 
         ##### Fine-tune the model #####
         # define the data collator
         if self.classifier == "cell":
-            data_collator = DataCollatorForCellClassification(
-                token_dictionary=self.gene_token_dict
-            )
+            data_collator = DataCollatorForCellClassification()
         elif self.classifier == "gene":
-            data_collator = DataCollatorForGeneClassification(
-                token_dictionary=self.gene_token_dict
-            )
+            data_collator = DataCollatorForGeneClassification()
 
         # define function to initiate model
         def model_init():
-            model = pu.load_model(
-                self.model_type,
-                num_classes,
-                model_directory,
-                "train",
-                quantize=self.quantize,
-            )
+            model = pu.load_model(model_type, num_classes, model_directory, "train")
 
             if self.freeze_layers is not None:
                 def_freeze_layers = self.freeze_layers
@@ -1091,8 +968,7 @@ class Classifier:
                     for param in module.parameters():
                         param.requires_grad = False
 
-            if self.quantize is False:
-                model = model.to("cuda:0")
+            model = model.to("cuda:0")
             return model
 
         # create the trainer
@@ -1142,7 +1018,6 @@ class Classifier:
                 metric="eval_macro_f1",
                 metric_columns=["loss", "eval_loss", "eval_accuracy", "eval_macro_f1"],
             ),
-            storage_path=output_directory,
         )
 
         return trainer
@@ -1205,13 +1080,11 @@ 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,
-        )
+        if self.classifier == "cell":
+            model_type = "CellClassifier"
+        elif self.classifier == "gene":
+            model_type = "GeneClassifier"
+        model = pu.load_model(model_type, num_classes, model_directory, "train")
 
         def_training_args, def_freeze_layers = cu.get_default_train_args(
             model, self.classifier, train_data, output_directory
@@ -1241,13 +1114,9 @@ class Classifier:
         ##### Fine-tune the model #####
         # define the data collator
         if self.classifier == "cell":
-            data_collator = DataCollatorForCellClassification(
-                token_dictionary=self.gene_token_dict
-            )
+            data_collator = DataCollatorForCellClassification()
         elif self.classifier == "gene":
-            data_collator = DataCollatorForGeneClassification(
-                token_dictionary=self.gene_token_dict
-            )
+            data_collator = DataCollatorForGeneClassification()
 
         # create the trainer
         trainer = Trainer(
@@ -1369,13 +1238,11 @@ 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,
-        )
+        if self.classifier == "cell":
+            model_type = "CellClassifier"
+        elif self.classifier == "gene":
+            model_type = "GeneClassifier"
+        model = pu.load_model(model_type, num_classes, model_directory, "eval")
 
         # evaluate the model
         result = self.evaluate_model(

geneformer/classifier_utils.py

@@ -1,6 +1,4 @@
-import json
 import logging
-import os
 import random
 from collections import Counter, defaultdict
 
@@ -8,7 +6,6 @@ import numpy as np
 import pandas as pd
 from scipy.stats import chisquare, ranksums
 from sklearn.metrics import accuracy_score, f1_score
-from sklearn.model_selection import StratifiedKFold, train_test_split
 
 from . import perturber_utils as pu
 
@@ -136,98 +133,64 @@ 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,
-):
-    # generate cross-validation splits
-    train_data = prep_gene_classifier_split(
-        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,
-    )
-    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, train_index, eval_index, max_ncells, iteration_num, num_proc
 ):
     # generate cross-validation splits
     targets = np.array(targets)
     labels = np.array(labels)
-    targets_subset = targets[index]
-    labels_subset = labels[index]
-    label_dict_subset = dict(zip(targets_subset, labels_subset))
+    targets_train, targets_eval = targets[train_index], targets[eval_index]
+    labels_train, labels_eval = labels[train_index], labels[eval_index]
+    label_dict_train = dict(zip(targets_train, labels_train))
+    label_dict_eval = dict(zip(targets_eval, labels_eval))
 
     # function to filter by whether contains train or eval labels
-    def if_contains_subset_label(example):
-        a = targets_subset
+    def if_contains_train_label(example):
+        a = targets_train
+        b = example["input_ids"]
+        return not set(a).isdisjoint(b)
+
+    def if_contains_eval_label(example):
+        a = targets_eval
         b = example["input_ids"]
         return not set(a).isdisjoint(b)
 
     # filter dataset for examples containing classes for this split
-    logger.info(f"Filtering data for {subset_name} genes in split {iteration_num}")
-    subset_data = data.filter(if_contains_subset_label, num_proc=num_proc)
+    logger.info(f"Filtering training data for genes in split {iteration_num}")
+    train_data = data.filter(if_contains_train_label, num_proc=num_proc)
     logger.info(
-        f"Filtered {round((1-len(subset_data)/len(data))*100)}%; {len(subset_data)} remain\n"
+        f"Filtered {round((1-len(train_data)/len(data))*100)}%; {len(train_data)} remain\n"
+    )
+    logger.info(f"Filtering evalation data for genes in split {iteration_num}")
+    eval_data = data.filter(if_contains_eval_label, num_proc=num_proc)
+    logger.info(
+        f"Filtered {round((1-len(eval_data)/len(data))*100)}%; {len(eval_data)} remain\n"
     )
-
-    # 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
-        )
 
     # subsample to max_ncells
-    subset_data = downsample_and_shuffle(subset_data, max_ncells, None, None)
+    train_data = downsample_and_shuffle(train_data, max_ncells, None, None)
+    eval_data = downsample_and_shuffle(eval_data, max_ncells, None, None)
 
     # relabel genes for this split
-    def subset_classes_to_ids(example):
+    def train_classes_to_ids(example):
         example["labels"] = [
-            label_dict_subset.get(token_id, -100) for token_id in example["input_ids"]
+            label_dict_train.get(token_id, -100) for token_id in example["input_ids"]
         ]
         return example
 
-    subset_data = subset_data.map(subset_classes_to_ids, num_proc=num_proc)
+    def eval_classes_to_ids(example):
+        example["labels"] = [
+            label_dict_eval.get(token_id, -100) for token_id in example["input_ids"]
+        ]
+        return example
 
-    return subset_data
+    train_data = train_data.map(train_classes_to_ids, num_proc=num_proc)
+    eval_data = eval_data.map(eval_classes_to_ids, num_proc=num_proc)
 
+    return train_data, eval_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):
     targets = np.array(targets)
     labels = np.array(labels)
     label_dict_train = dict(zip(targets, labels))
@@ -245,11 +208,6 @@ def prep_gene_classifier_all_data(
         f"Filtered {round((1-len(train_data)/len(data))*100)}%; {len(train_data)} remain\n"
     )
 
-    if balance is True:
-        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)
 
@@ -265,145 +223,6 @@ def prep_gene_classifier_all_data(
     return train_data
 
 
-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:
-        # 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_dict_subset_orig = label_dict_subset.copy()
-        # balance gene sets
-        max_ntrials = 25
-        boost = 1
-        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_ngenes = len(label0_genes)
-                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:
-                    # 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
-        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_ngenes = len(label1_genes)
-                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:
-                    # 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
-        ):
-            label_ratio_0to1 = label_ratio_0to1_orig
-            label_dict_subset_new = label_dict_subset_orig
-        logger.warning(
-            f"Gene sets were not able to be balanced within 0.8-1.25 fold after {max_ntrials} trials. Imbalance level: {label_ratio_0to1}\n"
-        )
-        return filter_data_balanced_genes(subset_data, label_dict_subset_new, num_proc)
-
-
-def count_genes_for_balancing(subset_data, label_dict_subset, num_proc):
-    def count_targets(example):
-        labels = [
-            label_dict_subset.get(token_id, -100) for token_id in example["input_ids"]
-        ]
-        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)]
-        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"]])
-
-    subset_data = subset_data.remove_columns("labels_counts")
-
-    return label0_counts, label1_counts
-
-
-def filter_data_balanced_genes(subset_data, label_dict_subset, num_proc):
-    # function to filter by whether contains labels
-    def if_contains_subset_label(example):
-        a = list(label_dict_subset.keys())
-        b = example["input_ids"]
-        return not set(a).isdisjoint(b)
-
-    # filter dataset for examples containing classes for this split
-    logger.info("Filtering data for balanced genes")
-    subset_data_len_orig = len(subset_data)
-    subset_data = subset_data.filter(if_contains_subset_label, num_proc=num_proc)
-    logger.info(
-        f"Filtered {round((1-len(subset_data)/subset_data_len_orig)*100)}%; {len(subset_data)} remain\n"
-    )
-
-    return subset_data, label_dict_subset
-
-
 def balance_attr_splits(
     data,
     attr_to_split,
@@ -490,7 +309,7 @@ def balance_attr_splits(
                     exp_counts[cat] * sum(obs) / sum(exp_counts.values())
                     for cat in all_categ
                 ]
-                pval = chisquare(f_obs=obs, f_exp=exp).pvalue
+                chisquare(f_obs=obs, f_exp=exp).pvalue
                 train_attr_counts = str(obs_counts).strip("Counter(").strip(")")
                 eval_attr_counts = str(exp_counts).strip("Counter(").strip(")")
                 df_vals += [train_attr_counts, eval_attr_counts, pval]
@@ -604,45 +423,3 @@ def get_default_train_args(model, classifier, data, output_dir):
     training_args.update(default_training_args)
 
     return training_args, freeze_layers
-
-
-def load_best_model(directory, model_type, num_classes, mode="eval"):
-    file_dict = dict()
-    for subdir, dirs, files in os.walk(directory):
-        for file in files:
-            if file.endswith("result.json"):
-                with open(f"{subdir}/{file}", "rb") as fp:
-                    result_json = json.load(fp)
-                file_dict[f"{subdir}"] = result_json["eval_macro_f1"]
-    file_df = pd.DataFrame(
-        {"dir": file_dict.keys(), "eval_macro_f1": file_dict.values()}
-    )
-    model_superdir = (
-        "run-"
-        + file_df.iloc[file_df["eval_macro_f1"].idxmax()]["dir"]
-        .split("_objective_")[2]
-        .split("_")[0]
-    )
-
-    for subdir, dirs, files in os.walk(f"{directory}/{model_superdir}"):
-        for file in files:
-            if file.endswith("model.safetensors"):
-                model = pu.load_model(model_type, num_classes, f"{subdir}", mode)
-    return model
-
-
-class StratifiedKFold3(StratifiedKFold):
-    def split(self, targets, labels, test_ratio=0.5, groups=None):
-        s = super().split(targets, labels, groups)
-        for train_indxs, test_indxs in s:
-            if test_ratio == 0:
-                yield train_indxs, test_indxs, None
-            else:
-                labels_test = np.array(labels)[test_indxs]
-                valid_indxs, test_indxs = train_test_split(
-                    test_indxs,
-                    stratify=labels_test,
-                    test_size=test_ratio,
-                    random_state=0,
-                )
-                yield train_indxs, valid_indxs, test_indxs

geneformer/collator_for_classification.py

@@ -1,22 +1,24 @@
 """
 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 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
 
+from .pretrainer import token_dictionary
+
 EncodedInput = List[int]
 logger = logging.get_logger(__name__)
 VERY_LARGE_INTEGER = int(
@@ -28,7 +30,6 @@ LARGE_INTEGER = int(
 
 # precollator functions
 
-
 class ExplicitEnum(Enum):
     """
     Enum with more explicit error message for missing values.
@@ -41,7 +42,6 @@ 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
@@ -54,6 +54,7 @@ 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,6 +66,7 @@ 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
@@ -76,41 +78,21 @@ class TensorType(ExplicitEnum):
     NUMPY = "np"
     JAX = "jax"
 
-
+    
 class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(mask_token="<mask>", pad_token="<pad>")
-
-        self.token_dictionary = kwargs.get("token_dictionary")
-        self.padding_side = "right"
-        self.model_input_names = ["input_ids"]
-        self._mask_token_id = self.token_dictionary.get("<mask>")
-        self._pad_token_id = self.token_dictionary.get("<pad>")
-        self._all_special_ids = [
-            self.token_dictionary.get("<mask>"),
-            self.token_dictionary.get("<pad>"),
-        ]
-
-    @property
-    def all_special_ids(self):
-        return self._all_special_ids
-
-    @property
-    def mask_token_id(self):
-        return self._mask_token_id
-
-    @property
-    def pad_token_id(self):
-        return self._pad_token_id
+    mask_token = "<mask>"
+    mask_token_id = token_dictionary.get("<mask>")
+    pad_token = "<pad>"
+    pad_token_id = token_dictionary.get("<pad>")
+    padding_side = "right"
+    all_special_ids = [
+        token_dictionary.get("<mask>"),
+        token_dictionary.get("<pad>")
+    ]
+    model_input_names = ["input_ids"]
 
     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
@@ -123,9 +105,7 @@ 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. "
@@ -153,9 +133,7 @@ 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):
@@ -195,9 +173,7 @@ 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."
@@ -210,24 +186,18 @@ 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.)` "
@@ -258,7 +228,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,
@@ -269,23 +239,29 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
         """
         Pad a single encoded input or a batch of encoded inputs up to predefined length or to the max sequence length
         in the batch.
+
         Padding side (left/right) padding token ids are defined at the tokenizer level (with ``self.padding_side``,
         ``self.pad_token_id`` and ``self.pad_token_type_id``)
+
         .. note::
+
             If the ``encoded_inputs`` passed are dictionary of numpy arrays, PyTorch tensors or TensorFlow tensors, the
             result will use the same type unless you provide a different tensor type with ``return_tensors``. In the
             case of PyTorch tensors, you will lose the specific device of your tensors however.
+
         Args:
             encoded_inputs (:class:`~transformers.BatchEncoding`, list of :class:`~transformers.BatchEncoding`, :obj:`Dict[str, List[int]]`, :obj:`Dict[str, List[List[int]]` or :obj:`List[Dict[str, List[int]]]`):
                 Tokenized inputs. Can represent one input (:class:`~transformers.BatchEncoding` or :obj:`Dict[str,
                 List[int]]`) or a batch of tokenized inputs (list of :class:`~transformers.BatchEncoding`, `Dict[str,
                 List[List[int]]]` or `List[Dict[str, List[int]]]`) so you can use this method during preprocessing as
                 well as in a PyTorch Dataloader collate function.
+
                 Instead of :obj:`List[int]` you can have tensors (numpy arrays, PyTorch tensors or TensorFlow tensors),
                 see the note above for the return type.
             padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
                  Select a strategy to pad the returned sequences (according to the model's padding side and padding
                  index) among:
+
                 * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a
                   single sequence if provided).
                 * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
@@ -296,14 +272,17 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
                 Maximum length of the returned list and optionally padding length (see above).
             pad_to_multiple_of (:obj:`int`, `optional`):
                 If set will pad the sequence to a multiple of the provided value.
+
                 This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability
                 >= 7.5 (Volta).
             return_attention_mask (:obj:`bool`, `optional`):
                 Whether to return the attention mask. If left to the default, will return the attention mask according
                 to the specific tokenizer's default, defined by the :obj:`return_outputs` attribute.
+
                 `What are attention masks? <../glossary.html#attention-mask>`__
             return_tensors (:obj:`str` or :class:`~transformers.tokenization_utils_base.TensorType`, `optional`):
                 If set, will return tensors instead of list of python integers. Acceptable values are:
+
                 * :obj:`'tf'`: Return TensorFlow :obj:`tf.constant` objects.
                 * :obj:`'pt'`: Return PyTorch :obj:`torch.Tensor` objects.
                 * :obj:`'np'`: Return Numpy :obj:`np.ndarray` objects.
@@ -312,13 +291,8 @@ 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:
@@ -412,7 +386,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,
@@ -420,15 +394,18 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
     ) -> dict:
         """
         Pad encoded inputs (on left/right and up to predefined length or max length in the batch)
+
         Args:
             encoded_inputs: Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`).
             max_length: maximum length of the returned list and optionally padding length (see below).
                 Will truncate by taking into account the special tokens.
             padding_strategy: PaddingStrategy to use for padding.
+
                 - PaddingStrategy.LONGEST Pad to the longest sequence in the batch
                 - PaddingStrategy.MAX_LENGTH: Pad to the max length (default)
                 - PaddingStrategy.DO_NOT_PAD: Do not pad
                 The tokenizer padding sides are defined in self.padding_side:
+
                     - 'left': pads on the left of the sequences
                     - 'right': pads on the right of the sequences
             pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value.
@@ -445,73 +422,46 @@ 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
@@ -535,15 +485,11 @@ 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.
@@ -567,15 +513,14 @@ class PrecollatorForGeneAndCellClassification(SpecialTokensMixin):
         if token is None:
             return None
 
-        return self.token_dictionary.get(token)
+        return token_dictionary.get(token)
 
     def __len__(self):
-        return len(self.token_dictionary)
+        return len(token_dictionary)
 
 
 # collator functions
 
-
 class DataCollatorForGeneClassification(DataCollatorForTokenClassification):
     """
     Data collator that will dynamically pad the inputs received, as well as the labels.
@@ -601,33 +546,25 @@ class DataCollatorForGeneClassification(DataCollatorForTokenClassification):
             The id to use when padding the labels (-100 will be automatically ignore by PyTorch loss functions).
     """
 
+    tokenizer = PrecollatorForGeneAndCellClassification()
     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=self.tokenizer,
             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,
@@ -637,31 +574,29 @@ 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_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 TOKEN_DICTIONARY_FILE
 from . import perturber_utils as pu
+from .tokenizer import TOKEN_DICTIONARY_FILE
 
 logger = logging.getLogger(__name__)
 
@@ -38,8 +38,6 @@ def get_embs(
     layer_to_quant,
     pad_token_id,
     forward_batch_size,
-    token_gene_dict,
-    special_token=False,
     summary_stat=None,
     silent=False,
 ):
@@ -49,8 +47,10 @@ def get_embs(
     if summary_stat is None:
         embs_list = []
     elif summary_stat is not None:
-        # get # of emb dims
-        emb_dims = pu.get_model_emb_dims(model)
+        # test embedding extraction for example cell and extract # emb dims
+        example = filtered_input_data.select([i for i in range(1)])
+        example.set_format(type="torch")
+        emb_dims = test_emb(model, example["input_ids"], layer_to_quant)
         if emb_mode == "cell":
             # initiate tdigests for # of emb dims
             embs_tdigests = [TDigest() for _ in range(emb_dims)]
@@ -67,27 +67,6 @@ def get_embs(
                 k: [TDigest() for _ in range(emb_dims)] for k in gene_set
             }
 
-    # Check if CLS and EOS token is present in the token dictionary
-    cls_present = any("<cls>" in value for value in token_gene_dict.values())
-    eos_present = any("<eos>" in value for value in token_gene_dict.values())
-    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()}
-        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"
-    elif emb_mode == "cell":
-        if cls_present:
-            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."
-            )
-
     overall_max_len = 0
 
     for i in trange(0, total_batch_length, forward_batch_size, leave=(not silent)):
@@ -113,14 +92,7 @@ def get_embs(
         embs_i = outputs.hidden_states[layer_to_quant]
 
         if emb_mode == "cell":
-            if cls_present:
-                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:
-                    mean_embs = pu.mean_nonpadding_embs(non_cls_embs, original_lens - 1)
-            else:
-                mean_embs = pu.mean_nonpadding_embs(embs_i, original_lens)
+            mean_embs = pu.mean_nonpadding_embs(embs_i, original_lens)
             if summary_stat is None:
                 embs_list.append(mean_embs)
             elif summary_stat is not None:
@@ -149,12 +121,6 @@ def get_embs(
                         accumulate_tdigests(
                             embs_tdigests_dict[int(k)], dict_h[k], emb_dims
                         )
-                    del embs_h
-                    del dict_h
-        elif emb_mode == "cls":
-            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
@@ -165,7 +131,7 @@ def get_embs(
         torch.cuda.empty_cache()
 
     if summary_stat is None:
-        if (emb_mode == "cell") or (emb_mode == "cls"):
+        if emb_mode == "cell":
             embs_stack = torch.cat(embs_list, dim=0)
         elif emb_mode == "gene":
             embs_stack = pu.pad_tensor_list(
@@ -243,6 +209,14 @@ def tdigest_median(embs_tdigests, emb_dims):
     return [embs_tdigests[i].percentile(50) for i in range(emb_dims)]
 
 
+def test_emb(model, example, layer_to_quant):
+    with torch.no_grad():
+        outputs = model(input_ids=example.to("cuda"))
+
+    embs_test = outputs.hidden_states[layer_to_quant]
+    return embs_test.size()[2]
+
+
 def label_cell_embs(embs, downsampled_data, emb_labels):
     embs_df = pd.DataFrame(embs.cpu().numpy())
     if emb_labels is not None:
@@ -278,7 +252,7 @@ def label_gene_embs(embs, downsampled_data, token_gene_dict):
     return embs_df
 
 
-def plot_umap(embs_df, emb_dims, label, output_file, kwargs_dict, seed=0):
+def plot_umap(embs_df, emb_dims, label, output_file, kwargs_dict):
     only_embs_df = embs_df.iloc[:, :emb_dims]
     only_embs_df.index = pd.RangeIndex(0, only_embs_df.shape[0], name=None).astype(str)
     only_embs_df.columns = pd.RangeIndex(0, only_embs_df.shape[1], name=None).astype(
@@ -288,27 +262,15 @@ def plot_umap(embs_df, emb_dims, label, output_file, kwargs_dict, seed=0):
     obs_dict = {"cell_id": list(only_embs_df.index), f"{label}": list(embs_df[label])}
     adata = anndata.AnnData(X=only_embs_df, obs=obs_dict, var=vars_dict)
     sc.tl.pca(adata, svd_solver="arpack")
-    sc.pp.neighbors(adata, random_state=seed)
-    sc.tl.umap(adata, random_state=seed)
+    sc.pp.neighbors(adata)
+    sc.tl.umap(adata)
     sns.set(rc={"figure.figsize": (10, 10)}, font_scale=2.3)
     sns.set_style("white")
-    default_kwargs_dict = {"size": 200}
+    default_kwargs_dict = {"palette": "Set2", "size": 200}
     if kwargs_dict is not None:
         default_kwargs_dict.update(kwargs_dict)
 
-    cats = set(embs_df[label])
-
-    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),
-        )
-        plt.show()
-        plt.savefig(output_file, bbox_inches="tight")
+    sc.pl.umap(adata, color=label, save=output_file, **default_kwargs_dict)
 
 
 def gen_heatmap_class_colors(labels, df):
@@ -384,8 +346,7 @@ def plot_heatmap(embs_df, emb_dims, label, output_file, kwargs_dict):
             bbox_to_anchor=(0.5, 1),
             facecolor="white",
         )
-    plt.show()
-    logger.info(f"Output file: {output_file}")
+
     plt.savefig(output_file, bbox_inches="tight")
 
 
@@ -393,7 +354,7 @@ class EmbExtractor:
     valid_option_dict = {
         "model_type": {"Pretrained", "GeneClassifier", "CellClassifier"},
         "num_classes": {int},
-        "emb_mode": {"cls", "cell", "gene"},
+        "emb_mode": {"cell", "gene"},
         "cell_emb_style": {"mean_pool"},
         "gene_emb_style": {"mean_pool"},
         "filter_data": {None, dict},
@@ -402,7 +363,6 @@ class EmbExtractor:
         "emb_label": {None, list},
         "labels_to_plot": {None, list},
         "forward_batch_size": {int},
-        "token_dictionary_file": {None, str},
         "nproc": {int},
         "summary_stat": {None, "mean", "median", "exact_mean", "exact_median"},
     }
@@ -411,7 +371,7 @@ class EmbExtractor:
         self,
         model_type="Pretrained",
         num_classes=0,
-        emb_mode="cls",
+        emb_mode="cell",
         cell_emb_style="mean_pool",
         gene_emb_style="mean_pool",
         filter_data=None,
@@ -422,7 +382,7 @@ class EmbExtractor:
         forward_batch_size=100,
         nproc=4,
         summary_stat=None,
-        token_dictionary_file=None,
+        token_dictionary_file=TOKEN_DICTIONARY_FILE,
     ):
         """
         Initialize embedding extractor.
@@ -434,11 +394,10 @@ class EmbExtractor:
         num_classes : int
             | If model is a gene or cell classifier, specify number of classes it was trained to classify.
             | For the pretrained Geneformer model, number of classes is 0 as it is not a classifier.
-        emb_mode : {"cls", "cell", "gene"}
-            | Whether to output CLS, cell, or gene embeddings.
-            | CLS embeddings are cell embeddings derived from the CLS token in the front of the rank value encoding.
-        cell_emb_style : {"mean_pool"}
-            | Method for summarizing cell embeddings if not using CLS token.
+        emb_mode : {"cell", "gene"}
+            | Whether to output cell or gene embeddings.
+        cell_emb_style : "mean_pool"
+            | Method for summarizing cell embeddings.
             | Currently only option is mean pooling of gene embeddings for given cell.
         gene_emb_style : "mean_pool"
             | Method for summarizing gene embeddings.
@@ -473,7 +432,6 @@ class EmbExtractor:
             | Non-exact recommended if encountering memory constraints while generating goal embedding positions.
             | Non-exact is slower but more memory-efficient.
         token_dictionary_file : Path
-            | Default is the Geneformer token dictionary
             | Path to pickle file containing token dictionary (Ensembl ID:token).
 
         **Examples:**
@@ -486,6 +444,7 @@ class EmbExtractor:
             ...         emb_mode="cell",
             ...         filter_data={"cell_type":["cardiomyocyte"]},
             ...         max_ncells=1000,
+            ...         max_ncells_to_plot=1000,
             ...         emb_layer=-1,
             ...         emb_label=["disease", "cell_type"],
             ...         labels_to_plot=["disease", "cell_type"])
@@ -502,7 +461,6 @@ class EmbExtractor:
         self.emb_layer = emb_layer
         self.emb_label = emb_label
         self.labels_to_plot = labels_to_plot
-        self.token_dictionary_file = token_dictionary_file
         self.forward_batch_size = forward_batch_size
         self.nproc = nproc
         if (summary_stat is not None) and ("exact" in summary_stat):
@@ -515,8 +473,6 @@ class EmbExtractor:
         self.validate_options()
 
         # load token dictionary (Ensembl IDs:token)
-        if self.token_dictionary_file is None:
-            token_dictionary_file = TOKEN_DICTIONARY_FILE
         with open(token_dictionary_file, "rb") as f:
             self.gene_token_dict = pickle.load(f)
 
@@ -532,7 +488,7 @@ class EmbExtractor:
                     continue
             valid_type = False
             for option in valid_options:
-                if (option in [int, list, dict, bool, str]) and isinstance(
+                if (option in [int, list, dict, bool]) and isinstance(
                     attr_value, option
                 ):
                     valid_type = True
@@ -606,14 +562,13 @@ class EmbExtractor:
         )
         layer_to_quant = pu.quant_layers(model) + self.emb_layer
         embs = get_embs(
-            model=model,
-            filtered_input_data=downsampled_data,
-            emb_mode=self.emb_mode,
-            layer_to_quant=layer_to_quant,
-            pad_token_id=self.pad_token_id,
-            forward_batch_size=self.forward_batch_size,
-            token_gene_dict=self.token_gene_dict,
-            summary_stat=self.summary_stat,
+            model,
+            downsampled_data,
+            self.emb_mode,
+            layer_to_quant,
+            self.pad_token_id,
+            self.forward_batch_size,
+            self.summary_stat,
         )
 
         if self.emb_mode == "cell":
@@ -627,8 +582,6 @@ class EmbExtractor:
             elif self.summary_stat is not None:
                 embs_df = pd.DataFrame(embs).T
                 embs_df.index = [self.token_gene_dict[token] for token in embs_df.index]
-        elif self.emb_mode == "cls":
-            embs_df = label_cell_embs(embs, downsampled_data, self.emb_label)
 
         # save embeddings to output_path
         if cell_state is None:
@@ -637,17 +590,13 @@ class EmbExtractor:
 
         if self.exact_summary_stat == "exact_mean":
             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:255].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:255].median(axis="rows"), columns=[self.exact_summary_stat]
             ).T
 
         if cell_state is not None:
@@ -800,15 +749,15 @@ class EmbExtractor:
             logger.error("Plotting UMAP requires 'labels_to_plot'. ")
             raise
 
-        if max_ncells_to_plot is not None:
-            if max_ncells_to_plot > self.max_ncells:
-                max_ncells_to_plot = self.max_ncells
-                logger.warning(
-                    "max_ncells_to_plot must be <= max_ncells. "
-                    f"Changing max_ncells_to_plot to {self.max_ncells}."
-                )
-            elif max_ncells_to_plot < self.max_ncells:
-                embs = embs.sample(max_ncells_to_plot, axis=0)
+        if max_ncells_to_plot > self.max_ncells:
+            max_ncells_to_plot = self.max_ncells
+            logger.warning(
+                "max_ncells_to_plot must be <= max_ncells. "
+                f"Changing max_ncells_to_plot to {self.max_ncells}."
+            )
+
+        if (max_ncells_to_plot is not None) and (max_ncells_to_plot < self.max_ncells):
+            embs = embs.sample(max_ncells_to_plot, axis=0)
 
         if self.emb_label is None:
             label_len = 0
@@ -830,11 +779,11 @@ class EmbExtractor:
                         f"not present in provided embeddings dataframe."
                     )
                     continue
-                output_prefix_label = output_prefix + f"_umap_{label}"
+                output_prefix_label = "_" + output_prefix + f"_umap_{label}"
                 output_file = (
                     Path(output_directory) / output_prefix_label
                 ).with_suffix(".pdf")
-                plot_umap(embs, emb_dims, label, output_file, kwargs_dict)
+                plot_umap(embs, emb_dims, label, output_prefix_label, kwargs_dict)
 
         if plot_style == "heatmap":
             for label in self.labels_to_plot:

geneformer/ensembl_mapping_dict_gc95M.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:0819bcbd869cfa14279449b037eb9ed1d09a91310e77bd1a19d927465030e95c
-size 3957652

geneformer/evaluation_utils.py

@@ -20,20 +20,20 @@ from sklearn.metrics import (
 )
 from tqdm.auto import trange
 
-from . import TOKEN_DICTIONARY_FILE
 from .emb_extractor import make_colorbar
+from .tokenizer import TOKEN_DICTIONARY_FILE
 
 logger = logging.getLogger(__name__)
 
+# load token dictionary (Ensembl IDs:token)
+with open(TOKEN_DICTIONARY_FILE, "rb") as f:
+    gene_token_dict = pickle.load(f)
+
 
 def preprocess_classifier_batch(cell_batch, max_len, label_name):
     if max_len is None:
         max_len = max([len(i) for i in cell_batch["input_ids"]])
 
-    # load token dictionary (Ensembl IDs:token)
-    with open(TOKEN_DICTIONARY_FILE, "rb") as f:
-        gene_token_dict = pickle.load(f)
-
     def pad_label_example(example):
         example[label_name] = np.pad(
             example[label_name],

geneformer/gene_dictionaries_30m/ensembl_mapping_dict_gc30M.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:eac0fb0b3007267871b6305ac0003ceba19d4f28d85686cb9067ecf142787869
-size 584125

geneformer/gene_dictionaries_30m/gene_median_dictionary_gc30M.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b3b589bb5ec75040d05fc44dd6bf0184cf87f3c362cf158d196a6ed3b7fe5f39
-size 940965

geneformer/gene_dictionaries_30m/token_dictionary_gc30M.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ab9dc40973fa5224d77b793e2fd114cacf3d08423ed9c4c49caf0ba9c7f218f1
-size 788424

geneformer/gene_median_dictionary_gc95M.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a51c53f6a771d64508dfaf61529df70e394c53bd20856926117ae5d641a24bf5
-size 1512661

geneformer/gene_name_id_dict_gc95M.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8b0fd0521406ed18b2e341ef0acb5f53aa1a62457a07ca5840e1c142f46dd326
-size 2038812

geneformer/in_silico_perturber.py

@@ -38,17 +38,19 @@ import logging
 import os
 import pickle
 from collections import defaultdict
+from typing import List
 
+import seaborn as sns
 import torch
-from datasets import Dataset, disable_progress_bars
-from multiprocess import set_start_method
+from datasets import Dataset
 from tqdm.auto import trange
 
-from . import TOKEN_DICTIONARY_FILE
 from . import perturber_utils as pu
 from .emb_extractor import get_embs
+from .tokenizer import TOKEN_DICTIONARY_FILE
+
+sns.set()
 
-disable_progress_bars()
 
 logger = logging.getLogger(__name__)
 
@@ -60,9 +62,9 @@ class InSilicoPerturber:
         "genes_to_perturb": {"all", list},
         "combos": {0, 1},
         "anchor_gene": {None, str},
-        "model_type": {"Pretrained", "GeneClassifier", "CellClassifier", "MTLCellClassifier", "MTLCellClassifier-Quantized"},
+        "model_type": {"Pretrained", "GeneClassifier", "CellClassifier"},
         "num_classes": {int},
-        "emb_mode": {"cls", "cell", "cls_and_gene", "cell_and_gene"},
+        "emb_mode": {"cell", "cell_and_gene"},
         "cell_emb_style": {"mean_pool"},
         "filter_data": {None, dict},
         "cell_states_to_model": {None, dict},
@@ -70,7 +72,6 @@ class InSilicoPerturber:
         "max_ncells": {None, int},
         "cell_inds_to_perturb": {"all", dict},
         "emb_layer": {-1, 0},
-        "token_dictionary_file": {None, str},
         "forward_batch_size": {int},
         "nproc": {int},
     }
@@ -94,8 +95,7 @@ class InSilicoPerturber:
         emb_layer=-1,
         forward_batch_size=100,
         nproc=4,
-        token_dictionary_file=None,
-        clear_mem_ncells=1000,
+        token_dictionary_file=TOKEN_DICTIONARY_FILE,
     ):
         """
         Initialize in silico perturber.
@@ -130,16 +130,16 @@ class InSilicoPerturber:
             | ENSEMBL ID of gene to use as anchor in combination perturbations.
             | For example, if combos=1 and anchor_gene="ENSG00000148400":
             |     anchor gene will be perturbed in combination with each other gene.
-        model_type : {"Pretrained", "GeneClassifier", "CellClassifier", "MTLCellClassifier", "MTLCellClassifier-Quantized"}
-            | Whether model is the pretrained Geneformer or a fine-tuned gene, cell, or multitask cell classifier (+/- 8bit quantization).
+        model_type : {"Pretrained", "GeneClassifier", "CellClassifier"}
+            | Whether model is the pretrained Geneformer or a fine-tuned gene or cell classifier.
         num_classes : int
             | If model is a gene or cell classifier, specify number of classes it was trained to classify.
             | For the pretrained Geneformer model, number of classes is 0 as it is not a classifier.
-        emb_mode : {"cls", "cell", "cls_and_gene","cell_and_gene"}
-            | Whether to output impact of perturbation on CLS token, cell, and/or gene embeddings.
+        emb_mode : {"cell", "cell_and_gene"}
+            | Whether to output impact of perturbation on cell and/or gene embeddings.
             | Gene embedding shifts only available as compared to original cell, not comparing to goal state.
         cell_emb_style : "mean_pool"
-            | Method for summarizing cell embeddings if not using CLS token.
+            | Method for summarizing cell embeddings.
             | Currently only option is mean pooling of gene embeddings for given cell.
         filter_data : None, dict
             | Default is to use all input data for in silico perturbation study.
@@ -184,13 +184,7 @@ class InSilicoPerturber:
             | Number of CPU processes to use.
         token_dictionary_file : Path
             | Path to pickle file containing token dictionary (Ensembl ID:token).
-        clear_mem_ncells : int
-            | Clear memory every n cells.
         """
-        try:
-            set_start_method("spawn")
-        except RuntimeError:
-            pass
 
         self.perturb_type = perturb_type
         self.perturb_rank_shift = perturb_rank_shift
@@ -222,32 +216,14 @@ class InSilicoPerturber:
         self.emb_layer = emb_layer
         self.forward_batch_size = forward_batch_size
         self.nproc = nproc
-        self.token_dictionary_file = token_dictionary_file
-        self.clear_mem_ncells = clear_mem_ncells
 
         self.validate_options()
 
         # load token dictionary (Ensembl IDs:token)
-        if self.token_dictionary_file is None:
-            token_dictionary_file = TOKEN_DICTIONARY_FILE
         with open(token_dictionary_file, "rb") as f:
             self.gene_token_dict = pickle.load(f)
-        self.token_gene_dict = {v: k for k, v in self.gene_token_dict.items()}
 
         self.pad_token_id = self.gene_token_dict.get("<pad>")
-        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."
-                )
-                raise
-            self.special_token = False
 
         if self.anchor_gene is None:
             self.anchor_token = None
@@ -305,7 +281,7 @@ class InSilicoPerturber:
                         continue
             valid_type = False
             for option in valid_options:
-                if (option in [bool, int, list, dict, str]) and isinstance(
+                if (option in [bool, int, list, dict]) and isinstance(
                     attr_value, option
                 ):
                     valid_type = True
@@ -451,45 +427,16 @@ class InSilicoPerturber:
         filtered_input_data = pu.load_and_filter(
             self.filter_data, self.nproc, input_data_file
         )
-
-        # 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
-            ):
-                logger.error(
-                    "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
-                ):
-                    logger.error(
-                        "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)
 
         if self.perturb_group is True:
-            if (self.special_token) and ("cls" in self.emb_mode):
-                self.isp_perturb_set_special(
-                    model, filtered_input_data, layer_to_quant, output_path_prefix
-                )
-            else:
-                self.isp_perturb_set(
-                    model, filtered_input_data, layer_to_quant, output_path_prefix
-                )
+            self.isp_perturb_set(
+                model, filtered_input_data, layer_to_quant, output_path_prefix
+            )
         else:
-            if (self.special_token) and ("cls" in self.emb_mode):
-                self.isp_perturb_all_special(
-                    model, filtered_input_data, layer_to_quant, output_path_prefix
-                )
-            else:
-                self.isp_perturb_all(
-                    model, filtered_input_data, layer_to_quant, output_path_prefix
-                )
+            self.isp_perturb_all(
+                model, filtered_input_data, layer_to_quant, output_path_prefix
+            )
 
     def apply_additional_filters(self, filtered_input_data):
         # additional filtering of input data dependent on isp mode
@@ -552,9 +499,7 @@ 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)
                 example["n_overflow"] = pu.calc_n_overflow(
                     self.max_len,
                     example["length"],
@@ -575,7 +520,6 @@ 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"]
@@ -608,7 +552,6 @@ class InSilicoPerturber:
                     layer_to_quant,
                     self.pad_token_id,
                     self.forward_batch_size,
-                    token_gene_dict=self.token_gene_dict,
                     summary_stat=None,
                     silent=True,
                 )
@@ -628,7 +571,6 @@ class InSilicoPerturber:
                     layer_to_quant,
                     self.pad_token_id,
                     self.forward_batch_size,
-                    token_gene_dict=self.token_gene_dict,
                     summary_stat=None,
                     silent=True,
                 )
@@ -728,6 +670,8 @@ class InSilicoPerturber:
                 cos_sims_dict = self.update_perturbation_dictionary(
                     cos_sims_dict,
                     cos_sims_data,
+                    filtered_input_data,
+                    indices_to_perturb,
                     gene_list,
                 )
             else:
@@ -736,6 +680,8 @@ class InSilicoPerturber:
                     cos_sims_dict[state] = self.update_perturbation_dictionary(
                         cos_sims_dict[state],
                         cos_sims_data[state],
+                        filtered_input_data,
+                        indices_to_perturb,
                         gene_list,
                     )
             del minibatch
@@ -757,264 +703,6 @@ class InSilicoPerturber:
                 f"{output_path_prefix}_gene_embs_dict_{self.tokens_to_perturb}",
             )
 
-    def isp_perturb_set_special(
-        self,
-        model,
-        filtered_input_data: Dataset,
-        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
-            indices_to_perturb = [
-                example_input_ids.index(token) if token in example_input_ids else None
-                for token in self.tokens_to_perturb
-            ]
-            indices_to_perturb = [
-                item for item in indices_to_perturb if item is not None
-            ]
-            if len(indices_to_perturb) > 0:
-                example["perturb_index"] = indices_to_perturb
-            else:
-                # -100 indicates tokens to overexpress are not present in rank value encoding
-                example["perturb_index"] = [-100]
-            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["n_overflow"] = pu.calc_n_overflow(
-                    self.max_len,
-                    example["length"],
-                    self.tokens_to_perturb,
-                    indices_to_perturb,
-                )
-            return example
-
-        total_batch_length = len(filtered_input_data)
-        if self.cell_states_to_model is None:
-            cos_sims_dict = defaultdict(list)
-        else:
-            cos_sims_dict = {
-                state: defaultdict(list)
-                for state in pu.get_possible_states(self.cell_states_to_model)
-            }
-
-        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"]
-            )
-            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)
-
-        # iterate through batches
-        for i in trange(0, total_batch_length, self.forward_batch_size):
-            max_range = min(i + self.forward_batch_size, total_batch_length)
-            inds_select = [i for i in range(i, max_range)]
-
-            minibatch = filtered_input_data.select(inds_select)
-            perturbation_batch = perturbed_data.select(inds_select)
-
-            ##### CLS Embedding Mode #####
-            if self.emb_mode == "cls":
-                indices_to_perturb = perturbation_batch["perturb_index"]
-
-                original_cls_emb = get_embs(
-                    model,
-                    minibatch,
-                    "cls",
-                    layer_to_quant,
-                    self.pad_token_id,
-                    self.forward_batch_size,
-                    token_gene_dict=self.token_gene_dict,
-                    summary_stat=None,
-                    silent=True,
-                )
-
-                perturbation_cls_emb = get_embs(
-                    model,
-                    perturbation_batch,
-                    "cls",
-                    layer_to_quant,
-                    self.pad_token_id,
-                    self.forward_batch_size,
-                    token_gene_dict=self.token_gene_dict,
-                    summary_stat=None,
-                    silent=True,
-                )
-
-                # 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",
-                )
-
-                # 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,
-                    )
-                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,
-                        )
-
-            ##### CLS and Gene Embedding Mode #####
-            elif self.emb_mode == "cls_and_gene":
-                full_original_emb = get_embs(
-                    model,
-                    minibatch,
-                    "gene",
-                    layer_to_quant,
-                    self.pad_token_id,
-                    self.forward_batch_size,
-                    self.token_gene_dict,
-                    summary_stat=None,
-                    silent=True,
-                )
-                indices_to_perturb = perturbation_batch["perturb_index"]
-                # remove indices that were perturbed
-                original_emb = pu.remove_perturbed_indices_set(
-                    full_original_emb,
-                    self.perturb_type,
-                    indices_to_perturb,
-                    self.tokens_to_perturb,
-                    minibatch["length"],
-                )
-                full_perturbation_emb = get_embs(
-                    model,
-                    perturbation_batch,
-                    "gene",
-                    layer_to_quant,
-                    self.pad_token_id,
-                    self.forward_batch_size,
-                    self.token_gene_dict,
-                    summary_stat=None,
-                    silent=True,
-                )
-
-                # 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, :
-                    ]
-                elif self.perturb_type == "delete":
-                    perturbation_emb = full_perturbation_emb[
-                        :, 1 : max(perturbation_batch["length"]) - 1, :
-                    ]
-
-                n_perturbation_genes = perturbation_emb.size()[1]
-
-                gene_cos_sims = pu.quant_cos_sims(
-                    perturbation_emb,
-                    original_emb,
-                    self.cell_states_to_model,
-                    self.state_embs_dict,
-                    emb_mode="gene",
-                )
-
-                # get cls emb
-                original_cls_emb = full_original_emb[:, 0, :]
-                perturbation_cls_emb = full_perturbation_emb[:, 0, :]
-
-                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",
-                )
-
-                # get cosine similarities in gene embeddings
-                # since getting gene embeddings, need gene names
-
-                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,
-                ]
-                gene_list = [
-                    [g for g in genes if g not in genes_to_exclude][
-                        :n_perturbation_genes
-                    ]
-                    for genes in gene_list
-                ]
-
-                for cell_i, genes in enumerate(gene_list):
-                    for gene_j, affected_gene in enumerate(genes):
-                        if len(self.genes_to_perturb) > 1:
-                            tokens_to_perturb = tuple(self.tokens_to_perturb)
-                        else:
-                            tokens_to_perturb = self.tokens_to_perturb[0]
-
-                        # fill in the gene cosine similarities
-                        try:
-                            stored_gene_embs_dict[
-                                (tokens_to_perturb, affected_gene)
-                            ].append(gene_cos_sims[cell_i, gene_j].item())
-                        except KeyError:
-                            stored_gene_embs_dict[
-                                (tokens_to_perturb, affected_gene)
-                            ] = gene_cos_sims[cell_i, gene_j].item()
-
-                if self.cell_states_to_model is None:
-                    cos_sims_dict = self.update_perturbation_dictionary(
-                        cos_sims_dict,
-                        cls_cos_sims,
-                        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,
-                        )
-                del full_original_emb
-                del original_emb
-                del full_perturbation_emb
-                del perturbation_emb
-                del gene_cos_sims
-
-            del original_cls_emb
-            del perturbation_cls_emb
-            del cls_cos_sims
-            del minibatch
-            del perturbation_batch
-
-            torch.cuda.empty_cache()
-
-        pu.write_perturbation_dictionary(
-            cos_sims_dict,
-            f"{output_path_prefix}_cell_embs_dict_{self.tokens_to_perturb}",
-        )
-
-        if self.emb_mode == "cls_and_gene":
-            pu.write_perturbation_dictionary(
-                stored_gene_embs_dict,
-                f"{output_path_prefix}_gene_embs_dict_{self.tokens_to_perturb}",
-            )
-
     def isp_perturb_all(
         self,
         model,
@@ -1033,10 +721,8 @@ class InSilicoPerturber:
 
         if self.emb_mode == "cell_and_gene":
             stored_gene_embs_dict = defaultdict(list)
-
-        num_inds_perturbed = 1 + self.combos
-        for h in trange(len(filtered_input_data)):
-            example_cell = filtered_input_data.select([h])
+        for i in trange(len(filtered_input_data)):
+            example_cell = filtered_input_data.select([i])
             full_original_emb = get_embs(
                 model,
                 example_cell,
@@ -1044,30 +730,16 @@ class InSilicoPerturber:
                 layer_to_quant,
                 self.pad_token_id,
                 self.forward_batch_size,
-                self.token_gene_dict,
                 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
+            gene_list = example_cell["input_ids"][0][:]
             if self.anchor_token is not None:
                 for token in self.anchor_token:
                     gene_list.remove(token)
-            # index 0 is not overexpressed so remove
-            if self.perturb_type == "overexpress":
-                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,
@@ -1078,459 +750,147 @@ class InSilicoPerturber:
                 self.nproc,
             )
 
-            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
-
-                full_perturbation_emb = get_embs(
-                    model,
-                    perturbation_minibatch,
-                    "gene",
-                    layer_to_quant,
-                    self.pad_token_id,
-                    self.forward_batch_size,
-                    self.token_gene_dict,
-                    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"
-                ):
-                    original_emb_minibatch = pu.make_comparison_batch(
-                        full_original_emb, indices_to_perturb_mini, perturb_group=False
-                    )
-                    gene_cos_sims = pu.quant_cos_sims(
-                        perturbation_emb,
-                        original_emb_minibatch,
-                        self.cell_states_to_model,
-                        self.state_embs_dict,
-                        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,
-                        self.cell_states_to_model,
-                        self.state_embs_dict,
-                        emb_mode="cell",
-                    )
-                    del perturbation_cell_emb
-
-                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]
-                        ):
-                            try:
-                                stored_gene_embs_dict[
-                                    (perturbed_gene, affected_gene)
-                                ].append(gene_cos_sims[perturbation_i, gene_j].item())
-                            except KeyError:
-                                stored_gene_embs_dict[
-                                    (perturbed_gene, affected_gene)
-                                ] = gene_cos_sims[perturbation_i, gene_j].item()
+            full_perturbation_emb = get_embs(
+                model,
+                perturbation_batch,
+                "gene",
+                layer_to_quant,
+                self.pad_token_id,
+                self.forward_batch_size,
+                summary_stat=None,
+                silent=True,
+            )
 
-                del full_perturbation_emb
+            num_inds_perturbed = 1 + self.combos
+            # need to remove overexpressed gene to quantify cosine shifts
+            if self.perturb_type == "overexpress":
+                perturbation_emb = full_perturbation_emb[:, num_inds_perturbed:, :]
+                gene_list = gene_list[
+                    num_inds_perturbed:
+                ]  # index 0 is not overexpressed
 
-                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(
-                        cos_sims_dict,
-                        cos_sims_data,
-                        gene_list_mini,
-                    )
-                else:
-                    cos_sims_data = cell_cos_sims
-                    for state in cos_sims_dict.keys():
-                        cos_sims_dict[state] = self.update_perturbation_dictionary(
-                            cos_sims_dict[state],
-                            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:
-                    pu.write_perturbation_dictionary(
-                        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,
-                            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
-                    if self.cell_states_to_model is None:
-                        cos_sims_dict = defaultdict(list)
-                    else:
-                        cos_sims_dict = {
-                            state: defaultdict(list)
-                            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()
+            elif self.perturb_type == "delete":
+                perturbation_emb = full_perturbation_emb
 
-            pu.write_perturbation_dictionary(
-                cos_sims_dict,
-                f"{output_path_prefix}_dict_cell_embs_{h}batch{pickle_batch}",
+            original_batch = pu.make_comparison_batch(
+                full_original_emb, indices_to_perturb, perturb_group=False
             )
 
-            if self.emb_mode == "cell_and_gene":
-                pu.write_perturbation_dictionary(
-                    stored_gene_embs_dict,
-                    f"{output_path_prefix}_dict_gene_embs_{h}batch{pickle_batch}",
+            if self.cell_states_to_model is None or self.emb_mode == "cell_and_gene":
+                gene_cos_sims = pu.quant_cos_sims(
+                    perturbation_emb,
+                    original_batch,
+                    self.cell_states_to_model,
+                    self.state_embs_dict,
+                    emb_mode="gene",
                 )
-
-            pickle_batch = -1
-            if self.cell_states_to_model is None:
-                cos_sims_dict = defaultdict(list)
-            else:
-                cos_sims_dict = {
-                    state: defaultdict(list)
-                    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)
-
-            # clear memory between cells
-            del perturbation_batch
-            del full_original_emb
             if self.cell_states_to_model is not None:
-                del original_cell_emb
-            torch.cuda.empty_cache()
-
-    def isp_perturb_all_special(
-        self,
-        model,
-        filtered_input_data: Dataset,
-        layer_to_quant: int,
-        output_path_prefix: str,
-    ):
-        pickle_batch = -1
-        if self.cell_states_to_model is None:
-            cos_sims_dict = defaultdict(list)
-        else:
-            cos_sims_dict = {
-                state: defaultdict(list)
-                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)
-
-        num_inds_perturbed = 1 + self.combos
-        for h in trange(len(filtered_input_data)):
-            example_cell = filtered_input_data.select([h])
-
-            # get original example cell cls and/or gene embs for comparison
-            if self.emb_mode == "cls":
-                original_cls_emb = get_embs(
-                    model,
-                    example_cell,
-                    "cls",
-                    layer_to_quant,
-                    self.pad_token_id,
-                    self.forward_batch_size,
-                    self.token_gene_dict,
-                    summary_stat=None,
-                    silent=True,
+                original_cell_emb = pu.compute_nonpadded_cell_embedding(
+                    full_original_emb, "mean_pool"
                 )
-            elif self.emb_mode == "cls_and_gene":
-                full_original_emb = get_embs(
-                    model,
-                    example_cell,
-                    "gene",
-                    layer_to_quant,
-                    self.pad_token_id,
-                    self.forward_batch_size,
-                    self.token_gene_dict,
-                    summary_stat=None,
-                    silent=True,
+                perturbation_cell_emb = pu.compute_nonpadded_cell_embedding(
+                    full_perturbation_emb, "mean_pool"
                 )
-                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][:]
 
-            # need to remove special tokens
-            for token in [self.cls_token_id, self.eos_token_id]:
-                gene_list.remove(token)
-            # need to remove the anchor gene
-            if self.anchor_token is not None:
-                for token in self.anchor_token:
-                    gene_list.remove(token)
-            # index 0 is not overexpressed so remove
-            if self.perturb_type == "overexpress":
-                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_special(
-                example_cell,
-                self.perturb_type,
-                self.tokens_to_perturb,
-                self.anchor_token,
-                self.combos,
-                self.nproc,
-            )
-
-            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)]
+                cell_cos_sims = pu.quant_cos_sims(
+                    perturbation_cell_emb,
+                    original_cell_emb,
+                    self.cell_states_to_model,
+                    self.state_embs_dict,
+                    emb_mode="cell",
                 )
-                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
-                    perturbation_cls_emb = get_embs(
-                        model,
-                        perturbation_minibatch,
-                        "cls",
-                        layer_to_quant,
-                        self.pad_token_id,
-                        self.forward_batch_size,
-                        self.token_gene_dict,
-                        summary_stat=None,
-                        silent=True,
-                    )
 
-                    # Calculate 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",
-                    )
+            if self.emb_mode == "cell_and_gene":
+                # remove perturbed index for gene list
+                perturbed_gene_dict = {
+                    gene: gene_list[:i] + gene_list[i + 1 :]
+                    for i, gene in enumerate(gene_list)
+                }
 
-                    if self.cell_states_to_model is None:
-                        cos_sims_dict = self.update_perturbation_dictionary(
-                            cos_sims_dict,
-                            cls_cos_sims,
-                            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":
-                    full_perturbation_emb = get_embs(
-                        model,
-                        perturbation_minibatch,
-                        "gene",
-                        layer_to_quant,
-                        self.pad_token_id,
-                        self.forward_batch_size,
-                        self.token_gene_dict,
-                        summary_stat=None,
-                        silent=True,
-                    )
+                for perturbation_i, perturbed_gene in enumerate(gene_list):
+                    for gene_j, affected_gene in enumerate(
+                        perturbed_gene_dict[perturbed_gene]
+                    ):
+                        try:
+                            stored_gene_embs_dict[
+                                (perturbed_gene, affected_gene)
+                            ].append(gene_cos_sims[perturbation_i, gene_j].item())
+                        except KeyError:
+                            stored_gene_embs_dict[
+                                (perturbed_gene, affected_gene)
+                            ] = gene_cos_sims[perturbation_i, gene_j].item()
 
-                    # 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()
-                        )
-                    elif self.perturb_type == "delete":
-                        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
+            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(
+                    cos_sims_dict,
+                    cos_sims_data,
+                    filtered_input_data,
+                    indices_to_perturb,
+                    gene_list,
+                )
+            else:
+                cos_sims_data = cell_cos_sims
+                for state in cos_sims_dict.keys():
+                    cos_sims_dict[state] = self.update_perturbation_dictionary(
+                        cos_sims_dict[state],
+                        cos_sims_data[state],
+                        filtered_input_data,
+                        indices_to_perturb,
+                        gene_list,
                     )
 
-                    original_emb_minibatch = (
-                        original_emb_minibatch[:, 1:-1, :].clone().detach()
-                    )
-                    gene_cos_sims = pu.quant_cos_sims(
-                        perturbation_emb,
-                        original_emb_minibatch,
-                        self.cell_states_to_model,
-                        self.state_embs_dict,
-                        emb_mode="gene",
+            # save dict to disk every 100 cells
+            if i % 100 == 0:
+                pu.write_perturbation_dictionary(
+                    cos_sims_dict,
+                    f"{output_path_prefix}_dict_cell_embs_1Kbatch{pickle_batch}",
+                )
+                if self.emb_mode == "cell_and_gene":
+                    pu.write_perturbation_dictionary(
+                        stored_gene_embs_dict,
+                        f"{output_path_prefix}_dict_gene_embs_1Kbatch{pickle_batch}",
                     )
 
-                    for perturbation_i, perturbed_gene in enumerate(gene_list_mini):
-                        for gene_j, affected_gene in enumerate(
-                            perturbed_gene_dict[perturbed_gene]
-                        ):
-                            try:
-                                stored_gene_embs_dict[
-                                    (perturbed_gene, affected_gene)
-                                ].append(gene_cos_sims[perturbation_i, gene_j].item())
-                            except KeyError:
-                                stored_gene_embs_dict[
-                                    (perturbed_gene, affected_gene)
-                                ] = gene_cos_sims[perturbation_i, gene_j].item()
+            # reset and clear memory every 1000 cells
+            if i % 1000 == 0:
+                pickle_batch += 1
+                if self.cell_states_to_model is None:
+                    cos_sims_dict = defaultdict(list)
+                else:
+                    cos_sims_dict = {
+                        state: defaultdict(list)
+                        for state in pu.get_possible_states(self.cell_states_to_model)
+                    }
 
-                    # get cls emb
-                    perturbation_cls_emb = (
-                        full_perturbation_emb[:, 0, :].clone().detach()
-                    )
+                if self.emb_mode == "cell_and_gene":
+                    stored_gene_embs_dict = defaultdict(list)
 
-                    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",
-                    )
+                torch.cuda.empty_cache()
 
-                    if self.cell_states_to_model is None:
-                        cos_sims_dict = self.update_perturbation_dictionary(
-                            cos_sims_dict,
-                            cls_cos_sims,
-                            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 original_emb_minibatch
-                    del full_perturbation_emb
-                    del perturbation_emb
-                    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:
-                    pu.write_perturbation_dictionary(
-                        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,
-                            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
-                    if self.cell_states_to_model is None:
-                        cos_sims_dict = defaultdict(list)
-                    else:
-                        cos_sims_dict = {
-                            state: defaultdict(list)
-                            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_1Kbatch{pickle_batch}"
+        )
 
+        if self.emb_mode == "cell_and_gene":
             pu.write_perturbation_dictionary(
-                cos_sims_dict,
-                f"{output_path_prefix}_dict_cell_embs_{h}batch{pickle_batch}",
+                stored_gene_embs_dict,
+                f"{output_path_prefix}_dict_gene_embs_1Kbatch{pickle_batch}",
             )
 
-            if self.emb_mode == "cls_and_gene":
-                pu.write_perturbation_dictionary(
-                    stored_gene_embs_dict,
-                    f"{output_path_prefix}_dict_gene_embs_{h}batch{pickle_batch}",
-                )
-
-            pickle_batch = -1
-            if self.cell_states_to_model is None:
-                cos_sims_dict = defaultdict(list)
-            else:
-                cos_sims_dict = {
-                    state: defaultdict(list)
-                    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)
-
-            # clear memory between cells
-            del perturbation_batch
-            del original_cls_emb
-            if self.emb_mode == "cls_and_gene":
-                del full_original_emb
-            torch.cuda.empty_cache()
-
     def update_perturbation_dictionary(
         self,
         cos_sims_dict: defaultdict,
         cos_sims_data: torch.Tensor,
+        filtered_input_data: Dataset,
+        indices_to_perturb: List[List[int]],
         gene_list=None,
     ):
         if gene_list is not None and cos_sims_data.shape[0] != len(gene_list):
             logger.error(
                 f"len(cos_sims_data.shape[0]) != len(gene_list). \n \
-                            {cos_sims_data.shape[0]=}.\n \
-                            {len(gene_list)=}."
+                            cos_sims_data.shape[0] = {cos_sims_data.shape[0]}.\n \
+                            len(gene_list) = {len(gene_list)}."
             )
             raise
 

geneformer/in_silico_perturber_stats.py

@@ -37,8 +37,10 @@ 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 .tokenizer import TOKEN_DICTIONARY_FILE
+
+GENE_NAME_ID_DICTIONARY_FILE = Path(__file__).parent / "gene_name_id_dict.pkl"
 
 logger = logging.getLogger(__name__)
 
@@ -114,7 +116,6 @@ def read_dictionaries(
                             state_dict[state_value][key] += new_dict[key]
                         except KeyError:
                             state_dict[state_value][key] = new_dict[key]
-
     if not file_found:
         logger.error(
             "No raw data for processing found within provided directory. "
@@ -191,69 +192,34 @@ def get_impact_component(test_value, gaussian_mixture_model):
 
 
 # aggregate data for single perturbation in multiple cells
-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),
-                ),
-                :,
-            ]
-        else:
-            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
-
-    if gene_ids_df.empty:
-        logger.error("genes_to_perturb not found in data.")
-        raise
-
-    tokens = gene_ids_df["Gene"]
-    symbols = gene_ids_df["Gene_name"]
-
-    for token, symbol in zip(tokens, symbols):
-        cos_shift_data = []
-        for dict_i in dict_list:
-            cos_shift_data += dict_i.get((token, "cell_emb"), [])
-
-        df = pd.DataFrame(columns=names)
-        df["Cosine_sim"] = cos_shift_data
-        df["Gene"] = symbol
-        cos_sims_full_dfs.append(df)
+def isp_aggregate_grouped_perturb(cos_sims_df, dict_list):
+    names = ["Cosine_shift"]
+    cos_sims_full_df = pd.DataFrame(columns=names)
 
-    return pd.concat(cos_sims_full_dfs)
+    cos_shift_data = []
+    token = cos_sims_df["Gene"][0]
+    for dict_i in dict_list:
+        cos_shift_data += dict_i.get((token, "cell_emb"), [])
+    cos_sims_full_df["Cosine_shift"] = cos_shift_data
+    return cos_sims_full_df
 
 
 def find(variable, x):
     try:
         if x in variable:  # Test if variable is iterable and contains x
             return True
-        elif x == variable:
-            return True
     except (ValueError, TypeError):
         return x == variable  # Test if variable is x if non-iterable
 
 
 def isp_aggregate_gene_shifts(
-    cos_sims_df, dict_list, gene_token_id_dict, gene_id_name_dict, token_dtype
+    cos_sims_df, dict_list, gene_token_id_dict, gene_id_name_dict
 ):
     cos_shift_data = dict()
     for i in trange(cos_sims_df.shape[0]):
         token = cos_sims_df["Gene"][i]
         for dict_i in dict_list:
-            if token_dtype == "nontuple":
-                affected_pairs = [k for k, v in dict_i.items() if k[0] == token]
-            else:
-                affected_pairs = [k for k, v in dict_i.items() if find(k[0], token)]
+            affected_pairs = [k for k, v in dict_i.items() if find(k[0], token)]
             for key in affected_pairs:
                 if key in cos_shift_data.keys():
                     cos_shift_data[key] += dict_i.get(key, [])
@@ -266,11 +232,11 @@ def isp_aggregate_gene_shifts(
     cos_sims_full_df = pd.DataFrame()
     cos_sims_full_df["Perturbed"] = [k[0] for k, v in cos_data_mean.items()]
     cos_sims_full_df["Gene_name"] = [
-        cos_sims_df[cos_sims_df["Gene"] == k[0]]["Gene_name"].item()
+        cos_sims_df[cos_sims_df["Gene"] == k[0]]["Gene_name"][0]
         for k, v in cos_data_mean.items()
     ]
     cos_sims_full_df["Ensembl_ID"] = [
-        cos_sims_df[cos_sims_df["Gene"] == k[0]]["Ensembl_ID"].item()
+        cos_sims_df[cos_sims_df["Gene"] == k[0]]["Ensembl_ID"][0]
         for k, v in cos_data_mean.items()
     ]
 
@@ -282,15 +248,15 @@ def isp_aggregate_gene_shifts(
     cos_sims_full_df["Affected_Ensembl_ID"] = [
         gene_token_id_dict.get(token, np.nan) for token in cos_sims_full_df["Affected"]
     ]
-    cos_sims_full_df["Cosine_sim_mean"] = [v[0] for k, v in cos_data_mean.items()]
-    cos_sims_full_df["Cosine_sim_stdev"] = [v[1] for k, v in cos_data_mean.items()]
+    cos_sims_full_df["Cosine_shift_mean"] = [v[0] for k, v in cos_data_mean.items()]
+    cos_sims_full_df["Cosine_shift_stdev"] = [v[1] for k, v in cos_data_mean.items()]
     cos_sims_full_df["N_Detections"] = [v[2] for k, v in cos_data_mean.items()]
 
     specific_val = "cell_emb"
     cos_sims_full_df["temp"] = list(cos_sims_full_df["Affected"] == specific_val)
-    # reorder so cell embs are at the top and all are subordered by magnitude of cosine sim
+    # reorder so cell embs are at the top and all are subordered by magnitude of cosine shift
     cos_sims_full_df = cos_sims_full_df.sort_values(
-        by=(["temp", "Cosine_sim_mean"]), ascending=[False, True]
+        by=(["temp", "Cosine_shift_mean"]), ascending=[False, False]
     ).drop("temp", axis=1)
 
     return cos_sims_full_df
@@ -681,7 +647,7 @@ class InSilicoPerturberStats:
         cell_states_to_model=None,
         pickle_suffix="_raw.pickle",
         token_dictionary_file=TOKEN_DICTIONARY_FILE,
-        gene_name_id_dictionary_file=ENSEMBL_DICTIONARY_FILE,
+        gene_name_id_dictionary_file=GENE_NAME_ID_DICTIONARY_FILE,
     ):
         """
         Initialize in silico perturber stats generator.
@@ -700,7 +666,7 @@ class InSilicoPerturberStats:
             | Default is assuming genes_to_perturb in isp experiment was "all" (each gene in each cell).
             | Otherwise, may provide a list of ENSEMBL IDs of genes perturbed as a group all together.
         combos : {0,1,2}
-            | Whether genex perturbed in isp experiment were perturbed individually (0), in pairs (1), or in triplets (2).
+            | Whether to perturb genes individually (0), in pairs (1), or in triplets (2).
         anchor_gene : None, str
             | ENSEMBL ID of gene to use as anchor in combination perturbations or in testing effect on downstream genes.
             | For example, if combos=1 and anchor_gene="ENSG00000136574":
@@ -948,11 +914,11 @@ class InSilicoPerturberStats:
         |     1: within impact component; 0: not within impact component
         | "Impact_component_percent": percent of cells in which given perturbation was modeled to be within impact component
 
-        | In case of aggregating data / gene shifts:
+        | In case of aggregating gene shifts:
         | "Perturbed": ID(s) of gene(s) being perturbed
         | "Affected": ID of affected gene or "cell_emb" indicating the impact on the cell embedding as a whole
-        | "Cosine_sim_mean": mean of cosine similarity of cell or affected gene in original vs. perturbed
-        | "Cosine_sim_stdev": standard deviation of cosine similarity of cell or affected gene in original vs. perturbed
+        | "Cosine_shift_mean": mean of cosine shift of modeled perturbation on affected gene or cell
+        | "Cosine_shift_stdev": standard deviation of cosine shift of modeled perturbation on affected gene or cell
         """
 
         if self.mode not in [
@@ -1052,30 +1018,14 @@ class InSilicoPerturberStats:
             )
 
         elif self.mode == "aggregate_data":
-            cos_sims_df = isp_aggregate_grouped_perturb(
-                cos_sims_df_initial, dict_list, self.genes_perturbed
-            )
+            cos_sims_df = isp_aggregate_grouped_perturb(cos_sims_df_initial, dict_list)
 
         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
-                ]
-                if all(tuple_types):
-                    token_dtype = "tuple"
-                elif not any(tuple_types):
-                    token_dtype = "nontuple"
-                else:
-                    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,
             )
 
         # save perturbation stats to output_path

geneformer/mtl/__init__.py

@@ -1 +0,0 @@
-# ruff: noqa: F401

geneformer/mtl/collators.py

@@ -1,76 +0,0 @@
-# imports
-import torch
-import pickle
-from ..collator_for_classification import DataCollatorForGeneClassification
-from .. import TOKEN_DICTIONARY_FILE
-
-"""Geneformer collator for multi-task cell classification."""
-
-class DataCollatorForMultitaskCellClassification(DataCollatorForGeneClassification):
-    class_type = "cell"
-
-    @staticmethod
-    def load_token_dictionary():
-        with open(TOKEN_DICTIONARY_FILE, 'rb') as f:
-            return pickle.load(f)
-
-    def __init__(self, *args, **kwargs) -> None:
-        # Load the token dictionary
-        token_dictionary = self.load_token_dictionary()
-        # Use the loaded token dictionary
-        super().__init__(token_dictionary=token_dictionary, *args, **kwargs)
-
-    def _prepare_batch(self, features):
-        # Process inputs as usual
-        batch = self.tokenizer.pad(
-            features,
-            class_type=self.class_type,
-            padding=self.padding,
-            max_length=self.max_length,
-            pad_to_multiple_of=self.pad_to_multiple_of,
-            return_tensors="pt",
-        )
-
-        # Check if labels are present
-        if "label" in features[0]:
-            # Initialize labels dictionary for all tasks
-            labels = {task: [] for task in features[0]["label"].keys()}
-            # Populate labels for each task
-            for feature in features:
-                for task, label in feature["label"].items():
-                    labels[task].append(label)
-
-            # Convert label lists to tensors, handling dictionaries appropriately
-            for task in labels:
-                if isinstance(labels[task][0], (list, torch.Tensor)):
-                    dtype = torch.long
-                    labels[task] = torch.tensor(labels[task], dtype=dtype)
-                elif isinstance(labels[task][0], dict):
-                    # Handle dict specifically if needed
-                    pass  # Resolve nested data structure
-
-            # Update the batch to include task-specific labels
-            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()
-            }
-
-        return batch
-
-    def __call__(self, features):
-        batch = self._prepare_batch(features)
-        for k, v in batch.items():
-            if torch.is_tensor(v):
-                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()
-                }
-            else:
-                batch[k] = torch.tensor(v, dtype=torch.int64)
-        return batch

geneformer/mtl/data.py

@@ -1,150 +0,0 @@
-import os
-
-from .collators import DataCollatorForMultitaskCellClassification
-from .imports import *
-
-
-def load_and_preprocess_data(dataset_path, config, is_test=False, dataset_type=""):
-    try:
-        dataset = load_from_disk(dataset_path)
-
-        task_names = [f"task{i+1}" for i in range(len(config["task_columns"]))]
-        task_to_column = dict(zip(task_names, config["task_columns"]))
-        config["task_names"] = task_names
-
-        if not is_test:
-            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)}"
-                    )
-
-        label_mappings = {}
-        task_label_mappings = {}
-        cell_id_mapping = {}
-        num_labels_list = []
-
-        # Load or create task label mappings
-        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)
-                }
-                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
-
-            # Save the task label mappings as a pickle file
-            with open(f"{config['results_dir']}/task_label_mappings.pkl", "wb") as f:
-                pickle.dump(task_label_mappings, f)
-        else:
-            # 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_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
-            )
-
-            # Use index-based cell ID for internal tracking
-            transformed_record["cell_id"] = idx
-
-            if not is_test:
-                # Prepare labels
-                label_dict = {}
-                for task, column in task_to_column.items():
-                    label_value = record[column]
-                    label_index = task_label_mappings[task][label_value]
-                    label_dict[task] = label_index
-                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_dataset.append(transformed_record)
-
-        return transformed_dataset, cell_id_mapping, num_labels_list
-    except KeyError as e:
-        print(f"Missing configuration or dataset key: {e}")
-    except Exception as e:
-        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,
-    )
-
-
-def get_data_loader(preprocessed_dataset, batch_size):
-    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,
-    )
-    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

geneformer/mtl/eval_utils.py

@@ -1,88 +0,0 @@
-import pandas as pd
-
-from .imports import *  # noqa # isort:skip
-from .data import prepare_test_loader  # noqa # isort:skip
-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)
-
-    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)
-            _, 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()]
-                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()
-                    )
-                    task_pred_labels[task_name].append(pred_label)
-                    task_pred_probs[task_name].append(pred_prob)
-
-    # Save test predictions with cell IDs and probabilities to CSV
-    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]}
-        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])
-            )
-        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)
-    model_directory = os.path.join(config["model_save_path"], "GeneformerMultiTask")
-    hyperparams_path = os.path.join(model_directory, "hyperparameters.json")
-
-    # Load the saved best hyperparameters
-    with open(hyperparams_path, "r") as f:
-        best_hyperparams = json.load(f)
-
-    # Extract the task weights if present, otherwise set to None
-    task_weights = best_hyperparams.get("task_weights", None)
-    normalized_task_weights = task_weights if task_weights else []
-
-    # Print the loaded hyperparameters
-    print("Loaded hyperparameters:")
-    for param, value in best_hyperparams.items():
-        if param == "task_weights":
-            print(f"normalized_task_weights: {value}")
-        else:
-            print(f"{param}: {value}")
-
-    best_model_path = os.path.join(model_directory, "pytorch_model.bin")
-    best_model = GeneformerMultiTask(
-        config["pretrained_path"],
-        num_labels_list,
-        dropout_rate=best_hyperparams["dropout_rate"],
-        use_task_weights=config["use_task_weights"],
-        task_weights=normalized_task_weights,
-    )
-    best_model.load_state_dict(torch.load(best_model_path))
-    best_model.to(device)
-
-    evaluate_test_dataset(best_model, device, test_loader, cell_id_mapping, config)
-    print("Evaluation completed.")

geneformer/mtl/imports.py

@@ -1,43 +0,0 @@
-import functools
-import gc
-import json
-import os
-import pickle
-import sys
-import warnings
-from enum import Enum
-from itertools import chain
-from typing import Dict, List, Optional, Union
-
-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,
-    DataCollatorForTokenClassification,
-    SpecialTokensMixin,
-    get_cosine_schedule_with_warmup,
-    get_linear_schedule_with_warmup,
-    get_scheduler,
-)
-from transformers.utils import logging, to_py_obj
-
-from .collators import DataCollatorForMultitaskCellClassification
-
-# local modules
-from .data import get_data_loader, preload_and_process_data
-from .model import GeneformerMultiTask
-from .optuna_utils import create_optuna_study
-from .utils import save_model

geneformer/mtl/model.py

@@ -1,121 +0,0 @@
-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
-
-    def forward(self, hidden_states):
-        attention_scores = torch.matmul(hidden_states, self.attention_weights)
-        attention_scores = torch.softmax(attention_scores, dim=1)
-        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,
-    ):
-        super(GeneformerMultiTask, self).__init__()
-        self.config = BertConfig.from_pretrained(pretrained_path)
-        self.bert = BertModel(self.config)
-        self.num_labels_list = num_labels_list
-        self.use_task_weights = use_task_weights
-        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)
-        )
-
-        # 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.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)
-            nn.init.zeros_(head.bias)
-
-    def forward(self, input_ids, attention_mask, labels=None):
-        try:
-            outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
-        except Exception as e:
-            raise RuntimeError(f"Error during BERT forward pass: {e}")
-
-        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.dropout(pooled_output)
-        except Exception as e:
-            raise RuntimeError(f"Error during pooling and dropout: {e}")
-
-        total_loss = 0
-        logits = []
-        losses = []
-
-        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}"
-                )
-
-            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)
-                    )
-                    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}"
-                    )
-
-        return total_loss, logits, losses if labels is not None else logits

geneformer/mtl/optuna_utils.py

@@ -1,27 +0,0 @@
-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"
-        )
-
-
-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"
-    )
-
-    # callback and TensorBoard callback
-    callbacks = [
-        lambda study, trial: save_trial_callback(study, trial, trials_result_path),
-        tensorboard_callback,
-    ]
-
-    study.optimize(objective, n_trials=n_trials, callbacks=callbacks)
-    return study

geneformer/mtl/train.py

@@ -1,380 +0,0 @@
-import os
-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, get_layer_freeze_range
-
-
-def set_seed(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(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."
-        )
-
-
-def create_model(config, num_labels_list, device):
-    model = GeneformerMultiTask(
-        config["pretrained_path"],
-        num_labels_list,
-        dropout_rate=config["dropout_rate"],
-        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"],
-    )
-    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"],
-    )
-    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
-        )
-    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,
-        )
-
-    return optimizer, scheduler
-
-
-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"]
-        ]
-
-        loss, _, _ = model(input_ids, attention_mask, labels)
-        loss.backward()
-
-        if config["gradient_clipping"]:
-            torch.nn.utils.clip_grad_norm_(model.parameters(), config["max_grad_norm"])
-
-        optimizer.step()
-        scheduler.step()
-
-        writer.add_scalar(
-            "Training Loss", loss.item(), epoch * len(train_loader) + batch_idx
-        )
-        if config.get("use_wandb", False):
-            import wandb
-
-            wandb.log({"Training Loss": loss.item()})
-
-        # Update progress bar
-        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
-    task_true_labels = {task_name: [] for task_name in config["task_names"]}
-    task_pred_labels = {task_name: [] for task_name in config["task_names"]}
-    task_pred_probs = {task_name: [] for task_name in config["task_names"]}
-
-    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"]
-            ]
-            loss, logits, _ = model(input_ids, attention_mask, labels)
-            val_loss += loss.item()
-
-            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()
-                    pred_label = torch.argmax(logits[i][sample_idx], dim=-1).item()
-                    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)
-
-    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}"
-        )
-        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
-        )
-
-
-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]}
-        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])
-            )
-        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,
-):
-    set_seed(config["seed"])
-    initialize_wandb(config)
-
-    model = create_model(config, num_labels_list, device)
-    total_steps = len(train_loader) * config["epochs"]
-    optimizer, scheduler = setup_optimizer_and_scheduler(model, config, total_steps)
-
-    log_dir = os.path.join(config["tensorboard_log_dir"], "manual_run")
-    writer = SummaryWriter(log_dir=log_dir)
-
-    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}"})
-
-    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
-    )
-
-    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,
-):
-    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", [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))
-        ]
-        weight_sum = sum(config["task_weights"])
-        config["task_weights"] = [
-            weight / weight_sum for weight in config["task_weights"]
-        ]
-    else:
-        config["task_weights"] = None
-
-    # Dynamic range for max_layers_to_freeze
-    freeze_range = get_layer_freeze_range(config["pretrained_path"])
-    config["max_layers_to_freeze"] = trial.suggest_int(
-        "max_layers_to_freeze",
-        freeze_range["min"],
-        freeze_range["max"]
-    )
-
-    model = create_model(config, num_labels_list, device)
-    total_steps = len(train_loader) * config["epochs"]
-    optimizer, scheduler = setup_optimizer_and_scheduler(model, config, total_steps)
-
-    log_dir = os.path.join(config["tensorboard_log_dir"], f"trial_{trial.number}")
-    writer = SummaryWriter(log_dir=log_dir)
-
-    for epoch in range(config["epochs"]):
-        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
-    )
-    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,
-    )
-
-    trial.set_user_attr("model_state_dict", model.state_dict())
-    trial.set_user_attr("task_weights", config["task_weights"])
-
-    trial.report(val_loss, config["epochs"])
-
-    if trial.should_prune():
-        raise optuna.TrialPruned()
-
-    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.finish()
-
-    return val_loss

geneformer/mtl/train_utils.py

@@ -1,161 +0,0 @@
-import random
-
-from .data import get_data_loader, preload_and_process_data
-from .imports import *
-from .model import GeneformerMultiTask
-from .train import objective, train_model
-from .utils import save_model
-
-
-def set_seed(seed):
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_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"])
-
-    # Print the manual hyperparameters being used
-    print("\nManual hyperparameters being used:")
-    for key, value in config["manual_hyperparameters"].items():
-        print(f"{key}: {value}")
-    print()  # Add an empty line for better readability
-
-    # Use the manual hyperparameters
-    for key, value in config["manual_hyperparameters"].items():
-        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,
-    )
-
-    print(f"\nValidation loss with manual hyperparameters: {val_loss}")
-
-    # Save the trained model
-    model_save_directory = os.path.join(
-        config["model_save_path"], "GeneformerMultiTask"
-    )
-    save_model(trained_model, model_save_directory)
-
-    # Save the hyperparameters
-    hyperparams_to_save = {
-        **config["manual_hyperparameters"],
-        "dropout_rate": config["dropout_rate"],
-        "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"],
-    }
-    hyperparams_path = os.path.join(model_save_directory, "hyperparameters.json")
-    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"])
-
-    if config["use_manual_hyperparameters"]:
-        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,
-            train_loader=train_loader,
-            val_loader=val_loader,
-            train_cell_id_mapping=train_cell_id_mapping,
-            val_cell_id_mapping=val_cell_id_mapping,
-            num_labels_list=num_labels_list,
-            config=config,
-            device=device,
-        )
-
-        study = optuna.create_study(
-            direction="minimize",  # Minimize validation loss
-            study_name=config["study_name"],
-            # storage=config["storage"],
-            load_if_exists=True,
-        )
-
-        study.optimize(objective_with_config_and_data, n_trials=config["n_trials"])
-
-        # After finding the best trial
-        best_params = study.best_trial.params
-        best_task_weights = study.best_trial.user_attrs["task_weights"]
-        print("Saving the best model and its hyperparameters...")
-
-        # Saving model as before
-        best_model = GeneformerMultiTask(
-            config["pretrained_path"],
-            num_labels_list,
-            dropout_rate=best_params["dropout_rate"],
-            use_task_weights=config["use_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()
-        }
-
-        # 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"
-        )
-        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:
-            json.dump({**best_params, "task_weights": best_task_weights}, f)
-        print(f"Best hyperparameters and task weights saved to {hyperparams_path}")

geneformer/mtl/utils.py

@@ -1,129 +0,0 @@
-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
-    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_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")
-        )
-    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")
-        accuracy = accuracy_score(true_labels, pred_labels)
-        task_metrics[task_name] = {"f1": f1, "accuracy": accuracy}
-    return task_metrics
-
-
-def calculate_combined_f1(combined_labels, combined_preds):
-    # Initialize the LabelEncoder
-    le = LabelEncoder()
-
-    # Fit and transform combined labels and predictions to numerical values
-    le.fit(combined_labels + combined_preds)
-    encoded_true_labels = le.transform(combined_labels)
-    encoded_pred_labels = le.transform(combined_preds)
-
-    # Print out the mapping for sanity check
-    print("\nLabel Encoder Mapping:")
-    for index, class_label in enumerate(le.classes_):
-        print(f"'{class_label}': {index}")
-
-    # Calculate accuracy
-    accuracy = accuracy_score(encoded_true_labels, encoded_pred_labels)
-
-    # Calculate F1 Macro score
-    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"
-#     if not os.path.exists(new_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")
-#     )
-
-#     # Initialize a new BERT model without the classification heads
-#     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")
-#     }
-
-#     # 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,
-#     )
-
-#     print(f"Model without classification heads saved to {new_model_save_directory}")
-
-
-def get_layer_freeze_range(pretrained_path):
-    """
-    Dynamically determines the number of layers to freeze based on the model depth from its configuration.
-    Args:
-        pretrained_path (str): Path to the pretrained model directory or model identifier.
-    Returns:
-        dict: A dictionary with 'min' and 'max' keys indicating the range of layers to freeze.
-    """
-    if pretrained_path:
-        config = AutoConfig.from_pretrained(pretrained_path)
-        total_layers = config.num_hidden_layers
-        return {"min": 0, "max": total_layers - 1}
-    else:
-        return {"min": 0, "max": 0}