Merge pull request #1556 from chenmoneygithub/fix-saving

Fix signature saving at Predict saving

Author: okhat

dspy/predict/predict.py

@@ -1,16 +1,15 @@
+import logging
 import random
+from functools import lru_cache
 
 from pydantic import BaseModel
 
 import dsp
 from dspy.predict.parameter import Parameter
-from dspy.primitives.program import Module
-
 from dspy.primitives.prediction import Prediction
+from dspy.primitives.program import Module
 from dspy.signatures.signature import ensure_signature, signature_to_template
 
-import logging
-from functools import lru_cache
 
 @lru_cache(maxsize=None)
 def warn_once(msg: str):
@@ -31,7 +30,13 @@ def reset(self):
         self.train = []
         self.demos = []
 
-    def dump_state(self, save_verbose=False):
+    def dump_state(self, save_verbose=None):
+        if save_verbose:
+            logging.warning(
+                "`save_verbose` is deprecated and will be removed in DSPy 2.6.0 release. Currently `save_verbose` "
+                "does nothing."
+            )
+
         state_keys = ["lm", "traces", "train"]
         state = {k: getattr(self, k) for k in state_keys}
 
@@ -45,33 +50,47 @@ def dump_state(self, save_verbose=False):
 
             state["demos"].append(demo)
 
-        # If `save_verbose` save all field metadata as well.
-        if save_verbose:
-            fields = []
-            for field_key in self.signature.fields.keys():
-                field_metadata = self.signature.fields[field_key]
-                fields.append({
-                    "name":  field_key,
-                    "field_type": field_metadata.json_schema_extra["__dspy_field_type"],
-                    "description": field_metadata.json_schema_extra["desc"],
-                    "prefix": field_metadata.json_schema_extra["prefix"]
-                })
-            state["fields"] = fields
-        
-        # Cache the signature instructions and the last field's name.
-        *_, last_key = self.signature.fields.keys()
-        state["signature_instructions"] = self.signature.instructions
-        state["signature_prefix"] = self.signature.fields[last_key].json_schema_extra["prefix"]
-
-        # Some special stuff for CoT.
+        state["signature"] = self.signature.dump_state()
+        # `extended_signature` is a special field for `Predict`s like CoT.
         if hasattr(self, "extended_signature"):
-            # Cache the signature instructions and the last field's name.
-            state["extended_signature_instructions"] = self.extended_signature.instructions
-            state["extended_signature_prefix"] = self.extended_signature.fields[last_key].json_schema_extra['prefix']
+            state["extended_signature"] = self.extended_signature.dump_state()
 
         return state
 
-    def load_state(self, state):
+    def load_state(self, state, use_legacy_loading=False):
+        """Load the saved state of a `Predict` object.
+
+        Args:
+            state (dict): The saved state of a `Predict` object.
+            use_legacy_loading (bool): Whether to use the legacy loading method. Only use it when you are loading a
+                saved state from a version of DSPy prior to v2.5.3.
+        """
+        if use_legacy_loading:
+            self._load_state_legacy(state)
+            return
+        if "signature" not in state:
+            # Check if the state is from a version of DSPy prior to v2.5.3.
+            raise ValueError(
+                "The saved state is from a version of DSPy prior to v2.5.3. Please use `use_legacy_loading=True` to "
+                "load the state."
+            )
+
+        excluded_keys = ["signature", "extended_signature"]
+        for name, value in state.items():
+            # `excluded_keys` are fields that go through special handling.
+            if name not in excluded_keys:
+                setattr(self, name, value)
+
+        self.signature = self.signature.load_state(state["signature"])
+
+        if "extended_signature" in state:
+            self.extended_signature.load_state(state["extended_signature"])
+
+    def _load_state_legacy(self, state):
+        """Legacy state loading for backwards compatibility.
+
+        This method is used to load the saved state of a `Predict` object from a version of DSPy prior to v2.5.3.
+        """
         for name, value in state.items():
             setattr(self, name, value)
 
@@ -84,7 +103,7 @@ def load_state(self, state):
             prefix = state["signature_prefix"]
             *_, last_key = self.signature.fields.keys()
             self.signature = self.signature.with_updated_fields(last_key, prefix=prefix)
-        
+
         # Some special stuff for CoT.
         if "extended_signature_instructions" in state:
             instructions = state["extended_signature_instructions"]
@@ -95,6 +114,7 @@ def load_state(self, state):
             *_, last_key = self.extended_signature.fields.keys()
             self.extended_signature = self.extended_signature.with_updated_fields(last_key, prefix=prefix)
 
+
     def __call__(self, **kwargs):
         return self.forward(**kwargs)
 

dspy/primitives/module.py

@@ -152,17 +152,17 @@ def dump_state(self, save_verbose):
         print(self.named_parameters())
         return {name: param.dump_state(save_verbose) for name, param in self.named_parameters()}
 
-    def load_state(self, state):
+    def load_state(self, state, use_legacy_loading=False):
         for name, param in self.named_parameters():
-            param.load_state(state[name])
+            param.load_state(state[name], use_legacy_loading=use_legacy_loading)
 
     def save(self, path, save_field_meta=False):
         with open(path, "w") as f:
             f.write(ujson.dumps(self.dump_state(save_field_meta), indent=2))
 
