Custom tasks modeling

[IlyasMoutawwakil]

What does this PR do?

This PR adds custom tasks modeling through a new class OVModelForCustomTasks which allows using models that are exported with arbitrary inputs/outputs.

Fixes #611

how to use

For example to use google/vit-base-patch16-224 model for image-classification while also getting all the hidden_states:

• First we need to custom export the model using main_export and model_kwargs={"output_hidden_states": True} which is passed to the model at inference time, and custom_export_configs which is a dictionary of custom export configs that tells the exporter how to handle inputs/outputs (e.g. what to name them).

from transformers import AutoConfig

from optimum.exporters.onnx.model_configs import ViTOnnxConfig
from optimum.exporters.openvino import main_export


custom_task = "hidden_states"  # "attentions"
model_id = "google/vit-base-patch16-224"
base_task = "image-classification"

output = f"vit_with_{custom_task}"
model_kwargs = {f"output_{custom_task}": True}
repo_id = f"IlyasMoutawwakil/vit-with-{custom_task}"


class ViTOnnxConfigWithHiddenStates(ViTOnnxConfig):
    @property
    def outputs(self):
        common_outputs = {}
        common_outputs["logits"] = {0: "batch_size"}

        for i in range(self._normalized_config.num_hidden_layers + 1):
            common_outputs[f"hidden_states.{i}"] = {0: "batch_size"}

        return common_outputs


config = AutoConfig.from_pretrained(model_id)
constom_export_config = ViTOnnxConfigWithHiddenStates(config, task=base_task)


main_export(
    model_name_or_path=output,
    custom_export_configs={"model": constom_export_config},
    model_kwargs=model_kwargs,
    task=base_task,
    output=output,
)

• Then we can use OVModeForCustomTasks:

import requests
from PIL import Image
from transformers import AutoImageProcessor

from optimum.intel.openvino import OVModelForCustomTasks

image_processor = AutoImageProcessor.from_pretrained(model_id)
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)
inputs = image_processor(images=image, return_tensors="pt")

ov_model = OVModelForCustomTasks.from_pretrained(output)
ov_outputs = ov_model(**inputs)
print(ov_outputs.keys()) # dict_keys(['logits', 'hidden_states'])

ov_model.push_to_hub(save_directory=output, repository_id=repo_id)

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]

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.

[echarlaix]

Looks great, thanks a lot @IlyasMoutawwakil ! Might be worth adding a test with a customconfig as well, what do you think ?

[IlyasMoutawwakil]

@echarlaix I added an openvino variant of bert-with-pooler in IlyasMoutawwakil/sbert-all-MiniLM-L6-v2-with-pooler (can move it to optimum if you prefer), and added a custom export test that creates it

[echarlaix]

Looks great, thanks a lot @IlyasMoutawwakil !