recursive json convert chatgpt

scripts/evalexperts.py

@@ -547,27 +547,32 @@ def save_results(results: Dict[str, Any], expert_stats: Dict[str, Any], args) ->
     if args.model_type == "custom":
         metadata["routing_type"] = "top-k (default)"
     
-    def convert_for_json(obj):
-        """Recursively convert objects to JSON-serializable formats."""
+        # Recursive conversion function
+    def recursive_convert(obj):
         if isinstance(obj, (np.integer, np.floating)):
             return int(obj) if isinstance(obj, np.integer) else float(obj)
         elif isinstance(obj, np.ndarray):
             return obj.tolist()
-        elif isinstance(obj, (torch.Tensor, torch.dtype)):
-            return str(obj) if isinstance(obj, torch.dtype) else obj.tolist()
-        elif isinstance(obj, (dict, list, tuple, str, int, float, bool, type(None))):
+        elif isinstance(obj, torch.Tensor):
+            return obj.cpu().tolist()
+        elif isinstance(obj, torch.dtype):
+            return str(obj)
+        elif isinstance(obj, dict):
+            return {k: recursive_convert(v) for k, v in obj.items()}
+        elif isinstance(obj, (list, tuple)):
+            return [recursive_convert(v) for v in obj]
+        elif isinstance(obj, (int, float, str, bool)) or obj is None:
             return obj
         else:
             return str(obj)
-    
-    # Convert all data to JSON-serializable format
-    serializable_results = {
+
+    # Convert everything
+    serializable_results = recursive_convert({
         "metadata": metadata,
         "task_results": results,
-        "expert_statistics": {
-            k: convert_for_json(v) for k, v in expert_stats.items()
-        }
-    }
+        "expert_statistics": expert_stats
+    })
+
     
     # Save to file
     with open(output_path, 'w') as f: