Add config generation for models evaluated via ONNXRT

1 files changed, 39 insertions, 0 deletions

diff --git a/config-model/src/main/java/com/yahoo/searchdefinition/OnnxModels.java b/config-model/src/main/java/com/yahoo/searchdefinition/OnnxModels.java
new file mode 100644
index 00000000000..87663ac79a3
--- /dev/null
+++ b/config-model/src/main/java/com/yahoo/searchdefinition/OnnxModels.java
@@ -0,0 +1,39 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.searchdefinition;
+
+import com.yahoo.vespa.model.AbstractService;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * ONNX models tied to a search definition.
+ *
+ * @author lesters
+ */
+public class OnnxModels {
+
+    private final Map<String, OnnxModel> models = new HashMap<>();
+
+    public void add(OnnxModel model) {
+        model.validate();
+        String name = model.getName();
+        models.put(name, model);
+    }
+
+    public OnnxModel get(String name) {
+        return models.get(name);
+    }
+
+    public Map<String, OnnxModel> asMap() {
+        return Collections.unmodifiableMap(models);
+    }
+
+    /** Initiate sending of these models to some services over file distribution */
+    public void sendTo(Collection<? extends AbstractService> services) {
+        models.values().forEach(model -> model.sendTo(services));
+    }
+
+}