initial: weights + modeling code + lean config

train.py

@@ -4,9 +4,12 @@ import sys
 from dataclasses import dataclass
 from typing import Dict, List, Optional
 
+import numpy as np
 import torch
-from datasets import load_dataset
+from datasets import DatasetDict, load_dataset
+from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
 from transformers import EarlyStoppingCallback, Trainer, TrainingArguments
+from transformers.trainer_utils import EvalPrediction
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 from modeling_virtual_cell import VirtualCellPatientConfig, VirtualCellPatientModel
@@ -27,11 +30,60 @@ class PatientCollator:
         }
 
 
+def _patient_predictions(logits: np.ndarray, entity_ids: np.ndarray):
+    """Average softmax probabilities across augmented views, one row per patient."""
+    entity_ids = np.asarray(entity_ids).astype(str)
+    unique = np.unique(entity_ids)
+    agg = []
+    for eid in unique:
+        views = logits[entity_ids == eid]
+        exp = np.exp(views - np.max(views, axis=1, keepdims=True))
+        agg.append(np.mean(exp / exp.sum(axis=1, keepdims=True), axis=0))
+    return np.array(agg), unique
+
+
+def _clf_metrics(y_true: np.ndarray, y_pred: np.ndarray, prefix: str) -> Dict[str, float]:
+    return {
+        f"{prefix}accuracy":  accuracy_score(y_true, y_pred),
+        f"{prefix}f1_macro":  f1_score(y_true, y_pred, average="macro", zero_division=0),
+        f"{prefix}precision": precision_score(y_true, y_pred, average="macro", zero_division=0),
+        f"{prefix}recall":    recall_score(y_true, y_pred, average="macro", zero_division=0),
+    }
+
+
+def compute_metrics(eval_pred: EvalPrediction) -> Dict[str, float]:
+    logits_with_entity = eval_pred.predictions  # (N, num_classes + 1)
+    logits     = logits_with_entity[:, :-1]
+    entity_ids = logits_with_entity[:, -1].astype(int)
+    labels     = eval_pred.label_ids
+
+    metrics = _clf_metrics(labels, np.argmax(logits, axis=1), "per_view/")
+
+    patient_preds, unique_entities = _patient_predictions(logits, entity_ids)
+    patient_labels = np.array([
+        labels[np.where(entity_ids == int(eid))[0][0]]
+        for eid in unique_entities
+    ])
+    metrics.update(_clf_metrics(patient_labels, np.argmax(patient_preds, axis=1), "patient/"))
+
+    return metrics
+
+
 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 prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
+        entity_id = inputs.pop("entity_id")
+        loss, logits, labels = super().prediction_step(
+            model, inputs, prediction_loss_only, ignore_keys=ignore_keys
+        )
+        if logits is not None:
+            entity_col = entity_id.float().unsqueeze(1).to(logits.device)
+            logits = torch.cat([logits, entity_col], dim=1)
+        return loss, logits, labels
+
 
 def parse_args():
     p = argparse.ArgumentParser()
@@ -53,6 +105,7 @@ def parse_args():
     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("--prefetch_factor",    type=int,   default=2)
     p.add_argument("--wandb_project",      default=None)
     p.add_argument("--run_name",           default=None)
 
@@ -62,7 +115,10 @@ def parse_args():
 def main():
     args = parse_args()
 
-    ds = load_dataset(args.dataset_path)
+    if os.path.isdir(args.dataset_path):
+        ds = DatasetDict.load_from_disk(args.dataset_path)
+    else:
+        ds = load_dataset(args.dataset_path, num_proc=args.num_workers)
     train_ds = ds["train"]
     val_ds: Optional[object] = ds.get("validation")
 
@@ -108,6 +164,9 @@ def main():
         report_to="wandb" if args.wandb_project else "none",
         run_name=args.run_name,
         dataloader_num_workers=args.num_workers,
+        dataloader_prefetch_factor=args.prefetch_factor if args.num_workers > 0 else None,
+        dataloader_persistent_workers=args.num_workers > 0,
+        dataloader_pin_memory=True,
         remove_unused_columns=False,
     )
 
@@ -119,6 +178,7 @@ def main():
         train_dataset=train_ds,
         eval_dataset=val_ds,
         data_collator=PatientCollator(),
+        compute_metrics=compute_metrics if has_val else None,
         callbacks=callbacks,
     )
 

wandb/debug-internal.log

@@ -0,0 +1,13 @@
+{"time":"2026-05-03T17:12:14.512438+03:00","level":"INFO","msg":"stream: starting","core version":"0.21.0"}
+{"time":"2026-05-03T17:12:15.049447+03:00","level":"INFO","msg":"stream: created new stream","id":"h9m78x54"}
+{"time":"2026-05-03T17:12:15.049472+03:00","level":"INFO","msg":"stream: started","id":"h9m78x54"}
+{"time":"2026-05-03T17:12:15.049488+03:00","level":"INFO","msg":"writer: Do: started","stream_id":"h9m78x54"}
+{"time":"2026-05-03T17:12:15.049533+03:00","level":"INFO","msg":"sender: started","stream_id":"h9m78x54"}
+{"time":"2026-05-03T17:12:15.049551+03:00","level":"INFO","msg":"handler: started","stream_id":"h9m78x54"}
+{"time":"2026-05-03T17:12:15.531811+03:00","level":"ERROR","msg":"git repo not found","error":"repository does not exist"}
+{"time":"2026-05-03T17:14:51.01985+03:00","level":"INFO","msg":"stream: closing","id":"h9m78x54"}
+{"time":"2026-05-03T17:14:51.643326+03:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"}
+{"time":"2026-05-03T17:14:51.997995+03:00","level":"INFO","msg":"sender: closed","stream_id":"h9m78x54"}
+{"time":"2026-05-03T17:14:51.998011+03:00","level":"INFO","msg":"handler: closed","stream_id":"h9m78x54"}
+{"time":"2026-05-03T17:14:51.998039+03:00","level":"INFO","msg":"writer: Close: closed","stream_id":"h9m78x54"}
+{"time":"2026-05-03T17:14:51.998606+03:00","level":"INFO","msg":"stream: closed","id":"h9m78x54"}

wandb/debug.log

