Extract OmegaConf resolver registration into shared utility

predict.py

@@ -5,6 +5,9 @@
 from hydra.utils import instantiate
 import torch
 from muonly.callbacks import PredictionWriter
+from muonly.utils.resolvers import register_resolvers
+
+register_resolvers()
 
 
 def run(ckpt_file_path: Path, gpu_id: int):

src/muonly/utils/resolvers.py

@@ -0,0 +1,26 @@
+import secrets
+
+from coolname.impl import generate_slug
+from omegaconf import OmegaConf
+
+
+def register_resolvers() -> None:
+    """Register custom OmegaConf resolvers used across scripts."""
+    OmegaConf.register_new_resolver(
+        "slug",
+        lambda pattern=2: generate_slug(pattern=pattern),
+        use_cache=True,
+        replace=True,
+    )
+
+    OmegaConf.register_new_resolver(
+        name="len",
+        resolver=len,
+        replace=True,
+    )
+
+    OmegaConf.register_new_resolver(
+        name="randbits",
+        resolver=secrets.randbits,
+        replace=True,
+    )

train.py

@@ -2,21 +2,19 @@
 import logging
 from pathlib import Path
 import warnings
-import secrets
 import os
 from aim.pytorch_lightning import AimLogger
 import hydra
 from hydra.utils import instantiate
 import torch
 from lightning.pytorch import LightningDataModule, Trainer
 from lightning import seed_everything
-from omegaconf import DictConfig
-from omegaconf import OmegaConf
-from coolname.impl import generate_slug
+from omegaconf import DictConfig, OmegaConf
 from tqdm import TqdmExperimentalWarning
 import mplhep as mh
 from muonly.nn.utils import init_params
 from muonly.utils.logging import log_everything
+from muonly.utils.resolvers import register_resolvers
 
 mh.style.use("CMS")
 
@@ -32,22 +30,7 @@
 
 warnings.filterwarnings("ignore", category=TqdmExperimentalWarning)
 
-OmegaConf.register_new_resolver(
-    "slug",
-    lambda pattern=2: generate_slug(pattern=pattern),
-    use_cache=True,
-    replace=True,
-)
-
-OmegaConf.register_new_resolver(
-    name="len",
-    resolver=len,
-)
-
-OmegaConf.register_new_resolver(
-    name="randbits",
-    resolver=secrets.randbits,
-)
+register_resolvers()
 
 
 @hydra.main(