Merge pull request #90 from WUR-AI/develop

Inference model, encoder wrapper and more

Author: gabrieletijunaityte

configs/data/yield_africa_KEN.yaml

@@ -0,0 +1,29 @@
+_target_: src.data.base_datamodule.BaseDataModule
+
+dataset:
+  _target_: src.data.yield_africa_dataset.YieldAfricaDataset
+  data_dir: ${paths.data_dir}
+  modalities:
+    coords: {}
+  use_target_data: true
+  use_features: true
+  use_aux_data: none
+  csv_name: null
+  seed: ${seed}
+  cache_dir: ${paths.cache_dir}
+  # Country/year filters — set to a list to restrict, null to include all.
+  # countries and years select only the listed values;
+  # exclude_countries and exclude_years drop the listed values.
+  countries: ["KEN"]
+  years: [2017, 2018, 2019, 2020]
+  exclude_countries: null
+  exclude_years: null
+
+batch_size: 64
+num_workers: 0
+pin_memory: false
+
+split_mode: "random"
+train_val_test_split: [0.7, 0.15, 0.15]
+save_split: false
+seed: ${seed}

configs/data/yield_africa_tessera_KEN.yaml

@@ -0,0 +1,32 @@
+_target_: src.data.base_datamodule.BaseDataModule
+
+dataset:
+  _target_: src.data.yield_africa_dataset.YieldAfricaDataset
+  data_dir: ${paths.data_dir}
+  modalities:
+    tessera:
+      # size must match the tile_size used when running the preprocessing script.
+      # Default: 9 pixels (set by yield_africa_tessera_preprocess.py --tile_size).
+      size: 9
+      format: npy
+      # year is intentionally omitted: yield_africa fetches per-record year tiles
+      # via the preprocessing script rather than a single bulk-year download.
+  use_target_data: true
+  use_features: true
+  use_aux_data: none
+  csv_name: null
+  seed: ${seed}
+  cache_dir: ${paths.cache_dir}
+  countries: ["KEN"]
+  years: [2017, 2018, 2019, 2020]
+  exclude_countries: null
+  exclude_years: null
+
+batch_size: 64
+num_workers: 0
+pin_memory: false
+
+split_mode: "random"
+train_val_test_split: [0.7, 0.15, 0.15]
+save_split: false
+seed: ${seed}

configs/eval.yaml

@@ -1,18 +1,43 @@
 # @package _global_
 
+# specify here default configuration
+# order of defaults determines the order in which configs override each other
 defaults:
   - _self_
-  - data: mnist # choose datamodule with `test_dataloader()` for evaluation
-  - model: mnist
-  - logger: null
-  - trainer: default
-  - paths: default
+  - data: butterfly_coords
+  - model: predictive_geoclip
+  - callbacks: default
+  - logger: ${oc.env:LOGGER,wandb}
+  - trainer: ${oc.env:TRAINER_PROFILE,default}
+  - paths: ${oc.env:STORAGE_MODE,local}
   - extras: default
   - hydra: default
+  - metrics: butterfly_predictive
 
+  # experiment configs allow for version control of specific hyperparameters
+  # e.g. best hyperparameters for given model and datamodule
+  - experiment: null
+
+  # config for hyperparameter optimization
+  - hparams_search: null
+
+  # optional local config for machine/user specific settings
+  # it's optional since it doesn't need to exist and is excluded from version control
+  - optional local: default
+
+  # debugging config (enable through command line, e.g. `python train.py debug=default)
+  - debug: null
+
+# task name, determines output directory path
 task_name: "eval"
 
+# tags to help you identify your experiments
+# you can overwrite this in experiment configs
+# overwrite from command line with `python train.py tags="[first_tag, second_tag]"`
 tags: ["dev"]
 
+# seed for random number generators in pytorch, numpy and python.random
+seed: 12345
+
 # passing checkpoint path is necessary for evaluation
-ckpt_path: ???
+ckpt_path: aether/logs/train/runs/2026-03-15_13-15-08/checkpoints/epoch_098.ckpt

configs/experiment/yield_africa_tabular_KEN.yaml

@@ -0,0 +1,23 @@
+# @package _global_
+
+defaults:
+  - override /model: yield_tabular_reg
+  - override /data: yield_africa_KEN
+  - override /metrics: yield_africa_regression
+
+tags: ["yield_africa", "tabular_only", "regression", "Kenya"]
+seed: 12345
+
+trainer:
+  min_epochs: 1
+  max_epochs: 150
+
+data:
+  batch_size: 64
+
+logger:
+  wandb:
+    tags: ${tags}
+    group: "yield_africa"
+  aim:
+    experiment: "yield_africa"

