initial: weights + modeling code + lean config

modeling_virtual_cell.py

@@ -1,3 +1,18 @@
+"""
+Virtual Cell Patient Model — HuggingFace release.
+
+Architecture: PaSCient (Cui et al., 2025). ConvergeBio contribution: training
+recipe, data scale, and model parameters.
+
+Usage:
+    from transformers import AutoModel
+    model = AutoModel.from_pretrained(
+        "ConvergeBio/virtual-cell-patient", trust_remote_code=True
+    )
+    # input_ids: [batch, num_cells, num_genes]  float32 log-normalized expression
+    out = model(input_ids=x)   # out.logits: [batch, num_classes]
+"""
+
 from typing import List, Optional
 
 import torch
@@ -79,6 +94,8 @@ class MLP(nn.Module):
 
 
 class MLPCellEmbedder(nn.Module):
+    # Thin wrapper that preserves the .encoder attribute name required
+    # for state-dict key compatibility with the checkpoint.
     def __init__(
         self,
         n_genes: int,

train.py

@@ -0,0 +1,130 @@
+import argparse
+import os
+import sys
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import torch
+from datasets import load_dataset
+from transformers import EarlyStoppingCallback, Trainer, TrainingArguments
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from modeling_virtual_cell import VirtualCellPatientConfig, VirtualCellPatientModel
+
+
+@dataclass
+class PatientCollator:
+    def __call__(self, features: List[Dict]) -> Dict[str, torch.Tensor]:
+        return {
+            "input_ids": torch.stack([
+                torch.tensor(f["input_ids"], dtype=torch.float32) for f in features
+            ]),
+            "attention_mask": torch.stack([
+                torch.tensor(f["attention_mask"], dtype=torch.bool) for f in features
+            ]),
+            "labels":    torch.tensor([f["labels"]    for f in features], dtype=torch.long),
+            "entity_id": torch.tensor([f["entity_id"] for f in features], dtype=torch.long),
+        }
+
+
+class PatientTrainer(Trainer):
+    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
+        outputs = model(**inputs)
+        return (outputs.loss, outputs) if return_outputs else outputs.loss
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+
+    p.add_argument("--dataset_path", required=True,
+                   help="HF dataset ID or local path with train (and optionally validation) splits")
+    p.add_argument("--model_name_or_path", default="ConvergeBio/virtual-cell-patient")
+    p.add_argument("--hf_token",           default=None)
+    p.add_argument("--output_dir",         default="./vc_output")
+    p.add_argument("--from_scratch",       action="store_true")
+    p.add_argument("--freeze_embedder",    action="store_true")
+    p.add_argument("--num_classes",        type=int,   default=None)
+    p.add_argument("--num_train_epochs",   type=int,   default=15)
+    p.add_argument("--per_device_train_batch_size", type=int,   default=32)
+    p.add_argument("--per_device_eval_batch_size",  type=int,   default=32)
+    p.add_argument("--learning_rate",      type=float, default=1e-4)
+    p.add_argument("--weight_decay",       type=float, default=0.05)
+    p.add_argument("--warmup_ratio",       type=float, default=0.1)
+    p.add_argument("--lr_scheduler_type",             default="cosine")
+    p.add_argument("--patience",           type=int,   default=5)
+    p.add_argument("--num_workers",        type=int,   default=4)
+    p.add_argument("--wandb_project",      default=None)
+    p.add_argument("--run_name",           default=None)
+
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    ds = load_dataset(args.dataset_path)
+    train_ds = ds["train"]
+    val_ds: Optional[object] = ds.get("validation")
+
+    hf_kwargs = {"trust_remote_code": True}
+    if args.hf_token:
+        hf_kwargs["token"] = args.hf_token
+
+    config = VirtualCellPatientConfig.from_pretrained(args.model_name_or_path, **hf_kwargs)
+    if args.num_classes is not None:
+        config.num_classes = args.num_classes
+        config.id2label    = {str(i): str(i) for i in range(args.num_classes)}
+        config.label2id    = {str(i): i       for i in range(args.num_classes)}
+
+    if args.from_scratch:
+        model = VirtualCellPatientModel(config)
+    else:
+        model = VirtualCellPatientModel.from_pretrained(
+            args.model_name_or_path, config=config, **hf_kwargs
+        )
+
+    if args.freeze_embedder:
+        for param in model.patient_embedder.parameters():
+            param.requires_grad = False
+
+    if args.wandb_project:
+        os.environ["WANDB_PROJECT"] = args.wandb_project
+
+    has_val = val_ds is not None
+    training_args = TrainingArguments(
+        output_dir=args.output_dir,
+        num_train_epochs=args.num_train_epochs,
+        per_device_train_batch_size=args.per_device_train_batch_size,
+        per_device_eval_batch_size=args.per_device_eval_batch_size,
+        learning_rate=args.learning_rate,
+        weight_decay=args.weight_decay,
+        warmup_ratio=args.warmup_ratio,
+        lr_scheduler_type=args.lr_scheduler_type,
+        eval_strategy="epoch" if has_val else "no",
+        save_strategy="epoch",
+        load_best_model_at_end=has_val,
+        metric_for_best_model="eval_loss" if has_val else None,
+        greater_is_better=False,
+        report_to="wandb" if args.wandb_project else "none",
+        run_name=args.run_name,
+        dataloader_num_workers=args.num_workers,
+        remove_unused_columns=False,
+    )
+
+    callbacks = [EarlyStoppingCallback(args.patience)] if has_val else []
+
+    trainer = PatientTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_ds,
+        eval_dataset=val_ds,
+        data_collator=PatientCollator(),
+        callbacks=callbacks,
+    )
+
+    trainer.train()
+    trainer.save_model(args.output_dir)
+
+
+if __name__ == "__main__":
+    main()