Add option for modifying chunk size for anndata tokenizer

geneformer/tokenizer.py

@@ -11,18 +11,16 @@ Optional col (cell) attributes: any other cell metadata can be passed on to the
 Usage:
   from geneformer import TranscriptomeTokenizer
   tk = TranscriptomeTokenizer({"cell_type": "cell_type", "organ_major": "organ_major"}, nproc=4)
-  tk.tokenize_data("loom_data_directory", "output_directory", "output_prefix")
+  tk.tokenize_data("data_directory", "output_directory", "output_prefix")
 """
 
 from __future__ import annotations
-from typing import Literal
-import pickle
-from pathlib import Path
 
 import logging
-
+import pickle
 import warnings
-warnings.filterwarnings("ignore", message=".*The 'nopython' keyword.*")
+from pathlib import Path
+from typing import Literal
 
 import anndata as ad
 import loompy as lp
@@ -30,6 +28,7 @@ import numpy as np
 import scipy.sparse as sp
 from datasets import Dataset
 
+warnings.filterwarnings("ignore", message=".*The 'nopython' keyword.*")
 logger = logging.getLogger(__name__)
 
 GENE_MEDIAN_FILE = Path(__file__).parent / "gene_median_dictionary.pkl"
@@ -61,6 +60,7 @@ class TranscriptomeTokenizer:
         self,
         custom_attr_name_dict=None,
         nproc=1,
+        chunk_size=512,
         gene_median_file=GENE_MEDIAN_FILE,
         token_dictionary_file=TOKEN_DICTIONARY_FILE,
     ):
@@ -75,6 +75,8 @@ class TranscriptomeTokenizer:
             Values are the names of the attributes in the dataset.
         nproc : int
             Number of processes to use for dataset mapping.
+        chunk_size: int = 512
+            Chunk size for anndata tokenizer.
         gene_median_file : Path
             Path to pickle file containing dictionary of non-zero median
             gene expression values across Genecorpus-30M.
@@ -87,6 +89,9 @@ class TranscriptomeTokenizer:
         # number of processes for dataset mapping
         self.nproc = nproc
 
+        # chunk size for anndata tokenizer
+        self.chunk_size = chunk_size
+
         # load dictionary of gene normalization factors
         # (non-zero median value of expression across Genecorpus-30M)
         with open(gene_median_file, "rb") as f:
@@ -111,11 +116,11 @@ class TranscriptomeTokenizer:
         use_generator: bool = False,
     ):
         """
-        Tokenize .loom files in loom_data_directory and save as tokenized .dataset in output_directory.
+        Tokenize .loom files in data_directory and save as tokenized .dataset in output_directory.
 
         Parameters
         ----------
-        loom_data_directory : Path
+        data_directory : Path
             Path to directory containing loom files or anndata files
         output_directory : Path
             Path to directory where tokenized data will be saved as .dataset
@@ -129,7 +134,9 @@ class TranscriptomeTokenizer:
         tokenized_cells, cell_metadata = self.tokenize_files(
             Path(data_directory), file_format
         )
-        tokenized_dataset = self.create_dataset(tokenized_cells, cell_metadata, use_generator=use_generator)
+        tokenized_dataset = self.create_dataset(
+            tokenized_cells, cell_metadata, use_generator=use_generator
+        )
 
         output_path = (Path(output_directory) / output_prefix).with_suffix(".dataset")
         tokenized_dataset.save_to_disk(output_path)
@@ -140,7 +147,9 @@ class TranscriptomeTokenizer:
         tokenized_cells = []
         if self.custom_attr_name_dict is not None:
             cell_attr = [attr_key for attr_key in self.custom_attr_name_dict.keys()]
-            cell_metadata = {attr_key: [] for attr_key in self.custom_attr_name_dict.values()}
+            cell_metadata = {
+                attr_key: [] for attr_key in self.custom_attr_name_dict.values()
+            }
 
         # loops through directories to tokenize .loom files
         file_found = 0
@@ -155,17 +164,20 @@ class TranscriptomeTokenizer:
             tokenized_cells += file_tokenized_cells
             if self.custom_attr_name_dict is not None:
                 for k in cell_attr:
-                    cell_metadata[self.custom_attr_name_dict[k]] += file_cell_metadata[k]
+                    cell_metadata[self.custom_attr_name_dict[k]] += file_cell_metadata[
+                        k
+                    ]
             else:
                 cell_metadata = None
 
         if file_found == 0:
             logger.error(
-                f"No .{file_format} files found in directory {data_directory}.")
+                f"No .{file_format} files found in directory {data_directory}."
+            )
             raise
         return tokenized_cells, cell_metadata
 
