run pre-commit

piperwolters · piperwolters · commit 41a0325339c4 · 2026-05-16T00:13:45.000Z
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -54,13 +54,10 @@ repos:
           - -s
           - B101,B311,B614
 
-  - repo: local
+  - repo: https://github.com/econchick/interrogate
+    rev: 1.7.0
     hooks:
       - id: interrogate
-        name: interrogate
-        language: system
-        entry: uv run interrogate
-        types: [python]
         args:
           - --ignore-init-method
           - --ignore-init-module
diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/__init__.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/__init__.py
@@ -1,5 +1,7 @@
 """OlmoEarth Pretrain v1 model package."""
 
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.olmoearth_pretrain_v1 import OlmoEarthPretrain_v1
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.olmoearth_pretrain_v1 import (
+    OlmoEarthPretrain_v1,
+)
 
 __all__ = ["OlmoEarthPretrain_v1"]
diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/data/normalize.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/data/normalize.py
@@ -6,7 +6,9 @@
 
 import numpy as np
 
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.constants import ModalitySpec
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.constants import (
+    ModalitySpec,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -18,7 +20,8 @@ def load_computed_config() -> dict[str, dict]:
     and std keys.
     """
     with (
-        files("olmoearth_pretrain_minimal.olmoearth_pretrain_v1.data.norm_configs") / "computed.json"
+        files("olmoearth_pretrain_minimal.olmoearth_pretrain_v1.data.norm_configs")
+        / "computed.json"
     ).open() as f:
         return json.load(f)
 
diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_patch_embed.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_patch_embed.py
@@ -12,7 +12,9 @@
 from einops import rearrange
 from torch import Tensor
 
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.constants import ModalitySpec
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.constants import (
+    ModalitySpec,
+)
 
 logger = logging.getLogger(__name__)
 
diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/flexi_vit.py
@@ -11,14 +11,6 @@
 from torch import Tensor, nn
 from torch.distributed.fsdp import fully_shard
 
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.config import Config
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.constants import (
-    BASE_GSD,
-    Modality,
-    ModalitySpec,
-    get_modality_specs_from_names,
-)
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.datatypes import MaskedOlmoEarthSample, MaskValue
 from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.attention import Block
 from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.encodings import (
     get_1d_sincos_pos_encoding,
@@ -29,8 +21,23 @@
     FlexiPatchEmbed,
     FlexiPatchReconstruction,
 )
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.tokenization import TokenizationConfig
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.utils import get_cumulative_sequence_lengths
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.tokenization import (
+    TokenizationConfig,
+)
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.utils import (
+    get_cumulative_sequence_lengths,
+)
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.config import Config
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.constants import (
+    BASE_GSD,
+    Modality,
+    ModalitySpec,
+    get_modality_specs_from_names,
+)
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.datatypes import (
+    MaskedOlmoEarthSample,
+    MaskValue,
+)
 
 logger = logging.getLogger(__name__)
 
diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/latent_mim.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/latent_mim.py
@@ -14,10 +14,15 @@
     register_fsdp_forward_method,
 )
 
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.config import Config
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.datatypes import MaskedOlmoEarthSample
 from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.flexi_vit import TokensAndMasks
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.utils import DistributedMixins, unpack_encoder_output
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.utils import (
+    DistributedMixins,
+    unpack_encoder_output,
+)
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.config import Config
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.datatypes import (
+    MaskedOlmoEarthSample,
+)
 
 logger = logging.getLogger(__name__)
 
diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/tokenization.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/nn/tokenization.py
@@ -24,7 +24,10 @@
 
 from dataclasses import dataclass, field
 
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.constants import Modality, ModalitySpec
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.constants import (
+    Modality,
+    ModalitySpec,
+)
 
 
 @dataclass
diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/olmoearth_pretrain_v1.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/olmoearth_pretrain_v1.py
@@ -9,9 +9,14 @@
 
 import torch
 
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.flexi_vit import (
+    EncoderConfig,
+    PredictorConfig,
+)
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.latent_mim import (
+    LatentMIMConfig,
+)
 from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.constants import Modality
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.flexi_vit import EncoderConfig, PredictorConfig
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.nn.latent_mim import LatentMIM, LatentMIMConfig
 
 # Model size configurations matching the official OlmoEarth v1 models
 MODEL_SIZE_CONFIGS = {
diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/utils/config.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/utils/config.py
@@ -50,11 +50,17 @@ def _resolve_class(cls, class_name: str) -> type | None:
         # Map old package paths to new ones for compatibility
         # Handle both "helios" (old name) and "olmoearth_pretrain" package names
         if class_name.startswith("helios."):
-            class_name = class_name.replace("helios.", "olmoearth_pretrain_minimal.olmoearth_pretrain_v1.", 1)
+            class_name = class_name.replace(
+                "helios.", "olmoearth_pretrain_minimal.olmoearth_pretrain_v1.", 1
+            )
             # Fix common typos in config files
             class_name = class_name.replace("flexihelios", "flexi_vit")
         elif class_name.startswith("olmoearth_pretrain."):
-            class_name = class_name.replace("olmoearth_pretrain.", "olmoearth_pretrain_minimal.olmoearth_pretrain_v1.", 1)
+            class_name = class_name.replace(
+                "olmoearth_pretrain.",
+                "olmoearth_pretrain_minimal.olmoearth_pretrain_v1.",
+                1,
+            )
 
         *modules, cls_name = class_name.split(".")
         module_name = ".".join(modules)
@@ -95,10 +101,20 @@ def _clean_data(cls, data: Any) -> Any:
                             # Try to resolve as Config using from_dict
                             if cls.CLASS_NAME_FIELD in value:
                                 nested_class_name = value[cls.CLASS_NAME_FIELD]
-                                nested_resolved_cls = cls._resolve_class(nested_class_name)
-                                if nested_resolved_cls is not None and is_dataclass(nested_resolved_cls):
-                                    nested_dict = {k: v for k, v in value.items() if k != cls.CLASS_NAME_FIELD}
-                                    valid_kwargs[key] = cast("type[_StandaloneConfig]", nested_resolved_cls).from_dict(nested_dict)
+                                nested_resolved_cls = cls._resolve_class(
+                                    nested_class_name
+                                )
+                                if nested_resolved_cls is not None and is_dataclass(
+                                    nested_resolved_cls
+                                ):
+                                    nested_dict = {
+                                        k: v
+                                        for k, v in value.items()
+                                        if k != cls.CLASS_NAME_FIELD
+                                    }
+                                    valid_kwargs[key] = cast(
+                                        "type[_StandaloneConfig]", nested_resolved_cls
+                                    ).from_dict(nested_dict)
                                 else:
                                     raise ValueError(
                                         f"Could not resolve nested config class '{nested_class_name}' for field '{key}'"
@@ -158,8 +174,15 @@ def from_dict(
                 class_name = cleaned[cls.CLASS_NAME_FIELD]
                 resolved_cls = cls._resolve_class(class_name)
                 if resolved_cls is not None and is_dataclass(resolved_cls):
-                    config_dict = {k: v for k, v in cleaned.items() if k != cls.CLASS_NAME_FIELD}
-                    return cast("type[_StandaloneConfig]", resolved_cls).from_dict(config_dict)
+                    config_dict = {
+                        k: v for k, v in cleaned.items() if k != cls.CLASS_NAME_FIELD
+                    }
+                    return cast(
+                        C,
+                        cast("type[_StandaloneConfig]", resolved_cls).from_dict(
+                            config_dict
+                        ),
+                    )
                 else:
                     raise ValueError(
                         f"Could not resolve class '{class_name}' from _CLASS_ field. "
diff --git a/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/utils/datatypes.py b/olmoearth_pretrain_minimal/olmoearth_pretrain_v1/utils/datatypes.py
@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import TYPE_CHECKING, Any, NamedTuple
+from typing import Any, NamedTuple
 
 import torch
 
diff --git a/tests/test_model_equivalence.py b/tests/test_model_equivalence.py
@@ -1,12 +1,16 @@
 """Tests for model loading functionality."""
 
+import olmoearth_pretrain.model_loader as op_loader
 import pytest
+import torch
+from olmoearth_pretrain.datatypes import (
+    MaskedOlmoEarthSample as opMaskedOlmoEarthSample,
+)
 
 import olmoearth_pretrain_minimal as opm
-import olmoearth_pretrain.model_loader as op_loader
-from olmoearth_pretrain.datatypes import MaskedOlmoEarthSample as opMaskedOlmoEarthSample
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.datatypes import MaskedOlmoEarthSample as opm_MaskedOlmoEarthSample
-import torch
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.datatypes import (
+    MaskedOlmoEarthSample as opm_MaskedOlmoEarthSample,
+)
 
 
 @pytest.mark.slow
@@ -38,7 +42,9 @@ def test_nano_model_with_weights_equivalence() -> None:
     opm_output = opm_model(opm_sample, patch_size=patch_size)
 
     # now lets do the same for an olmoearth_pretrain model
-    op_model = op_loader.load_model_from_id(op_loader.ModelID.OLMOEARTH_V1_NANO, load_weights=True)
+    op_model = op_loader.load_model_from_id(
+        op_loader.ModelID.OLMOEARTH_V1_NANO, load_weights=True
+    )
     assert op_model is not None
     op_param_count = sum(p.numel() for p in op_model.parameters())
     assert op_param_count == opm_param_count
diff --git a/tests/test_model_loader.py b/tests/test_model_loader.py
@@ -1,12 +1,20 @@
 """Tests for model loading functionality."""
 
-import pytest
 from pathlib import Path
 
-from olmoearth_pretrain_minimal import ModelID, load_model_from_id, OlmoEarthPretrain_v1, load_model_from_path
-from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.datatypes import MaskedOlmoEarthSample
+import pytest
 import torch
 
+from olmoearth_pretrain_minimal import (
+    ModelID,
+    OlmoEarthPretrain_v1,
+    load_model_from_id,
+    load_model_from_path,
+)
+from olmoearth_pretrain_minimal.olmoearth_pretrain_v1.utils.datatypes import (
+    MaskedOlmoEarthSample,
+)
+
 ARTIFACTS = Path(__file__).parent / "artifacts"
 assert ARTIFACTS.exists()
 
@@ -37,6 +45,7 @@ def test_load_nano_model_no_weights() -> None:
     sample = MaskedOlmoEarthSample(**masked_sample_dict)
     _ = model(sample, patch_size=patch_size)
 
+
 def test_load_tiny_model_no_weights() -> None:
     """Test loading tiny model without weights."""
     model = load_model_from_id(ModelID.OLMOEARTH_V1_TINY, load_weights=False)
@@ -111,7 +120,7 @@ def test_direct_initialization() -> None:
 def test_direct_initialization_all_sizes() -> None:
     """Test direct model initialization for all model sizes."""
     for model_size in ["nano", "tiny", "base", "large"]:
-        model = OlmoEarthPretrain_v1(model_size=model_size)
+        model = OlmoEarthPretrain_v1(model_size=model_size)  # type: ignore[arg-type]
         assert model is not None
         param_count = sum(p.numel() for p in model.parameters())
         assert param_count > 0
@@ -120,7 +129,8 @@ def test_direct_initialization_all_sizes() -> None:
 def test_invalid_model_size() -> None:
     """Test that invalid model size raises an error."""
     with pytest.raises(ValueError, match="Invalid model_size"):
-        OlmoEarthPretrain_v1(model_size="invalid")
+        OlmoEarthPretrain_v1(model_size="invalid")  # type: ignore[arg-type]
+
 
 def test_load_v1_1_config() -> None:
     """Test loading nano model from config."""
