Add support for sentence similarity

xadupre · xadupre · commit dc00d0786c33 · 2025-04-15T15:07:58.000+02:00
diff --git a/_doc/api/tasks/automatic_speech_recognition.rst b/_doc/api/tasks/automatic_speech_recognition.rst
@@ -1,6 +1,6 @@
 
-onnx_diagnostic.export.automatic_speech_recognition
-===================================================
+onnx_diagnostic.tasks.automatic_speech_recognition
+==================================================
 
 .. automodule:: onnx_diagnostic.tasks.automatic_speech_recognition
     :members:
diff --git a/_doc/api/tasks/fill_mask.rst b/_doc/api/tasks/fill_mask.rst
@@ -1,6 +1,6 @@
 
-onnx_diagnostic.export.fill_mask
-================================
+onnx_diagnostic.tasks.fill_mask
+===============================
 
 .. automodule:: onnx_diagnostic.tasks.fill_mask
     :members:
diff --git a/_doc/api/tasks/image_classification.rst b/_doc/api/tasks/image_classification.rst
@@ -1,6 +1,6 @@
 
-onnx_diagnostic.export.image_classification
-===========================================
+onnx_diagnostic.tasks.image_classification
+==========================================
 
 .. automodule:: onnx_diagnostic.tasks.image_classification
     :members:
diff --git a/_doc/api/tasks/index.rst b/_doc/api/tasks/index.rst
@@ -9,6 +9,7 @@ onnx_diagnostic.tasks
     fill_mask
     image_classification
     image_text_to_text
+    sentence_similarity
     text_classification
     text_generation
     text2text_generation
diff --git a/_doc/api/tasks/sentence_similarity.rst b/_doc/api/tasks/sentence_similarity.rst
@@ -0,0 +1,7 @@
+
+onnx_diagnostic.tasks.sentence_similarity
+=========================================
+
+.. automodule:: onnx_diagnostic.tasks.sentence_similarity
+    :members:
+    :no-undoc-members:
diff --git a/_doc/api/tasks/text2text_generation.rst b/_doc/api/tasks/text2text_generation.rst
@@ -1,6 +1,6 @@
 
-onnx_diagnostic.export.text2text_generation
-===========================================
+onnx_diagnostic.tasks.text2text_generation
+==========================================
 
 .. automodule:: onnx_diagnostic.tasks.text2text_generation
     :members:
diff --git a/_doc/api/tasks/text_classification.rst b/_doc/api/tasks/text_classification.rst
@@ -1,6 +1,6 @@
 
-onnx_diagnostic.export.text_classification
-==========================================
+onnx_diagnostic.tasks.text_classification
+=========================================
 
 .. automodule:: onnx_diagnostic.tasks.text_classification
     :members:
diff --git a/_doc/api/tasks/text_generation.rst b/_doc/api/tasks/text_generation.rst
@@ -1,6 +1,6 @@
 
-onnx_diagnostic.export.text_generation
-======================================
+onnx_diagnostic.tasks.text_generation
+=====================================
 
 .. automodule:: onnx_diagnostic.tasks.text_generation
     :members:
diff --git a/_doc/api/tasks/zero_shot_image_classification.rst b/_doc/api/tasks/zero_shot_image_classification.rst
@@ -1,6 +1,6 @@
 
-onnx_diagnostic.export.zero_shot_image_classification
-=====================================================
+onnx_diagnostic.tasks.zero_shot_image_classification
+====================================================
 
 .. automodule:: onnx_diagnostic.tasks.zero_shot_image_classification
     :members:
diff --git a/_unittests/ut_tasks/test_tasks.py b/_unittests/ut_tasks/test_tasks.py
@@ -9,7 +9,6 @@ class TestTasks(ExtTestCase):
     @hide_stdout()
     def test_text2text_generation(self):
         mid = "sshleifer/tiny-marian-en-de"
-        # mid = "Salesforce/codet5-small"
         data = get_untrained_model_with_inputs(mid, verbose=1)
         self.assertIn((data["size"], data["n_weights"]), [(473928, 118482)])
         model, inputs = data["model"], data["inputs"]
@@ -85,7 +84,6 @@ def test_automatic_speech_recognition(self):
     @hide_stdout()
     def test_imagetext2text_generation(self):
         mid = "HuggingFaceM4/tiny-random-idefics"
-        # mid = "Salesforce/codet5-small"
         data = get_untrained_model_with_inputs(mid, verbose=1)
         self.assertIn((data["size"], data["n_weights"]), [(12742888, 3185722)])
         model, inputs = data["model"], data["inputs"]
