""" Module: tokenization.py This module provides a tokenization pipeline for preprocessed single-cell RNA sequencing (scRNA-seq) data. It converts gene expression data stored in AnnData format into tokenized sequences that can be used for downstream machine learning tasks, such as masked language modeling or classification. Main Features: - Tokenizes gene expression data into integer tokens using a custom GeneTokenizer. - Supports additional biological annotations (e.g., disease, tissue, cell type, sex). - Handles both top-k and random gene selection for tokenization. - Configurable via JSON-based hyperparameters or TokenizationArgs objects. - Saves tokenized data in Hugging Face Dataset format for efficient processing. Dependencies: - anndata, numpy, torch, datasets, tqdm Usage: - Run this script as a standalone program with a configuration file specifying the hyperparameters. - Import the `tokenize` function and call it with the data path, metadata path, and tokenization arguments. """ import gc import os import json import random import shutil from argparse import ArgumentParser from typing import Union import anndata as ad import numpy as np import torch from datasets import Dataset, load_from_disk from tqdm import tqdm from teddy.tokenizer.gene_tokenizer import GeneTokenizer from teddy.tokenizer.tokenization_args import TokenizationArgs ############################################################################### # Updated Functions ############################################################################### def _bin_values(vals_list, tokenization_args, no_sorting=False): """ Bins expression values into specified bins, assigning bin 0 to non-expressed genes when `include_zero_genes` is True. no_sorting=False => "positional chunk" approach for topk-sorted arrays - provided data_processing is expected to be sorted through topk (input expression values). no_sorting=True => simple bucketize approach ignoring the topk order - provided data_processing is not sorted (labels). """ binned_vals = [] for vals in vals_list: if isinstance(vals, np.ndarray): vals = torch.tensor(vals) vals_to_bin = vals # Original binning approach if not no_sorting: # "positional chunk" approach from the original code num_repetitions = max(1, len(vals_to_bin) // tokenization_args.bins) bin_pattern = torch.arange(0, tokenization_args.bins).unsqueeze(1).repeat(1, num_repetitions).flatten() # slice or pad to match the length of vals_to_bin if len(bin_pattern) > len(vals_to_bin): bin_pattern = bin_pattern[-len(vals_to_bin) :] else: extra = len(vals_to_bin) - len(bin_pattern) if extra > 0: bin_pattern = torch.cat([torch.zeros(extra), bin_pattern]) bin_pattern = bin_pattern.flip(0) binned_vals.append(bin_pattern) else: if len(vals_to_bin) > 0: bin_edges = torch.linspace(vals_to_bin.min(), vals_to_bin.max(), steps=tokenization_args.bins + 1) binned_non_zero_vals = torch.bucketize(vals_to_bin, bin_edges) binned_non_zero_vals = torch.clamp(binned_non_zero_vals, min=1) binned_tensor = binned_non_zero_vals.float() binned_vals.append(binned_tensor) else: binned_tensor = torch.zeros_like(vals_to_bin, dtype=torch.float) binned_vals.append(binned_tensor) return binned_vals def _rank_continuous(vals, tokenization_args): """ Ranks gene expression values in the range [-1, 1]. """ if isinstance(vals, np.ndarray): vals = torch.tensor(vals) if len(vals) > 0: ranked_vals = torch.linspace(-1, 1, steps=len(vals)).flip(0) else: ranked_vals = vals return ranked_vals def _prepare_tokenizer_args(tokenization_args: Union[dict, TokenizationArgs]): """ Prepares and validates tokenization arguments, ensuring reproducibility by setting random seeds if specified. """ if isinstance(tokenization_args, dict): load_dir = tokenization_args["load_dir"] save_dir = tokenization_args["save_dir"] token_args_obj = TokenizationArgs(**tokenization_args) else: # It's already TokenizationArgs load_dir = tokenization_args.load_dir save_dir = tokenization_args.save_dir token_args_obj = tokenization_args # If a random seed is specified, set it for reproducibility if token_args_obj.gene_seed is not None: random.seed(token_args_obj.gene_seed) np.random.seed(token_args_obj.gene_seed) torch.manual_seed(token_args_obj.gene_seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(token_args_obj.gene_seed) return token_args_obj, load_dir, save_dir def _check_genes_in_tokenizer(data: ad.AnnData, gene_id_column: str, tokenizer: GeneTokenizer): """ Checks if the genes in the dataset are present in the tokenizer's vocabulary. """ if gene_id_column == "index": gene_index = data.var.index else: gene_index = data.var[gene_id_column] # Check membership in vocab gene_in_vocab = np.where([g in tokenizer.vocab for g in gene_index])[0] coding_genes = gene_index[gene_in_vocab] ratio = len(gene_in_vocab) / len(data.var) if ratio < 0.1: raise OSError( f"Only {ratio:.2%} of gene IDs found in tokenizer vocab. " "Check gene_id_column or vocab mismatch." ) return gene_in_vocab, coding_genes, ratio def _build_batch_tensors(X_batch: torch.Tensor, token_array: torch.Tensor, token_args, data=None, obs_indices=None): """ Build topk or random subsets for each row in X_batch (batch_size x num_genes). Return gene_list, vals_list, labels_list. """ batch_size = X_batch.shape[0] seq_tokens = token_args.max_seq_len - 1 if token_args.add_cls else token_args.max_seq_len # If random_genes => pick random subset then topk that subset if token_args.random_genes: random_indices = torch.stack([torch.randperm(X_batch.shape[1])[:seq_tokens] for _ in range(batch_size)]) random_vals = torch.gather(X_batch, 1, random_indices) top_vals, rel_indices = torch.topk( random_vals, k=min(seq_tokens, random_vals.shape[1]), largest=True, sorted=True ) # Convert rel_indices => absolute indices top_indices = torch.gather(random_indices, 1, rel_indices) else: # normal topk top_vals, top_indices = torch.topk(X_batch, k=min(seq_tokens, X_batch.shape[1]), largest=True, sorted=True) gene_ids = token_array[top_indices] # If add_cls => prepend a CLS token if token_args.add_cls: cls_col = torch.tensor(token_args.cls_token_id).repeat(batch_size, 1) gene_ids = torch.cat([cls_col, gene_ids], dim=1) ones_col = torch.ones(batch_size, 1, dtype=top_vals.dtype) top_vals = torch.cat([ones_col, top_vals], dim=1) labels_list = None return gene_ids, top_vals, labels_list, None ############################################################################### # Main tokenize function ############################################################################### def tokenize(data_path: str, metadata_path: str, tokenization_args: Union[dict, TokenizationArgs]): """ Tokenizes gene expression data stored in AnnData format. Args: data_path (str): Path to the AnnData file containing preprocessed gene expression data. metadata_path (str): Path to the metadata file in JSON format. tokenization_args (Union[dict, TokenizationArgs]): Configuration for tokenization. """ token_args, load_dir, save_dir = _prepare_tokenizer_args(tokenization_args) # 1) Load GeneTokenizer tokenizer = GeneTokenizer.from_pretrained(token_args.tokenizer_name_or_path) if token_args.cls_token_id is None: token_args.cls_token_id = tokenizer.cls_token_id # 2) Load AnnData data = ad.read_h5ad(data_path) if "processed" not in data.layers: raise ValueError(f"Missing 'processed' layer in {data_path}") # 3) Genes in vocab gene_in_vocab, coding_genes, ratio = _check_genes_in_tokenizer(data, token_args.gene_id_column, tokenizer) print(f"{ratio:.2%} of genes found in tokenizer vocab") # 5) Build token array for these genes token_array = torch.tensor(tokenizer.encode(coding_genes.tolist(), add_special_tokens=False)) # 6) Convert processed layer to dense X_matrix = data.layers["processed"].toarray() # 7) Prepare final dictionary => HF Dataset all_data = {"gene_ids": [], "values": []} BATCH_SIZE = 512 n_obs = data.shape[0] for start_idx in tqdm(range(0, n_obs, BATCH_SIZE), desc="Tokenizing in batches"): end_idx = min(start_idx + BATCH_SIZE, n_obs) obs_indices = np.arange(start_idx, end_idx) X_batch = torch.tensor(X_matrix[obs_indices, :][:, gene_in_vocab], dtype=torch.float) gene_ids_batch, vals_batch, labels_batch, decoder_vals_batch = _build_batch_tensors( X_batch, token_array, token_args, data=None, obs_indices=None, ) final_gene_list = [] final_vals_list = [] final_labels_list = [] if "decoder_values" in data.layers: final_decoder_vals_list = [] # Filter out zero if needed # or keep them for row_idx in range(len(gene_ids_batch)): g_row = gene_ids_batch[row_idx] v_row = vals_batch[row_idx] if labels_batch is not None: lb_row = labels_batch[row_idx] else: lb_row = None if decoder_vals_batch is not None: dec_v_row = decoder_vals_batch[row_idx] else: dec_v_row = None if not token_args.include_zero_genes: nonzero_mask = v_row != 0 g_row = g_row[nonzero_mask] v_row = v_row[nonzero_mask] if lb_row is not None: lb_row = lb_row[nonzero_mask] if dec_v_row is not None: dec_v_row = dec_v_row[nonzero_mask] final_gene_list.append(g_row) final_vals_list.append(v_row) final_labels_list.append(lb_row) if "decoder_values" in data.layers: final_decoder_vals_list.append(dec_v_row) # If we do binning or rank => apply them if token_args.bins and token_args.continuous_rank: raise ValueError("Should not use bins and continuous_rank simultaneously.") if token_args.bins: # possibly do no_sorting if we are binning "labels" # we only do "no_sorting=True" for labels, but let's keep it simple for now final_vals_list = _bin_values(final_vals_list, token_args, no_sorting=False) elif token_args.continuous_rank: for i, vals in enumerate(final_vals_list): final_vals_list[i] = _rank_continuous(vals, token_args) # Add to all_data for row_idx in range(len(final_gene_list)): all_data["gene_ids"].append(final_gene_list[row_idx].tolist()) all_data["values"].append(final_vals_list[row_idx].tolist()) if token_args.label_column: all_data["labels"] = data.obs[token_args.label_column].cat.codes.values.tolist() # bio_annotations if token_args.bio_annotations: with open(token_args.disease_mapping) as f: disease_mapping = json.load(f) with open(token_args.tissue_mapping) as f: tissue_mapping = json.load(f) with open(token_args.cell_mapping) as f: cell_mapping = json.load(f) with open(token_args.sex_mapping) as f: sex_mapping = json.load(f) if "disease" not in data.obs.columns: data.obs["disease"] = "normal" if "tissue" not in data.obs.columns: data.obs["tissue"] = "cultured cell" if "sex" not in data.obs.columns: data.obs["sex"] = "unknown" if "cell_type" not in data.obs.columns: data.obs["cell_type"] = "unknown" mapped_diseases = [disease_mapping[k] for k in data.obs["disease"].tolist()] mapped_tissues = [tissue_mapping[k] for k in data.obs["tissue"].tolist()] mapped_cell_types = [cell_mapping[k] for k in data.obs["cell_type"].tolist()] mapped_sexes = [sex_mapping[k] for k in data.obs["sex"].tolist()] all_data["disease"] = tokenizer.encode(mapped_diseases, add_special_tokens=False) all_data["tissue"] = tokenizer.encode(mapped_tissues, add_special_tokens=False) all_data["cell_type"] = tokenizer.encode(mapped_cell_types, add_special_tokens=False) all_data["sex"] = tokenizer.encode(mapped_sexes, add_special_tokens=False) if token_args.add_disease_annotation: # We override "labels" with "disease" tokens all_data["labels"] = all_data["disease"] del data gc.collect() dataset = Dataset.from_dict(all_data) num_samples = len(dataset) if token_args.max_shard_samples: num_shards = num_samples // min(token_args.max_shard_samples, num_samples) else: num_shards = 1 # Compute the path of data_path relative to load_dir relative_data_path = os.path.relpath(data_path, load_dir) relative_metadata_path = os.path.relpath(metadata_path, load_dir) # Remove the ".h5ad" extension from data_path if desired no_extension_data_path = os.path.splitext(relative_data_path)[0] # Reconstruct the final paths under save_dir save_tokenized_data_path = os.path.join(save_dir, no_extension_data_path) save_metadata_path = os.path.join(save_dir, relative_metadata_path) dataset.save_to_disk(save_tokenized_data_path, num_shards=num_shards) shutil.copy(metadata_path, save_metadata_path) ############################################################################### # A simple shard function ############################################################################### def shard_hf_dataset(data_path: str, metadata_path: str, tokenization_args: Union[dict, TokenizationArgs]): """ Shards a Hugging Face Dataset into smaller chunks for efficient storage and processing. """ if isinstance(tokenization_args, dict): load_dir = tokenization_args["load_dir"] save_dir = tokenization_args["save_dir"] token_args_obj = TokenizationArgs(**tokenization_args) else: load_dir = tokenization_args.load_dir save_dir = tokenization_args.save_dir token_args_obj = tokenization_args all_data = load_from_disk(data_path) num_samples = len(all_data) if token_args_obj.max_shard_samples: num_shards = num_samples // min(token_args_obj.max_shard_samples, num_samples) else: num_shards = 1 save_tokenized_data_path = data_path.replace(load_dir, save_dir) save_metadata_path = metadata_path.replace(load_dir, save_dir) all_data.save_to_disk(save_tokenized_data_path, num_shards=num_shards) shutil.copy(metadata_path, save_metadata_path) ############################################################################### # Main block ############################################################################### if __name__ == "__main__": parser = ArgumentParser(description="Tokenize an AnnData file for downstream ML tasks.") parser.add_argument( "--data_path", type=str, required=True, help="Path to the .h5ad file containing the preprocessed scRNA-seq data." ) parser.add_argument( "--metadata_path", type=str, required=True, help="Path to the JSON file containing metadata." ) parser.add_argument( "--config_path", type=str, required=True, help="Path to the JSON file specifying tokenization hyperparameters." ) args = parser.parse_args() # Load tokenization arguments from JSON with open(args.config_path, "r") as f: tokenization_args = json.load(f) # Call the tokenize function tokenize( data_path=args.data_path, metadata_path=args.metadata_path, tokenization_args=tokenization_args )