@@ -0,0 +1,25 @@
+2026-05-03 17:12:13,856 INFO    MainThread:63423 [wandb_setup.py:_flush():80] Current SDK version is 0.21.0
+2026-05-03 17:12:13,857 INFO    MainThread:63423 [wandb_setup.py:_flush():80] Configure stats pid to 63423
+2026-05-03 17:12:13,857 INFO    MainThread:63423 [wandb_setup.py:_flush():80] Loading settings from /Users/daniellemillersayag/.config/wandb/settings
+2026-05-03 17:12:13,857 INFO    MainThread:63423 [wandb_setup.py:_flush():80] Loading settings from /Users/daniellemillersayag/Documents/vcell/paper/hf-release/wandb/settings
+2026-05-03 17:12:13,857 INFO    MainThread:63423 [wandb_setup.py:_flush():80] Loading settings from environment variables
+2026-05-03 17:12:13,857 INFO    MainThread:63423 [wandb_init.py:setup_run_log_directory():703] Logging user logs to /Users/daniellemillersayag/Documents/vcell/paper/hf-release/wandb/run-20260503_171213-h9m78x54/logs/debug.log
+2026-05-03 17:12:13,857 INFO    MainThread:63423 [wandb_init.py:setup_run_log_directory():704] Logging internal logs to /Users/daniellemillersayag/Documents/vcell/paper/hf-release/wandb/run-20260503_171213-h9m78x54/logs/debug-internal.log
+2026-05-03 17:12:13,858 INFO    MainThread:63423 [wandb_init.py:init():830] calling init triggers
+2026-05-03 17:12:13,858 INFO    MainThread:63423 [wandb_init.py:init():835] wandb.init called with sweep_config: {}
+config: {'_wandb': {'code_path': 'code/train.py'}}
+2026-05-03 17:12:13,858 INFO    MainThread:63423 [wandb_init.py:init():871] starting backend
+2026-05-03 17:12:14,495 INFO    MainThread:63423 [wandb_init.py:init():874] sending inform_init request
+2026-05-03 17:12:14,511 INFO    MainThread:63423 [wandb_init.py:init():882] backend started and connected
+2026-05-03 17:12:14,513 INFO    MainThread:63423 [wandb_init.py:init():953] updated telemetry
+2026-05-03 17:12:14,513 INFO    MainThread:63423 [wandb_init.py:init():977] communicating run to backend with 90.0 second timeout
+2026-05-03 17:12:15,529 INFO    MainThread:63423 [wandb_init.py:init():1029] starting run threads in backend
+2026-05-03 17:12:15,651 INFO    MainThread:63423 [wandb_run.py:_console_start():2458] atexit reg
+2026-05-03 17:12:15,651 INFO    MainThread:63423 [wandb_run.py:_redirect():2306] redirect: wrap_raw
+2026-05-03 17:12:15,651 INFO    MainThread:63423 [wandb_run.py:_redirect():2375] Wrapping output streams.
+2026-05-03 17:12:15,651 INFO    MainThread:63423 [wandb_run.py:_redirect():2398] Redirects installed.
+2026-05-03 17:12:15,652 INFO    MainThread:63423 [wandb_init.py:init():1075] run started, returning control to user process
+2026-05-03 17:12:15,653 INFO    MainThread:63423 [wandb_run.py:_config_callback():1363] config_cb None None {'return_dict': True, 'output_hidden_states': False, 'output_attentions': False, 'torchscript': False, 'torch_dtype': 'float32', 'use_bfloat16': False, 'tf_legacy_loss': False, 'pruned_heads': {}, 'tie_word_embeddings': True, 'chunk_size_feed_forward': 0, 'is_encoder_decoder': False, 'is_decoder': False, 'cross_attention_hidden_size': None, 'add_cross_attention': False, 'tie_encoder_decoder': False, 'max_length': 20, 'min_length': 0, 'do_sample': False, 'early_stopping': False, 'num_beams': 1, 'num_beam_groups': 1, 'diversity_penalty': 0.0, 'temperature': 1.0, 'top_k': 50, 'top_p': 1.0, 'typical_p': 1.0, 'repetition_penalty': 1.0, 'length_penalty': 1.0, 'no_repeat_ngram_size': 0, 'encoder_no_repeat_ngram_size': 0, 'bad_words_ids': None, 'num_return_sequences': 1, 'output_scores': False, 'return_dict_in_generate': False, 'forced_bos_token_id': None, 'forced_eos_token_id': None, 'remove_invalid_values': False, 'exponential_decay_length_penalty': None, 'suppress_tokens': None, 'begin_suppress_tokens': None, 'architectures': ['VirtualCellPatientModel'], 'finetuning_task': None, 'id2label': {0: 'oncological', 1: 'immune_inflammatory', 2: 'neurological', 3: 'metabolic_vascular', 4: 'gastrointestinal', 5: 'respiratory', 6: 'epithelial_barrier', 7: 'sensory_specialized', 8: 'healthy_control', 9: 'other'}, 'label2id': {'oncological': 0, 'immune_inflammatory': 1, 'neurological': 2, 'metabolic_vascular': 3, 'gastrointestinal': 4, 'respiratory': 5, 'epithelial_barrier': 6, 'sensory_specialized': 7, 'healthy_control': 8, 'other': 9}, 'tokenizer_class': None, 'prefix': None, 'bos_token_id': None, 'pad_token_id': None, 'eos_token_id': None, 'sep_token_id': None, 'decoder_start_token_id': None, 'task_specific_params': None, 'problem_type': None, '_name_or_path': '/Users/daniellemillersayag/Documents/vcell/paper/hf-release', '_attn_implementation_autoset': True, 'transformers_version': '4.51.3', 'model_type': 'virtual_cell_patient', 'auto_map': {'AutoConfig': 'modeling_virtual_cell.VirtualCellPatientConfig', 'AutoModel': 'modeling_virtual_cell.VirtualCellPatientModel'}, 'n_genes': 18301, 'embed_dim': 512, 'hidden_dim': [4096, 1024], 'dropout': 0.1, 'residual': False, 'activation': 'prelu', 'attention_hidden_dim': 512, 'num_classes': 10, 'classifier_dropout': 0.1, 'output_dir': '/tmp/vc_smoke_test_wandb', 'overwrite_output_dir': False, 'do_train': False, 'do_eval': True, 'do_predict': False, 'eval_strategy': 'epoch', 'prediction_loss_only': False, 'per_device_train_batch_size': 4, 'per_device_eval_batch_size': 4, 'per_gpu_train_batch_size': None, 'per_gpu_eval_batch_size': None, 'gradient_accumulation_steps': 1, 'eval_accumulation_steps': None, 'eval_delay': 0, 'torch_empty_cache_steps': None, 'learning_rate': 0.0001, 'weight_decay': 0.05, 'adam_beta1': 0.9, 'adam_beta2': 0.999, 'adam_epsilon': 1e-08, 'max_grad_norm': 1.0, 'num_train_epochs': 2, 'max_steps': -1, 'lr_scheduler_type': 'cosine', 'lr_scheduler_kwargs': {}, 'warmup_ratio': 0.1, 'warmup_steps': 0, 'log_level': 'passive', 'log_level_replica': 'warning', 'log_on_each_node': True, 'logging_dir': '/tmp/vc_smoke_test_wandb/runs/May03_17-12-12_Mac.lan', 'logging_strategy': 'steps', 'logging_first_step': False, 'logging_steps': 500, 'logging_nan_inf_filter': True, 'save_strategy': 'epoch', 'save_steps': 500, 'save_total_limit': None, 'save_safetensors': True, 'save_on_each_node': False, 'save_only_model': False, 'restore_callback_states_from_checkpoint': False, 'no_cuda': False, 'use_cpu': False, 'use_mps_device': False, 'seed': 42, 'data_seed': None, 'jit_mode_eval': False, 'use_ipex': False, 'bf16': False, 'fp16': False, 'fp16_opt_level': 'O1', 'half_precision_backend': 'auto', 'bf16_full_eval': False, 'fp16_full_eval': False, 'tf32': None, 'local_rank': 0, 'ddp_backend': None, 'tpu_num_cores': None, 'tpu_metrics_debug': False, 'debug': [], 'dataloader_drop_last': False, 'eval_steps': None, 'dataloader_num_workers': 2, 'dataloader_prefetch_factor': 2, 'past_index': -1, 'run_name': 'smoke-test', 'disable_tqdm': False, 'remove_unused_columns': False, 'label_names': None, 'load_best_model_at_end': True, 'metric_for_best_model': 'eval_loss', 'greater_is_better': False, 'ignore_data_skip': False, 'fsdp': [], 'fsdp_min_num_params': 0, 'fsdp_config': {'min_num_params': 0, 'xla': False, 'xla_fsdp_v2': False, 'xla_fsdp_grad_ckpt': False}, 'tp_size': 0, 'fsdp_transformer_layer_cls_to_wrap': None, 'accelerator_config': {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}, 'deepspeed': None, 'label_smoothing_factor': 0.0, 'optim': 'adamw_torch', 'optim_args': None, 'adafactor': False, 'group_by_length': False, 'length_column_name': 'length', 'report_to': ['wandb'], 'ddp_find_unused_parameters': None, 'ddp_bucket_cap_mb': None, 'ddp_broadcast_buffers': None, 'dataloader_pin_memory': True, 'dataloader_persistent_workers': True, 'skip_memory_metrics': True, 'use_legacy_prediction_loop': False, 'push_to_hub': False, 'resume_from_checkpoint': None, 'hub_model_id': None, 'hub_strategy': 'every_save', 'hub_token': '<HUB_TOKEN>', 'hub_private_repo': None, 'hub_always_push': False, 'gradient_checkpointing': False, 'gradient_checkpointing_kwargs': None, 'include_inputs_for_metrics': False, 'include_for_metrics': [], 'eval_do_concat_batches': True, 'fp16_backend': 'auto', 'push_to_hub_model_id': None, 'push_to_hub_organization': None, 'push_to_hub_token': '<PUSH_TO_HUB_TOKEN>', 'mp_parameters': '', 'auto_find_batch_size': False, 'full_determinism': False, 'torchdynamo': None, 'ray_scope': 'last', 'ddp_timeout': 1800, 'torch_compile': False, 'torch_compile_backend': None, 'torch_compile_mode': None, 'include_tokens_per_second': False, 'include_num_input_tokens_seen': False, 'neftune_noise_alpha': None, 'optim_target_modules': None, 'batch_eval_metrics': False, 'eval_on_start': False, 'use_liger_kernel': False, 'eval_use_gather_object': False, 'average_tokens_across_devices': False}
+2026-05-03 17:12:15,654 INFO    MainThread:63423 [wandb_config.py:__setitem__():154] [no run ID] config set model/num_parameters = 79963661 - <bound method Run._config_callback of <wandb.sdk.wandb_run.Run object at 0x14e776850>>
+2026-05-03 17:12:15,654 INFO    MainThread:63423 [wandb_run.py:_config_callback():1363] config_cb model/num_parameters 79963661 None
+2026-05-03 17:14:51,017 INFO    MsgRouterThr:63423 [mailbox.py:close():129] [no run ID] Closing mailbox, abandoning 1 handles.

