Add Geneformer tokenizer and updated model card

README.md

@@ -1,17 +1,14 @@
 # Geneformer
 Geneformer is a 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](manuscript_link) for details.
--->
 
 # Model Description
 Geneformer is transformer model pretrained on a [Genecorpus-30M](dataset_link), a pretraining corpus comprised of ~30 million single cell transcriptomes from a broad range of human tissues. 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. 
 
 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](manuscript_link). -->
-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. Fine-tuning Geneformer towards a diverse panel of downstream tasks relevant to chromatin and network dynamics using limited task-specific data demonstrated that Geneformer consistently boosted predictive accuracy. Applied to disease modeling with limited patient data, Geneformer identified candidate therapeutic targets. Overall, Geneformer represents an invaluable pretrained model from which fine-tuning towards a broad range of downstream applications can be pursued to accelerate discovery of key network regulators and candidate therapeutic targets.
+We detail applications and results in [our manuscript](manuscript_link). 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. Fine-tuning Geneformer towards a diverse panel of downstream tasks relevant to chromatin and network dynamics using limited task-specific data demonstrated that Geneformer consistently boosted predictive accuracy. Applied to disease modeling with limited patient data, Geneformer identified candidate therapeutic targets. Overall, Geneformer represents an invaluable pretrained model from which fine-tuning towards a broad range of downstream applications can be pursued to accelerate discovery of key network regulators and candidate therapeutic targets.
 
 # Application
-The pretrained Geneformer model can be used directly, for example for in silico deletion analysis, but is best used by fine-tuning towards the relevant downstream task, such as gene or cell state classification.
+The pretrained Geneformer model can be used directly, for example for in silico deletion analysis, but is best used by fine-tuning towards the relevant downstream task, such as gene or cell state classification.

geneformer/tokenizer.py

