1 files changed, 26 insertions, 30 deletions

diff --git a/linguistics/src/main/java/com/yahoo/language/significance/impl/DefaultSignificanceModelRegistry.java b/linguistics/src/main/java/com/yahoo/language/significance/impl/DefaultSignificanceModelRegistry.java
index 1be1d3f13b5..72874c15d9e 100644
--- a/linguistics/src/main/java/com/yahoo/language/significance/impl/DefaultSignificanceModelRegistry.java
+++ b/linguistics/src/main/java/com/yahoo/language/significance/impl/DefaultSignificanceModelRegistry.java
@@ -1,20 +1,21 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.language.significance.impl;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.yahoo.component.annotation.Inject;
 import com.yahoo.language.Language;
 import com.yahoo.language.significance.SignificanceModel;
 import com.yahoo.language.significance.SignificanceModelRegistry;
 import com.yahoo.search.significance.config.SignificanceConfig;
 
+import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.nio.file.Path;
 import java.util.EnumMap;
-import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.function.Supplier;
 
-import static com.yahoo.yolean.Exceptions.uncheck;
 /**
  * Default implementation of {@link SignificanceModelRegistry}.
  * This implementation loads models lazily and caches them.
@@ -24,24 +25,35 @@ import static com.yahoo.yolean.Exceptions.uncheck;
 public class DefaultSignificanceModelRegistry implements SignificanceModelRegistry {
 
     private final Map<Language, SignificanceModel> models;
+
     @Inject
-    public DefaultSignificanceModelRegistry(SignificanceConfig cfg) { this(new Builder(cfg)); }
-    private DefaultSignificanceModelRegistry(Builder b) {
+    public DefaultSignificanceModelRegistry(SignificanceConfig cfg) {
         this.models = new EnumMap<>(Language.class);
-        b.models.forEach((language, path) -> {
-            models.put(language,
-                    uncheck(() -> new DefaultSignificanceModel(path)));
-        });
+        for (var model : cfg.model()) {
+           addModel(model.path());
+        }
     }
 
-    public DefaultSignificanceModelRegistry(HashMap<Language, Path> map) {
+    public DefaultSignificanceModelRegistry(List<Path> models) {
         this.models = new EnumMap<>(Language.class);
-        map.forEach((language, path) -> {
-            models.put(language,
-                    uncheck(() -> new DefaultSignificanceModel(path)));
-        });
+        for (var path : models) {
+            addModel(path);
+        }
     }
 
+    public void addModel(Path path) {
+        ObjectMapper objectMapper = new ObjectMapper();
+        try {
+            SignificanceModelFile file = objectMapper.readValue(path.toFile(), SignificanceModelFile.class);
+            for (var pair : file.languages().entrySet()) {
+                this.models.put(
+                        Language.fromLanguageTag(pair.getKey()),
+                        new DefaultSignificanceModel(pair.getValue(), file.id()));
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException("Failed to load model from " + path, e);
+        }
+    }
 
     @Override
     public Optional<SignificanceModel> getModel(Language language) {
@@ -51,20 +63,4 @@ public class DefaultSignificanceModelRegistry implements SignificanceModelRegist
         }
         return Optional.of(models.get(language));
     }
-
-
-    public static final class Builder {
-        private final Map<Language, Path> models = new EnumMap<>(Language.class);
-
-        public Builder() {}
-        public Builder(SignificanceConfig cfg) {
-            for (var model : cfg.model()) {
-                addModel(Language.fromLanguageTag(model.language()), model.path());
-            }
-        }
-
-        public Builder addModel(Language lang, Path path) { models.put(lang, path); return this; }
-        public DefaultSignificanceModelRegistry build() { return new DefaultSignificanceModelRegistry(this); }
-    }
-
 }