diff --git a/examples/recipes/intfloat_multilingual-e5-small/cpu/cpu/feature-extraction_fp16_config.json b/examples/recipes/intfloat_multilingual-e5-small/cpu/cpu/feature-extraction_fp16_config.json new file mode 100644 index 000000000..705ba52c2 --- /dev/null +++ b/examples/recipes/intfloat_multilingual-e5-small/cpu/cpu/feature-extraction_fp16_config.json @@ -0,0 +1,91 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 250037 + ] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 2 + ] + }, + { + "name": "token_type_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 2 + ] + } + ], + "output_tensors": [ + { + "name": "last_hidden_state" + } + ], + "compatibility": { + "transformers_attention": "eager" + } + }, + "optim": { + "clamp_constant_values": true + }, + "quant": { + "mode": "fp16", + "samples": 10, + "calibration_method": "minmax", + "weight_type": "uint8", + "activation_type": "uint8", + "per_channel": false, + "symmetric": false, + "weight_symmetric": null, + "activation_symmetric": null, + "save_calibration": false, + "distribution": "uniform", + "seed": null, + "calibration_load_path": null, + "calibration_save_path": null, + "op_types_to_quantize": null, + "nodes_to_exclude": null, + "task": "feature-extraction", + "model_id": "intfloat/multilingual-e5-small", + "model_type": "bert", + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "compile": null, + "loader": { + "task": "feature-extraction", + "model_class": "AutoModel", + "model_type": "bert" + } +} diff --git a/examples/recipes/intfloat_multilingual-e5-small/cpu/cpu/feature-extraction_fp32_config.json b/examples/recipes/intfloat_multilingual-e5-small/cpu/cpu/feature-extraction_fp32_config.json new file mode 100644 index 000000000..7df1f4555 --- /dev/null +++ b/examples/recipes/intfloat_multilingual-e5-small/cpu/cpu/feature-extraction_fp32_config.json @@ -0,0 +1,69 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 250037 + ] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 2 + ] + }, + { + "name": "token_type_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 2 + ] + } + ], + "output_tensors": [ + { + "name": "last_hidden_state" + } + ], + "compatibility": { + "transformers_attention": "eager" + } + }, + "optim": { + "clamp_constant_values": true + }, + "quant": null, + "compile": null, + "loader": { + "task": "feature-extraction", + "model_class": "AutoModel", + "model_type": "bert" + } +} diff --git a/examples/recipes/intfloat_multilingual-e5-small/cpu/cpu/sentence-similarity_fp16_config.json b/examples/recipes/intfloat_multilingual-e5-small/cpu/cpu/sentence-similarity_fp16_config.json new file mode 100644 index 000000000..1ff23afc6 --- /dev/null +++ b/examples/recipes/intfloat_multilingual-e5-small/cpu/cpu/sentence-similarity_fp16_config.json @@ -0,0 +1,91 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 250037 + ] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 2 + ] + }, + { + "name": "token_type_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 2 + ] + } + ], + "output_tensors": [ + { + "name": "last_hidden_state" + } + ], + "compatibility": { + "transformers_attention": "eager" + } + }, + "optim": { + "clamp_constant_values": true + }, + "quant": { + "mode": "fp16", + "samples": 10, + "calibration_method": "minmax", + "weight_type": "uint8", + "activation_type": "uint8", + "per_channel": false, + "symmetric": false, + "weight_symmetric": null, + "activation_symmetric": null, + "save_calibration": false, + "distribution": "uniform", + "seed": null, + "calibration_load_path": null, + "calibration_save_path": null, + "op_types_to_quantize": null, + "nodes_to_exclude": null, + "task": "sentence-similarity", + "model_id": "intfloat/multilingual-e5-small", + "model_type": "bert", + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "compile": null, + "loader": { + "task": "sentence-similarity", + "model_class": "AutoModel", + "model_type": "bert" + } +} diff --git a/examples/recipes/intfloat_multilingual-e5-small/cpu/cpu/sentence-similarity_fp32_config.json b/examples/recipes/intfloat_multilingual-e5-small/cpu/cpu/sentence-similarity_fp32_config.json new file mode 100644 index 000000000..c52f3420c --- /dev/null +++ b/examples/recipes/intfloat_multilingual-e5-small/cpu/cpu/sentence-similarity_fp32_config.json @@ -0,0 +1,69 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 250037 + ] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 2 + ] + }, + { + "name": "token_type_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 2 + ] + } + ], + "output_tensors": [ + { + "name": "last_hidden_state" + } + ], + "compatibility": { + "transformers_attention": "eager" + } + }, + "optim": { + "clamp_constant_values": true + }, + "quant": null, + "compile": null, + "loader": { + "task": "sentence-similarity", + "model_class": "AutoModel", + "model_type": "bert" + } +} diff --git a/src/winml/modelkit/models/winml/feature_extraction.py b/src/winml/modelkit/models/winml/feature_extraction.py index df3444159..7e4978370 100644 --- a/src/winml/modelkit/models/winml/feature_extraction.py +++ b/src/winml/modelkit/models/winml/feature_extraction.py @@ -15,6 +15,7 @@ from collections import OrderedDict from typing import Any +import numpy as np from transformers.utils.generic import ModelOutput from .base import WinMLPreTrainedModel @@ -58,7 +59,31 @@ def forward(self, **kwargs: Any) -> ModelOutput: order. Tensors keep their native rank (no unsqueeze); downstream pooling handles 1-D and 2-D after raw[0]. """ - outputs = self._run_inference(self._format_inputs(**kwargs)) + inputs = self._format_inputs(**kwargs) + if "token_type_ids" not in inputs and "input_ids" in inputs: + input_names = self.io_config.get("input_names", []) + if "token_type_ids" in input_names: + input_index = input_names.index("token_type_ids") + input_types = self.io_config.get("input_types", []) + input_shapes = self.io_config.get("input_shapes", []) + if input_index < len(input_types) and input_index < len(input_shapes): + required_shape = input_shapes[input_index] + actual_shape = inputs["input_ids"].shape + shape_matches = len(required_shape) == len(actual_shape) and all( + not isinstance(dimension, int) + or dimension <= 0 + or dimension == actual_dimension + for dimension, actual_dimension in zip( + required_shape, actual_shape, strict=True + ) + ) + if shape_matches: + inputs["token_type_ids"] = np.zeros( + actual_shape, + dtype=np.dtype(input_types[input_index]), + ) + + outputs = self._run_inference(inputs) # WinMLEncoderDecoderModel expects its encoder sub-component to expose # hidden states as "last_hidden_state". Alias the primary output when an # encoder ONNX graph named it otherwise (e.g. "encoder_hidden_states"). diff --git a/tests/unit/models/auto/test_feature_extraction.py b/tests/unit/models/auto/test_feature_extraction.py index e813eac50..90bc10722 100644 --- a/tests/unit/models/auto/test_feature_extraction.py +++ b/tests/unit/models/auto/test_feature_extraction.py @@ -15,6 +15,7 @@ from unittest.mock import MagicMock import numpy as np +import pytest import torch from transformers.utils import ModelOutput @@ -27,6 +28,8 @@ def create_mock_model(): mock_session = MagicMock() mock_session.io_config = { "input_names": ["input_ids", "attention_mask", "token_type_ids"], + "input_types": [np.dtype("int32"), np.dtype("int32"), np.dtype("int32")], + "input_shapes": [[1, 8], [1, 8], [1, 8]], "output_names": ["last_hidden_state"], } mock_session.run.return_value = { @@ -89,13 +92,69 @@ def test_optional_inputs_forwarded(self): assert "attention_mask" in call_kwargs assert "token_type_ids" in call_kwargs - def test_none_inputs_excluded(self): + def test_missing_token_type_ids_synthesized_with_required_shape_and_dtype(self): model = create_mock_model() - model.forward(input_ids=torch.ones(1, 8, dtype=torch.long)) + model.forward( + input_ids=torch.ones(1, 8, dtype=torch.long), + attention_mask=torch.ones(1, 8, dtype=torch.long), + ) + + call_kwargs = model._session.run.call_args[0][0] + np.testing.assert_array_equal(call_kwargs["token_type_ids"], np.zeros((1, 8))) + assert call_kwargs["token_type_ids"].dtype == np.int32 + + def test_provided_token_type_ids_preserved(self): + model = create_mock_model() + provided = torch.tensor([[0, 0, 0, 0, 1, 1, 1, 1]], dtype=torch.int64) + + model.forward( + input_ids=torch.ones(1, 8, dtype=torch.long), + attention_mask=torch.ones(1, 8, dtype=torch.long), + token_type_ids=provided, + ) call_kwargs = model._session.run.call_args[0][0] - assert "attention_mask" not in call_kwargs - assert "token_type_ids" not in call_kwargs + np.testing.assert_array_equal(call_kwargs["token_type_ids"], provided.numpy()) + assert call_kwargs["token_type_ids"].dtype == provided.numpy().dtype + + def test_unrelated_missing_required_input_still_fails(self): + from winml.modelkit.session.session import WinMLSession + + model = create_mock_model() + model._session.run.side_effect = lambda inputs: ( + WinMLSession._validate_inputs(model._session, inputs) + ) + + with pytest.raises(ValueError, match="attention_mask"): + model.forward(input_ids=torch.ones(1, 8, dtype=torch.long)) + + def test_incompatible_static_token_type_shape_is_not_hidden(self): + from winml.modelkit.session.session import WinMLSession + + model = create_mock_model() + model._session.io_config["input_shapes"][2] = [1, 16] + model._session.run.side_effect = lambda inputs: ( + WinMLSession._validate_inputs(model._session, inputs) + ) + + with pytest.raises(ValueError, match="token_type_ids"): + model.forward( + input_ids=torch.ones(1, 8, dtype=torch.long), + attention_mask=torch.ones(1, 8, dtype=torch.long), + ) + + def test_sentence_similarity_feature_path_synthesizes_each_single_segment(self): + model = create_mock_model() + + for token_id in (1, 2): + model.forward( + input_ids=torch.full((1, 8), token_id, dtype=torch.long), + attention_mask=torch.ones(1, 8, dtype=torch.long), + ) + + assert model._session.run.call_count == 2 + for call in model._session.run.call_args_list: + np.testing.assert_array_equal(call.args[0]["token_type_ids"], np.zeros((1, 8))) class TestForwardPreservesOnnxOutputNames: