Repository: stanfordnlp/dspy

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
                )
  ```


Discussion:
Yes, this is reasonable — but no current optimizers optimize the little parts that are not saved. 
Thanks, I'm pretty sure optimize_signature for TypedPredictor does.

e.g. After optimizing this signature :
```
generate_related_questions.generate_related_questions.predictor.predictor = Predict(RelatedQuestionsSignature(query, topics -> related_questions
    instructions='\n    Generate related questions for a given query and topics\n    '
    query = Field(annotation=str required=True json_schema_extra={'__dspy_field_type': 'input', 'prefix': 'Query:', 'desc': '${query}'})
    topics = Field(annotation=str required=True json_schema_extra={'__dspy_field_type': 'input', 'prefix': 'Topics:', 'desc': '${topics}'})
    related_questions = Field(annotation=List[str] required=True json_schema_extra={'desc': 'A list of 4 related questions to the input question', '__dspy_field_type': 'output', 'prefix': 'Related Questions:'})
))
```
Becomes:
```
generate_related_questions.generate_related_questions.predictor.predictor = Predict(StringSignature(query, topics -> related_questions
    instructions='You are a professor of mathematics. Generate related questions for a given query and topics. Use lots of creativity and be as clear as possible. I really need your help!'
    query = Field(annotation=str required=True json_schema_extra={'__dspy_field_type': 'input', 'prefix': 'Main Query:', 'desc': '${query}'})
    topics = Field(annotation=str required=True json_schema_extra={'__dspy_field_type': 'input', 'prefix': 'Associated Topics:', 'desc': '${topics}'})
    related_questions = Field(annotation=List[str] required=True json_schema_extra={'desc': 'A list of 4 creative and clear questions related to the input query.', '__dspy_field_type': 'output', 'prefix': 'Related Questions:'})
))
```




Thank you @GFarnon ! I guess I reached the point where I don't know every little detail in all optimizers anymore. Ok yes this is a change we should merge. Would be so kind as to open a PR?
🫡
I'm very inexperienced so it's likely something I'm doing wrong, but this has tripped me up a lot.

The docs tell me that a saved program "contains all the parameters and steps in the source program", but it seems that the reality is more like: "running an optimiser and saving its output will save the specific output of that optimiser, not the whole program".
@chenmoneygithub Wanna take a look at this if you get a chance? Seems relevant to your current efforts.
Sure! I will revisit what states should be saved at `dump_state` call. 



Patch:
diff --git a/dspy/predict/predict.py b/dspy/predict/predict.py
index 8bd42699d5..a029a11712 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 @@ 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)
 
diff --git a/dspy/primitives/module.py b/dspy/primitives/module.py
index 5df2b308df..1dee9cbb19 100644
--- a/dspy/primitives/module.py
+++ b/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):
diff --git a/dspy/signatures/signature.py b/dspy/signatures/signature.py
index 16c2dc66da..38f546f638 100644
--- a/dspy/signatures/signature.py
+++ b/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)
\ 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 b04e087e8b..8d0121eedd 100644
--- a/tests/predict/test_predict.py
+++ b/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")
diff --git a/tests/signatures/test_signature.py b/tests/signatures/test_signature.py
index 8fb9c5ff24..bc63db1ac3 100644
--- a/tests/signatures/test_signature.py
+++ b/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:"