wandb/run-20260503_171213-h9m78x54/files/code/train.py

@@ -0,0 +1,190 @@
+import argparse
+import os
+import sys
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+from datasets import DatasetDict, load_dataset
+from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
+from transformers import EarlyStoppingCallback, Trainer, TrainingArguments
+from transformers.trainer_utils import EvalPrediction
+
+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),
+        }
+
+
+def _patient_predictions(logits: np.ndarray, entity_ids: np.ndarray):
+    """Average softmax probabilities across augmented views, one row per patient."""
+    entity_ids = np.asarray(entity_ids).astype(str)
+    unique = np.unique(entity_ids)
+    agg = []
+    for eid in unique:
+        views = logits[entity_ids == eid]
+        exp = np.exp(views - np.max(views, axis=1, keepdims=True))
+        agg.append(np.mean(exp / exp.sum(axis=1, keepdims=True), axis=0))
+    return np.array(agg), unique
+
+
+def _clf_metrics(y_true: np.ndarray, y_pred: np.ndarray, prefix: str) -> Dict[str, float]:
+    return {
+        f"{prefix}accuracy":  accuracy_score(y_true, y_pred),
+        f"{prefix}f1_macro":  f1_score(y_true, y_pred, average="macro", zero_division=0),
+        f"{prefix}precision": precision_score(y_true, y_pred, average="macro", zero_division=0),
+        f"{prefix}recall":    recall_score(y_true, y_pred, average="macro", zero_division=0),
+    }
+
+
+def compute_metrics(eval_pred: EvalPrediction) -> Dict[str, float]:
+    logits_with_entity = eval_pred.predictions  # (N, num_classes + 1)
+    logits     = logits_with_entity[:, :-1]
+    entity_ids = logits_with_entity[:, -1].astype(int)
+    labels     = eval_pred.label_ids
+
+    metrics = _clf_metrics(labels, np.argmax(logits, axis=1), "per_view/")
+
+    patient_preds, unique_entities = _patient_predictions(logits, entity_ids)
+    patient_labels = np.array([
+        labels[np.where(entity_ids == int(eid))[0][0]]
+        for eid in unique_entities
+    ])
+    metrics.update(_clf_metrics(patient_labels, np.argmax(patient_preds, axis=1), "patient/"))
+
+    return metrics
+
+
+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 prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
+        entity_id = inputs.pop("entity_id")
+        loss, logits, labels = super().prediction_step(
+            model, inputs, prediction_loss_only, ignore_keys=ignore_keys
+        )
+        if logits is not None:
+            entity_col = entity_id.float().unsqueeze(1).to(logits.device)
+            logits = torch.cat([logits, entity_col], dim=1)
+        return loss, logits, labels
+
+
+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("--prefetch_factor",    type=int,   default=2)
+    p.add_argument("--wandb_project",      default=None)
+    p.add_argument("--run_name",           default=None)
+
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    if os.path.isdir(args.dataset_path):
+        ds = DatasetDict.load_from_disk(args.dataset_path)
+    else:
+        ds = load_dataset(args.dataset_path, num_proc=args.num_workers)
+    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,
+        dataloader_prefetch_factor=args.prefetch_factor if args.num_workers > 0 else None,
+        dataloader_persistent_workers=args.num_workers > 0,
+        dataloader_pin_memory=True,
+        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(),
+        compute_metrics=compute_metrics if has_val else None,
+        callbacks=callbacks,
+    )
+
+    trainer.train()
+    trainer.save_model(args.output_dir)
+
+
+if __name__ == "__main__":
+    main()

wandb/run-20260503_171213-h9m78x54/files/config.yaml

