data:
  name: PerturbationDataModule
  kwargs:
    toml_config_path: /dev/shm/cd4pb/cd4.toml
    embed_key: X_hvg
    output_space: gene
    pert_rep: onehot
    basal_rep: sample
    num_workers: 0
    pin_memory: true
    n_basal_samples: 1
    downsample: null
    downsample_cells: null
    additional_obs: []
    basal_mapping_strategy: random
    should_yield_control_cells: true
    batch_col: donor
    pert_col: gene
    cell_type_key: condition
    control_pert: NTC
    map_controls: true
    perturbation_features_file: null
    store_raw_basal: false
    int_counts: false
    barcode: true
    use_consecutive_loading: false
    val_subsample_fraction: 1.0
  output_dir: null
  debug: true
model:
  name: state_halt
  checkpoint: null
  device: cuda
  kwargs:
    cell_set_len: 64
    blur: 0.05
    hidden_dim: 336
    loss: energy
    confidence_token: true
    n_encoder_layers: 1
    n_decoder_layers: 1
    predict_residual: true
    softplus: true
    freeze_pert_backbone: false
    transformer_decoder: false
    finetune_vci_decoder: false
    residual_decoder: false
    batch_encoder: true
    use_batch_token: false
    mask_attn: false
    n_refine_rounds: 6
    halt_tau: 0.05
    halt_alpha: 0.5
    halt_beta: 0.1
    halt_gamma: 0.1
    halt_warmup_steps: 2000
    halt_magnitude_free: true
    confidence_weight: 0.01
    confidence_target_scale: 10.0
    use_effect_gating_token: false
    distributional_loss: energy
    init_from: null
    mmd_num_chunks: 1
    randomize_mmd_chunks: false
    llm_name: null
    transformer_backbone_key: llama
    transformer_backbone_kwargs:
      bidirectional_attention: true
      max_position_embeddings: 64
      hidden_size: 336
      intermediate_size: 3072
      num_hidden_layers: 8
      num_attention_heads: 12
      num_key_value_heads: 12
      head_dim: 28
      use_cache: false
      attention_dropout: 0.0
      hidden_dropout: 0.0
      layer_norm_eps: 1.0e-06
      pad_token_id: 0
      bos_token_id: 1
      eos_token_id: 2
      tie_word_embeddings: false
      rotary_dim: 0
      use_rotary_embeddings: false
    lora:
      enable: false
      r: 16
      alpha: 32
      dropout: 0.05
      bias: none
      target: auto
      adapt_mlp: false
      task_type: FEATURE_EXTRACTION
      merge_on_eval: false
training:
  wandb_track: true
  weight_decay: 0.0005
  batch_size: 64
  lr: 0.001
  max_steps: 15000
  train_seed: 42
  val_freq: 1000
  ckpt_every_n_steps: 1000
  gradient_clip_val: 10
  gradient_accumulation_steps: 1
  loss_fn: mse
  devices: 1
  strategy: auto
  use_mfu: false
  mfu_kwargs:
    available_flops: 60000000000000.0
    use_backward: true
    logging_interval: 10
    window_size: 2
  cumulative_flops_use_backward: true
wandb:
  entity: yashraj59-personal
  project: cd4_state_halt
  local_wandb_dir: ./wandb_logs
  tags: []
name: cd4pb_halt_v1
output_dir: /workspace/halt/actrun_state/cd4pb_out
use_wandb: true
overwrite: false
return_adatas: false
pred_adata_path: null
true_adata_path: null