fix up datasetmetadata

FAIR-Chem · Dec 18, 2024 · bffcb10 · bffcb10
1 parent 0124228
commit bffcb10
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/src/fairchem/core/datasets/base_dataset.py b/src/fairchem/core/datasets/base_dataset.py
@@ -34,8 +34,11 @@
 T_co = TypeVar("T_co", covariant=True)
 
 
-class DatasetMetadata(NamedTuple):
-    natoms: ArrayLike | None = None
+class DatasetMetadata:
+    def __init__(self, natoms: ArrayLike | None = None, **kwargs):
+        self.natoms = natoms
+        for key, value in kwargs.items():
+            setattr(self, key, value)
 
 
 class UnsupportedDatasetError(ValueError):
@@ -106,9 +109,10 @@ def _metadata(self) -> DatasetMetadata:
         metadata = DatasetMetadata(
             **{
                 field: np.concatenate([metadata[field] for metadata in metadata_npzs])
-                for field in DatasetMetadata._fields
+                for field in metadata_npzs[0].keys()
             }
         )
+
         assert np.issubdtype(
             metadata.natoms.dtype, np.integer
         ), f"Metadata natoms must be an integer type! not {metadata.natoms.dtype}"