enable cell filter for OOM dataset and training with new data

[lidanwu]

@jdpearce4 Nice work and thanks for making this repository public!

In this PR, I added a few features to enable running the model on Blackwell GPU chip and training/fine-tuning the existing model with new data under new assay token. Hope this would also be helpful to the community.

Below is a summary of changes.

• relax version requirements for dependencies packages to work with torch that supports Blackwell GPU.

• update AnnDatasetOOM to enable cell filtering under same control knobs (filter_outliers and min_expressed_genes) as AnnDataset.

• Add FileAwareBatchSampler and update AnnDatasetOOM to enable process mujltiple rows/cells from same file at once.

  • Together these 2 changes allow faster I/O when going through all cells with in large anndata objects under back-end mode. (3x faster in my test during training).
  • Currently this PR enables file-aware batching of process for training mode only. But one could enable it for inference mode if desired.

• Add train submodule with yaml config template to source code.

  • The intention is to allow fine-tuning existing model with new data under different assay token.
  • So the train is deigned to
    1. allow optional addition of new assay tokens via using alternative json file expanded_assay_vocab and initializing its embedding based on existing assay token (Assay Embedding Init section in train config)
    2. work with large datasets via new AnnDatasetOOM dataset and FileAwareBatchSampler (use_oom_dataloader and enable_file_aware_batching).
    3. support multi-stage training with different freezing policies (Freezing & Training Policy section in train config)

• The training capacity is also added as a subcommand for transcriptformer CLI

All pytest passed under torch==2.12.0 with cuda 13.0 environment. And I have verified both the inference and new training sub-comands are working expected for me.

transcriptformer CLI

usage: transcriptformer [-h] {inference,download,download-data,train} ...

TranscriptFormer command-line interface

positional arguments:
  {inference,download,download-data,train}
                        Command to run
    inference           Run inference with a TranscriptFormer model
    download            Download and extract TranscriptFormer model artifacts
    download-data       Download CellxGene Discover datasets by species
    train               Train or continue training with expanded assay vocab

options:
  -h, --help            show this help message and exit

transcriptformer train CLI

usage: transcriptformer train [-h] --checkpoint-dir CHECKPOINT_DIR --output-dir OUTPUT_DIR --train-file TRAIN_FILE [--val-file VAL_FILE] [--expanded-assay-vocab EXPANDED_ASSAY_VOCAB]
                              [--obs-assay-col OBS_ASSAY_COL] [--gene-col-name GENE_COL_NAME] [--filter-to-vocabs] [--filter-outliers FILTER_OUTLIERS] [--sort-genes] [--randomize-genes]
                              [--min-expressed-genes MIN_EXPRESSED_GENES] [--n-data-workers N_DATA_WORKERS] [--resume-artifact-dir RESUME_ARTIFACT_DIR] [--resume-mode {weights,lightning}]
                              [--batch-size BATCH_SIZE] [--num-workers NUM_WORKERS] [--max-epochs MAX_EPOCHS] [--precision PRECISION] [--num-gpus NUM_GPUS] [--num-nodes NUM_NODES]
                              [--device {auto,cpu,cuda,mps}] [--lr LR] [--weight-decay WEIGHT_DECAY] [--adam-beta1 ADAM_BETA1] [--adam-beta2 ADAM_BETA2] [--adam-eps ADAM_EPS]
                              [--warmup-ratio WARMUP_RATIO] [--min-lr-ratio MIN_LR_RATIO] [--gene-id-loss-weight GENE_ID_LOSS_WEIGHT] [--softplus-approx | --no-softplus-approx]
                              [--init-default-source INIT_DEFAULT_SOURCE] [--assay-init-map ASSAY_INIT_MAP] [--freeze-transformer]
                              [--unfreeze-last-n-transformer-blocks UNFREEZE_LAST_N_TRANSFORMER_BLOCKS] [--freeze-gene-embeddings] [--freeze-count-head] [--freeze-gene-head]
                              [--train-aux-only] [--shuffle-expressed-each-batch] [--clip-counts CLIP_COUNTS] [--normalize-to-scale NORMALIZE_TO_SCALE] [--use-raw]
                              [--remove-duplicate-genes] [--use-oom-dataloader] [--file-aware-batching | --no-file-aware-batching] [--oom-batches-per-file OOM_BATCHES_PER_FILE]
                              [--seed SEED]

Fine-tune TranscriptFormer with expanded assay tokens and optional freezing.
...

Example command to continue training with expanded assay vocab

transcriptformer train \
    --checkpoint-dir ./checkpoints/tf_sapiens \
    --output-dir  ./outputs/tf_assay_expanded \
    --train-file ./data/train01.h5ad --train-file ./data/train02.h5ad \
    --expanded-assay-vocab ./data/expanded_assay_vocab.json \
    --assay-init-map "new assay 1==10x 3' v3" --assay-init-map "new assay 2==10x 3' v3" \
    --train-aux-only 

Example command to use a config file for train directly

python -m transcriptformer.cli.train \
    --config-path ./conf \
    --config-name train_config.yaml