@@ -94,7 +92,6 @@ def test_imagetext2text_generation(self):
     @hide_stdout()
     def test_fill_mask(self):
         mid = "google-bert/bert-base-multilingual-cased"
-        # mid = "Salesforce/codet5-small"
         data = get_untrained_model_with_inputs(mid, verbose=1)
         self.assertIn((data["size"], data["n_weights"]), [(428383212, 107095803)])
         model, inputs = data["model"], data["inputs"]
@@ -103,12 +100,19 @@ def test_fill_mask(self):
     @hide_stdout()
     def test_text_classification(self):
         mid = "Intel/bert-base-uncased-mrpc"
-        # mid = "Salesforce/codet5-small"
         data = get_untrained_model_with_inputs(mid, verbose=1)
         self.assertIn((data["size"], data["n_weights"]), [(154420232, 38605058)])
         model, inputs = data["model"], data["inputs"]
         model(**inputs)
 
+    @hide_stdout()
+    def test_sentence_similary(self):
+        mid = "sentence-transformers/all-MiniLM-L6-v1"
+        data = get_untrained_model_with_inputs(mid, verbose=1)
+        self.assertIn((data["size"], data["n_weights"]), [(62461440, 15615360)])
+        model, inputs = data["model"], data["inputs"]
+        model(**inputs)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)
diff --git a/_unittests/ut_tasks/try_tasks.py b/_unittests/ut_tasks/try_tasks.py
@@ -236,6 +236,55 @@ def test_text_classification(self):
         encoded_input["input_ids"][0]
         tokenizer.convert_ids_to_tokens(encoded_input["input_ids"][0])
 
+    @never_test()
+    def test_sentence_similary(self):
+        # clear&&NEVERTEST=1 python _unittests/ut_tasks/try_tasks.py -k ce_sim
+        # https://huggingface.co/sentence-transformers/all-MiniLM-L6-v1
+
+        from transformers import AutoTokenizer, AutoModel
+        import torch
+        import torch.nn.functional as F
+
+        # Mean Pooling - Take attention mask into account for correct averaging
+        def mean_pooling(model_output, attention_mask):
+            token_embeddings = model_output[
+                0
+            ]  # First element of model_output contains all token embeddings
+            input_mask_expanded = (
+                attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+            )
+            return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
+                input_mask_expanded.sum(1), min=1e-9
+            )
+
+        # Sentences we want sentence embeddings for
+        sentences = ["This is an example sentence", "Each sentence is converted"]
+
+        # Load model from HuggingFace Hub
+        tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v1")
+        model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v1")
+
+        # Tokenize sentences
+        encoded_input = tokenizer(
+            sentences, padding=True, truncation=True, return_tensors="pt"
+        )
+
+        # Compute token embeddings
+        with torch.no_grad():
+            print()
+            print("-- inputs", string_type(encoded_input, with_shape=True, with_min_max=True))
+            model_output = model(**encoded_input)
+            print("-- outputs", string_type(model_output, with_shape=True, with_min_max=True))
+
+        # Perform pooling
+        sentence_embeddings = mean_pooling(model_output, encoded_input["attention_mask"])
+
+        # Normalize embeddings
+        sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
+
+        print("Sentence embeddings:")
+        print(sentence_embeddings)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)
diff --git a/_unittests/ut_xrun_doc/test_documentation_examples.py b/_unittests/ut_xrun_doc/test_documentation_examples.py
@@ -54,7 +54,10 @@ def run_test(self, fold: str, name: str, verbose=0) -> int:
                     # dot not installed, this part
                     # is tested in onnx framework
                     raise unittest.SkipTest(f"failed: {name!r} due to missing dot.")