@@ -0,0 +1,204 @@
+"""
+Geneformer tokenizer.
+
+Usage:
+  from geneformer.tokenizer import Tokenizer
+  tk = Tokenizer({"cell_type": "cell_type", "organ_major": "organ_major"}, nproc=4)
+  tk.tokenize_data("loom_data_directory", "output_directory", "output_prefix")
+"""
+
+import pickle
+from pathlib import Path
+
+import loompy as lp
+import numpy as np
+from datasets import Dataset
+
+GENE_MEDIAN_FILE = Path(__file__).parent / "gene_median_dictionary.pkl"
+TOKEN_DICTIONARY_FILE = Path(__file__).parent / "token_dictionary.pkl"
+
+
+def tokenize_cell(gene_vector, gene_tokens):
+    """
+    Convert normalized gene expression vector to tokenized rank value encoding.
+    """
+    # create array of gene vector with token indices
+    # mask undetected genes
+    nonzero_mask = np.nonzero(gene_vector)[0]
+    # sort by median-scaled gene values
+    sorted_indices = np.argsort(-gene_vector[nonzero_mask])
+    # tokenize
+    sentence_tokens = gene_tokens[nonzero_mask][sorted_indices]
+    return sentence_tokens
+
+
+class Tokenizer:
+    def __init__(
+        self,
+        custom_attr_name_dict,
+        nproc=1,
+        gene_median_file=GENE_MEDIAN_FILE,
+        token_dictionary_file=TOKEN_DICTIONARY_FILE,
+    ):
+        """
+        Initialize tokenizer.
+
+        Parameters
+        ----------
+        custom_attr_name_dict : dict
+            Dictionary of custom attributes to be added to the dataset.
+            Keys are the names of the attributes in the loom file.
+            Values are the names of the attributes in the dataset.
+        nproc : int
+            Number of processes to use for dataset mapping.
+        gene_median_file : Path
+            Path to pickle file containing dictionary of non-zero median
+            gene expression values across Genecorpus-30M.
+        token_dictionary_file : Path
+            Path to pickle file containing token dictionary (Ensembl IDs:token).
+        """
+        # dictionary of custom attributes {output dataset column name: input .loom column name}
+        self.custom_attr_name_dict = custom_attr_name_dict
+
+        # number of processes for dataset mapping
+        self.nproc = nproc
+
+        # load dictionary of gene normalization factors
+        # (non-zero median value of expression across Genecorpus-30M)
+        with open(gene_median_file, "rb") as f:
+            self.gene_median_dict = pickle.load(f)
+
+        # load token dictionary (Ensembl IDs:token)
+        with open(token_dictionary_file, "rb") as f:
+            self.gene_token_dict = pickle.load(f)
+
+        # gene keys for full vocabulary
+        self.gene_keys = list(self.gene_median_dict.keys())
+
+        # protein-coding and miRNA gene list dictionary for selecting .loom rows for tokenization
+        self.genelist_dict = dict(zip(self.gene_keys, [True] * len(self.gene_keys)))
+
+    def tokenize_data(self, loom_data_directory, output_directory, output_prefix):
+        """
+        Tokenize .loom files in loom_data_directory and save as tokenized .dataset in output_directory.
+
+        Parameters
+        ----------
+        loom_data_directory : Path
+            Path to directory containing loom files
+        output_directory : Path
+            Path to directory where tokenized data will be saved as .dataset
+        output_prefix : str
+            Prefix for output .dataset
+        """
+        tokenized_cells, cell_metadata = self.tokenize_files(loom_data_directory)
+        tokenized_dataset = self.create_dataset(tokenized_cells, cell_metadata)
+
+        output_path = (Path(output_directory) / output_prefix).with_suffix(".dataset")
+        tokenized_dataset.save_to_disk(output_path)
+
+    def tokenize_files(self, loom_data_directory):
+        tokenized_cells = []
+        cell_metadata = {attr_key: [] for attr_key in self.custom_attr_name_dict.keys()}
+
+        # loops through directories to tokenize .loom files
+        for loom_file_path in loom_data_directory.glob("*.loom"):
+            print(f"Tokenizing {loom_file_path}")
+            file_tokenized_cells, file_cell_metadata = self.tokenize_file(
+                loom_file_path
+            )
+            tokenized_cells += file_tokenized_cells
+            cell_metadata.update(file_cell_metadata)
+
+        return tokenized_cells, cell_metadata
+
+    def tokenize_file(self, loom_file_path):
+        file_cell_metadata = {
+            attr_key: [] for attr_key in self.custom_attr_name_dict.keys()
+        }
+
+        with lp.connect(str(loom_file_path)) as data:
+            # define coordinates of detected protein-coding or miRNA genes and vector of their normalization factors
+            coding_miRNA_loc = np.where(
+                [self.genelist_dict.get(i, False) for i in data.ra["ensembl_id"]]
+            )[0]
+            norm_factor_vector = np.array(
+                [
+                    self.gene_median_dict[i]
+                    for i in data.ra["ensembl_id"][coding_miRNA_loc]
+                ]
+            )
+            coding_miRNA_ids = data.ra["ensembl_id"][coding_miRNA_loc]
+            coding_miRNA_tokens = np.array(
+                [self.gene_token_dict[i] for i in coding_miRNA_ids]
+            )
+
+            # define coordinates of cells passing filters for inclusion (e.g. QC)
+            try:
+                data.ca["filter_pass"]
+            except NameError:
+                var_exists = False
+            else:
+                var_exists = True
+
+            if var_exists is True:
+                filter_pass_loc = np.where(
+                    [True if i == 1 else False for i in data.ca["filter_pass"]]
+                )[0]
+            elif var_exists is False:
+                print(
+                    f"{loom_file_path} has no column attribute 'filter_pass'; tokenizing all cells."
+                )
+                filter_pass_loc = np.array([i for i in range(data.shape[1])])
+
+            # scan through .loom files and tokenize cells
+            tokenized_cells = []
+            for (_ix, _selection, view) in data.scan(items=filter_pass_loc, axis=1):
+                # select subview with protein-coding and miRNA genes
+                subview = view.view[coding_miRNA_loc, :]
+
+                # normalize by total counts per cell and multiply by 10,000 to allocate bits to precision
+                # and normalize by gene normalization factors
+                subview_norm_array = (
+                    subview[:, :]
+                    / subview.ca.n_counts
+                    * 10_000
+                    / norm_factor_vector[:, None]
+                )
+                # tokenize subview gene vectors
+                tokenized_cells += [
+                    tokenize_cell(subview_norm_array[:, i], coding_miRNA_tokens)
+                    for i in range(subview_norm_array.shape[1])
+                ]
+
+                # add custom attributes for subview to dict
+                for k in file_cell_metadata.keys():
+                    file_cell_metadata[k] += subview.ca[k].tolist()
+
+        return tokenized_cells, file_cell_metadata
+
+    def create_dataset(self, tokenized_cells, cell_metadata):
+        # create dict for dataset creation
+        dataset_dict = {"input_ids": tokenized_cells}
+        dataset_dict.update(cell_metadata)
+
+        # create dataset
+        output_dataset = Dataset.from_dict(dataset_dict)
+
+        # truncate dataset
+        def truncate(example):
+            example["input_ids"] = example["input_ids"][0:2048]
+            return example
+
+        output_dataset_truncated = output_dataset.map(truncate, num_proc=self.nproc)
+
+        # measure lengths of dataset
+        def measure_length(example):
+            example["length"] = len(example["input_ids"])
+            return example
+
+        output_dataset_truncated_w_length = output_dataset_truncated.map(
+            measure_length, num_proc=self.nproc
+        )
+
+        return output_dataset_truncated_w_length