aboutsummaryrefslogtreecommitdiffstats
path: root/config-model/src/main/java/com/yahoo/schema/processing/OnnxModelTypeResolver.java
diff options
context:
space:
mode:
authorJon Bratseth <bratseth@gmail.com>2022-05-19 12:03:06 +0200
committerJon Bratseth <bratseth@gmail.com>2022-05-19 12:03:06 +0200
commit5c24dc5c9642a8d9ed70aee4c950fd0678a1ebec (patch)
treebd9b74bf00c832456f0b83c1b2cd7010be387d68 /config-model/src/main/java/com/yahoo/schema/processing/OnnxModelTypeResolver.java
parentf17c4fe7de4c55f5c4ee61897eab8c2f588d8405 (diff)
Rename the 'searchdefinition' package to 'schema'
Diffstat (limited to 'config-model/src/main/java/com/yahoo/schema/processing/OnnxModelTypeResolver.java')
-rw-r--r--config-model/src/main/java/com/yahoo/schema/processing/OnnxModelTypeResolver.java40
1 files changed, 40 insertions, 0 deletions
diff --git a/config-model/src/main/java/com/yahoo/schema/processing/OnnxModelTypeResolver.java b/config-model/src/main/java/com/yahoo/schema/processing/OnnxModelTypeResolver.java
new file mode 100644
index 00000000000..32229ea635b
--- /dev/null
+++ b/config-model/src/main/java/com/yahoo/schema/processing/OnnxModelTypeResolver.java
@@ -0,0 +1,40 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+package com.yahoo.schema.processing;
+
+import com.yahoo.config.application.api.DeployLogger;
+import com.yahoo.schema.OnnxModel;
+import com.yahoo.schema.RankProfile;
+import com.yahoo.schema.RankProfileRegistry;
+import com.yahoo.schema.Schema;
+import com.yahoo.vespa.model.container.search.QueryProfiles;
+import com.yahoo.vespa.model.ml.OnnxModelInfo;
+
+/**
+ * Processes every "onnx-model" element in the schema. Associates model type
+ * information by retrieving from either the ONNX model file directly or from
+ * preprocessed information in ZK. Adds missing input and output mappings
+ * (assigning default names).
+ *
+ * Must be processed before RankingExpressingTypeResolver.
+ *
+ * @author lesters
+ */
+public class OnnxModelTypeResolver extends Processor {
+
+ public OnnxModelTypeResolver(Schema schema, DeployLogger deployLogger, RankProfileRegistry rankProfileRegistry, QueryProfiles queryProfiles) {
+ super(schema, deployLogger, rankProfileRegistry, queryProfiles);
+ }
+
+ @Override
+ public void process(boolean validate, boolean documentsOnly) {
+ if (documentsOnly) return;
+ for (OnnxModel onnxModel : schema.declaredOnnxModels().values())
+ onnxModel.setModelInfo(OnnxModelInfo.load(onnxModel.getFileName(), schema.applicationPackage()));
+ for (RankProfile profile : rankProfileRegistry.rankProfilesOf(schema)) {
+ for (OnnxModel onnxModel : profile.declaredOnnxModels().values())
+ onnxModel.setModelInfo(OnnxModelInfo.load(onnxModel.getFileName(), schema.applicationPackage()));
+ }
+ }
+
+}