-    def tokenize_anndata(self, adata_file_path, target_sum=10_000, chunk_size=512):
+    def tokenize_anndata(self, adata_file_path, target_sum=10_000):
         adata = ad.read(adata_file_path, backed="r")
 
         if self.custom_attr_name_dict is not None:
@@ -195,9 +207,7 @@ class TranscriptomeTokenizer:
             var_exists = True
 
         if var_exists:
-            filter_pass_loc = np.where(
-                [i == 1 for i in adata.obs["filter_pass"]]
-            )[0]
+            filter_pass_loc = np.where([i == 1 for i in adata.obs["filter_pass"]])[0]
         elif not var_exists:
             print(
                 f"{adata_file_path} has no column attribute 'filter_pass'; tokenizing all cells."
@@ -206,12 +216,12 @@ class TranscriptomeTokenizer:
 
         tokenized_cells = []
 
-        for i in range(0, len(filter_pass_loc), chunk_size):
-            idx = filter_pass_loc[i:i+chunk_size]
+        for i in range(0, len(filter_pass_loc), self.chunk_size):
+            idx = filter_pass_loc[i : i + self.chunk_size]
 
-            n_counts = adata[idx].obs['n_counts'].values[:, None]
+            n_counts = adata[idx].obs["n_counts"].values[:, None]
             X_view = adata[idx, coding_miRNA_loc].X
-            X_norm = (X_view / n_counts * target_sum / norm_factor_vector)
+            X_norm = X_view / n_counts * target_sum / norm_factor_vector
             X_norm = sp.csr_matrix(X_norm)
 
             tokenized_cells += [
@@ -259,9 +269,7 @@ class TranscriptomeTokenizer:
                 var_exists = True
 
             if var_exists:
-                filter_pass_loc = np.where(
-                    [i == 1 for i in data.ca["filter_pass"]]
-                )[0]
+                filter_pass_loc = np.where([i == 1 for i in data.ca["filter_pass"]])[0]
             elif not var_exists:
                 print(
                     f"{loom_file_path} has no column attribute 'filter_pass'; tokenizing all cells."
@@ -270,7 +278,7 @@ class TranscriptomeTokenizer:
 
             # scan through .loom files and tokenize cells
             tokenized_cells = []
-            for (_ix, _selection, view) in data.scan(items=filter_pass_loc, axis=1):
+            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, :]
 
@@ -297,7 +305,13 @@ class TranscriptomeTokenizer:
 
         return tokenized_cells, file_cell_metadata
 
-    def create_dataset(self, tokenized_cells, cell_metadata, use_generator=False, keep_uncropped_input_ids=False):
+    def create_dataset(
+        self,
+        tokenized_cells,
+        cell_metadata,
+        use_generator=False,
+        keep_uncropped_input_ids=False,
+    ):
         print("Creating dataset.")
         # create dict for dataset creation
         dataset_dict = {"input_ids": tokenized_cells}
@@ -306,30 +320,28 @@ class TranscriptomeTokenizer:
 
         # create dataset
         if use_generator:
+
             def dict_generator():
                 for i in range(len(tokenized_cells)):
                     yield {k: dataset_dict[k][i] for k in dataset_dict.keys()}
+
             output_dataset = Dataset.from_generator(dict_generator, num_proc=self.nproc)
         else:
             output_dataset = Dataset.from_dict(dataset_dict)
-            
+
         def format_cell_features(example):
             # Store original uncropped input_ids in separate feature
             if keep_uncropped_input_ids:
-                example['input_ids_uncropped'] = example['input_ids']
-                example['length_uncropped'] = len(example['input_ids'])
+                example["input_ids_uncropped"] = example["input_ids"]
+                example["length_uncropped"] = len(example["input_ids"])
 
             # Truncate/Crop input_ids to size 2,048
-            example['input_ids'] = example['input_ids'][0:2048]
-            example['length'] = len(example['input_ids'])
+            example["input_ids"] = example["input_ids"][0:2048]
+            example["length"] = len(example["input_ids"])
 
             return example
 
         output_dataset_truncated = output_dataset.map(
-            format_cell_features,
-            num_proc=self.nproc
+            format_cell_features, num_proc=self.nproc
         )
         return output_dataset_truncated
-
-
-