@@ -0,0 +1,516 @@
+_attn_implementation_autoset:
+    value: true
+_name_or_path:
+    value: /Users/daniellemillersayag/Documents/vcell/paper/hf-release
+_wandb:
+    value:
+        cli_version: 0.21.0
+        code_path: code/train.py
+        e:
+            fv6s7853m72kjtsdphyqhm5sm6sgz3ly:
+                apple: {}
+                args:
+                    - --dataset_path
+                    - /Users/daniellemillersayag/Documents/vcell/paper/example_dataset
+                    - --model_name_or_path
+                    - /Users/daniellemillersayag/Documents/vcell/paper/hf-release
+                    - --output_dir
+                    - /tmp/vc_smoke_test_wandb
+                    - --num_train_epochs
+                    - "2"
+                    - --per_device_train_batch_size
+                    - "4"
+                    - --per_device_eval_batch_size
+                    - "4"
+                    - --num_workers
+                    - "2"
+                    - --patience
+                    - "5"
+                    - --wandb_project
+                    - virtual-cell-patient
+                    - --run_name
+                    - smoke-test
+                codePath: train.py
+                codePathLocal: train.py
+                cpu_count: 11
+                cpu_count_logical: 11
+                disk:
+                    /:
+                        total: "994662584320"
+                        used: "276313182208"
+                email: danielle.miller@converge-bio.com
+                executable: /Users/daniellemillersayag/Documents/Repos/virtual-cell/venv/bin/python
+                host: Mac.lan
+                memory:
+                    total: "38654705664"
+                os: macOS-26.3.1-arm64-arm-64bit
+                program: /Users/daniellemillersayag/Documents/vcell/paper/hf-release/train.py
+                python: CPython 3.11.10
+                root: /Users/daniellemillersayag/Documents/vcell/paper/hf-release
+                startedAt: "2026-05-03T14:12:13.849133Z"
+                writerId: fv6s7853m72kjtsdphyqhm5sm6sgz3ly
+        m:
+            - "1": train/global_step
+              "6":
+                - 3
+              "7": []
+            - "2": '*'
+              "5": 1
+              "6":
+                - 1
+              "7": []
+        python_version: 3.11.10
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 12
+                - 49
+                - 51
+                - 53
+                - 71
+            "2":
+                - 1
+                - 5
+                - 11
+                - 12
+                - 49
+                - 51
+                - 53
+                - 71
+            "3":
+                - 7
+                - 13
+                - 19
+                - 62
+                - 66
+            "4": 3.11.10
+            "5": 0.21.0
+            "6": 4.51.3
+            "9":
+                "1": transformers_trainer
+            "12": 0.21.0
+            "13": darwin-arm64
+accelerator_config:
+    value:
+        dispatch_batches: null
+        even_batches: true
+        gradient_accumulation_kwargs: null
+        non_blocking: false
+        split_batches: false
+        use_seedable_sampler: true
+activation:
+    value: prelu
+adafactor:
+    value: false
+adam_beta1:
+    value: 0.9
+adam_beta2:
+    value: 0.999
+adam_epsilon:
+    value: 1e-08
+add_cross_attention:
+    value: false
+architectures:
+    value:
+        - VirtualCellPatientModel
+attention_hidden_dim:
+    value: 512
+auto_find_batch_size:
+    value: false
+auto_map:
+    value:
+        AutoConfig: modeling_virtual_cell.VirtualCellPatientConfig
+        AutoModel: modeling_virtual_cell.VirtualCellPatientModel
+average_tokens_across_devices:
+    value: false
+bad_words_ids:
+    value: null
+batch_eval_metrics:
+    value: false
+begin_suppress_tokens:
+    value: null
+bf16:
+    value: false
+bf16_full_eval:
+    value: false
+bos_token_id:
+    value: null
+chunk_size_feed_forward:
+    value: 0
+classifier_dropout:
+    value: 0.1
+cross_attention_hidden_size:
+    value: null
+data_seed:
+    value: null
+dataloader_drop_last:
+    value: false
+dataloader_num_workers:
+    value: 2
+dataloader_persistent_workers:
+    value: true
+dataloader_pin_memory:
+    value: true
+dataloader_prefetch_factor:
+    value: 2
+ddp_backend:
+    value: null
+ddp_broadcast_buffers:
+    value: null
+ddp_bucket_cap_mb:
+    value: null
+ddp_find_unused_parameters:
+    value: null
+ddp_timeout:
+    value: 1800
+debug:
+    value: []
+decoder_start_token_id:
+    value: null
+deepspeed:
+    value: null
+disable_tqdm:
+    value: false
+diversity_penalty:
+    value: 0
+do_eval:
+    value: true
+do_predict:
+    value: false
+do_sample:
+    value: false
+do_train:
+    value: false
+dropout:
+    value: 0.1
+early_stopping:
+    value: false
+embed_dim:
+    value: 512
+encoder_no_repeat_ngram_size:
+    value: 0
+eos_token_id:
+    value: null
+eval_accumulation_steps:
+    value: null
+eval_delay:
+    value: 0
+eval_do_concat_batches:
+    value: true
+eval_on_start:
+    value: false
+eval_steps:
+    value: null
+eval_strategy:
+    value: epoch
+eval_use_gather_object:
+    value: false
+exponential_decay_length_penalty:
+    value: null
+finetuning_task:
+    value: null
+forced_bos_token_id:
+    value: null
+forced_eos_token_id:
+    value: null
+fp16:
+    value: false
+fp16_backend:
+    value: auto
+fp16_full_eval:
+    value: false
+fp16_opt_level:
+    value: O1
+fsdp:
+    value: []
+fsdp_config:
+    value:
+        min_num_params: 0
+        xla: false
+        xla_fsdp_grad_ckpt: false
+        xla_fsdp_v2: false
+fsdp_min_num_params:
+    value: 0
+fsdp_transformer_layer_cls_to_wrap:
+    value: null
+full_determinism:
+    value: false
+gradient_accumulation_steps:
+    value: 1
+gradient_checkpointing:
+    value: false
+gradient_checkpointing_kwargs:
+    value: null
+greater_is_better:
+    value: false
+group_by_length:
+    value: false
+half_precision_backend:
+    value: auto
+hidden_dim:
+    value:
+        - 4096
+        - 1024
+hub_always_push:
+    value: false
+hub_model_id:
+    value: null
+hub_private_repo:
+    value: null
+hub_strategy:
+    value: every_save
+hub_token:
+    value: <HUB_TOKEN>
+id2label:
+    value:
+        "0": oncological
+        "1": immune_inflammatory
+        "2": neurological
+        "3": metabolic_vascular
+        "4": gastrointestinal
+        "5": respiratory
+        "6": epithelial_barrier
+        "7": sensory_specialized
+        "8": healthy_control
+        "9": other
+ignore_data_skip:
+    value: false
+include_for_metrics:
+    value: []
+include_inputs_for_metrics:
+    value: false
+include_num_input_tokens_seen:
+    value: false
+include_tokens_per_second:
+    value: false
+is_decoder:
+    value: false
+is_encoder_decoder:
+    value: false
+jit_mode_eval:
+    value: false
+label_names:
+    value: null
+label_smoothing_factor:
+    value: 0
+label2id:
+    value:
+        epithelial_barrier: 6
+        gastrointestinal: 4
+        healthy_control: 8
+        immune_inflammatory: 1
+        metabolic_vascular: 3
+        neurological: 2
+        oncological: 0
+        other: 9
+        respiratory: 5
+        sensory_specialized: 7
+learning_rate:
+    value: 0.0001
+length_column_name:
+    value: length
+length_penalty:
+    value: 1
+load_best_model_at_end:
+    value: true
+local_rank:
+    value: 0
+log_level:
+    value: passive
+log_level_replica:
+    value: warning
+log_on_each_node:
+    value: true
+logging_dir:
+    value: /tmp/vc_smoke_test_wandb/runs/May03_17-12-12_Mac.lan
+logging_first_step:
+    value: false
+logging_nan_inf_filter:
+    value: true
+logging_steps:
+    value: 500
+logging_strategy:
+    value: steps
+lr_scheduler_type:
+    value: cosine
+max_grad_norm:
+    value: 1
+max_length:
+    value: 20
+max_steps:
+    value: -1
+metric_for_best_model:
+    value: eval_loss
+min_length:
+    value: 0
+model/num_parameters:
+    value: 79963661
+model_type:
+    value: virtual_cell_patient
+mp_parameters:
+    value: ""
+n_genes:
+    value: 18301
+neftune_noise_alpha:
+    value: null
+no_cuda:
+    value: false
+no_repeat_ngram_size:
+    value: 0
+num_beam_groups:
+    value: 1
+num_beams:
+    value: 1
+num_classes:
+    value: 10
+num_return_sequences:
+    value: 1
+num_train_epochs:
+    value: 2
+optim:
+    value: adamw_torch
+optim_args:
+    value: null
+optim_target_modules:
+    value: null
+output_attentions:
+    value: false
+output_dir:
+    value: /tmp/vc_smoke_test_wandb
+output_hidden_states:
+    value: false
+output_scores:
+    value: false
+overwrite_output_dir:
+    value: false
+pad_token_id:
+    value: null
+past_index:
+    value: -1
+per_device_eval_batch_size:
+    value: 4
+per_device_train_batch_size:
+    value: 4
+per_gpu_eval_batch_size:
+    value: null
+per_gpu_train_batch_size:
+    value: null
+prediction_loss_only:
+    value: false
+prefix:
+    value: null
+problem_type:
+    value: null
+push_to_hub:
+    value: false
+push_to_hub_model_id:
+    value: null
+push_to_hub_organization:
+    value: null
+push_to_hub_token:
+    value: <PUSH_TO_HUB_TOKEN>
+ray_scope:
+    value: last
+remove_invalid_values:
+    value: false
+remove_unused_columns:
+    value: false
+repetition_penalty:
+    value: 1
+report_to:
+    value:
+        - wandb
+residual:
+    value: false
+restore_callback_states_from_checkpoint:
+    value: false
+resume_from_checkpoint:
+    value: null
+return_dict:
+    value: true
+return_dict_in_generate:
+    value: false
+run_name:
+    value: smoke-test
+save_on_each_node:
+    value: false
+save_only_model:
+    value: false
+save_safetensors:
+    value: true
+save_steps:
+    value: 500
+save_strategy:
+    value: epoch
+save_total_limit:
+    value: null
+seed:
+    value: 42
+sep_token_id:
+    value: null
+skip_memory_metrics:
+    value: true
+suppress_tokens:
+    value: null
+task_specific_params:
+    value: null
+temperature:
+    value: 1
+tf_legacy_loss:
+    value: false
+tf32:
+    value: null
+tie_encoder_decoder:
+    value: false
+tie_word_embeddings:
+    value: true
+tokenizer_class:
+    value: null
+top_k:
+    value: 50
+top_p:
+    value: 1
+torch_compile:
+    value: false
+torch_compile_backend:
+    value: null
+torch_compile_mode:
+    value: null
+torch_dtype:
+    value: float32
+torch_empty_cache_steps:
+    value: null
+torchdynamo:
+    value: null
+torchscript:
+    value: false
+tp_size:
+    value: 0
+tpu_metrics_debug:
+    value: false
+tpu_num_cores:
+    value: null
+transformers_version:
+    value: 4.51.3
+typical_p:
+    value: 1
+use_bfloat16:
+    value: false
+use_cpu:
+    value: false
+use_ipex:
+    value: false
+use_legacy_prediction_loop:
+    value: false
+use_liger_kernel:
+    value: false
+use_mps_device:
+    value: false
+warmup_ratio:
+    value: 0.1
+warmup_steps:
+    value: 0
+weight_decay:
+    value: 0.05

