Rename loom tokenizer function and modify example notebook for adata.

examples/tokenizing_scRNAseq_data.ipynb

@@ -1,31 +1,31 @@
 {
  "cells": [
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "a91bca46-c056-4784-8c6c-b0f5d3f33496",
    "metadata": {
     "tags": []
    },
    "source": [
-    "## Tokenizing .loom single cell RNA-seq data to rank value encoding .dataset format"
+    "## Tokenizing .loom or .h5ad single cell RNA-seq data to rank value encoding .dataset format"
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "350e6252-b783-494b-9767-f087eb868a15",
    "metadata": {},
    "source": [
-    "#### Input data is a directory with .loom files containing raw counts from single cell RNAseq data, including all genes detected in the transcriptome without feature selection. \n",
+    "#### Input data is a directory with .loom or .h5ad files containing raw counts from single cell RNAseq data, including all genes detected in the transcriptome without feature selection. The input file type is specified by the argument file_format in the tokenize_data function.\n",
     "\n",
-    "#### Genes should be labeled with Ensembl IDs (row attribute \"ensembl_id\"), which provide a unique identifer for conversion to tokens. Other forms of gene annotations (e.g. gene names) can be converted to Ensembl IDs via Ensembl Biomart. Cells should be labeled with the total read count in the cell (column attribute \"n_counts\") to be used for normalization.\n",
+    "#### The discussion below references the .loom file format, but the analagous labels are required for .h5ad files, just that they will be column instead of row attributes and vice versa due to the transposed format of the two file types.\n",
+    "\n",
+    "#### Genes should be labeled with Ensembl IDs (loom row attribute \"ensembl_id\"), which provide a unique identifer for conversion to tokens. Other forms of gene annotations (e.g. gene names) can be converted to Ensembl IDs via Ensembl Biomart. Cells should be labeled with the total read count in the cell (loom column attribute \"n_counts\") to be used for normalization.\n",
     "\n",
     "#### No cell metadata is required, but custom cell attributes may be passed onto the tokenized dataset by providing a dictionary of custom attributes to be added, which is formatted as loom_col_attr_name : desired_dataset_col_attr_name. For example, if the original .loom dataset has column attributes \"cell_type\" and \"organ_major\" and one would like to retain these attributes as labels in the tokenized dataset with the new names \"cell_type\" and \"organ\", respectively, the following custom attribute dictionary should be provided: {\"cell_type\": \"cell_type\", \"organ_major\": \"organ\"}. \n",
     "\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, one can use the relevant tools (such as Anndata tools) to convert the file to a .loom format prior to running the transcriptome tokenizer."
+    "#### 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."
    ]
   },
   {
@@ -45,8 +45,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tk = TranscriptomeTokenizer({\"cell_type\": \"cell_type\", \"organ_major\": \"organ_major\"}, nproc=4)\n",
-    "tk.tokenize_data(\"loom_data_directory\", \"output_directory\", \"output_prefix\", file_format=\"loom\")"
+    "tk = TranscriptomeTokenizer({\"cell_type\": \"cell_type\", \"organ_major\": \"organ\"}, nproc=16)\n",
+    "tk.tokenize_data(\"loom_data_directory\", \n",
+    "                 \"output_directory\", \n",
+    "                 \"output_prefix\", \n",
+    "                 file_format=\"loom\")"
    ]
   }
  ],

geneformer/tokenizer.py

@@ -146,7 +146,7 @@ class TranscriptomeTokenizer:
         file_found = 0
         # loops through directories to tokenize .loom or .h5ad files
         tokenize_file_fn = (
-            self.tokenize_file if file_format == "loom" else self.tokenize_anndata
+            self.tokenize_loom if file_format == "loom" else self.tokenize_anndata
         )
         for file_path in data_directory.glob("*.{}".format(file_format)):
             file_found = 1
@@ -209,8 +209,8 @@ class TranscriptomeTokenizer:
         for i in range(0, len(filter_pass_loc), chunk_size):
             idx = filter_pass_loc[i:i+chunk_size]
 
-            X_view = adata[idx, coding_miRNA_loc].X
             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 = sp.csr_matrix(X_norm)
 
@@ -228,7 +228,7 @@ class TranscriptomeTokenizer:
 
         return tokenized_cells, file_cell_metadata
 
-    def tokenize_file(self, loom_file_path, target_sum=10_000):
+    def tokenize_loom(self, loom_file_path, target_sum=10_000):
         if self.custom_attr_name_dict is not None:
             file_cell_metadata = {
                 attr_key: [] for attr_key in self.custom_attr_name_dict.keys()
@@ -298,7 +298,7 @@ class TranscriptomeTokenizer:
         return tokenized_cells, file_cell_metadata
 
     def create_dataset(self, tokenized_cells, cell_metadata, use_generator=False):
-        print("Creating dataset...")
+        print("Creating dataset.")
         # create dict for dataset creation
         dataset_dict = {"input_ids": tokenized_cells}
         if self.custom_attr_name_dict is not None:
@@ -329,4 +329,4 @@ class TranscriptomeTokenizer:
             measure_length, num_proc=self.nproc
         )
 
-        return output_dataset_truncated_w_length
+        return output_dataset_truncated_w_length