Repository: stanfordnlp/dspy
Pull request #677 (2024-03-19 12:30:05)

Fix save and load state of program with typed predictors
This simple PR aims to fix the `Pydantic model is not JSON serializable` error when saving the state of a compiled program that uses `TypedPredictors`.

However, this fix assumes that the demos of the predictors don't need to keep the same types after saving and loading their states. In older versions of DSPy the type was preserved, but after this commit it changed: https://github.com/stanfordnlp/dspy/commit/b9516c45dbcdedbed6b71dab7cd4a6c1024e0dcf

If that is not expected... then I don't how we would load the state to the pipeline. The JSON cannot know about the Pydantic models. A solution would be to pickle everything. Unfortunately I wasn't able to pickle the pipeline either, because it crashes on the dynamically created Pydantic models of the augmented TypedCoTs signatures.


Patch:
diff --git a/dspy/predict/predict.py b/dspy/predict/predict.py
index 37fff10c..17ae5672 100644
--- a/dspy/predict/predict.py
+++ b/dspy/predict/predict.py
@@ -1,5 +1,7 @@
 import random
 
+from pydantic import BaseModel
+
 import dsp
 from dspy.predict.parameter import Parameter
 from dspy.primitives.prediction import Prediction
@@ -20,9 +22,19 @@ class Predict(Parameter):
         self.demos = []
 
     def dump_state(self):
-        state_keys = ["lm", "traces", "train", "demos"]
+        state_keys = ["lm", "traces", "train"]
         state = {k: getattr(self, k) for k in state_keys}
 
+        state["demos"] = []
+        for demo in self.demos:
+            demo = demo.copy()
+
+            for field in demo:
+                if isinstance(demo[field], BaseModel):
+                    demo[field] = demo[field].model_dump_json()
+
+            state["demos"].append(demo)
+
         # Cache the signature instructions and the last field's name.
         state["signature_instructions"] = self.signature.instructions
 


Test patch:
diff --git a/tests/predict/test_predict.py b/tests/predict/test_predict.py
index c0407b93..c701b380 100644
--- a/tests/predict/test_predict.py
+++ b/tests/predict/test_predict.py
@@ -1,8 +1,10 @@
 import dspy
-from dspy import Predict, Signature
+from dspy import Predict, Signature, TypedPredictor
 from dspy.utils.dummies import DummyLM
 import copy
 import textwrap
+import pydantic
+import ujson
 
 
 def test_initialization_with_string_signature():
@@ -26,7 +28,7 @@ def test_reset_method():
     assert predict_instance.demos == []
 
 
-def test_dump_and_load_state():
+def test_lm_after_dump_and_load_state():
     predict_instance = Predict("input -> output")
     predict_instance.lm = "lm_state"
     dumped_state = predict_instance.dump_state()
@@ -53,7 +55,7 @@ def test_call_method():
     )
 
 
-def test_dump_load_state():
+def test_instructions_after_dump_and_load_state():
     predict_instance = Predict(Signature("input -> output", "original instructions"))
     dumped_state = predict_instance.dump_state()
     new_instance = Predict(Signature("input -> output", "new instructions"))
@@ -61,6 +63,78 @@ def test_dump_load_state():
     assert new_instance.signature.instructions == "original instructions"
 
 
+def test_demos_after_dump_and_load_state():
+    class TranslateToEnglish(dspy.Signature):
+        """Translate content from a language to English."""
+
+        content: str = dspy.InputField()
+        language: str = dspy.InputField()
+        translation: str = dspy.OutputField()
+
+    original_instance = Predict(TranslateToEnglish)
+    original_instance.demos = [
+        dspy.Example(
+            content="¿Qué tal?",
+            language="SPANISH",
+            translation="Hello there",
+        ).with_inputs("content", "language"),
+    ]
+
+    dumped_state = original_instance.dump_state()
+    assert len(dumped_state["demos"]) == len(original_instance.demos)
+    assert dumped_state["demos"][0]["content"] == original_instance.demos[0].content
+
+    saved_state = ujson.dumps(dumped_state)
+    loaded_state = ujson.loads(saved_state)
+
+    new_instance = Predict(TranslateToEnglish)
+    new_instance.load_state(loaded_state)
+    assert len(new_instance.demos) == len(original_instance.demos)
+    # Demos don't need to keep the same types after saving and loading the state.
+    assert new_instance.demos[0]["content"] == original_instance.demos[0].content
+
+
+def test_typed_demos_after_dump_and_load_state():
+    class TypedTranslateToEnglish(dspy.Signature):
+        """Translate content from a language to English."""
+
+        class Input(pydantic.BaseModel):
+            content: str
+            language: str
+
+        class Output(pydantic.BaseModel):
+            translation: str
+
+        input: Input = dspy.InputField()
+        output: Output = dspy.OutputField()
+
+    original_instance = TypedPredictor(TypedTranslateToEnglish).predictor
+    original_instance.demos = [
+        dspy.Example(
+            input=TypedTranslateToEnglish.Input(
+                content="¿Qué tal?",
+                language="SPANISH",
+            ),
+            output=TypedTranslateToEnglish.Output(
+                translation="Hello there",
+            ),
+        ).with_inputs("input"),
+    ]
+
+    dumped_state = original_instance.dump_state()
+    assert len(dumped_state["demos"]) == len(original_instance.demos)
+    assert dumped_state["demos"][0]["input"] == original_instance.demos[0].input.model_dump_json()
+
+    saved_state = ujson.dumps(dumped_state)
+    loaded_state = ujson.loads(saved_state)
+
+    new_instance = TypedPredictor(TypedTranslateToEnglish).predictor
+    new_instance.load_state(loaded_state)
+    assert len(new_instance.demos) == len(original_instance.demos)
+    # Demos don't need to keep the same types after saving and loading the state.
+    assert new_instance.demos[0]["input"] == original_instance.demos[0].input.model_dump_json()
+
+
 def test_forward_method():
     program = Predict("question -> answer")
     dspy.settings.configure(lm=DummyLM([]))
@@ -132,7 +206,8 @@ def test_output_only():
     dspy.settings.configure(lm=lm)
     assert predictor().output == "short answer"
 
-    assert lm.get_convo(-1) == textwrap.dedent("""\
+    assert lm.get_convo(-1) == textwrap.dedent(
+        """\
         Given the fields , produce the fields `output`.
         
         ---
@@ -143,4 +218,5 @@ def test_output_only():
         
         ---
         
-        Output: short answer""")
+        Output: short answer"""
+    )