Model Load

[fireking77]

Hi Team!

Framework version: dspy-ai-2.5.12

I'm encountering an issue when trying to load an optimized MiproV2 model and use it in my project. I have optimized my initial model using MiproV2 and saved it as a JSON file. However, when I load the model, it seems that the original model signature is being used, not the optimized one.
Despite following the example given in Issue #617,
I’m still unable to get it to work as expected.

Here's what I did:

Optimized my model with MiproV2 and saved it as mipro_optimized_20241017_100830.json.
Tried to load the model with the following approach:
python
Copy code
entity_extractor = EntityExtractor()
entity_extractor.load("mipro_optimized_20241017_100830.json")

I did some debuging and In dspy/primitives/module.py, I can see the load_state method:

def load_state(self, state, use_legacy_loading=False):
    for name, param in self.named_parameters():
        if isinstance(param, BaseModule):
            param.load_state(state[name], use_legacy_loading=use_legacy_loading)
        else:
            # `use_legacy_loading` is only applicable for BaseModule instances.
            param.load_state(state[name])

In my debuger, after this loop I saw the correct self.signature which came from the JSON manifest.
So It appears to correctly process the JSON file, but when returning to its invoker(the main load method):

def load(self, path, use_legacy_loading=False):
    with open(path) as f:
        self.load_state(ujson.loads(f.read()), use_legacy_loading=use_legacy_loading)

The original model is loaded (so the self.signature is the original one), not the optimized version (which was loaded in load_state).

So no suprise when I invoke the language model, it still uses the original signature rather than the optimized one from the JSON manifest.

Steps to reproduce:
entity_extractor = EntityExtractor() # My DSPy program
entity_extractor.load("mipro_optimized_20241017_100830.json")

Question:

1. Am I missing something in the model loading process? Is there a specific step needed to properly load and utilize the optimized model?
2. I use a TypedPredictor—is Mipro optimizing the signature for that? Should the optimized instructions be used as part of the system message in the LLM call?

Any help or guidance would be much appreciated. Thanks!

Darvi

[okhat]

Hey @fireking77 ! Thanks a lot for opening this.

It's possible that this is a bug in saving or loading. Let's figure out which one it is. Can you check the JSON file? Does it have your optimized instruction and/or examples? Check for extended_signature and signature.

[fireking77]

Hi @okhat !

Can you check the JSON file? Does it have your optimized instruction and/or examples?
Yes it have the optimisation part is perfect. This is why I asked this question :)

Check for extended_signature and signature.
I can't see the extended_signature in the JSON, only the signature:

