extend sentence transformers onnx test

fxmarty · fxmarty · commit 87dc5631d275 · 2024-02-28T11:02:05.000+01:00
diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py
@@ -332,7 +332,10 @@
 
 PYTORCH_SENTENCE_TRANSFORMERS_MODEL = {
     "clip": "sentence-transformers/clip-ViT-B-32",
-    "transformer": "sentence-transformers/all-MiniLM-L6-v2",
+    "transformer": {
+        "sentence-transformers/all-MiniLM-L6-v2": ["feature-extraction", "sentence-similarity"],
+        "fxmarty/tiny-dummy-mistral-sentence-transformer": ["feature-extraction", "sentence-similarity"],
+    },
 }
 
 

Original file line number	Diff line number	Diff line change
`@@ -332,7 +332,10 @@`
`332`	`332`
`333`	`333`	`PYTORCH_SENTENCE_TRANSFORMERS_MODEL = {`
`334`	`334`	`"clip": "sentence-transformers/clip-ViT-B-32",`
`335`		`- "transformer": "sentence-transformers/all-MiniLM-L6-v2",`
	`335`	`+ "transformer": {`
	`336`	`+ "sentence-transformers/all-MiniLM-L6-v2": ["feature-extraction", "sentence-similarity"],`
	`337`	`+ "fxmarty/tiny-dummy-mistral-sentence-transformer": ["feature-extraction", "sentence-similarity"],`
	`338`	`+ },`
`336`	`339`	`}`
`337`	`340`
`338`	`341`