diff options
author | Lester Solbakken <lesters@oath.com> | 2023-02-18 12:55:41 +0100 |
---|---|---|
committer | Lester Solbakken <lesters@oath.com> | 2023-02-18 12:55:41 +0100 |
commit | cbd165fe6af4d9d3f40316e78b800621700db943 (patch) | |
tree | 50797824edc53d9fa123a2f88a597cf2cdd4c2de | |
parent | 21e0c4913dd0bc88dfec3016d4552f57fc0e7c4b (diff) |
Add model ids for generator models
-rw-r--r-- | config-model/src/main/java/com/yahoo/vespa/model/container/xml/ModelIdResolver.java | 25 |
1 files changed, 21 insertions, 4 deletions
diff --git a/config-model/src/main/java/com/yahoo/vespa/model/container/xml/ModelIdResolver.java b/config-model/src/main/java/com/yahoo/vespa/model/container/xml/ModelIdResolver.java index 0abd7212017..9d282f0850b 100644 --- a/config-model/src/main/java/com/yahoo/vespa/model/container/xml/ModelIdResolver.java +++ b/config-model/src/main/java/com/yahoo/vespa/model/container/xml/ModelIdResolver.java @@ -3,6 +3,9 @@ package com.yahoo.vespa.model.container.xml; import com.yahoo.text.XML; import org.w3c.dom.Element; + +import java.util.Collections; +import java.util.HashMap; import java.util.Map; import java.util.stream.Collectors; @@ -14,10 +17,24 @@ import java.util.stream.Collectors; */ public class ModelIdResolver { - private static final Map<String, String> providedModels = - Map.of("minilm-l6-v2", "https://data.vespa.oath.cloud/onnx_models/sentence_all_MiniLM_L6_v2.onnx", - "mpnet-base-v2", "https://data.vespa.oath.cloud/onnx_models/sentence-all-mpnet-base-v2.onnx", - "bert-base-uncased", "https://data.vespa.oath.cloud/onnx_models/bert-base-uncased-vocab.txt"); + private static Map<String, String> setupProvidedModels() { + Map<String, String> models = new HashMap<>(); + models.put("minilm-l6-v2", "https://data.vespa.oath.cloud/onnx_models/sentence_all_MiniLM_L6_v2.onnx"); + models.put("mpnet-base-v2", "https://data.vespa.oath.cloud/onnx_models/sentence-all-mpnet-base-v2.onnx"); + models.put("bert-base-uncased", "https://data.vespa.oath.cloud/onnx_models/bert-base-uncased-vocab.txt"); + models.put("flan-t5-vocab", "https://data.vespa.oath.cloud/onnx_models/flan-t5-spiece.model"); + models.put("flan-t5-small-encoder", "https://data.vespa.oath.cloud/onnx_models/flan-t5-small-encoder-model.onnx"); + models.put("flan-t5-small-decoder", "https://data.vespa.oath.cloud/onnx_models/flan-t5-small-decoder-model.onnx"); + models.put("flan-t5-base-encoder", "https://data.vespa.oath.cloud/onnx_models/flan-t5-base-encoder-model.onnx"); + models.put("flan-t5-base-decoder", "https://data.vespa.oath.cloud/onnx_models/flan-t5-base-decoder-model.onnx"); + models.put("flan-t5-large-encoder", "https://data.vespa.oath.cloud/onnx_models/flan-t5-large-encoder-model.onnx"); + models.put("flan-t5-large-decoder", "https://data.vespa.oath.cloud/onnx_models/flan-t5-large-decoder-model.onnx"); + models.put("flan-t5-xl-encoder", "https://data.vespa.oath.cloud/onnx_models/flan-t5-xl-encoder-model.onnx"); + models.put("flan-t5-xl-decoder", "https://data.vespa.oath.cloud/onnx_models/flan-t5-xl-decoder-model.onnx"); + return Collections.unmodifiableMap(models); + } + + private static final Map<String, String> providedModels = setupProvidedModels(); /** * Finds any config values of type 'model' below the given config element and |