enable cell filter for OOM dataset and training with new data

.gitignore

@@ -180,3 +180,10 @@ figures/
 # hydra
 .hydra/
 new_env/
+
+data/
+checkpoints/
+scripts/
+examples/
+outputs/
+conf/

README.md

@@ -65,7 +65,7 @@ uv pip install -e .
 ### Requirements
 
 Transcriptformer has the following core dependencies:
-- PyTorch (<=2.5.1, as 2.6.0+ may cause pickle errors)
+- PyTorch (>=2.5.1, <2.13) # original says PyTorch (<=2.5.1, as 2.6.0+ may cause pickle errors)
 - PyTorch Lightning
 - anndata
 - scanpy

conf/train_config.yaml

@@ -0,0 +1,88 @@
+defaults:
+  - _self_  # Load the current config file
+
+train:
+  # === Checkpointing & Output ===
+  checkpoint_dir: null         # Path to source checkpoint directory (required)
+  resume_artifact_dir: null    # Path to previous training output for resuming (optional)
+  resume_mode: weights         # 'weights' (load weights only) or 'lightning' (resume full state)
+  output_dir: ./train_outputs/tf_assay_expanded  # Output directory for new artifacts
+
+  checkpoint:
+    save_top_k: 1              # Number of best checkpoints to keep (by val/total_loss)
+    save_last: true            # Always save the last checkpoint
+
+  live_metrics:
+    enabled: true              # Export per-epoch summaries and live plots while training
+    every_n_epochs: 1          # Refresh metrics/plots every N epochs
+
+  # === Data ===
+  train_files: []              # List of training .h5ad files
+  val_files: []                # List of validation .h5ad files
+  expanded_assay_vocab: null   # Path to expanded assay vocab JSON (optional)
+  obs_assay_col: assay         # Required AnnData.obs column used to create aux tokens for model forward
+  obs_keys: null               # Optional list of extra AnnData.obs fields to carry in batch.obs (metadata only)
+
+  data_config:
+    _target_: transcriptformer.data.dataclasses.DataConfig
+    gene_col_name: ensembl_id      # Column name in AnnData.var containing gene identifiers
+    clip_counts: 30.0              # Maximum count value (higher values will be clipped), Null means inherit from checkpoint model config
+    filter_to_vocabs: true         # Whether to filter genes to only those in the vocabulary, Null means inherit from checkpoint model config
+    filter_outliers: 3.0          # Standard deviation threshold for filtering outlier cells (0.0 = no filtering), Null means inherit from checkpoint model config
+    normalize_to_scale: 0.0        # Scale factor for count normalization (0 = no normalization)
+    sort_genes: null               # Whether to sort genes by expression level, Null means inherit from checkpoint model config
+    randomize_genes: true          # Whether to randomize gene order, Null means inherit from checkpoint model config
+    min_expressed_genes: 200      # Minimum number of expressed genes required per cell, Null means inherit from checkpoint model config
+    use_raw: "auto"                 # Whether to use .raw.X (True), .X (False), or auto-detect (auto/null), Null means inherit from checkpoint model config
+    remove_duplicate_genes: false  # Remove duplicate genes instead of raising
+    n_data_workers: 4              # Data workers used by data module config
+
+  # === Dataloader ===
+  batch_size: 8                # Batch size
+  num_workers: 4               # DataLoader workers
+  use_oom_dataloader: true     # Use backed OOM dataset + file-aware batching (lower memory, mostly file-local batches)
+  enable_file_aware_batching: true # Training-only and DDP-safe: keeps file-local OOM batches; set false to use Lightning DistributedSampler path
+  oom_batches_per_file: 1      # In OOM mode, number of consecutive batches drawn from one file before interleaving
+
+
+  # === Training Loop ===
+  # Recommended: Aim for a global batch size of 4–5 million tokens (examples * seq_len * num_gpus * num_nodes * accumulate_grad_batches)
+  max_epochs: 5                # Number of epochs
+  precision: 16-mixed          # Precision (16-mixed, 32, bf16, etc.)
+  device: auto                 # Preferred device policy ('auto', 'cpu', 'cuda', 'mps'), same resolution path as inference
+  num_gpus: 1                  # Requested GPU count (1, >1, or -1 for all available on each node)
+  num_nodes: 1                 # Number of nodes (for DDP)
+  log_every_n_steps: 10        # Logging frequency
+  seed: 42                     # Random seed
+  gradient_clip_val: 1.0       # Max norm for gradient clipping (improves stability)
+  accumulate_grad_batches: 6   # Number of steps to accumulate gradients before optimizer step
+
+  # === Optimizer & Scheduler ===
+  lr: 5.5e-5                   # Learning rate, try 1.5e-4 if train from scratch
+  weight_decay: 0.05           # AdamW weight decay, try 0.1 if train from scratch
+  adam_beta1: 0.9              # AdamW beta1
+  adam_beta2: 0.95             # AdamW beta2
+  adam_eps: 1e-8               # AdamW epsilon
+  warmup_ratio: 0.1            # LR warmup ratio
+  min_lr_ratio: 0.18            # Minimum LR as fraction of initial, cap at 1e-5
+
+  # === Loss Weights ===
+  loss_config:
+    _target_: transcriptformer.data.dataclasses.LossConfig
+    gene_id_loss_weight: 1.0    # Gene ID prediction loss weight
+    softplus_approx: true       # Use softplus approximation in loss path
+
+  # === Assay Embedding Init ===
+  init_default_source: unknown  # Default source for new assay tokens
+  assay_init_map: []           # List of 'new=source' mappings
+
+  # === Freezing & Training Policy ===
+  freeze_transformer: false    # Freeze transformer encoder
+  unfreeze_last_n_transformer_blocks: 0 # If freeze_transformer=true, unfreeze only the last N encoder blocks (0 keeps full encoder frozen)
+  freeze_gene_embeddings: false # Freeze gene embeddings
+  freeze_count_head: false     # Freeze count decoder
+  freeze_gene_head: false      # Freeze gene decoder
+  train_aux_only: false        # Train only aux (assay) embeddings
+
+  # === Data Preprocessing ===
+  shuffle_expressed_each_batch: false # Shuffle expressed genes per batch

