work on graph

ecmwf · Jun 15, 2024 · 6aa56a8 · 6aa56a8
1 parent 9b46a6f
commit 6aa56a8
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/src/anemoi/inference/checkpoint/metadata/__init__.py b/src/anemoi/inference/checkpoint/metadata/__init__.py
@@ -351,12 +351,12 @@ def digraph(self, label_maker=lambda x: dict(label=x.kind)):
         for node, label in nodes.items():
             for k, v in label.items():
 
-                # if isinstance(v, str) and v.startswith("<") and v.endswith(">"):
-                #     # Keep HTML labels as is
-                #     pass
-                # else:
-                #     # Use json.dumps to escape special characters
-                label[k] = json.dumps(v)
+                if isinstance(v, str) and v.startswith("<") and v.endswith(">"):
+                    # Keep HTML labels as i
+                    label[k] = "<" + json.dumps(v[1:-1])[1:-1] + ">"
+                else:
+                    # Use json.dumps to escape special characters
+                    label[k] = json.dumps(v)
 
             label = " ".join([f"{k}={v}" for k, v in label.items()])