-                if "We couldn't connect to 'https://huggingface.co'" in st:
+                if (
+                    "We couldn't connect to 'https://huggingface.co'" in st
+                    or "Cannot access content at: https://huggingface.co/" in st
+                ):
                     raise unittest.SkipTest(f"Connectivity issues due to\n{err}")
                 raise AssertionError(  # noqa: B904
                     "Example '{}' (cmd: {} - exec_prefix='{}') "
diff --git a/onnx_diagnostic/tasks/__init__.py b/onnx_diagnostic/tasks/__init__.py
@@ -4,6 +4,7 @@
     fill_mask,
     image_classification,
     image_text_to_text,
+    sentence_similarity,
     text_classification,
     text_generation,
     text2text_generation,
@@ -15,6 +16,7 @@
     fill_mask,
     image_classification,
     image_text_to_text,
+    sentence_similarity,
     text_classification,
     text_generation,
     text2text_generation,
diff --git a/onnx_diagnostic/tasks/sentence_similarity.py b/onnx_diagnostic/tasks/sentence_similarity.py
@@ -0,0 +1,67 @@
+from typing import Any, Callable, Dict, Optional, Tuple
+import torch
+from ..helpers.config_helper import update_config, check_hasattr
+
+__TASK__ = "sentence-similarity"
+
+
+def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+    """Reduces a model size."""
+    check_hasattr(config, "num_attention_heads", "num_hidden_layers")
+    kwargs = dict(
+        num_hidden_layers=min(config.num_hidden_layers, 2),
+        num_attention_heads=min(config.num_attention_heads, 4),
+    )
+    update_config(config, kwargs)
+    return kwargs
+
+
+def get_inputs(
+    model: torch.nn.Module,
+    config: Optional[Any],
+    batch_size: int,
+    sequence_length: int,
+    dummy_max_token_id: int,
+    **kwargs,  # unused
+):
+    """
+    Generates inputs for task ``sentence-similarity``.
+    Example:
+
+    ::
+
+        input_ids:T7s1x13[101,72654:A16789.23076923077],
+        token_type_ids:T7s1x13[0,0:A0.0],
+        attention_mask:T7s1x13[1,1:A1.0])
+    """
+    batch = torch.export.Dim("batch", min=1, max=1024)
+    seq_length = "seq_length"
+    shapes = {
+        "input_ids": {0: batch, 1: seq_length},
+        "token_type_ids": {0: batch, 1: seq_length},
+        "attention_mask": {0: batch, 1: seq_length},
+    }
+    inputs = dict(
+        input_ids=torch.randint(0, dummy_max_token_id, (batch_size, sequence_length)).to(
+            torch.int64
+        ),
+        token_type_ids=torch.zeros((batch_size, sequence_length)).to(torch.int64),
+        attention_mask=torch.ones((batch_size, sequence_length)).to(torch.int64),
+    )
+    return dict(inputs=inputs, dynamic_shapes=shapes)
+
+
+def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:
+    """
+    Inputs kwargs.
+
+    If the configuration is None, the function selects typical dimensions.
+    """
+    if config is not None:
+        check_hasattr(config, "vocab_size")
+    kwargs = dict(
+        batch_size=2,
+        sequence_length=30,
+        dummy_max_token_id=31999 if config is None else (config.vocab_size - 1),
+    )
+    return kwargs, get_inputs
diff --git a/onnx_diagnostic/tasks/text_classification.py b/onnx_diagnostic/tasks/text_classification.py
@@ -25,7 +25,7 @@ def get_inputs(
     **kwargs,  # unused
 ):
     """
-    Generates inputs for task ``fill-mask``.
+    Generates inputs for task ``text-classification``.
     Example:
 
     ::
diff --git a/onnx_diagnostic/torch_models/hghub/hub_data.py b/onnx_diagnostic/torch_models/hghub/hub_data.py
@@ -15,6 +15,7 @@
     BeitForImageClassification,image-classification
     BertForMaskedLM,fill-mask
     BertForSequenceClassification,text-classification
+    BertModel,sentence-similarity
     BigBirdModel,feature-extraction
     BlenderbotModel,feature-extraction
     BloomModel,feature-extraction
@@ -146,6 +147,7 @@
     "no-pipeline-tag",
     "object-detection",
     "reinforcement-learning",
+    "sentence-similarity",
     "text-classification",
     "text-generation",
     "text-to-audio",
diff --git a/onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py b/onnx_diagnostic/torch_models/hghub/hub_data_cached_configs.py
@@ -3502,3 +3502,31 @@ def _ccached_intel_bert_base_uncased_mrpc():
             "vocab_size": 30522,
         }
     )
+
+
+def _ccached_sentence_transformers_all_MiniLM_L6_v1():
+    "sentence-transformers/all-MiniLM-L6-v1"
+    return transformers.BertConfig(
+        **{
+            "_name_or_path": "nreimers/MiniLM-L6-H384-uncased",
+            "architectures": ["BertModel"],
+            "attention_probs_dropout_prob": 0.1,
+            "gradient_checkpointing": false,
+            "hidden_act": "gelu",
+            "hidden_dropout_prob": 0.1,
+            "hidden_size": 384,
+            "initializer_range": 0.02,
+            "intermediate_size": 1536,
+            "layer_norm_eps": 1e-12,
+            "max_position_embeddings": 512,
+            "model_type": "bert",
+            "num_attention_heads": 12,
+            "num_hidden_layers": 6,
+            "pad_token_id": 0,
+            "position_embedding_type": "absolute",
+            "transformers_version": "4.8.2",
+            "type_vocab_size": 2,
+            "use_cache": true,
+            "vocab_size": 30522,
+        }
+    )