wandb/run-20260503_171213-h9m78x54/files/output.log

@@ -0,0 +1,5 @@
+100%|██████████| 20/20 [02:14<00:00,  6.74s/it]
+                                             
+{'eval_loss': 2.8007757663726807, 'eval_per_view/accuracy': 0.3333333333333333, 'eval_per_view/f1_macro': 0.25, 'eval_per_view/precision': 0.25, 'eval_per_view/recall': 0.25, 'eval_patient/accuracy': 0.3333333333333333, 'eval_patient/f1_macro': 0.25, 'eval_patient/precision': 0.25, 'eval_patient/recall': 0.25, 'eval_runtime': 20.6865, 'eval_samples_per_second': 0.725, 'eval_steps_per_second': 0.193, 'epoch': 1.0}
+{'eval_loss': 3.5730626583099365, 'eval_per_view/accuracy': 0.3333333333333333, 'eval_per_view/f1_macro': 0.25, 'eval_per_view/precision': 0.25, 'eval_per_view/recall': 0.25, 'eval_patient/accuracy': 0.3333333333333333, 'eval_patient/f1_macro': 0.25, 'eval_patient/precision': 0.25, 'eval_patient/recall': 0.25, 'eval_runtime': 21.0117, 'eval_samples_per_second': 0.714, 'eval_steps_per_second': 0.19, 'epoch': 2.0}
+{'train_runtime': 137.2132, 'train_samples_per_second': 0.583, 'train_steps_per_second': 0.146, 'train_loss': 0.570319652557373, 'epoch': 2.0}

wandb/run-20260503_171213-h9m78x54/files/requirements.txt

