Merge pull request stanfordnlp#1370 from stanfordnlp/extend-save-info

Save field metadata in `module.save()` with `add_field_meta` feature flag

docs/docs/building-blocks/6-optimizers.md

@@ -122,6 +122,8 @@ optimized_program.save(YOUR_SAVE_PATH)
 
 The resulting file is in plain-text JSON format. It contains all the parameters and steps in the source program. You can always read it and see what the optimizer generated.
 
+You can add `save_field_meta` to additionally save the list of fields with the keys, `name`, `field_type`, `description`, and `prefix` with: `optimized_program.save(YOUR_SAVE_PATH, save_field_meta=True).
+
 ### Loading a program
 
 To load a program from a file, you can instantiate an object from that class and then call the load method on it.

dspy/predict/predict.py

@@ -23,7 +23,7 @@ def reset(self):
         self.train = []
         self.demos = []
 
-    def dump_state(self):
+    def dump_state(self, save_field_meta=False):
         state_keys = ["lm", "traces", "train"]
         state = {k: getattr(self, k) for k in state_keys}
 
@@ -37,6 +37,19 @@ def dump_state(self):
 
             state["demos"].append(demo)
 
+        # If `save_field_meta` save all field metadata as well.
+        if save_field_meta:
+            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
@@ -196,10 +209,8 @@ def new_generate(lm, signature, example, max_depth=6, **kwargs):
 
     return completions
 
-
-
 # TODO: get some defaults during init from the context window?
 # # TODO: FIXME: Hmm, I guess expected behavior is that contexts can
 # affect execution. Well, we need to determine whether context dominates, __init__ demoninates, or forward dominates.
 # Generally, unless overwritten, we'd see n=None, temperature=None.
-# That will eventually mean we have to learn them.
+# That will eventually mean we have to learn them.

dspy/primitives/module.py

@@ -112,9 +112,9 @@ def reset_copy(self):
 
         return obj
 
-    def dump_state(self):
+    def dump_state(self, save_field_meta):
         print(self.named_parameters())
-        return {name: param.dump_state() for name, param in self.named_parameters()}
+        return {name: param.dump_state(save_field_meta) for name, param in self.named_parameters()}
 
     def load_state(self, state):
         for name, param in self.named_parameters():
@@ -127,9 +127,9 @@ def load_state(self, state):
             #     else:
             #         raise
 
-    def save(self, path):
+    def save(self, path, save_field_meta=False):
         with open(path, "w") as f:
-            f.write(ujson.dumps(self.dump_state(), indent=2))
+            f.write(ujson.dumps(self.dump_state(save_field_meta), indent=2))
 
     def load(self, path):
         with open(path) as f:
@@ -147,4 +147,4 @@ def postprocess_parameter_name(name, value):
     if name == "_predict":
         return "self"
 
-    return name
+    return name

dspy/primitives/program.py

@@ -77,4 +77,4 @@ def activate_assertions(self, handler=backtrack_handler, **handler_args):
 def set_attribute_by_name(obj, name, value):
     magicattr.set(obj, name, value)
 
-Program = Module
+Program = Module

dspy/signatures/signature.py

@@ -423,4 +423,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)