configs/experiment/yield_africa_tessera_fusion_KEN.yaml

@@ -0,0 +1,38 @@
+# @package _global_
+# configs/experiment/yield_africa_tessera_fusion_loco.yaml
+# TESSERA + tabular fusion model evaluated with leave-one-country-out split.
+# Default held-out country: KEN (largest, most representative test set).
+#
+# Requires:
+#   1. TESSERA tiles pre-fetched:
+#        python src/data_preprocessing/yield_africa_tessera_preprocess.py --data_dir <data_dir>
+#   2. LOCO split files pre-generated:
+#        python src/data_preprocessing/yield_africa_loco_splits.py --data_dir <data_dir>
+#
+# To evaluate on a different held-out country:
+#   python src/train.py experiment=yield_africa_tessera_fusion_loco \
+#     data.saved_split_file_name=split_loco_RWA.pth
+
+defaults:
+  - override /model: yield_tessera_fusion_reg
+  - override /data: yield_africa_tessera_KEN
+  - override /metrics: yield_africa_regression
+
+tags: ["yield_africa", "tessera_fusion", "regression", "Kenya"]
+seed: 12345
+
+trainer:
+  min_epochs: 1
+  max_epochs: 150
+
+data:
+  batch_size: 64
+  dataset:
+    use_features: true
+
+logger:
+  wandb:
+    tags: ${tags}
+    group: "yield_africa"
+  aim:
+    experiment: "yield_africa"

configs/experiment/yield_africa_tessera_reg_KEN.yaml

@@ -0,0 +1,27 @@
+# @package _global_
+# configs/experiment/yield_africa_tessera_reg.yaml
+# Variant: TESSERA spatial encoder only (no tabular features).
+# Requires tiles pre-fetched by:
+#   python src/data_preprocessing/yield_africa_tessera_preprocess.py --data_dir <data_dir>
+
+defaults:
+  - override /model: yield_tessera_reg
+  - override /data: yield_africa_tessera_KEN
+  - override /metrics: yield_africa_regression
+
+tags: ["yield_africa", "tessera_only", "regression", "Kenya"]
+seed: 12345
+
+trainer:
+  min_epochs: 1
+  max_epochs: 150
+
+data:
+  batch_size: 64
+
+logger:
+  wandb:
+    tags: ${tags}
+    group: "yield_africa"
+  aim:
+    experiment: "yield_africa"

configs/inference.yaml

@@ -0,0 +1,25 @@
+defaults:
+  - _self_
+  - paths: ${oc.env:STORAGE_MODE,local}
+  - extras: default
+  - hydra: default
+
+# task name, determines output directory path
+task_name: "inference"
+tags: ["dev"]
+
+# If set, inference.py will load this merged checkpoint directly.
+#inference_ckpt_path: ${paths.log_dir}${task_name}/2026-04-07_12-56-16/inference_model.ckpt
+
+# If `inference_ckpt_path` is not set, stitch the inference model from:
+# - predictive ckpt (provides prediction_head weights)
+# - alignment ckpt (provides text_encoder weights + geo_encoder)
+predictive_ckpt_path: ${paths.log_dir}train/runs/2026-04-02_15-54-53/checkpoints/epoch_000.ckpt
+alignment_ckpt_path: ${paths.log_dir}train/runs/2026-04-02_15-40-03/checkpoints/epoch_000.ckpt
+
+# If set, inference.py will save a merged inference checkpoint you can reload with
+# `inference_ckpt_path`.
+save_inference_ckpt_path: ${paths.log_dir}${task_name}/${now:%Y-%m-%d}_${now:%H-%M-%S}/inference_model.ckpt
+#save_inference_ckpt_path: null
+
+training_order: ["alignment_model", "prediction_model"]

configs/paths/shared.yaml

@@ -9,7 +9,7 @@ data_dir: ${oc.env:DATA_DIR,${oc.env:SHARED_ROOT,${paths.root_dir}}/data}
 cache_dir: ${oc.env:SHARED_CACHE,${paths.data_dir}/cache}
 
 # path to logging directory
-log_dir: ${oc.env:SHARED_ROOT}/logs/
+log_dir: ${oc.env:SHARED_ROOT,${paths.root_dir}}/logs/
 
 # path to output directory, created dynamically by hydra
 # path generation pattern is specified in `configs/hydra/local.yaml`
@@ -21,4 +21,4 @@ work_dir: ${hydra:runtime.cwd}
 
 # huggingface cache directory
 # can be overridden via HF_HOME environment variable
-huggingface_cache: ${oc.env:HF_HOME,oc.env:SHARED_CACHE/huggingface}
+huggingface_cache: ${oc.env:HF_HOME,${oc.env:SHARED_CACHE}/huggingface}