Merge pull request #236 from allenai/henryh/helios-to-olmoearth_pretrain

Update to latest commit of Olmoearth_pretrain

Author: Hgherzog

.github/workflows/build_test.yaml

@@ -55,17 +55,17 @@ jobs:
           sudo rm -rf /opt/ghc
           sudo rm -rf /usr/local/share/boost
 
-      # Add deploy key and clone Helios repository.
-      - name: Start ssh-agent and add helios deploy key
+      # Add deploy key and clone olmoearth_pretrain repository.
+      - name: Start ssh-agent and add olmoearth_pretrain deploy key
         uses: webfactory/[email protected]
         with:
           ssh-private-key: ${{ secrets.DEPLOY_KEY_FOR_HELIOS_CLONE }}
-      - name: Clone helios and update requirements-extra.txt.
+      - name: Clone olmoearth_pretrain and update requirements-extra.txt.
         run: |
           mkdir docker_build
-          git clone [email protected]:allenai/helios.git docker_build/helios
-          git -C docker_build/helios checkout f0a63f190b0f99d9c503249daf7e3e47bbd4792a
-          echo "helios @ /opt/rslearn_projects/docker_build/helios/" >> requirements-extra.txt
+          git clone [email protected]:allenai/olmoearth_pretrain.git docker_build/olmoearth_pretrain
+          git -C docker_build/olmoearth_pretrain checkout 0cd82b4784bc2f246a23f6da98a9bab27761c1ba
+          echo "olmoearth_pretrain @ /opt/rslearn_projects/docker_build/olmoearth_pretrain/" >> requirements-extra.txt
 
       # Same thing for olmoearth_run repository.
       - name: Start ssh-agent and add olmoearth_run deploy key

requirements-helios.txt

@@ -1 +1 @@
-helios @ git+https://github.com/allenai/helios.git@f0a63f190b0f99d9c503249daf7e3e47bbd4792a
+olmoearth_pretrain @ git+https://github.com/allenai/olmoearth_pretrain.git@0cd82b4784bc2f246a23f6da98a9bab27761c1ba

rslp/helios/model.py

@@ -6,11 +6,11 @@
 
 import torch
 from einops import rearrange
-from helios.data.constants import Modality
-from helios.nn.flexihelios import Encoder, TokensAndMasks
-from helios.train.masking import MaskedHeliosSample, MaskValue
 from olmo_core.config import Config
 from olmo_core.distributed.checkpoint import load_model_and_optim_state
+from olmoearth_pretrain.data.constants import Modality
+from olmoearth_pretrain.nn.flexihelios import Encoder, TokensAndMasks
+from olmoearth_pretrain.train.masking import MaskedOlmoEarthSample, MaskValue
 from upath import UPath
 
 from rslp.log_utils import get_logger
@@ -53,7 +53,7 @@ def __init__(
             selector: an optional sequence of attribute names or list indices to select
                 the sub-module that should be applied on the input images.
             forward_kwargs: additional arguments to pass to forward pass besides the
-                MaskedHeliosSample.
+                 MaskedOlmoEarthSample.
             random_initialization: whether to skip loading the checkpoint so the
                 weights are randomly initialized. In this case, the checkpoint is only
                 used to define the model architecture.
@@ -148,7 +148,7 @@ def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
         timestamps[:, :, 2] = 2024  # year
         kwargs["timestamps"] = timestamps
 
-        sample = MaskedHeliosSample(**kwargs)
+        sample = MaskedOlmoEarthSample(**kwargs)
 
         # Decide context based on self.autocast_dtype.
         if self.autocast_dtype is None:

rslp/helios/norm.py

@@ -3,7 +3,7 @@
 import json
 from typing import Any
 
-from helios.data.normalize import load_computed_config
+from olmoearth_pretrain.data.normalize import load_computed_config
 from rslearn.train.transforms.transform import Transform
 
 from rslp.log_utils import get_logger