geneformer.tokenizer

Geneformer tokenizer.

Input data:

Required format: raw counts scRNAseq data without feature selection as .loom or anndata file.

Required row (gene) attribute: “ensembl_id”; Ensembl ID for each gene.

Required col (cell) attribute: “n_counts”; total read counts in that cell.

Optional col (cell) attribute: “filter_pass”; binary indicator of whether cell should be tokenized based on user-defined filtering criteria.

Optional col (cell) attributes: any other cell metadata can be passed on to the tokenized dataset as a custom attribute dictionary as shown below.

Usage:

>>> from geneformer import TranscriptomeTokenizer
>>> tk = TranscriptomeTokenizer({"cell_type": "cell_type", "organ_major": "organ"}, nproc=4)
>>> tk.tokenize_data("data_directory", "output_directory", "output_prefix")

Description:

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.

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.

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.

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”}.

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.

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.

OF NOTE: Use model_version to auto-select settings for model version other than current default. For V1 model series (original Geneformer pretrained in 2021 on ~30M cells), one must use correct corresponding token dictionary and gene median file, set special_token to False, and set model_input_size to 2048. This argument enables auto-selection of these settings. (For V2 model series, special_token must be True and model_input_size is 4096.)

class TranscriptomeTokenizer(custom_attr_name_dict=None, nproc=1, chunk_size=512, model_input_size=4096, special_token=True, collapse_gene_ids=True, model_version='V2', gene_median_file=FILEPATH, token_dictionary_file=FILEPATH, gene_mapping_file=FILEPATH)

Initialize tokenizer.

Parameters:

custom_attr_name_dictNone, 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.

nprocint

Number of processes to use for dataset mapping.

chunk_sizeint = 512

Chunk size for anndata tokenizer.

model_input_sizeint = 4096

Max input size of model to truncate input to.

For the V1 model series, should be 2048. For the V2 model series, should be 4096.

special_tokenbool = True

Adds CLS token before and EOS token after rank value encoding.

For the V1 model series, should be False. For the V2 model series, should be True.

collapse_gene_idsbool = True

Whether to collapse gene IDs based on gene mapping dictionary.

model_versionstr

To auto-select settings for model version other than current default.

Current options: V1: models pretrained on ~30M cells, V2: models pretrained on ~104M cells

gene_median_filePath

Path to pickle file containing dictionary of non-zero median

gene expression values across Genecorpus.

token_dictionary_filePath

Path to pickle file containing token dictionary (Ensembl IDs:token).

gene_mapping_fileNone, Path

Path to pickle file containing dictionary for collapsing gene IDs.

tokenize_data(data_directory: Path | str, output_directory: Path | str, output_prefix: str, file_format: Literal['loom', 'h5ad'] = 'loom', use_generator: bool = False)

Tokenize .loom files in data_directory and save as tokenized .dataset in output_directory.

Parameters:

data_directoryPath

Path to directory containing loom files or anndata files

output_directoryPath

Path to directory where tokenized data will be saved as .dataset

output_prefixstr

Prefix for output .dataset

file_formatstr

Format of input files. Can be “loom” or “h5ad”.

use_generatorbool

Whether to use generator or dict for tokenization.