Add plants training dataset and eval (#10)

* Refactor evals

* Refactor label column

* Update batch size

* Refactor evals metrics

* Add plants config

* Update configs

* Fix eval dataset ordering

Author: gonzalobenegas

configs/data/default.yaml

@@ -23,12 +23,18 @@ max_val_lm_samples: null # Maximum number of samples for LM validation (null = u
 seed: 42
 
 # Evaluation datasets (optional)
-# Set evals: null to disable all evals, or configure specific evals below
-evals:
-  traitgym_mendelian_promoter:
-    dataset_name: songlab/TraitGym
-    dataset_config: mendelian_traits
-    genome_url: https://ftp.ensembl.org/pub/release-115/fasta/homo_sapiens/dna/Homo_sapiens.GRCh38.dna_sm.toplevel.fa.gz
-    genome_path: data/Homo_sapiens.GRCh38.dna_sm.toplevel.fa.gz
-    window_size: 512
-    batch_size: 128
+# Set evals: null to disable all evals, or configure specific evals in dataset-specific configs
+# Example structure:
+# evals:
+#   - name: eval_name
+#     dataset_name: songlab/TraitGym
+#     dataset_config: mendelian_traits
+#     split: test # Dataset split to load (default: "test")
+#     genome_url: https://ftp.ensembl.org/...
+#     filter_name: traitgym_promoter # Filter from EVAL_FILTERS registry (default: "none")
+#     window_size: 512
+#     batch_size: 128
+#     label_column: label # Column to preserve as labels (default: "label")
+#     transform: minus # Transform to apply to raw LLR: minus, identity, abs (default: identity)
+#     metrics: [auprc] # Metrics to compute: auprc, auroc, spearman, pearson (default: [auprc])
+evals: null

configs/data/gpn_animal_promoter.yaml

@@ -4,3 +4,17 @@ defaults:
 batch_size: 2048 # Total effective batch size
 per_device_batch_size: 128 # Batch size per device (adjust based on GPU memory)
 num_workers: 8
+
+# Evaluation datasets
+evals:
+  - name: traitgym_mendelian_promoter
+    dataset_name: songlab/TraitGym
+    dataset_config: mendelian_traits
+    split: test # Dataset split to load (default: "test")
+    genome_url: https://ftp.ensembl.org/pub/release-115/fasta/homo_sapiens/dna/Homo_sapiens.GRCh38.dna_sm.toplevel.fa.gz
+    filter_name: traitgym_promoter
+    window_size: 512
+    batch_size: 128
+    label_column: label # Column to preserve as labels (default: "label")
+    transform: minus # Transform to apply to raw LLR: minus, identity, abs (default: identity)
+    metrics: [auprc] # Metrics to compute: auprc, auroc, spearman, pearson (default: [auprc])

configs/data/plants.yaml

@@ -0,0 +1,24 @@
+defaults:
+  - default
+
+# Training dataset: Angiosperm 16 genomes
+dataset_name: kuleshov-group/Angiosperm_16_genomes
+
+# Batch size configuration
+batch_size: 2048 # Total effective batch size
+per_device_batch_size: 128 # Batch size per device (adjust based on GPU memory)
+num_workers: 8
+
+# Evaluation datasets
+evals:
+  - name: maize_af
+    dataset_name: plantcad/maize-allele-frequency
+    dataset_config: null
+    split: validation
+    genome_url: https://ftp.ensemblgenomes.ebi.ac.uk/pub/plants/release-62/fasta/zea_mays/dna/Zea_mays.Zm-B73-REFERENCE-NAM-5.0.dna_sm.toplevel.fa.gz
+    filter_name: none
+    window_size: 512
+    batch_size: 128
+    label_column: AF # Allele frequency column
+    transform: identity # No transform for regression (default: identity)
+    metrics: [pearson, spearman] # Correlation metrics for regression task

configs/experiment/clm_transformer_small.yaml

@@ -3,6 +3,7 @@
 # Short training run with small Transformer encoder for quick testing
 
 defaults:
+  - override /data: plants
   - override /model: clm_transformer_small
 
 logger:

configs/model/bert_bytenet_small.yaml

@@ -1,6 +1,7 @@
 _target_: glm_experiments.models.lm_lit_module.MLMLitModule
 
 soft_masked_weight: ${data.soft_masked_weight}
+evals: ${data.evals}
 
 net:
   _target_: glm_experiments.models.components.lm.MLM

configs/model/clm_transformer_base.yaml

@@ -1,6 +1,7 @@
 _target_: glm_experiments.models.lm_lit_module.CLMLitModule
 
 soft_masked_weight: ${data.soft_masked_weight}
+evals: ${data.evals}
 
 net:
   _target_: glm_experiments.models.components.lm.CLM

configs/model/clm_transformer_small.yaml

@@ -1,6 +1,7 @@
 _target_: glm_experiments.models.lm_lit_module.CLMLitModule
 
 soft_masked_weight: ${data.soft_masked_weight}
+evals: ${data.evals}
 
 net:
   _target_: glm_experiments.models.components.lm.CLM

configs/model/gpn_animal_promoter.yaml

@@ -1,6 +1,7 @@
 _target_: glm_experiments.models.lm_lit_module.MLMLitModule
 
 soft_masked_weight: ${data.soft_masked_weight}
+evals: ${data.evals}
 
 net:
   _target_: glm_experiments.models.components.lm.MLM

configs/model/mlm_transformer_base.yaml

@@ -1,6 +1,7 @@
 _target_: glm_experiments.models.lm_lit_module.MLMLitModule
 
 soft_masked_weight: ${data.soft_masked_weight}
+evals: ${data.evals}
 
 net:
   _target_: glm_experiments.models.components.lm.MLM

configs/model/mlm_transformer_small.yaml

@@ -1,6 +1,7 @@
 _target_: glm_experiments.models.lm_lit_module.MLMLitModule
 
 soft_masked_weight: ${data.soft_masked_weight}
+evals: ${data.evals}
 
 net:
   _target_: glm_experiments.models.components.lm.MLM