"signature": {
      "instructions": "Extract and analyze entities from input descriptions related to fleet management and security systems.......<continues>

I installed the following packages to monitor the LLM calls, and it’s sending my original prompt, not the one shown above:

arize-phoenix
openinference-instrumentation-dspy

If you can tell me what I should debug, I’m ready to work with you on this!
By the way, aside from this hiccup, the framework works great and is very easy to use.

Thanks for all the time and effort you’ve put into this!
Darvi

[okhat]

Hey @fireking77 ! Thank you, this is super helpful!

So the JSON file has the proper, optimized instructions? But they're just not getting loaded?

Or are they not getting saved?

[okhat]

adding @chenmoneygithub re: saving a dspy.TypedPredictor

[fireking77]

Hey @fireking77 ! Thank you, this is super helpful!

So the JSON file has the proper, optimized instructions? But they're just not getting loaded?

Or are they not getting saved?

Not getting loaded, I think it is saved properly...I copied the part above from that JSON which was created after the mipro optimisation like this:

    from dspy.teleprompt import MIPROv2

    MAX_LABELED_DEMO = 16
    MAX_BOOTSTRAP_DEMO = 6
    NUM_OF_CANDIDATES = 20
    NUM_OF_TRIALS = 45

    teleprompter = MIPROv2(
        metric=metric,
        prompt_model=prompt_model,
        task_model=task_model,
        num_threads=3,
        num_candidates=NUM_OF_CANDIDATES,
        init_temperature=0.75,
        max_labeled_demos=MAX_LABELED_DEMO,
        max_bootstrapped_demos=MAX_BOOTSTRAP_DEMO,
        verbose=False,
    )

    # Optimize program
    print(f"Optimizing program with MIPRO...")
    optimized_program = teleprompter.compile(
        student=entity_extractor.deepcopy(),
        trainset=train_set,
        valset=validation_set,
        num_trials=NUM_OF_TRIALS,
        minibatch=False,
        # minibatch_size=15,
        # minibatch_full_eval_steps=10,
        max_labeled_demos=MAX_LABELED_DEMO,
        max_bootstrapped_demos=MAX_BOOTSTRAP_DEMO,
        requires_permission_to_run=False,
    )

    optimized_program.save("_data/entity_extractor_model.json")

I copied the whole thing from the doc, because after you guys created this nice interface it was not a rocket science! :D

actually this is my program:

class EntityExtractor(dspy.Module):
    def __init__(self):
        super().__init__()

        self.extractor = dspy.TypedPredictor(ExtractEntityFromDescription)

    def forward(self, input: ExtractEntityFromDescriptionInput):
        pred = self.extractor(input=input)
        return pred

[chenmoneygithub]

The original model is loaded (so the self.signature is the original one), not the optimized version (which was loaded in load_state).

@fireking77 could you elaborate a bit about this? load_state is invoked by this load, so once the state loading is successful it shouldn't be overwritten.

[fireking77]

demo1: original instruction - before init/load

demo2: optimized / expected instruction - after load

demo3: original instruction again before we return so the self object should be overrided with the info from the model JSON but we lost it somewhere - we can see on the predictor, but the signature is original again.

and when I check the prompt that was sent to the LLM, it's the original one. The issue might be elsewhere, and I may have confused you by mentioning these debug info, but it clearly does not send the calls with the optimized signature.

[chenmoneygithub]

@fireking77 I just realized that it's weird that MiproV2 is changing your signature instead of extended_signature. @okhat Omar I think we expect only the extended_signature to be changed by MiproV2?

For this specific loading, I am still looking into it.

[chenmoneygithub]

I tried to manually change the signature in json file, and I am able to load it back successfully:

@fireking77 Could you share how you define the DSPy model? We cannot reproduce the error on our side.

[chenmoneygithub]

I think I found the issue:

class TypedPredictor(dspy.Module):
    def __init__(self, signature, instructions=None, *, max_retries=3, wrap_json=False, explain_errors=False):
        """Like dspy.Predict, but enforces type annotations in the signature.

        Args:
            signature: The signature of the module. Can use type annotations.
            instructions: A description of what the model should do.
            max_retries: The number of times to retry the prediction if the output is invalid.
            wrap_json: If True, json objects in the input will be wrapped in ```json ... ```
            explain_errors: If True, the model will try to explain the errors it encounters.
        """
        super().__init__()
        self.signature = ensure_signature(signature, instructions)
        self.predictor = dspy.Predict(signature, _parse_values=False)
        self.max_retries = max_retries
        self.wrap_json = wrap_json
        self.explain_errors = explain_errors

@fireking77 In our TypedPredictor we store this self.signature, but saving/loading only handles this self.predictor. That means when you are loading the saved file, self.signature will diverge from self.predictor.signature. We have two options here:

• Add saving/loading logic for TypedPredictor.
• Remove the self.signature in TypePredictor.

@okhat Omar please let me know which one you prefer.

[okhat]

@chenmoneygithub We can remove the self.signature but it may break other stuff. How much code is saving this signature?

FYI: TypedPredictors will be deleted in 2.6 since dspy.Predict will replace their functionality/

[fireking77]

@chenmoneygithub , @okhat ,

I think I found the issue:
In our TypedPredictor we store this self.signature, but saving/loading only handles this self.predictor. That means when you are loading the saved file, self.signature will diverge from self.predictor.signature.

@chenmoneygithub Thanks so much for identifying this.
Omar mentioned that the entire TypedPredictor will be removed from the codebase with the upcoming Predict module.

In the meantime, how should I work around this issue? Should I use Typed CoT? Should I install the bleeding edge version of dspy from the repo? 😄
What do you suggest to help me move forward?
Actually, I manually added the optimized signature to my signature definition, and it worked! 😄 I noticed better performance. But how can I automate this until version 2.6?

Thanks,
Darvi

[okhat]

Hey @fireking77 ! I see a few paths.

The most temporary fix: If your predictor needs to be a TypedPredictor (i.e., dspy.Predict doesn't cut it), if your compile run is cached, you can recreate the same optimized program by compiling every time. It takes a few seconds and costs $0 if your cache is still there.

Beyond that, if you'd like to help us implement the fix "Add saving/loading logic for TypedPredictor" that would be awesome. The other path is to wait for @chenmoneygithub to do it, which may mean a Monday/Tuesday release.

Let me know what you think. Happy to help with any of these options.

[fireking77]

@okhat Omar,

For now, I'm manually changing the signature, but I plan to deep dive into the framework once I'm able/would be able to contribute a fix I will contribute help. Until then, the best I can do for the sake of everybody :D is some debugging and testing. So there's no rush on my end.

May I ask another question? In the next release, you mentioned removing this module. I believe the base Predict module will become the typed predictor, so it should be able to populate an object like this one, correct?

class ExtractEntityFromDescriptionOutput(BaseModel):
    entity_hu: str = Field(
        description="The extracted entity in Hungarian, refactored, lowercased, and cleaned from the input description, without any special Hungarian characters (only ASCII characters)."
    )
    entity_en: str = Field(
        description="The English translation of the extracted Hungarian entity."
    )
    is_inverted: bool = Field(
        description="Boolean flag indicating if the input is connected in an inverted way (for example, True if 'fordított', 'ford', 'for', or 'f' is present in the input description)."
    )
    categories: str = Field(
        description="It contains the categories, separated by '|', to which the extracted entity belongs."
    )
    review: bool = Field(
        description="Boolean flag indicating if the entity cannot be accurately identified or there's low confidence in the extraction."
    )

    def get_categories(self) -> Set[EntityCategory]:
        return set([EntityCategory(category.strip()) for category in self.categories.split("|")])

Thanks, guys! I really appreciate all the effort.
Darvi

[okhat]

Thank you @fireking77 ! Yes, we won't remove TypedPredictor until Predict can reliably handle the Pydantic model you showed :D