pyproject.toml

@@ -23,15 +23,15 @@ dependencies = [
   "hydra-core==1.3.2",
   "numpy==2.2.6",
   "pandas==2.2.2",
-  "psutil==5.9",
+  "psutil>=5.9,<8",
   "pynvml==12",
   "pytest==8.4",
   "pytorch-lightning==2.5.1",
   "requests==2.32.3",
   "scanpy==1.11.2",
   "scipy==1.15.3",
   "timeout-decorator==0.5",
-  "torch==2.5.1",
+  "torch>=2.5.1,<2.13",
 ]
 
 optional-dependencies.build = [

src/transcriptformer/cli/__init__.py

@@ -69,6 +69,7 @@
 from omegaconf import OmegaConf
 
 from transcriptformer.model.inference import run_inference
+from transcriptformer.train.engine import run_train_from_dict, setup_runtime_for_training
 
 # Suppress annoying warnings
 warnings.filterwarnings("ignore", category=FutureWarning, module="anndata")
@@ -284,6 +285,102 @@ def setup_download_data_parser(subparsers):
     )
 
 
+def setup_train_parser(subparsers):
+    """Setup the parser for the train command."""
+    parser = subparsers.add_parser(
+        "train",
+        help="Train or continue training with expanded assay vocab",
+        description="Fine-tune TranscriptFormer with expanded assay tokens and optional freezing.",
+    )
+
+    parser.add_argument("--checkpoint-dir", required=True, help="Base artifact directory with config.json/model_weights.pt")
+    parser.add_argument("--output-dir", required=True, help="Output artifact directory")
+    parser.add_argument("--train-file", action="append", required=True, help="Training .h5ad file (repeatable)")
+    parser.add_argument("--val-file", action="append", default=[], help="Validation .h5ad file (repeatable)")
+
+    parser.add_argument("--expanded-assay-vocab", help="Expanded assay_vocab.json path")
+    parser.add_argument("--obs-assay-col", default="assay", help="obs assay column")
+    parser.add_argument("--gene-col-name", default="ensembl_id", help="adata.var gene ID column")
+    parser.add_argument("--filter-to-vocabs", action="store_true", default=True, help="filter genes to vocabulary")
+    parser.add_argument("--filter-outliers", type=float, default=0.0, help="cell outlier filtering threshold")
+    parser.add_argument("--sort-genes", action="store_true", help="sort genes by expression")
+    parser.add_argument("--randomize-genes", action="store_true", help="randomize gene order")
+    parser.add_argument("--min-expressed-genes", type=int, default=0, help="minimum expressed genes per cell")
+    parser.add_argument("--n-data-workers", type=int, default=4, help="DataConfig n_data_workers value")
+
+    parser.add_argument("--resume-artifact-dir", default=None, help="Resume from previous output artifact directory")
+    parser.add_argument(
+        "--resume-mode",
+        default="weights",
+        choices=["weights", "lightning"],
+        help="Resume mode: weights (supports runtime policy changes) or lightning (resume optimizer/scheduler state)",
+    )
+
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument("--max-epochs", type=int, default=5)
+    parser.add_argument("--precision", default="16-mixed")
+    parser.add_argument(
+        "--num-gpus",
+        type=int,
+        default=1,
+        help="Number of GPUs per node to use for training (1=single, -1=all available, >1=specific number)",
+    )
+    parser.add_argument("--num-nodes", type=int, default=1)
+    parser.add_argument(
+        "--device",
+        default="auto",
+        choices=["auto", "cpu", "cuda", "mps"],
+        help="Preferred device policy for training",
+    )
+
+    parser.add_argument("--lr", type=float, default=5.5e-5)
+    parser.add_argument("--weight-decay", type=float, default=0.05)
+    parser.add_argument("--adam-beta1", type=float, default=0.9)
+    parser.add_argument("--adam-beta2", type=float, default=0.95)
+    parser.add_argument("--adam-eps", type=float, default=1e-8)
+    parser.add_argument("--warmup-ratio", type=float, default=0.1)
+    parser.add_argument("--min-lr-ratio", type=float, default=0.1)
+
+    parser.add_argument("--gene-id-loss-weight", type=float, default=1.0)
+    parser.add_argument("--softplus-approx", action=argparse.BooleanOptionalAction, default=True)
+
+    parser.add_argument("--init-default-source", default="unknown")
+    parser.add_argument("--assay-init-map", action="append", default=[], help="new_assay=source_assay mapping")
+
+    parser.add_argument("--freeze-transformer", action="store_true")
+    parser.add_argument(
+        "--unfreeze-last-n-transformer-blocks",
+        type=int,
+        default=0,
+        help="When --freeze-transformer is set, unfreeze only the last N encoder blocks (0 keeps full transformer frozen)",
+    )
+    parser.add_argument("--freeze-gene-embeddings", action="store_true")
+    parser.add_argument("--freeze-count-head", action="store_true")
+    parser.add_argument("--freeze-gene-head", action="store_true")
+    parser.add_argument("--train-aux-only", action="store_true")
+
+    parser.add_argument("--shuffle-expressed-each-batch", action="store_true")
+    parser.add_argument("--clip-counts", type=float, default=30.0)
+    parser.add_argument("--normalize-to-scale", type=float, default=0.0)
+    parser.add_argument("--use-raw", action="store_true")
+    parser.add_argument("--remove-duplicate-genes", action="store_true")
+    parser.add_argument("--use-oom-dataloader", action="store_true")
+    parser.add_argument(
+        "--file-aware-batching",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Keep OOM batches mostly file-local; disable to fall back to Lightning DistributedSampler batching",
+    )
+    parser.add_argument(
+        "--oom-batches-per-file",
+        type=int,
+        default=1,
+        help="In OOM mode, number of consecutive batches to draw from one file before interleaving",
+    )
+    parser.add_argument("--seed", type=int, default=42)
+
+
 def run_inference_cli(args):
     """Run inference using command line arguments."""
     # Only print logo if not in distributed mode (avoids duplicates)