@@ -0,0 +1,243 @@
+propcache==0.3.2
+soupsieve==2.7
+fsspec==2026.4.0
+contourpy==1.3.2
+arrow==1.3.0
+s3fs==2026.4.0
+threadpoolctl==3.6.0
+uri-template==1.3.0
+rfc3339-validator==0.1.4
+pydantic_core==2.33.2
+sparse==0.17.0
+flax==0.10.7
+pexpect==4.9.0
+argcomplete==3.6.3
+optax==0.2.5
+transformers==4.51.3
+fqdn==1.5.1
+llvmlite==0.44.0
+jupyter_core==5.8.1
+idna==3.10
+babel==2.17.0
+traitlets==5.14.3
+aioitertools==0.13.0
+urllib3==2.5.0
+jupyter_server==2.16.0
+nest-asyncio==1.6.0
+pyro-ppl==1.9.1
+debugpy==1.8.14
+kiwisolver==1.4.8
+ptyprocess==0.7.0
+jaxlib==0.6.2
+isoduration==20.11.0
+nox==2026.2.9
+certifi==2025.7.14
+pytz==2025.2
+narwhals==1.47.0
+toml==0.10.2
+legacy-api-wrap==1.4.1
+dependency-groups==1.3.1
+joblib==1.5.1
+protobuf==6.31.1
+typing-inspection==0.4.1
+multidict==6.6.3
+virtualenv==21.2.0
+overrides==7.7.0
+parso==0.8.4
+webencodings==0.5.1
+tinycss2==1.4.0
+jupyterlab==4.4.4
+array-api-compat==1.12.0
+docrep==0.3.2
+hf-xet==1.1.5
+pydeck==0.9.1
+ipython==9.4.0
+beautifulsoup4==4.13.4
+wandb==0.21.0
+h5py==3.14.0
+mudata==0.3.2
+platformdirs==4.3.8
+wrapt==2.1.2
+opt_einsum==3.4.0
+anyio==4.9.0
+defusedxml==0.7.1
+ipywidgets==8.1.7
+pip==24.0
+attrs==25.3.0
+pandas==2.3.1
+natsort==8.4.0
+async-lru==2.0.5
+blinker==1.9.0
+tenacity==9.1.2
+notebook==7.4.4
+markdown-it-py==3.0.0
+seaborn==0.13.2
+gseapy==1.1.12
+GitPython==3.1.44
+pyparsing==3.2.3
+pyzmq==27.0.0
+python-pptx==1.0.2
+jupyter-console==6.6.3
+jedi==0.19.2
+pytest==8.4.1
+charset-normalizer==3.4.2
+numpyro==0.18.0
+requests==2.32.4
+torchmetrics==1.7.4
+typing_extensions==4.14.1
+jupyter==1.1.1
+numba==0.61.2
+patsy==1.0.1
+aiohttp==3.12.14
+treescope==0.1.9
+jupyter_client==8.6.3
+distlib==0.4.0
+pynndescent==0.5.13
+asttokens==3.0.0
+tqdm==4.67.1
+matplotlib==3.10.3
+pandocfilters==1.5.1
+prometheus_client==0.22.1
+json5==0.12.0
+huggingface-hub==0.33.4
+fastjsonschema==2.21.1
+jsonpointer==3.0.0
+tzdata==2025.2
+ipython_pygments_lexers==1.1.1
+appnope==0.1.4
+lightning==2.5.2
+numpy==1.26.4
+jax==0.6.2
+httpcore==1.0.9
+filelock==3.25.2
+decorator==5.2.1
+msgpack==1.1.1
+cffi==1.17.1
+captum==0.8.0
+executing==2.2.0
+nbformat==5.10.4
+jupyterlab_widgets==3.0.15
+yarl==1.20.1
+setuptools==65.5.0
+umap-learn==0.5.9.post2
+aiobotocore==3.6.0
+stack-data==0.6.3
+jmespath==1.1.0
+tensorboard==2.19.0
+multiprocess==0.70.16
+Werkzeug==3.1.3
+jsonschema==4.24.0
+xxhash==3.5.0
+nbconvert==7.16.6
+referencing==0.36.2
+regex==2024.11.6
+absl-py==2.3.1
+sentry-sdk==2.33.1
+Send2Trash==1.8.3
+jupyter-lsp==2.2.5
+python-dotenv==1.2.2
+scvi-tools==1.3.2
+nbclient==0.10.2
+h11==0.16.0
+gitdb==4.0.12
+sniffio==1.3.1
+simplejson==3.20.1
+psutil==7.0.0
+fonttools==4.58.5
+rpds-py==0.26.0
+mdurl==0.1.2
+magika==0.6.3
+networkx==3.5
+python-dateutil==2.9.0.post0
+colorlog==6.10.1
+mpmath==1.3.0
+jupyterlab_pygments==0.3.0
+mistune==3.1.3
+torch==2.5.1
+anndata==0.11.4
+wcwidth==0.2.13
+streamlit==1.51.0
+markdownify==1.2.2
+scikit-learn==1.7.0
+tokenizers==0.21.2
+jupyter-events==0.12.0
+prompt_toolkit==3.0.51
+botocore==1.43.0
+aiosignal==1.4.0
+grpcio==1.73.1
+plotly==6.2.0
+toolz==1.0.0
+click==8.2.1
+lightning-utilities==0.14.3
+packaging==25.0
+jupyterlab_server==2.27.3
+argon2-cffi==25.1.0
+webcolors==24.11.1
+jsonschema-specifications==2025.4.1
+pycparser==2.22
+cycler==0.12.1
+Jinja2==3.1.6
+tornado==6.5.1
+session-info2==0.1.2
+dill==0.3.8
+comm==0.2.2
+multipledispatch==1.0.0
+pure_eval==0.2.3
+pydantic==2.11.7
+flatbuffers==25.12.19
+pluggy==1.6.0
+Pygments==2.19.2
+etils==1.13.0
+rfc3986-validator==0.1.1
+python-discovery==1.2.1
+aiohappyeyeballs==2.6.1
+python-json-logger==3.3.0
+terminado==0.18.1
+xgboost==3.1.1
+types-python-dateutil==2.9.0.20250708
+sympy==1.13.1
+argon2-cffi-bindings==21.2.0
+xlsxwriter==3.2.9
+PyYAML==6.0.2
+httpx==0.28.1
+humanize==4.12.3
+lxml==6.1.0
+rich==14.0.0
+matplotlib-inline==0.1.7
+smmap==5.0.2
+matplotlib-venn==1.1.2
+safetensors==0.5.3
+xarray==2025.7.1
+pillow==11.3.0
+ml_collections==1.1.0
+tensorboard-data-server==0.7.2
+pytorch-lightning==2.5.2
+pyro-api==0.1.2
+scipy==1.15.3
+jupyter_server_terminals==0.5.3
+bleach==6.2.0
+orbax-checkpoint==0.11.19
+ml_dtypes==0.5.1
+altair==5.5.0
+tensorstore==0.1.76
+iniconfig==2.1.0
+ipykernel==6.29.5
+zipp==3.23.0
+annotated-types==0.7.0
+scanpy==1.11.3
+datasets==3.2.0
+widgetsnbextension==4.0.14
+Markdown==3.8.2
+six==1.17.0
+importlib_resources==6.5.2
+chex==0.1.89
+pyarrow==20.0.0
+markitdown==0.1.5
+statsmodels==0.14.5
+cachetools==6.2.1
+notebook_shim==0.2.4
+frozenlist==1.7.0
+onnxruntime==1.25.1
+accelerate==1.1.1
+websocket-client==1.8.0
+MarkupSafe==3.0.2

wandb/run-20260503_171213-h9m78x54/files/wandb-metadata.json

@@ -0,0 +1,47 @@
+{
+  "os":  "macOS-26.3.1-arm64-arm-64bit",
+  "python":  "CPython 3.11.10",
+  "startedAt":  "2026-05-03T14:12:13.849133Z",
+  "args":  [
+    "--dataset_path",
+    "/Users/daniellemillersayag/Documents/vcell/paper/example_dataset",
+    "--model_name_or_path",
+    "/Users/daniellemillersayag/Documents/vcell/paper/hf-release",
+    "--output_dir",
+    "/tmp/vc_smoke_test_wandb",
+    "--num_train_epochs",
+    "2",
+    "--per_device_train_batch_size",
+    "4",
+    "--per_device_eval_batch_size",
+    "4",
+    "--num_workers",
+    "2",
+    "--patience",
+    "5",
+    "--wandb_project",
+    "virtual-cell-patient",
+    "--run_name",
+    "smoke-test"
+  ],
+  "program":  "/Users/daniellemillersayag/Documents/vcell/paper/hf-release/train.py",
+  "codePath":  "train.py",
+  "codePathLocal":  "train.py",
+  "email":  "danielle.miller@converge-bio.com",
+  "root":  "/Users/daniellemillersayag/Documents/vcell/paper/hf-release",
+  "host":  "Mac.lan",
+  "executable":  "/Users/daniellemillersayag/Documents/Repos/virtual-cell/venv/bin/python",
+  "cpu_count":  11,
+  "cpu_count_logical":  11,
+  "disk":  {
+    "/":  {
+      "total":  "994662584320",
+      "used":  "276313182208"
+    }
+  },
+  "memory":  {
+    "total":  "38654705664"
+  },
+  "apple":  {},
+  "writerId":  "fv6s7853m72kjtsdphyqhm5sm6sgz3ly"
+}

wandb/run-20260503_171213-h9m78x54/files/wandb-summary.json

@@ -0,0 +1 @@
+{"train/epoch":2,"eval/per_view/f1_macro":0.25,"train_loss":0.570319652557373,"eval/patient/accuracy":0.3333333333333333,"eval/samples_per_second":0.714,"train_steps_per_second":0.146,"_wandb":{"runtime":155},"train/global_step":20,"_timestamp":1.7778176704983711e+09,"_runtime":155,"train_samples_per_second":0.583,"eval/per_view/accuracy":0.3333333333333333,"eval/patient/precision":0.25,"eval/per_view/precision":0.25,"eval/patient/f1_macro":0.25,"total_flos":3.5121959039064e+17,"eval/loss":3.5730626583099365,"train_runtime":137.2132,"eval/per_view/recall":0.25,"eval/runtime":21.0117,"_step":2,"eval/patient/recall":0.25,"eval/steps_per_second":0.19}

wandb/run-20260503_171213-h9m78x54/logs/debug-core.log

