Custom model export to onnx-runtime
I modified the BertEmbeddins, BertModel and BertForTokenClassification to accept additional feature (whether token in capitalized or not), in pure transformers it all works, but I am struggling with implementing the export of this custom model (so I can optimize it with optimum and get an inference speed up)register_for_onnx = TasksManager.create_register("onnx")

@register_for_onnx("custom-bert", "token-classification")
class CustomOnnxConfig(TextEncoderOnnxConfig):
    # Specifies how to normalize the BertConfig, this is needed to access common attributes
    # during dummy input generation.
    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
    # Sets the absolute tolerance to when validating the exported ONNX model against the
    # reference model.
    ATOL_FOR_VALIDATION = 1e-4


    def inputs(self) -> Dict[str, Dict[int, str]]:
        if self.task == "multiple-choice":
            dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"}
        else:
            dynamic_axis = {0: "batch_size", 1: "sequence_length"}
        return {
            "input_ids": dynamic_axis,
            "attention_mask": dynamic_axis,
            "capitalization_ids": dynamic_axis,
            "token_type_ids": dynamic_axis,
        }



base_model = CustomBertForTokenClassification.from_pretrained("my-checkpoint")

onnx_path = Path("model.onnx")Here I do not understand what to do nextbase_model.configreturns BertConfig, which I think I need to overwrite with the custom config I created in the previous step.Can you please help me?
2023-05-04T15:45:36Z
First, I can see that your new model will have a new input. We need aDummyInputGeneratorthat can handle this. So you could try something like:from optimum.utils import DummyTextInputGenerator

class MyDummyTextInputenerator(DummyTextInputGenerator):
    SUPPORTED_INPUT_NAMES = (                                                                                                                                                               
        "input_ids",                                                                                                                                                                        
        "attention_mask",                                                                                                                                                                   
        "token_type_ids",
        "capitalization_ids"
      )

class CustomOnnxConfig(TextEncoderOnnxConfig):
    # Specifies how to normalize the BertConfig, this is needed to access common attributes
    # during dummy input generation.
    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
    DUMMY_INPUT_GENERATOR_CLASSES = (MyDummyTextInputenerator,)
    # Sets the absolute tolerance to when validating the exported ONNX model against the
    # reference model.
    ATOL_FOR_VALIDATION = 1e-4

    @property
    def inputs(self) -> Dict[str, Dict[int, str]]:
        if self.task == "multiple-choice":
            dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"}
        else:
            dynamic_axis = {0: "batch_size", 1: "sequence_length"}
        return {
            "input_ids": dynamic_axis,
            "attention_mask": dynamic_axis,
            "capitalization_ids": dynamic_axis,
            "token_type_ids": dynamic_axis,
        }Since anOnnxConfigalready exists forbert, the register method will not do anything.If you want to be able to overwrite existing register configurations you can do that:register_for_onnx = TasksManager.create_register("onnx", overwrite_existing=True)Another approach is to do it programmatically: you can first export your model using Python code and then optimize it using the CLI.from pathlib import Path
from optimum.exporters.onnx import export

export(
    model,
    CustomOnnxConfig(model.config),
    Path("my_custom_model_onnx"),
)
2023-05-04T16:53:30Z
Thank you a lot! I am sorry, I still have issuesfrom optimum.utils import DummyTextInputGenerator
from optimum.exporters.onnx.config import OnnxConfig


class MyDummyTextInputenerator(DummyTextInputGenerator):
    SUPPORTED_INPUT_NAMES = (                                                                                                                                                               
        "input_ids",                                                                                                                                                                        
        "attention_mask",                                                                                                                                                                   
        "token_type_ids",
        "capitalization_ids"
      )
    
class TextEncoderOnnxConfig(OnnxConfig):
    # Describes how to generate the dummy inputs.
    DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator,)


register_for_onnx = TasksManager.create_register("onnx", overwrite_existing=True)

@register_for_onnx("bert", "token-classification")
class CustomOnnxConfig(TextEncoderOnnxConfig):
    # Specifies how to normalize the BertConfig, this is needed to access common attributes
    # during dummy input generation.
    NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
    DUMMY_INPUT_GENERATOR_CLASSES = (MyDummyTextInputenerator,)
    # Sets the absolute tolerance to when validating the exported ONNX model against the
    # reference model.
    ATOL_FOR_VALIDATION = 1e-4

    @property
    def inputs(self) -> Dict[str, Dict[int, str]]:
        if self.task == "multiple-choice":
            dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"}
        else:
            dynamic_axis = {0: "batch_size", 1: "sequence_length"}
        return {
            "input_ids": dynamic_axis,
            "attention_mask": dynamic_axis,
            "capitalization_ids": dynamic_axis,
            "token_type_ids": dynamic_axis,
        }

model = CustomBertForTokenClassification.from_pretrained("my-checkpoint")

from pathlib import Path
from optimum.exporters.onnx import export

export(
    model,
    CustomOnnxConfig(model.config),
    Path("model.onnx"),
)I get errorRuntimeError: number of output names provided (1) exceeded number of outputs (0)
2023-05-05T09:41:20Z
I think this might be linked tothis issue.Could you try installing optimum from sources and re-run your script?
2023-05-05T10:44:47Z
Did installation like this:pip install transformers[onnx]==4.28.1
python -m pip install git+https://github.com/huggingface/optimum.gitBut the error stayed sameRuntimeError: number of output names provided (1) exceeded number of outputs (0)
2023-05-09T18:35:18Z
Hi Maiia,I saw your messagehere, thanks for creating a notebook, it really helped me to try things out.You do not need to register anything since you are not using theTasksManagerin the end. Creating yourOnnxConfigand using theexportfunction is enough.But because you are not using theTasksManager, you have to instantiate theOnnxConfigmanually, andyou need to specify the name of the task when doing so, otherwise it will infere the task to be the default one(BertModeland notBertForTokenClassification).So to be able to export your model the fix is actually easy:export(
    model,
    CustomOnnxConfig(model.config, task="token-classification"),
    Path("model.onnx"),
)
2023-05-09T19:02:13Z
michaelbenayoun:hereThank you a lot! You helped me so, so much!Last question - if I need to reload it (asOrtModelForTokenClassificationto optimize withORTOptimizer) I probably need to save a modified config? (If a save original config and doreloaded_model = ORTModelForTokenClassification.from_pretrained("onnx")
inputs = {k: torch.zeros([2, 16], dtype=torch.long) for k in ["input_ids", "attention_mask", "capitalization_ids", "token_type_ids"]}
reloaded_model(**inputs)I get an error aboutValueError: Model requires 4 inputs. Input Feed contains 3Or maybe I just need to modifyOrtModelForTokenClassificationclass and it will fix the issue…
2023-05-09T19:16:34Z
ModifiedOrtForTokenClassification- and everything works perfectly now!Thank you so, so, so much for patience and kindness and spending time on a beginner!And for wonderful work which Huggingface gives people! Thank you!!!