-    def load(self, path):
+    def load(self, path, use_legacy_loading=False):
         with open(path) as f:
-            self.load_state(ujson.loads(f.read()))
+            self.load_state(ujson.loads(f.read()), use_legacy_loading=use_legacy_loading)
 
 
 def postprocess_parameter_name(name, value):

dspy/signatures/signature.py

@@ -96,6 +96,10 @@ def signature(cls) -> str:
     def instructions(cls) -> str:
         return inspect.cleandoc(getattr(cls, "__doc__", ""))
 
+    @instructions.setter
+    def instructions(cls, instructions: str) -> None:
+        setattr(cls, "__doc__", instructions)
+
     def with_instructions(cls, instructions: str) -> Type["Signature"]:
         return Signature(cls.fields, instructions)
 
@@ -159,6 +163,28 @@ def insert(cls, index: int, name: str, field, type_: Type = None) -> Type["Signa
         new_fields = dict(input_fields + output_fields)
         return Signature(new_fields, cls.instructions)
 
+    def dump_state(cls):
+        state = {"instructions": cls.instructions, "fields": []}
+        for field in cls.fields:
+            state["fields"].append(
+                {
+                    "prefix": cls.fields[field].json_schema_extra["prefix"],
+                    "description": cls.fields[field].json_schema_extra["desc"],
+                }
+            )
+
+        return state
+
+    def load_state(cls, state):
+        signature_copy = Signature(deepcopy(cls.fields), cls.instructions)
+
+        signature_copy.instructions = state["instructions"]
+        for field, saved_field in zip(signature_copy.fields.values(), state["fields"]):
+            field.json_schema_extra["prefix"] = saved_field["prefix"]
+            field.json_schema_extra["desc"] = saved_field["description"]
+
+        return signature_copy
+
     def equals(cls, other) -> bool:
         """Compare the JSON schema of two Pydantic models."""
         if not isinstance(other, type) or not issubclass(other, BaseModel):
@@ -423,4 +449,4 @@ def infer_prefix(attribute_name: str) -> str:
         else:
             title_cased_words.append(word.capitalize())
 
-    return " ".join(title_cased_words)
+    return " ".join(title_cased_words)

tests/predict/test_predict.py

@@ -126,6 +126,27 @@ class Output(pydantic.BaseModel):
     # 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_signature_fields_after_dump_and_load_state(tmp_path):
+    class CustomSignature(dspy.Signature):
+        """I am just an instruction."""
+        sentence = dspy.InputField(desc="I am an innocent input!")
+        sentiment = dspy.OutputField()
+
+    file_path = tmp_path / "tmp.json"
+    original_instance = Predict(CustomSignature)
+    original_instance.save(file_path)
+
+    class CustomSignature2(dspy.Signature):
+        """I am not a pure instruction."""
+        sentence = dspy.InputField(desc="I am a malicious input!")
+        sentiment = dspy.OutputField(desc="I am a malicious output!", prefix="I am a prefix!")
+
+    new_instance = Predict(CustomSignature2)
+    assert new_instance.signature.dump_state() != original_instance.signature.dump_state()
+    # After loading, the fields should be the same.
+    new_instance.load(file_path)
+    assert new_instance.signature.dump_state() == original_instance.signature.dump_state()
+
 
 def test_forward_method():
     program = Predict("question -> answer")

tests/signatures/test_signature.py

@@ -238,3 +238,44 @@ class SignatureFour(Signature):
         assert "input4" in SignatureTwo.input_fields
     assert "input1" in SignatureOne.input_fields
     assert "input2" in SignatureTwo.input_fields
+
+
+def test_dump_and_load_state():
+    class CustomSignature(dspy.Signature):
+        """I am just an instruction."""
+
+        sentence = dspy.InputField(desc="I am an innocent input!")
+        sentiment = dspy.OutputField()
+
+    state = CustomSignature.dump_state()
+    expected = {
+        "instructions": "I am just an instruction.",
+        "fields": [
+            {
+                "prefix": "Sentence:",
+                "description": "I am an innocent input!",
+            },
+            {
+                "prefix": "Sentiment:",
+                "description": "${sentiment}",
+            },
+        ],
+    }
+    assert state == expected
+
+    class CustomSignature2(dspy.Signature):
+        """I am a malicious instruction."""
+
+        sentence = dspy.InputField(desc="I am an malicious input!")
+        sentiment = dspy.OutputField()
+
+    assert CustomSignature2.dump_state() != expected
+    # Overwrite the state with the state of CustomSignature.
+    loaded_signature = CustomSignature2.load_state(state)
+    assert loaded_signature.instructions == "I am just an instruction."
+    # After `load_state`, the state should be the same as CustomSignature.
+    assert loaded_signature.dump_state() == expected
+    # CustomSignature2 should not have been modified.
+    assert CustomSignature2.instructions == "I am a malicious instruction."
+    assert CustomSignature2.fields["sentence"].json_schema_extra["desc"] == "I am an malicious input!"
+    assert CustomSignature2.fields["sentiment"].json_schema_extra["prefix"] == "Sentiment:"