each loss component is now a separate tag, to avoid character buffer …

…overflow on wandb.
choderalab · Dec 3, 2024 · 8c34641 · 8c34641
1 parent b8bd255
commit 8c34641
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 4 deletions.
diff --git a/modelforge/curation/scripts/curate_spice114_openff.py b/modelforge/curation/scripts/curate_spice114_openff.py
@@ -99,7 +99,7 @@ def main():
 
     # We'll want to provide some simple means of versioning
     # if we make updates to either the underlying dataset, curation modules, or parameters given to the code
-    version = "1"
+    version = "2"
     # version of the dataset to curate
     version_select = f"v_0"
 

diff --git a/modelforge/curation/spice_1_openff_curation.py b/modelforge/curation/spice_1_openff_curation.py
@@ -163,7 +163,7 @@ def _init_record_entries_series(self):
             "dft_total_force": "series_atom",
             "formation_energy": "series_mol",
             "mbis_charges": "series_atom",
-            "scf_dipole": "series_atom",
+            "scf_dipole": "series_mol",
         }
 
     # we will use the retry package to allow us to resume download if we lose connection to the server

diff --git a/modelforge/train/training.py b/modelforge/train/training.py
@@ -1701,6 +1701,8 @@ def setup_logger(self) -> pL.loggers.Logger:
                     name=experiment_name,
                 ),
             )
+
+            log.debug(f'tags: {self._generate_tags(["tensorboard"])}')
         elif self.training_parameter.experiment_logger.logger_name == "wandb":
             from modelforge.utils.io import check_import
 
@@ -1909,14 +1911,27 @@ def _generate_tags(self, tags: List[str]) -> List[str]:
         """Generates tags for the experiment."""
         import modelforge
 
+        try:
+            version = modelforge.__version__
+        except:
+            # for editable local install
+            from modelforge._version import __version__
+
+            version = __version__
+        losses = [
+            f"loss-{loss}"
+            for loss in self.training_parameter.loss_parameter.loss_components
+        ]
         tags.extend(
             [
-                str(modelforge.__version__),
+                str(version),
                 self.dataset_parameter.dataset_name,
                 self.potential_parameter.potential_name,
-                f"loss-{'-'.join(self.training_parameter.loss_parameter.loss_components)}",
+                # f"loss-{'-'.join(self.training_parameter.loss_parameter.loss_components)}",
             ]
         )
+        tags.extend(losses)
+
         return tags