add json mode

Author: chenmoneygithub

dspy/predict/predict.py

@@ -30,7 +30,7 @@ def reset(self):
         self.train = []
         self.demos = []
 
-    def dump_state(self):
+    def dump_state(self, json_mode=True):
         state_keys = ["traces", "train"]
         state = {k: getattr(self, k) for k in state_keys}
 
@@ -42,7 +42,7 @@ def dump_state(self):
                 # FIXME: Saving BaseModels as strings in examples doesn't matter because you never re-access as an object
                 demo[field] = serialize_object(demo[field])
 
-            if isinstance(demo, dict):
+            if isinstance(demo, dict) or not json_mode:
                 state["demos"].append(demo)
             else:
                 state["demos"].append(demo.toDict())

dspy/primitives/base_module.py

@@ -1,4 +1,5 @@
 import copy
+import json
 import logging
 from collections import deque
 from collections.abc import Generator
@@ -153,8 +154,8 @@ def reset_copy(self):
 
         return new_instance
 
-    def dump_state(self):
-        return {name: param.dump_state() for name, param in self.named_parameters()}
+    def dump_state(self, json_mode=True):
+        return {name: param.dump_state(json_mode=json_mode) for name, param in self.named_parameters()}
 
     def load_state(self, state):
         for name, param in self.named_parameters():
@@ -220,9 +221,9 @@ def save(self, path, save_program=False, modules_to_serialize=None):
 
             return
 
-        state = self.dump_state()
-        state["metadata"] = metadata
         if path.suffix == ".json":
+            state = self.dump_state()
+            state["metadata"] = metadata
             try:
                 with open(path, "w", encoding="utf-8") as f:
                     f.write(orjson.dumps(state, option=orjson.OPT_INDENT_2 | orjson.OPT_APPEND_NEWLINE).decode("utf-8"))
@@ -233,6 +234,8 @@ def save(self, path, save_program=False, modules_to_serialize=None):
                     "with `.pkl`, or saving the whole program by setting `save_program=True`."
                 )
         elif path.suffix == ".pkl":
+            state = self.dump_state(json_mode=False)
+            state["metadata"] = metadata
             with open(path, "wb") as f:
                 cloudpickle.dump(state, f)
         else: