teddy/data_processing/tokenization/README.md

The tokenize_for_model.py script is designed to tokenize gene expression data for use in our models. It takes in processesed the, applies various tokenization techniques, and prepares it for training or inference.

General Workflow

The script follows these main steps: 0. Load Tokenization Arguments: The script starts by loading the tokenization arguments from a configuration file or dictionary.

1. Load Gene Tokenizer: It loads a pre-trained gene tokenizer based on the provided tokenization arguments.
2. Load AnnData: The script reads the gene expression data from an AnnData file.
3. Check Genes in Tokenizer: It verifies that the genes in the dataset are present in the tokenizer's vocabulary.
4. Build Token Array: The script constructs a token array for the genes in the dataset.
5. Convert Processed Layer to Dense: It converts the processed layer of the AnnData object to a dense matrix.
6. Tokenize in Batches: The script processes the data in batches, applying tokenization and optional binning or ranking.
7. Save Tokenized Data: Finally, the script saves the tokenized data to disk.

Tokenization Arguments

The script uses several tokenization arguments to control its behavior. Here is an explanation of each argument and the steps they influence:

• max_seq_len
  • Description: Specifies the maximum sequence length for the tokenized data.
  • Impact: Determines the number of genes to include in each tokenized sequence (cell). If add_cls is enabled, the sequence length is reduced by one to accommodate the CLS token.
• add_cls
  • Description: Indicates whether to prepend a CLS token to each sequence.
  • Impact: If enabled, a CLS token is added to the beginning of each sequence, and the sequence length is adjusted accordingly.
• cls_token_id
  • Description: The token ID to use for the CLS token.
  • Impact: If add_cls is enabled, this token ID is used for the CLS token.
• random_genes
  • Description: Specifies whether to select a random subset of genes before applying top-k selection
  • Impact: If enabled, a random subset of genes is selected for each batch, and then the top-k values are determined from this subset.
• include_zero_genes
  • Description: Indicates whether to include zero-expression genes in the tokenized data.
  • Impact: If enabled, zero-expression genes are included in the tokenized sequences. Otherwise, they are filtered out.
• bins
  • Description: Specifies the number of bins to use for binning expression values.
  • Impact: If set, the script bins the expression values into the specified number of bins. This argument is only relevant for TEDDY-X.
• continuous_rank
  • Description: Indicates whether to rank expression values continuously.
  • Impact: If enabled, the script ranks the expression values in the range [-1, 1]. This argument is only relevant for TEDDY-X.
• gene_seed
  • Description: A random seed for reproducibility.
  • Impact: If set, the script uses this seed to ensure reproducible random operations.
• gene_id_column
  • Description: The column name in the AnnData object that contains gene IDs.
  • Impact: The script uses this column to identify genes from vocab in the dataset.
• label_column
  • Description: The column name in the AnnData object that contains classification labels
  • Impact: If set, the script adds these labels to the tokenized data.
• bio_annotations
  • Description: Indicates whether to add biological annotations to the tokenized data.
  • Impact: If enabled, the script adds annotations such as disease, tissue, cell type, and sex to the tokenized data.
• disease_mapping, tissue_mapping, cell_mapping, sex_mapping
  • Description: File paths to JSON files containing mappings for biological annotations.
  • Impact: The script uses these mappings to convert biological annotations to token IDs.
• add_disease_annotation
  • Description: Indicates whether to override labels with disease annotations.
  • Impact: If enabled, the script overrides the labels with disease annotations.
• max_shard_samples
  • Description: The maximum number of samples per shard when saving the tokenized data.
  • Impact: The script splits the tokenized data into shards with the specified maximum number of samples.