@@ -452,6 +549,68 @@ def run_download_data_cli(args):
         sys.exit(1)
 
 
+def run_train_cli(args):
+    """Run training from CLI args."""
+    setup_runtime_for_training()
+    cfg = {
+        "checkpoint_dir": args.checkpoint_dir,
+        "resume_artifact_dir": args.resume_artifact_dir,
+        "resume_mode": args.resume_mode,
+        "output_dir": args.output_dir,
+        "train_files": args.train_file,
+        "val_files": args.val_file,
+        "expanded_assay_vocab": args.expanded_assay_vocab,
+        "obs_assay_col": args.obs_assay_col,
+        "data_config": {
+            "gene_col_name": args.gene_col_name,
+            "clip_counts": args.clip_counts,
+            "filter_to_vocabs": args.filter_to_vocabs,
+            "filter_outliers": args.filter_outliers,
+            "normalize_to_scale": args.normalize_to_scale,
+            "sort_genes": args.sort_genes,
+            "randomize_genes": args.randomize_genes,
+            "min_expressed_genes": args.min_expressed_genes,
+            "use_raw": args.use_raw,
+            "remove_duplicate_genes": args.remove_duplicate_genes,
+            "n_data_workers": args.n_data_workers,
+        },
+        "batch_size": args.batch_size,
+        "num_workers": args.num_workers,
+        "max_epochs": args.max_epochs,
+        "precision": args.precision,
+        "device": args.device,
+        "num_gpus": args.num_gpus,
+        "num_nodes": args.num_nodes,
+        "lr": args.lr,
+        "weight_decay": args.weight_decay,
+        "adam_beta1": args.adam_beta1,
+        "adam_beta2": args.adam_beta2,
+        "adam_eps": args.adam_eps,
+        "warmup_ratio": args.warmup_ratio,
+        "min_lr_ratio": args.min_lr_ratio,
+        "loss_config": {
+            "gene_id_loss_weight": args.gene_id_loss_weight,
+            "softplus_approx": args.softplus_approx,
+        },
+        "init_default_source": args.init_default_source,
+        "assay_init_map": args.assay_init_map,
+        "freeze_transformer": args.freeze_transformer,
+        "unfreeze_last_n_transformer_blocks": args.unfreeze_last_n_transformer_blocks,
+        "freeze_gene_embeddings": args.freeze_gene_embeddings,
+        "freeze_count_head": args.freeze_count_head,
+        "freeze_gene_head": args.freeze_gene_head,
+        "train_aux_only": args.train_aux_only,
+        "shuffle_expressed_each_batch": args.shuffle_expressed_each_batch,
+        "use_oom_dataloader": args.use_oom_dataloader,
+        "enable_file_aware_batching": args.file_aware_batching,
+        "oom_batches_per_file": args.oom_batches_per_file,
+        "seed": args.seed,
+    }
+
+    result = run_train_from_dict(cfg)
+    print(f"Training complete. Artifacts saved to {result['output_dir']}")
+
+
 def main():
     """Main entry point for the CLI."""
     parser = argparse.ArgumentParser(
@@ -465,6 +624,7 @@ def main():
     setup_inference_parser(subparsers)
     setup_download_parser(subparsers)
     setup_download_data_parser(subparsers)
+    setup_train_parser(subparsers)
 
     # Parse arguments
     args = parser.parse_args()
@@ -480,6 +640,8 @@ def main():
         run_download_cli(args)
     elif args.command == "download-data":
         run_download_data_cli(args)
+    elif args.command == "train":
+        run_train_cli(args)
 
 
 if __name__ == "__main__":

src/transcriptformer/cli/conf/inference_config.yaml

@@ -25,7 +25,7 @@ model:
     emb_type: cell # Type of embeddings to extract: "cell" for mean-pooled cell embeddings or "cge" for contextual gene embeddings
     num_gpus: 1 # Number of GPUs to use for inference (1 = single GPU, -1 = all available GPUs, >1 = specific number)
     device: auto # Specific device to use: "auto" (best available), "cpu", "cuda", "mps"
-    use_oom_dataloader: false # Use OOM-safe map-style DataLoader with DistributedSampler
+    use_oom_dataloader: false # Use OOM-safe map-style DataLoader (Lightning handles distributed sampling in DDP)
 
   data_config:
     _target_: transcriptformer.data.dataclasses.DataConfig