Improved metrics config (#21)

Author: spietras

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/modules/basic.py

@@ -61,8 +61,14 @@ def label(cls) -> str:
     def config(cls) -> Dict[str, Any]:
         return {
             "type": "line",
-            "data": {"datasets": [{"data": []}]},
-            "options": {"parsing": {"xAxisKey": "step", "yAxisKey": "loss"}},
+            "data": {"datasets": [{"label": "Supervised Loss", "data": []}]},
+            "options": {
+                "parsing": {"xAxisKey": "step", "yAxisKey": "loss"},
+                "scales": {
+                    "x": {"title": {"text": "Step"}},
+                    "y": {"title": {"text": "Loss"}},
+                },
+            },
         }
 
 
@@ -79,8 +85,14 @@ def label(cls) -> str:
     def config(cls) -> Dict[str, Any]:
         return {
             "type": "line",
-            "data": {"datasets": [{"data": []}]},
-            "options": {"parsing": {"xAxisKey": "step", "yAxisKey": "score"}},
+            "data": {"datasets": [{"label": "Reinforced Score", "data": []}]},
+            "options": {
+                "parsing": {"xAxisKey": "step", "yAxisKey": "score"},
+                "scales": {
+                    "x": {"title": {"text": "Step"}},
+                    "y": {"title": {"text": "Score"}},
+                },
+            },
         }
 
 

kilroy_module_pytorch_py_sdk/src/kilroy_module_pytorch_py_sdk/modules/reward.py

@@ -66,8 +66,14 @@ def label(cls) -> str:
     def config(cls) -> Dict[str, Any]:
         return {
             "type": "line",
-            "data": {"datasets": [{"data": []}]},
-            "options": {"parsing": {"xAxisKey": "step", "yAxisKey": "loss"}},
+            "data": {"datasets": [{"label": "Supervised Loss", "data": []}]},
+            "options": {
+                "parsing": {"xAxisKey": "step", "yAxisKey": "loss"},
+                "scales": {
+                    "x": {"title": {"text": "Step"}},
+                    "y": {"title": {"text": "Loss"}},
+                },
+            },
         }
 
 
@@ -84,8 +90,14 @@ def label(cls) -> str:
     def config(cls) -> Dict[str, Any]:
         return {
             "type": "line",
-            "data": {"datasets": [{"data": []}]},
-            "options": {"parsing": {"xAxisKey": "step", "yAxisKey": "score"}},
+            "data": {"datasets": [{"label": "Reinforced Score", "data": []}]},
+            "options": {
+                "parsing": {"xAxisKey": "step", "yAxisKey": "score"},
+                "scales": {
+                    "x": {"title": {"text": "Step"}},
+                    "y": {"title": {"text": "Score"}},
+                },
+            },
         }
 
 
@@ -102,8 +114,14 @@ def label(cls) -> str:
     def config(cls) -> Dict[str, Any]:
         return {
             "type": "line",
-            "data": {"datasets": [{"data": []}]},
-            "options": {"parsing": {"xAxisKey": "step", "yAxisKey": "loss"}},
+            "data": {"datasets": [{"label": "Reward Model Loss", "data": []}]},
+            "options": {
+                "parsing": {"xAxisKey": "step", "yAxisKey": "loss"},
+                "scales": {
+                    "x": {"title": {"text": "Step"}},
+                    "y": {"title": {"text": "Loss"}},
+                },
+            },
         }
 
 
@@ -120,8 +138,16 @@ def label(cls) -> str:
     def config(cls) -> Dict[str, Any]:
         return {
             "type": "line",
-            "data": {"datasets": [{"data": []}]},
-            "options": {"parsing": {"xAxisKey": "step", "yAxisKey": "score"}},
+            "data": {
+                "datasets": [{"label": "Reward Model Score", "data": []}]
+            },
+            "options": {
+                "parsing": {"xAxisKey": "step", "yAxisKey": "score"},
+                "scales": {
+                    "x": {"title": {"text": "Step"}},
+                    "y": {"title": {"text": "Score"}},
+                },
+            },
         }