Add support for SigLIP

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")

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,

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",

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,

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",

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,
+ )

tests/onnxruntime/test_modeling.py

@@ -1987,6 +1987,7 @@ class ORTModelForFeatureExtractionIntegrationTest(ORTModelTestMixin):
  "mpnet",
  "roberta",
  "xlm_roberta",
+ "nomic-bert",
  ]
 
  FULL_GRID = {"model_arch": SUPPORTED_ARCHITECTURES}