transformers

transformers copied to clipboard

Feature request

Currently, when we use the save_pretrained function from this library the model signature used to save the model is the default one that only calls the model on the inputs, I would like to be able to provide a custom signature while using the save_pretrained function.

Motivation

Persisting models with custom signatures is quite important for models that target production setups, especially if they are going to be served with TF Serving.

I might be wrong but it seems that currently, the only way to save a Transformer model with a custom signature is by saving it using functions from the TF library, it would be very nice if the HF ecosystem could also support this feature.

Your contribution

I think this might be simple to implement and I would be happy to draft a PR if you think this could be a helpful feature.

cc @Rocketknight1 @gante

Hi @dimitreOliveira, that sounds like a great feature, and we'd be happy to accept that PR! We've been working on making our default signatures more general and usable, but this sounds like a good idea too. Are you planning to add a signatures argument that's passed through to model.save() when saved_model=True?

Hey @Rocketknight1 I am glad you liked the feature, I am happy to collaborate with the TF side of the lib.

Yes, my idea is to just add a signature parameters to the signatures function, that parameter would default to None and if that is the case we would just use self.serving as we already do, this way there would not be any relevant side-effect, and users could just create their custom signatures and pass it while saving. Looking at the code design it seems that this change would be compatible with all TF transformers models ; ) I have not looked yet to see if that would generate any issues with the tests, but if the plan is good I will work on the code during the weekend.

For context, the idea for this feature came to me while I was working on this repository, that also have a collection of custom signatures that range from text classification to text generation. Maybe this feature also works for the vision and speech models but I do not have a lot of experience with those, maybe later I could also take a look there.

@Rocketknight1 @gante you can find the draft PR above, let me know if it looks good, then I can finish the work, if needed, I can provide some examples of cool use cases using custom signatures with the models.