Skip to content

Add feature extraction support for image models#2255

Merged
IlyasMoutawwakil merged 3 commits into
mainfrom
ortmodel-image-feature-extraction
May 12, 2025
Merged

Add feature extraction support for image models#2255
IlyasMoutawwakil merged 3 commits into
mainfrom
ortmodel-image-feature-extraction

Conversation

@Abdennacer-Badaoui

Copy link
Copy Markdown
Member

What does this PR do?

This PR adds support for image models in ORTModelForFeatureExtraction, specifically addressing ViT and other visual models. The implementation modifies the forward method to dynamically handle both text and image inputs by detecting the presence of image-specific input keys like 'pixel_values'. This allows users to seamlessly use ORTModelForFeatureExtraction for both text and image feature extraction tasks without requiring a separate class structure. Tests have been updated to properly validate this functionality across all supported architectures, including Vision Transformer (ViT) models.

Fixes # (issue)
Fixes the issue : #2226

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you make sure to update the documentation with your changes?
  • Did you write any new necessary tests?

@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

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.

Comment thread optimum/onnxruntime/modeling_ort.py Outdated
Comment thread optimum/onnxruntime/modeling_ort.py Outdated
Comment thread optimum/onnxruntime/modeling_ort.py Outdated
Comment thread tests/onnxruntime/test_modeling.py Outdated
@IlyasMoutawwakil

Copy link
Copy Markdown
Member

No need to change the testing suite for this, as it's gonna be undergoing some refactoring

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants