diff --git a/optimum/commands/env.py b/optimum/commands/env.py index 9b76a9fb69..ee01ed0601 100644 --- a/optimum/commands/env.py +++ b/optimum/commands/env.py @@ -57,8 +57,8 @@ def run(self): "Platform": platform.platform(), "Python version": platform.python_version(), "Huggingface_hub version": huggingface_hub.__version__, - "PyTorch version (GPU?)": f"{pt_version} (cuda availabe: {pt_cuda_available})", - "Tensorflow version (GPU?)": f"{tf_version} (cuda availabe: {tf_cuda_available})", + "PyTorch version (GPU?)": f"{pt_version} (cuda available: {pt_cuda_available})", + "Tensorflow version (GPU?)": f"{tf_version} (cuda available: {tf_cuda_available})", } print("\nCopy-and-paste the text below in your GitHub issue:\n") diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py index 26202e889b..277f9f2cfc 100644 --- a/optimum/exporters/onnx/model_configs.py +++ b/optimum/exporters/onnx/model_configs.py @@ -116,6 +116,10 @@ def inputs(self) -> Dict[str, Dict[int, str]]: } +class NomicBertOnnxConfig(BertOnnxConfig): + DEFAULT_ONNX_OPSET = 14 + + class AlbertOnnxConfig(BertOnnxConfig): DEFAULT_ONNX_OPSET = 11 @@ -785,6 +789,19 @@ class ResNetOnnxConfig(ViTOnnxConfig): DEFAULT_ONNX_OPSET = 11 +class SigLIPVisionOnnxConfig(ViTOnnxConfig): + pass + + +class SigLIPVisionModelOnnxConfig(SigLIPVisionOnnxConfig): + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + return { + "last_hidden_state": {0: "batch_size"}, + "pooler_output": {0: "batch_size"}, + } + + class DetrOnnxConfig(ViTOnnxConfig): DEFAULT_ONNX_OPSET = 12 @@ -1004,6 +1021,10 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): return dummy_inputs +class SigLIPTextOnnxConfig(CLIPTextOnnxConfig): + pass + + class UNetOnnxConfig(VisionOnnxConfig): ATOL_FOR_VALIDATION = 1e-3 # The ONNX export of a CLIPText architecture, an other Stable Diffusion component, needs the Trilu @@ -1750,7 +1771,8 @@ class SpeechT5OnnxConfig(OnnxSeq2SeqConfigWithPast): # so we won't support for now. NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( hidden_size="hidden_size", - num_attention_heads="encoder_attention_heads", # TODO: bugged in case encoder and decoder have different number of heads + num_attention_heads="encoder_attention_heads", + # TODO: bugged in case encoder and decoder have different number of heads encoder_num_layers="encoder_layers", decoder_num_layers="decoder_layers", allow_new=True, diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py index f02f176923..cbfccfa285 100644 --- a/optimum/exporters/tasks.py +++ b/optimum/exporters/tasks.py @@ -871,6 +871,7 @@ class TasksManager: "text2text-generation-with-past", onnx="M2M100OnnxConfig", ), + "nomic-bert": supported_tasks_mapping("feature-extraction", onnx="NomicBertOnnxConfig"), "nystromformer": supported_tasks_mapping( "feature-extraction", "fill-mask", diff --git a/optimum/utils/normalized_config.py b/optimum/utils/normalized_config.py index 81207b7649..5ae5310ab1 100644 --- a/optimum/utils/normalized_config.py +++ b/optimum/utils/normalized_config.py @@ -259,6 +259,7 @@ class NormalizedConfigManager: "mpt": MPTNormalizedTextConfig, "mt5": T5LikeNormalizedTextConfig, "m2m-100": BartLikeNormalizedTextConfig, + "nomic-bert": NormalizedTextConfig, "nystromformer": NormalizedTextConfig, "opt": NormalizedTextConfig, "pegasus": BartLikeNormalizedTextConfig, diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py index a55c7a124d..25b4123071 100644 --- a/tests/exporters/exporters_utils.py +++ b/tests/exporters/exporters_utils.py @@ -251,6 +251,7 @@ "mpt": "mosaicml/mpt-7b", "mt5": "lewtun/tiny-random-mt5", # Not using google/mt5-small because it takes too much time for testing. "musicgen": "facebook/musicgen-small", + "nomic-bert": "nomic-ai/nomic-embed-text-v1.5", "nystromformer": "hf-internal-testing/tiny-random-NystromformerModel", "owlv2": "google/owlv2-base-patch16", "owlvit": "google/owlvit-base-patch32", diff --git a/tests/exporters/onnx/test_exporters_onnx_cli.py b/tests/exporters/onnx/test_exporters_onnx_cli.py index ed611ade04..6b57558279 100644 --- a/tests/exporters/onnx/test_exporters_onnx_cli.py +++ b/tests/exporters/onnx/test_exporters_onnx_cli.py @@ -730,3 +730,20 @@ def test_complex_synonyms(self): model.save_pretrained(tmpdir_in) main_export(model_name_or_path=tmpdir_in, output=tmpdir_out, task="text-classification") + + @slow + def test_custom_model(self): + with TemporaryDirectory() as tmpdirname: + out = subprocess.run( + f"python3 -m optimum.exporters.onnx --trust-remote-code --model nomic-ai/nomic-embed-text-v1.5 --task feature-extraction {tmpdirname}", + shell=True, + capture_output=True, + ) + self.assertFalse(out.returncode) + + with TemporaryDirectory() as tmpdirname: + out = subprocess.run( + f"python3 -m optimum.exporters.onnx --trust-remote-code --model nomic-ai/nomic-embed-text-v1.5 --task feature-extraction {tmpdirname}", + shell=True, + check=True, + ) diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py index 473ab4cf3b..5c449e72e5 100644 --- a/tests/onnxruntime/test_modeling.py +++ b/tests/onnxruntime/test_modeling.py @@ -1987,6 +1987,7 @@ class ORTModelForFeatureExtractionIntegrationTest(ORTModelTestMixin): "mpnet", "roberta", "xlm_roberta", + "nomic-bert", ] FULL_GRID = {"model_arch": SUPPORTED_ARCHITECTURES}