Improve DSPy saving

[chenmoneygithub]

We are reworking DSPy saving in order to make it more robust:

• Introduce whole model saving, users can specify state_only=False to save the model architecture along with the state. This is done via cloudpickle.
• Improve the state-only saving to support json and pickle, which resolves the problem that some demos are not json-serializable.
• Add unit test on both saving format.

From the user perspective, it means they have 3 journeys on saving a DSPy model.

Method 1: Save everything.

# Saving.
predict = dspy.Predict("question->answer")
optimizer = dspy.BootstrapFewShot(max_bootstrapped_demos=4, max_labeled_demos=4, max_rounds=5, metric=some_metric)
compiled_predict = optimizer.compile(predict, trainset=trainset)
compiled_predict.save("/tmp/my_model", save_program=True)   # `my_model` is a directory.

# Loading.
loaded_predict = dspy.load("/tmp/my_model")

Method 2 (our old approach): Save state with Json

# Saving.
predict = dspy.Predict("question->answer")
optimizer = dspy.BootstrapFewShot(max_bootstrapped_demos=4, max_labeled_demos=4, max_rounds=5, metric=some_metric)
compiled_predict = optimizer.compile(predict, trainset=trainset)
compiled_predict.save("/tmp/my_model/model.json", save_program=False)

# Loading.
loaded_predict = dspy.Predict("question->answer")
loaded_predict.load("/tmp/my_model/model.json")

Method 3: Save state with cloudpickle

# Saving.
predict = dspy.Predict("question->answer")
optimizer = dspy.BootstrapFewShot(max_bootstrapped_demos=4, max_labeled_demos=4, max_rounds=5, metric=some_metric)
compiled_predict = optimizer.compile(predict, trainset=trainset)
compiled_predict.save("/tmp/my_model/model.pkl", save_program=False)

# Loading.
loaded_predict = dspy.Predict("question->answer")
loaded_predict.load("/tmp/my_model/model.pkl")

[okhat]

Just thinking out loud: I don't like that we already have so many flags, since this is our opportunity to simplify save/load a lot, rather than to make it more complicated from a DX standpoint.

More actionable: I don't like that state_only is a "double negative". Turning it off does more work. It should probably be reversed, e.g. save_program=False by default.

Separately, do we need use_json? It seems to me that we use JSON iff save_program == False.

[chenmoneygithub]

that makes sense! We allow users to pickle the state when save_program == False and use_json = False, which handles the problem of serializing data like datetime.date.

What about this:

• I will remove the metadata arg, which is not useful now.
• I will remove the use_json arg, and use the file suffix (.json or .pkl) to determine the format for saving state.
• Replace state_only by save_program.

[okhat]

Does this method even need to exist?

[chenmoneygithub]

We can delete it since we have made our new plan for backward compatibility.