@@ -0,0 +1,14 @@
+{"time":"2026-05-03T17:12:14.46496+03:00","level":"INFO","msg":"main: starting server","port-filename":"/var/folders/rp/15xk3vwn341d11km1j04wfvm0000gn/T/tmpwtefwqry/port-63423.txt","pid":63423,"log-level":0,"disable-analytics":false,"shutdown-on-parent-exit":false,"enable-dcgm-profiling":false}
+{"time":"2026-05-03T17:12:14.46556+03:00","level":"INFO","msg":"server: will exit if parent process dies","ppid":63423}
+{"time":"2026-05-03T17:12:14.465548+03:00","level":"INFO","msg":"server: accepting connections","addr":{"Name":"/var/folders/rp/15xk3vwn341d11km1j04wfvm0000gn/T/wandb-63423-63522-633304008/socket","Net":"unix"}}
+{"time":"2026-05-03T17:12:14.49552+03:00","level":"INFO","msg":"connection: ManageConnectionData: new connection created","id":"1"}
+{"time":"2026-05-03T17:12:14.512233+03:00","level":"INFO","msg":"handleInformInit: received","streamId":"h9m78x54","id":"1"}
+{"time":"2026-05-03T17:12:15.049479+03:00","level":"INFO","msg":"handleInformInit: stream started","streamId":"h9m78x54","id":"1"}
+{"time":"2026-05-03T17:14:51.019456+03:00","level":"INFO","msg":"handleInformTeardown: server teardown initiated","id":"1"}
+{"time":"2026-05-03T17:14:51.019791+03:00","level":"INFO","msg":"server is shutting down"}
+{"time":"2026-05-03T17:14:51.01977+03:00","level":"INFO","msg":"connection: closing","id":"1"}
+{"time":"2026-05-03T17:14:51.019978+03:00","level":"INFO","msg":"connection: closed successfully","id":"1"}
+{"time":"2026-05-03T17:14:51.020159+03:00","level":"INFO","msg":"server: listener closed","addr":{"Name":"/var/folders/rp/15xk3vwn341d11km1j04wfvm0000gn/T/wandb-63423-63522-633304008/socket","Net":"unix"}}
+{"time":"2026-05-03T17:14:52.003896+03:00","level":"INFO","msg":"handleInformTeardown: server shutdown complete","id":"1"}
+{"time":"2026-05-03T17:14:52.003946+03:00","level":"INFO","msg":"connection: ManageConnectionData: connection closed","id":"1"}
+{"time":"2026-05-03T17:14:52.003972+03:00","level":"INFO","msg":"server is closed"}

wandb/run-20260503_171213-h9m78x54/logs/debug-internal.log

@@ -0,0 +1,13 @@
+{"time":"2026-05-03T17:12:14.512438+03:00","level":"INFO","msg":"stream: starting","core version":"0.21.0"}
+{"time":"2026-05-03T17:12:15.049447+03:00","level":"INFO","msg":"stream: created new stream","id":"h9m78x54"}
+{"time":"2026-05-03T17:12:15.049472+03:00","level":"INFO","msg":"stream: started","id":"h9m78x54"}
+{"time":"2026-05-03T17:12:15.049488+03:00","level":"INFO","msg":"writer: Do: started","stream_id":"h9m78x54"}
+{"time":"2026-05-03T17:12:15.049533+03:00","level":"INFO","msg":"sender: started","stream_id":"h9m78x54"}
+{"time":"2026-05-03T17:12:15.049551+03:00","level":"INFO","msg":"handler: started","stream_id":"h9m78x54"}
+{"time":"2026-05-03T17:12:15.531811+03:00","level":"ERROR","msg":"git repo not found","error":"repository does not exist"}
+{"time":"2026-05-03T17:14:51.01985+03:00","level":"INFO","msg":"stream: closing","id":"h9m78x54"}
+{"time":"2026-05-03T17:14:51.643326+03:00","level":"INFO","msg":"fileTransfer: Close: file transfer manager closed"}
+{"time":"2026-05-03T17:14:51.997995+03:00","level":"INFO","msg":"sender: closed","stream_id":"h9m78x54"}
+{"time":"2026-05-03T17:14:51.998011+03:00","level":"INFO","msg":"handler: closed","stream_id":"h9m78x54"}
+{"time":"2026-05-03T17:14:51.998039+03:00","level":"INFO","msg":"writer: Close: closed","stream_id":"h9m78x54"}
+{"time":"2026-05-03T17:14:51.998606+03:00","level":"INFO","msg":"stream: closed","id":"h9m78x54"}

wandb/run-20260503_171213-h9m78x54/logs/debug.log

