Fix symmetry & abspath issues

models/rfd3/src/rfd3/engine.py

@@ -21,11 +21,13 @@
 from rfd3.inference.datasets import (
     assemble_distributed_inference_loader_from_json,
 )
-from rfd3.inference.input_parsing import DesignInputSpecification
+from rfd3.inference.input_parsing import (
+    DesignInputSpecification,
+    ensure_input_is_abspath,
+)
 from rfd3.model.inference_sampler import SampleDiffusionConfig
 from rfd3.utils.inference import (
     ensure_inference_sampler_matches_design_spec,
-    ensure_input_is_abspath,
 )
 from rfd3.utils.io import (
     CIF_LIKE_EXTENSIONS,
@@ -391,9 +393,12 @@ def _multiply_specifications(
         design_specifications = {}
         for prefix, example_spec in inputs.items():
             # Record task name in the specification
-            if "extra" not in example_spec:
-                example_spec["extra"] = {}
-            example_spec["extra"]["task_name"] = prefix
+            if isinstance(example_spec, DesignInputSpecification):
+                example_spec.extra["task_name"] = prefix
+            else:
+                if "extra" not in example_spec:
+                    example_spec["extra"] = {}
+                example_spec["extra"]["task_name"] = prefix
 
             # ... Create n_batches for example
             for batch_id in range((n_batches) if exists(n_batches) else 1):

models/rfd3/src/rfd3/inference/datasets.py

@@ -14,8 +14,8 @@
 from omegaconf import DictConfig, OmegaConf
 from rfd3.inference.input_parsing import (
     DesignInputSpecification,
+    ensure_input_is_abspath,
 )
-from rfd3.utils.inference import ensure_input_is_abspath
 from torch.utils.data import (
     DataLoader,
     SequentialSampler,

models/rfd3/src/rfd3/inference/input_parsing.py

@@ -5,6 +5,7 @@
 import time
 import warnings
 from contextlib import contextmanager
+from os import PathLike
 from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
@@ -1121,3 +1122,35 @@ def accumulate_components(
     if atom_array_accum.bonds is None:
         atom_array_accum.bonds = BondList(atom_array_accum.array_length())
     return atom_array_accum
+
+
+def ensure_input_is_abspath(
+    args: Dict[str, DesignInputSpecification | dict], path: PathLike | None
+):
+    """
+    Ensures the input source is an absolute path if exists, if not it will convert
+
+    args:
+        spec: Inference specification for atom array
+        path: None or file to which the input is relative to.
+    """
+    if isinstance(args, str):
+        raise ValueError(
+            "Expected args to be a dictionary, got a string: {}. If you are using an input JSON ensure it contains dictionaries of arguments".format(
+                args
+            )
+        )
+    if "input" not in args or not exists(args["input"]):
+        return args
+    input = str(args["input"])
+    if not os.path.isabs(input):
+        if path is None:
+            raise ValueError(
+                "input path provided in input, but no path to resolve relative to (required)."
+            )
+        input = os.path.abspath(os.path.join(os.path.dirname(str(path)), input))
+        logger.info(
+            f"Input source path is relative, converted to absolute path: {input}"
+        )
+        args["input"] = input
+    return args

models/rfd3/src/rfd3/utils/inference.py

@@ -3,7 +3,6 @@
 """
 
 import logging
-import os
 from os import PathLike
 from typing import Dict
 
@@ -365,32 +364,6 @@ def inference_load_(
     return data
 
 
-def ensure_input_is_abspath(args: dict, path: PathLike | None):
-    """
-    Ensures the input source is an absolute path if exists, if not it will convert
-
-    args:
-        spec: Inference specification for atom array
-        path: None or file to which the input is relative to.
-    """
-    if isinstance(args, str):
-        raise ValueError(
-            "Expected args to be a dictionary, got a string: {}. If you are using an input JSON ensure it contains dictionaries of arguments".format(
-                args
-            )
-        )
-    if "input" not in args or not exists(args["input"]):
-        return args
-    input = args["input"]
-    if not os.path.isabs(input):
-        input = os.path.abspath(os.path.join(os.path.dirname(path), input))
-        ranked_logger.info(
-            f"Input source path is relative, converted to absolute path: {input}"
-        )
-        args["input"] = input
-    return args
-
-
 def ensure_inference_sampler_matches_design_spec(
     design_spec: dict, inference_sampler: dict | None = None
 ):
@@ -401,7 +374,10 @@ def ensure_inference_sampler_matches_design_spec(
         inference_sampler: Inference sampler dictionary
     """
     has_symmetry_specification = [
-        True if "symmetry" in item.keys() else False for item in design_spec.values()
+        True
+        if "symmetry" in item.keys() and item.get("symmetry") is not None
+        else False
+        for item in design_spec.values()
     ]
     if any(has_symmetry_specification):
         if (