Merge pull request #33171 from vespa-engine/arnej/detect-no-token-typ…

…e-ids detect if model does not use token_type_ids
vespa-engine · Jan 25, 2025 · 1fc09f8 · 1fc09f8
2 parents effc2d1 + 2e88fc4
commit 1fc09f8
Showing 1 changed file with 16 additions and 4 deletions.
diff --git a/model-integration/src/main/java/ai/vespa/embedding/huggingface/HuggingFaceEmbedder.java b/model-integration/src/main/java/ai/vespa/embedding/huggingface/HuggingFaceEmbedder.java
@@ -48,7 +48,6 @@ public HuggingFaceEmbedder(OnnxRuntime onnx, Embedder.Runtime runtime, HuggingFa
         this.runtime = runtime;
         inputIdsName = config.transformerInputIds();
         attentionMaskName = config.transformerAttentionMask();
-        tokenTypeIdsName = config.transformerTokenTypeIds();
         outputName = config.transformerOutput();
         normalize = config.normalize();
         prependQuery = config.prependQuery();
@@ -75,15 +74,29 @@ public HuggingFaceEmbedder(OnnxRuntime onnx, Embedder.Runtime runtime, HuggingFa
         onnxOpts.setExecutionMode(config.transformerExecutionMode().toString());
         onnxOpts.setThreads(config.transformerInterOpThreads(), config.transformerIntraOpThreads());
         evaluator = onnx.evaluatorOf(config.transformerModel().toString(), onnxOpts);
+        tokenTypeIdsName = detectTokenTypeIds(config, evaluator);
         validateModel();
     }
 
+    private static String detectTokenTypeIds(HuggingFaceEmbedderConfig config, OnnxEvaluator evaluator) {
+        String configured = config.transformerTokenTypeIds();
+        Map<String, TensorType> inputs = evaluator.getInputInfo();
+        if (inputs.size() < 3) {
+            // newer models have only 2 inputs (they do not use token type IDs)
+            return "";
+        } else {
+            // could detect fallback from inputs here, currently set as default in .def file
+            return configured;
+        }
+    }
+
     private void validateModel() {
         Map<String, TensorType> inputs = evaluator.getInputInfo();
         validateName(inputs, inputIdsName, "input");
         validateName(inputs, attentionMaskName, "input");
-        if (!tokenTypeIdsName.isEmpty()) validateName(inputs, tokenTypeIdsName, "input");
-
+        if (!tokenTypeIdsName.isEmpty()) {
+            validateName(inputs, tokenTypeIdsName, "input");
+        }
         Map<String, TensorType> outputs = evaluator.getOutputInfo();
         validateName(outputs, outputName, "output");
     }
@@ -250,4 +263,3 @@ protected record HFEmbeddingResult(IndexedTensor output, Tensor attentionMask, S
     protected record HFEmbedderCacheKey(String embedderId, Object embeddedValue) { }
 
 }
-