@@ -0,0 +1,25 @@
+2026-05-03 17:12:13,856 INFO    MainThread:63423 [wandb_setup.py:_flush():80] Current SDK version is 0.21.0
+2026-05-03 17:12:13,857 INFO    MainThread:63423 [wandb_setup.py:_flush():80] Configure stats pid to 63423
+2026-05-03 17:12:13,857 INFO    MainThread:63423 [wandb_setup.py:_flush():80] Loading settings from /Users/daniellemillersayag/.config/wandb/settings
+2026-05-03 17:12:13,857 INFO    MainThread:63423 [wandb_setup.py:_flush():80] Loading settings from /Users/daniellemillersayag/Documents/vcell/paper/hf-release/wandb/settings
+2026-05-03 17:12:13,857 INFO    MainThread:63423 [wandb_setup.py:_flush():80] Loading settings from environment variables
+2026-05-03 17:12:13,857 INFO    MainThread:63423 [wandb_init.py:setup_run_log_directory():703] Logging user logs to /Users/daniellemillersayag/Documents/vcell/paper/hf-release/wandb/run-20260503_171213-h9m78x54/logs/debug.log
+2026-05-03 17:12:13,857 INFO    MainThread:63423 [wandb_init.py:setup_run_log_directory():704] Logging internal logs to /Users/daniellemillersayag/Documents/vcell/paper/hf-release/wandb/run-20260503_171213-h9m78x54/logs/debug-internal.log
+2026-05-03 17:12:13,858 INFO    MainThread:63423 [wandb_init.py:init():830] calling init triggers
+2026-05-03 17:12:13,858 INFO    MainThread:63423 [wandb_init.py:init():835] wandb.init called with sweep_config: {}
+config: {'_wandb': {'code_path': 'code/train.py'}}
+2026-05-03 17:12:13,858 INFO    MainThread:63423 [wandb_init.py:init():871] starting backend
+2026-05-03 17:12:14,495 INFO    MainThread:63423 [wandb_init.py:init():874] sending inform_init request
+2026-05-03 17:12:14,511 INFO    MainThread:63423 [wandb_init.py:init():882] backend started and connected
+2026-05-03 17:12:14,513 INFO    MainThread:63423 [wandb_init.py:init():953] updated telemetry
+2026-05-03 17:12:14,513 INFO    MainThread:63423 [wandb_init.py:init():977] communicating run to backend with 90.0 second timeout
+2026-05-03 17:12:15,529 INFO    MainThread:63423 [wandb_init.py:init():1029] starting run threads in backend
+2026-05-03 17:12:15,651 INFO    MainThread:63423 [wandb_run.py:_console_start():2458] atexit reg
+2026-05-03 17:12:15,651 INFO    MainThread:63423 [wandb_run.py:_redirect():2306] redirect: wrap_raw
+2026-05-03 17:12:15,651 INFO    MainThread:63423 [wandb_run.py:_redirect():2375] Wrapping output streams.
+2026-05-03 17:12:15,651 INFO    MainThread:63423 [wandb_run.py:_redirect():2398] Redirects installed.
+2026-05-03 17:12:15,652 INFO    MainThread:63423 [wandb_init.py:init():1075] run started, returning control to user process
+2026-05-03 17:12:15,653 INFO    MainThread:63423 [wandb_run.py:_config_callback():1363] config_cb None None {'return_dict': True, 'output_hidden_states': False, 'output_attentions': False, 'torchscript': False, 'torch_dtype': 'float32', 'use_bfloat16': False, 'tf_legacy_loss': False, 'pruned_heads': {}, 'tie_word_embeddings': True, 'chunk_size_feed_forward': 0, 'is_encoder_decoder': False, 'is_decoder': False, 'cross_attention_hidden_size': None, 'add_cross_attention': False, 'tie_encoder_decoder': False, 'max_length': 20, 'min_length': 0, 'do_sample': False, 'early_stopping': False, 'num_beams': 1, 'num_beam_groups': 1, 'diversity_penalty': 0.0, 'temperature': 1.0, 'top_k': 50, 'top_p': 1.0, 'typical_p': 1.0, 'repetition_penalty': 1.0, 'length_penalty': 1.0, 'no_repeat_ngram_size': 0, 'encoder_no_repeat_ngram_size': 0, 'bad_words_ids': None, 'num_return_sequences': 1, 'output_scores': False, 'return_dict_in_generate': False, 'forced_bos_token_id': None, 'forced_eos_token_id': None, 'remove_invalid_values': False, 'exponential_decay_length_penalty': None, 'suppress_tokens': None, 'begin_suppress_tokens': None, 'architectures': ['VirtualCellPatientModel'], 'finetuning_task': None, 'id2label': {0: 'oncological', 1: 'immune_inflammatory', 2: 'neurological', 3: 'metabolic_vascular', 4: 'gastrointestinal', 5: 'respiratory', 6: 'epithelial_barrier', 7: 'sensory_specialized', 8: 'healthy_control', 9: 'other'}, 'label2id': {'oncological': 0, 'immune_inflammatory': 1, 'neurological': 2, 'metabolic_vascular': 3, 'gastrointestinal': 4, 'respiratory': 5, 'epithelial_barrier': 6, 'sensory_specialized': 7, 'healthy_control': 8, 'other': 9}, 'tokenizer_class': None, 'prefix': None, 'bos_token_id': None, 'pad_token_id': None, 'eos_token_id': None, 'sep_token_id': None, 'decoder_start_token_id': None, 'task_specific_params': None, 'problem_type': None, '_name_or_path': '/Users/daniellemillersayag/Documents/vcell/paper/hf-release', '_attn_implementation_autoset': True, 'transformers_version': '4.51.3', 'model_type': 'virtual_cell_patient', 'auto_map': {'AutoConfig': 'modeling_virtual_cell.VirtualCellPatientConfig', 'AutoModel': 'modeling_virtual_cell.VirtualCellPatientModel'}, 'n_genes': 18301, 'embed_dim': 512, 'hidden_dim': [4096, 1024], 'dropout': 0.1, 'residual': False, 'activation': 'prelu', 'attention_hidden_dim': 512, 'num_classes': 10, 'classifier_dropout': 0.1, 'output_dir': '/tmp/vc_smoke_test_wandb', 'overwrite_output_dir': False, 'do_train': False, 'do_eval': True, 'do_predict': False, 'eval_strategy': 'epoch', 'prediction_loss_only': False, 'per_device_train_batch_size': 4, 'per_device_eval_batch_size': 4, 'per_gpu_train_batch_size': None, 'per_gpu_eval_batch_size': None, 'gradient_accumulation_steps': 1, 'eval_accumulation_steps': None, 'eval_delay': 0, 'torch_empty_cache_steps': None, 'learning_rate': 0.0001, 'weight_decay': 0.05, 'adam_beta1': 0.9, 'adam_beta2': 0.999, 'adam_epsilon': 1e-08, 'max_grad_norm': 1.0, 'num_train_epochs': 2, 'max_steps': -1, 'lr_scheduler_type': 'cosine', 'lr_scheduler_kwargs': {}, 'warmup_ratio': 0.1, 'warmup_steps': 0, 'log_level': 'passive', 'log_level_replica': 'warning', 'log_on_each_node': True, 'logging_dir': '/tmp/vc_smoke_test_wandb/runs/May03_17-12-12_Mac.lan', 'logging_strategy': 'steps', 'logging_first_step': False, 'logging_steps': 500, 'logging_nan_inf_filter': True, 'save_strategy': 'epoch', 'save_steps': 500, 'save_total_limit': None, 'save_safetensors': True, 'save_on_each_node': False, 'save_only_model': False, 'restore_callback_states_from_checkpoint': False, 'no_cuda': False, 'use_cpu': False, 'use_mps_device': False, 'seed': 42, 'data_seed': None, 'jit_mode_eval': False, 'use_ipex': False, 'bf16': False, 'fp16': False, 'fp16_opt_level': 'O1', 'half_precision_backend': 'auto', 'bf16_full_eval': False, 'fp16_full_eval': False, 'tf32': None, 'local_rank': 0, 'ddp_backend': None, 'tpu_num_cores': None, 'tpu_metrics_debug': False, 'debug': [], 'dataloader_drop_last': False, 'eval_steps': None, 'dataloader_num_workers': 2, 'dataloader_prefetch_factor': 2, 'past_index': -1, 'run_name': 'smoke-test', 'disable_tqdm': False, 'remove_unused_columns': False, 'label_names': None, 'load_best_model_at_end': True, 'metric_for_best_model': 'eval_loss', 'greater_is_better': False, 'ignore_data_skip': False, 'fsdp': [], 'fsdp_min_num_params': 0, 'fsdp_config': {'min_num_params': 0, 'xla': False, 'xla_fsdp_v2': False, 'xla_fsdp_grad_ckpt': False}, 'tp_size': 0, 'fsdp_transformer_layer_cls_to_wrap': None, 'accelerator_config': {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}, 'deepspeed': None, 'label_smoothing_factor': 0.0, 'optim': 'adamw_torch', 'optim_args': None, 'adafactor': False, 'group_by_length': False, 'length_column_name': 'length', 'report_to': ['wandb'], 'ddp_find_unused_parameters': None, 'ddp_bucket_cap_mb': None, 'ddp_broadcast_buffers': None, 'dataloader_pin_memory': True, 'dataloader_persistent_workers': True, 'skip_memory_metrics': True, 'use_legacy_prediction_loop': False, 'push_to_hub': False, 'resume_from_checkpoint': None, 'hub_model_id': None, 'hub_strategy': 'every_save', 'hub_token': '<HUB_TOKEN>', 'hub_private_repo': None, 'hub_always_push': False, 'gradient_checkpointing': False, 'gradient_checkpointing_kwargs': None, 'include_inputs_for_metrics': False, 'include_for_metrics': [], 'eval_do_concat_batches': True, 'fp16_backend': 'auto', 'push_to_hub_model_id': None, 'push_to_hub_organization': None, 'push_to_hub_token': '<PUSH_TO_HUB_TOKEN>', 'mp_parameters': '', 'auto_find_batch_size': False, 'full_determinism': False, 'torchdynamo': None, 'ray_scope': 'last', 'ddp_timeout': 1800, 'torch_compile': False, 'torch_compile_backend': None, 'torch_compile_mode': None, 'include_tokens_per_second': False, 'include_num_input_tokens_seen': False, 'neftune_noise_alpha': None, 'optim_target_modules': None, 'batch_eval_metrics': False, 'eval_on_start': False, 'use_liger_kernel': False, 'eval_use_gather_object': False, 'average_tokens_across_devices': False}
+2026-05-03 17:12:15,654 INFO    MainThread:63423 [wandb_config.py:__setitem__():154] [no run ID] config set model/num_parameters = 79963661 - <bound method Run._config_callback of <wandb.sdk.wandb_run.Run object at 0x14e776850>>
+2026-05-03 17:12:15,654 INFO    MainThread:63423 [wandb_run.py:_config_callback():1363] config_cb model/num_parameters 79963661 None
+2026-05-03 17:14:51,017 INFO    MsgRouterThr:63423 [mailbox.py:close():129] [no run ID] Closing mailbox, abandoning 1 handles.