Upload config.py

configs/config.py

@@ -1,19 +1,63 @@
-### configs/config.py
-
-```python
-class Config:
-    model_name = "facebook/esm2_t33_650M_UR50D"
-    latent_dim = 1280  # Adjust based on ESM-2 latent dimension
-    optim = {"lr": 1e-4}
-    training = {
-        "ema": 0.999,
-        "epochs": 10,
-        "batch_size": 32,
-        "gpus": 8,
-        "precision": 16,  # Mixed precision training
-        "accumulate_grad_batches": 2,  # Gradient accumulation
-        "save_dir": "./checkpoints/",
-    }
-    data_path = "./data/"
-    T = 1000  # Number of diffusion steps
-    subs_masking = False
+MODE = "train" # train, ppl_eval, sample_eval
+DIFFUSION = "absorbing_state"
+BACKBONE = "dit"
+PARAMETERIZATION = "subs" # subs, d3pm, sedd
+TIME_CONDITIONING = False
+T = 1000
+SUBS_MASKING = False
+LATENT_DIM = 1280
+MODEL_NAME = 'facebook/esm2_t33_650M_UR50D'
+
+SEED = 1
+
+class Loader:
+    BATCH_SIZE = 8
+    DATA_PATH = "/workspace/a03-sgoel/MDpLM/data"
+    
+class Sampling:
+    PREDICTOR = "ddpm_cache"  # analytic, ddpm, ddpm_cache (recommended)
+    STEPS = 128
+    NOISE_REMOVAL = True
+    STRIDE_LENGTH = 1
+    NUM_STRIDES = 1
+    NUM_SAMPLING_BATCHES = 2
+
+class model:
+    hidden_size = 1280
+    cond_dim = 1280
+    n_heads = 8
+    n_blocks = 2
+    dropout = 0.5
+    length = 512
+
+class Training:
+    EMA = 0.9999
+    ANTITHETIC_SAMPLING = True
+    SAMPLING_EPS = 1e-3
+    ACCUMULATE_GRAD_BATCHES = 2
+    NUM_EPOCHS = 10
+    GRADIENT_CLIP_VAL = 1.0
+    PRECISION = 'bf16'
+    MAX_STEPS = 1,000,000
+    LOG_EVERY_N_STEPS: 10
+    GPUS = 2
+    SAVE_DIR = "/workspace/a03-sgoel/MDpLM/models"
+
+class Eval:
+    CHECKPOINT_PATH = "/workspace/a03-sgoel/MDpLM/checkpoints"
+    DISABLE_EMA = False
+    COMPUTE_GENERATIVE_PERPLEXITY = False
+    PERPLEXITY_BATCH_SIZE = 8
+    GENERATE_SAMPLES = True
+
+class Optim:
+    LR = 1e-4
+    BETA1 = 0.9
+    BETA2 = 0.999
+    EPS = 1e-8
+
+
+# finish this when training actually works
+class Wandb:
+    PROJECT = "MDpLM_shrey_test"
+    GROUP = "programmablebio"