Add support to export ColPali Model to ONNX#2074
Conversation
|
@fxmarty, @echarlaix, @JingyaHuang, @michaelbenayoun Are you open to merging this? |
|
Apologies for the delay @akshayballal95, could you add a test with a tiny random model like https://huggingface.co/hf-internal-testing/tiny-random-PaliGemmaForConditionalGeneration, can be added here https://github.com/huggingface/optimum/blob/main/tests/exporters/exporters_utils.py#L37 |
| class ColPaliModelPatcher(ModelPatcher): | ||
| def __init__( | ||
| self, | ||
| config: "OnnxConfig", | ||
| model: Union["PreTrainedModel", "TFPreTrainedModel"], | ||
| model_kwargs: Optional[Dict[str, Any]] = None, | ||
| ): | ||
| super().__init__(config, model, model_kwargs) | ||
|
|
||
| def patched_forward(input_ids=None, pixel_values=None, attention_mask=None, **kwargs): | ||
| outputs = self.orig_forward( | ||
| input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask, **kwargs | ||
| ) | ||
| return outputs | ||
|
|
||
| self.patched_forward = patched_forward |
There was a problem hiding this comment.
its needed because the original ColPali Model only takes **kwargs and no named arguments. This resulted in an error. This fixes that error
I have added the conversion test. It works fine locally. |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
This PR has been marked as stale because it has been open for 90 days with no activity. This thread will be automatically closed in 30 days if no further activity occurs. |
3 participants
What does this PR do?
This PR adds support for exporting the ColPali merged model to ONNX format. The model is based on the "pali gemma" model type, and thus, I have added it under the "feature-extraction" task. Do suggest if there is a better way to integrate this. If this looks fine with a few modifications, I can add support for the Paligemma text-generation task as well.
Before submitting
Who can review?
@fxmarty, @echarlaix, @JingyaHuang, @michaelbenayoun