Repository: stanfordnlp/dspy (python)

Issue:
Saving/Loading Method for signatures in dspy.Predict
Currently the dump_state / load_state method in dspy.Predict only saves signature_instructions and signature_prefix - the prefix for the final key in the signature:
```
def dump_state(self):
       
         ............
        # Cache the signature instructions and the last field's name.
        state["signature_instructions"] = self.signature.instructions

        *_, last_key = self.signature.fields.keys()
        state["signature_prefix"] = self.signature.fields[last_key].json_schema_extra["prefix"]

        return state
        
def load_state(self, state):
        for name, value in state.items():
            setattr(self, name, value)

        # Reconstruct the signature.
        if "signature_instructions" in state:
            instructions = state["signature_instructions"]
            self.signature = self.signature.with_instructions(instructions)

        if "signature_prefix" in state:
            prefix = state["signature_prefix"]
            *_, last_key = self.signature.fields.keys()
            self.signature = self.signature.with_updated_fields(last_key, prefix=prefix)
  ```
  
  But the optimizers also optimize the descriptions and prefixes of all the other fields in the signature, so we lose those when we save the state. Can we change to something like this?
  ```
  def dump_state(self):
        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

        field_details = {}
        for key, field in self.signature.fields.items():
            field_details[key] = {
                "prefix": field.json_schema_extra.get("prefix", ""),
                "desc": field.json_schema_extra.get("desc", "")
            }
        state["signature_field_details"] = field_details

        return state

    def load_state(self, state):
        for name, value in state.items():
            setattr(self, name, value)

        # Reconstruct the signature.
        if "signature_instructions" in state:
            instructions = state["signature_instructions"]
            self.signature = self.signature.with_instructions(instructions)

        if "signature_field_details" in state:
            field_details = state["signature_field_details"]
            for key, details in field_details.items():
                prefix = details.get("prefix", "")
                desc = details.get("desc", "")
                self.signature = self.signature.with_updated_fields(
                    key, prefix=prefix, desc=desc
                )
  ```

Pull request:
Fix signature saving at Predict saving
resolve #1045 

Change summary:

- Add a `dump_state`/`load_state` pair to Signature class for better modularization.
- Remove the `save_verbose` option, as it doesn't make sense to change input fields to output fields at loading time. As we are doing "weights-only" saving, the signatures' architecture should match, i.e., input fields and output fields' number and name.  Only changeable attributes should be included with saving.
- Save prefix for every field as requested in #1045.
- Added unit testing. 

Patch:
diff --git a/dspy/predict/predict.py b/dspy/predict/predict.py
index 8bd42699..a029a117 100644
--- a/dspy/predict/predict.py
+++ b/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 @@ class Predict(Module, Parameter):
         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 @@ class Predict(Module, Parameter):
 
             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 @@ class Predict(Module, Parameter):
             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 @@ class Predict(Module, Parameter):
             *_, 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)
 
diff --git a/dspy/primitives/module.py b/dspy/primitives/module.py
index 5df2b308..1dee9cbb 100644
--- a/dspy/primitives/module.py
+++ b/dspy/primitives/module.py
@@ -152,17 +152,17 @@ class BaseModule:
         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):
diff --git a/dspy/signatures/signature.py b/dspy/signatures/signature.py
index 16c2dc66..38f546f6 100644
--- a/dspy/signatures/signature.py
+++ b/dspy/signatures/signature.py
@@ -96,6 +96,10 @@ class SignatureMeta(type(BaseModel)):
     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 @@ class SignatureMeta(type(BaseModel)):
         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)
\ No newline at end of file
+    return " ".join(title_cased_words)


Test patch:
diff --git a/tests/predict/test_predict.py b/tests/predict/test_predict.py
index b04e087e..8d0121ee 100644
--- a/tests/predict/test_predict.py
+++ b/tests/predict/test_predict.py
@@ -126,6 +126,27 @@ def test_typed_demos_after_dump_and_load_state():
     # 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")
diff --git a/tests/signatures/test_signature.py b/tests/signatures/test_signature.py
index 8fb9c5ff..bc63db1a 100644
--- a/tests/signatures/test_signature.py
+++ b/tests/signatures/test_signature.py
@@ -238,3 +238,44 @@ def test_multiple_replaced_by_update_signatures():
         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:"