Proper sentence-transformers ONNX export support

[fxmarty]

As reported in #1519, simple mapping sentence-transformers to transformers library allows to use only a subset of sentence-transformers library.

This PR adds the support of the export of sentence_embedding for sentence-transformers models.

Examples:

optimum-cli export onnx -m sentence-transformers/clip-ViT-B-32-multilingual-v1 clip_vit_multilingual_onnx
optimum-cli export onnx -m sentence-transformers/all-MiniLM-L6-v2 minilm_onnx

[JingyaHuang]

Can those extra outputs be supported directly in transformers? Just find the changes a bit hacky, and this is causing errors in optimum neuron subpackage: aws-neuron/aws-neuron-sdk#808

[JingyaHuang]

standardizing model attributes made it a bit misleading for debug, eg. it was surprising to get sentence-transformers-transformer as model_type from a config where the model_type is marked as bert:

[fxmarty]

I doubt this is feasible, as sentence-transformers adds quite a few features on top of transformers. For example, the sentence embeddings.

standardizing model attributes made it a bit misleading for debug, eg. it was surprising to get sentence-transformers-transformer as model_type from a config where the model_type is marked as bert:

I don't find this to be too hacky, as the model_type bert refers to the bottleneck model in sentence-transformers' nn.Sequential. Some sentence-transformers models use the Transformer as bottleneck, some use a CLIPModel, and the export is different depending on the architecture:

optimum/optimum/exporters/onnx/model_configs.py

Lines 818 to 842 in fc214d4

	class SentenceTransformersTransformerOnnxConfig(TextEncoderOnnxConfig):
	NORMALIZED_CONFIG_CLASS = NormalizedTextConfig
	DEFAULT_ONNX_OPSET = 14 # Some bottleneck transformers models require a specific ONNX opset to be successfully exported. We put a rather high opset here for the export to work for all architectures.
	@property
	def inputs(self) -> Dict[str, Dict[int, str]]:
	return {
	"input_ids": {0: "batch_size", 1: "sequence_length"},
	"attention_mask": {0: "batch_size", 1: "sequence_length"},
	}
	@property
	def outputs(self) -> Dict[str, Dict[int, str]]:
	return {
	"token_embeddings": {0: "batch_size", 1: "sequence_length"},
	"sentence_embedding": {0: "batch_size"},
	}
	# we need to set output_attentions=True in the model input to avoid calling
	# torch.nn.functional.scaled_dot_product_attention that is not supported by the ONNX export
	# due to the op torch.nn.functional.multi_head_attention_forward used for WavLM
	def patch_model_for_export(
	self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
	) -> "ModelPatcher":
	return SentenceTransformersTransformerPatcher(self, model, model_kwargs=model_kwargs)

&

optimum/optimum/exporters/onnx/model_configs.py

Lines 872 to 883 in fc214d4

	class SentenceTransformersCLIPOnnxConfig(CLIPOnnxConfig):
	@property
	def outputs(self) -> Dict[str, Dict[int, str]]:
	return {
	"text_embeds": {0: "text_batch_size"},
	"image_embeds": {0: "image_batch_size"},
	}
	def patch_model_for_export(
	self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None
	) -> "ModelPatcher":
	return SentenceTransformersCLIPPatcher(self, model, model_kwargs=model_kwargs)

.

Happy to refactor if needed though

[JingyaHuang]

Why is it timm test? @fxmarty

[fxmarty]

It is a typo