Read and resolve constants in model evaluation

11 files changed, 185 insertions, 29 deletions

diff --git a/model-evaluation/pom.xml b/model-evaluation/pom.xml
index edb22c1b529..0b6a5d08155 100644
--- a/model-evaluation/pom.xml
+++ b/model-evaluation/pom.xml
@@ -40,6 +40,12 @@
     </dependency>
     <dependency>
       <groupId>com.yahoo.vespa</groupId>
+      <artifactId>searchcore</artifactId>
+      <version>${project.version}</version>
+      <scope>provided</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.yahoo.vespa</groupId>
       <artifactId>config</artifactId>
       <version>${project.version}</version>
       <scope>provided</scope>
diff --git a/model-evaluation/src/main/java/ai/vespa/models/evaluation/Constant.java b/model-evaluation/src/main/java/ai/vespa/models/evaluation/Constant.java
new file mode 100644
index 00000000000..e664693ab38
--- /dev/null
+++ b/model-evaluation/src/main/java/ai/vespa/models/evaluation/Constant.java
@@ -0,0 +1,27 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package ai.vespa.models.evaluation;
+
+import com.yahoo.tensor.Tensor;
+
+/**
+ * A named constant loaded from a file.
+ *
+ * This is immutable.
+ *
+ * @author bratseth
+ */
+class Constant {
+
+    private final String name;
+    private final Tensor value;
+
+    Constant(String name, Tensor value) {
+        this.name = name;
+        this.value = value;
+    }
+
+    public String name() { return name; }
+
+    public Tensor value() { return value; }
+
+}
diff --git a/model-evaluation/src/main/java/ai/vespa/models/evaluation/FunctionEvaluator.java b/model-evaluation/src/main/java/ai/vespa/models/evaluation/FunctionEvaluator.java
index 520986ffb77..e08b9f77d15 100644
--- a/model-evaluation/src/main/java/ai/vespa/models/evaluation/FunctionEvaluator.java
+++ b/model-evaluation/src/main/java/ai/vespa/models/evaluation/FunctionEvaluator.java
@@ -56,4 +56,6 @@ public class FunctionEvaluator {
         return function.getBody().evaluate(context).asTensor();
     }
 
+    LazyArrayContext context() { return context; }
+
 }
diff --git a/model-evaluation/src/main/java/ai/vespa/models/evaluation/LazyArrayContext.java b/model-evaluation/src/main/java/ai/vespa/models/evaluation/LazyArrayContext.java
index 2dcfd204077..beaa36b898f 100644
--- a/model-evaluation/src/main/java/ai/vespa/models/evaluation/LazyArrayContext.java
+++ b/model-evaluation/src/main/java/ai/vespa/models/evaluation/LazyArrayContext.java
@@ -8,6 +8,7 @@ import com.yahoo.searchlib.rankingexpression.Reference;
 import com.yahoo.searchlib.rankingexpression.evaluation.Context;
 import com.yahoo.searchlib.rankingexpression.evaluation.ContextIndex;
 import com.yahoo.searchlib.rankingexpression.evaluation.DoubleValue;
+import com.yahoo.searchlib.rankingexpression.evaluation.TensorValue;
 import com.yahoo.searchlib.rankingexpression.evaluation.Value;
 import com.yahoo.searchlib.rankingexpression.rule.CompositeNode;
 import com.yahoo.searchlib.rankingexpression.rule.ExpressionNode;
@@ -16,6 +17,7 @@ import com.yahoo.tensor.TensorType;
 
 import java.util.Arrays;
 import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -37,8 +39,11 @@ final class LazyArrayContext extends Context implements ContextIndex {
      *
      * @param expression the expression to create a context for
      */
-    LazyArrayContext(RankingExpression expression, Map<FunctionReference, ExpressionFunction> functions, Model model) {
-        this.indexedBindings = new IndexedBindings(expression, functions, this, model);
+    LazyArrayContext(RankingExpression expression,
+                     Map<FunctionReference, ExpressionFunction> functions,
+                     List<Constant> constants,
+                     Model model) {
+        this.indexedBindings = new IndexedBindings(expression, functions, constants, this, model);
     }
 
     /**
@@ -139,8 +144,10 @@ final class LazyArrayContext extends Context implements ContextIndex {
          */
         IndexedBindings(RankingExpression expression,
                         Map<FunctionReference, ExpressionFunction> functions,
+                        List<Constant> constants,
                         LazyArrayContext owner,
                         Model model) {
+            // 1. Determine and prepare bind targets
             Set<String> bindTargets = new LinkedHashSet<>();
             extractBindTargets(expression.getRoot(), functions, bindTargets);
 
@@ -150,9 +157,18 @@ final class LazyArrayContext extends Context implements ContextIndex {
             int i = 0;
             ImmutableMap.Builder<String, Integer> nameToIndexBuilder = new ImmutableMap.Builder<>();
             for (String variable : bindTargets)
-                nameToIndexBuilder.put(variable,i++);
+                nameToIndexBuilder.put(variable, i++);
             nameToIndex = nameToIndexBuilder.build();
 
+
+            // 2. Bind the bind targets
+            for (Constant constant : constants) {
+                String constantReference = "constant(" + constant.name() + ")";
+                Integer index = nameToIndex.get(constantReference);
+                if (index != null)
+                    values[index] = new TensorValue(constant.value());
+            }
+
             for (Map.Entry<FunctionReference, ExpressionFunction> function : functions.entrySet()) {
                 Integer index = nameToIndex.get(function.getKey().serialForm());
                 if (index != null) // Referenced in this, so bind it
@@ -170,7 +186,7 @@ final class LazyArrayContext extends Context implements ContextIndex {
                 extractBindTargets(functions.get(reference).getBody().getRoot(), functions, bindTargets);
             }
             else if (isConstant(node)) {
-                // Ignore
+                bindTargets.add(node.toString());
             }
             else if (node instanceof ReferenceNode) {
                 bindTargets.add(node.toString());
@@ -193,7 +209,7 @@ final class LazyArrayContext extends Context implements ContextIndex {
             if ( ! (node instanceof ReferenceNode)) return false;
 
             ReferenceNode reference = (ReferenceNode)node;
-            return reference.getName().equals("value") && reference.getArguments().size() == 1;
+            return reference.getName().equals("constant") && reference.getArguments().size() == 1;
         }
 
         Value get(int index) { return values[index]; }
diff --git a/model-evaluation/src/main/java/ai/vespa/models/evaluation/Model.java b/model-evaluation/src/main/java/ai/vespa/models/evaluation/Model.java
index 95eb923786d..3fb43d73187 100644
--- a/model-evaluation/src/main/java/ai/vespa/models/evaluation/Model.java
+++ b/model-evaluation/src/main/java/ai/vespa/models/evaluation/Model.java
@@ -36,11 +36,15 @@ public class Model {
 
     private final ExpressionOptimizer expressionOptimizer = new ExpressionOptimizer();
 
+    /** Programmatically create a model containing functions without constant of function references only */
     public Model(String name, Collection<ExpressionFunction> functions) {
-        this(name, functions, Collections.emptyMap());
+        this(name, functions, Collections.emptyMap(), Collections.emptyList());
     }
 
-    Model(String name, Collection<ExpressionFunction> functions, Map<FunctionReference, ExpressionFunction> referencedFunctions) {
+    Model(String name,
+          Collection<ExpressionFunction> functions,
+          Map<FunctionReference, ExpressionFunction> referencedFunctions,
+          List<Constant> constants) {
         // TODO: Optimize functions
         this.name = name;
         this.functions = ImmutableList.copyOf(functions);
@@ -48,7 +52,8 @@ public class Model {
         ImmutableMap.Builder<String, LazyArrayContext> contextBuilder = new ImmutableMap.Builder<>();
         for (ExpressionFunction function : functions) {
             try {
-                contextBuilder.put(function.getName(), new LazyArrayContext(function.getBody(), referencedFunctions, this));
+                contextBuilder.put(function.getName(),
+                                   new LazyArrayContext(function.getBody(), referencedFunctions, constants, this));
             }
             catch (RuntimeException e) {
                 throw new IllegalArgumentException("Could not prepare an evaluation context for " + function, e);
diff --git a/model-evaluation/src/main/java/ai/vespa/models/evaluation/ModelsEvaluator.java b/model-evaluation/src/main/java/ai/vespa/models/evaluation/ModelsEvaluator.java
index dacf20b7ef2..48c71b5a04a 100644
--- a/model-evaluation/src/main/java/ai/vespa/models/evaluation/ModelsEvaluator.java
+++ b/model-evaluation/src/main/java/ai/vespa/models/evaluation/ModelsEvaluator.java
@@ -5,6 +5,7 @@ import com.google.common.annotations.Beta;
 import com.google.common.collect.ImmutableMap;
 import com.yahoo.component.AbstractComponent;
 import com.yahoo.vespa.config.search.RankProfilesConfig;
+import com.yahoo.vespa.config.search.core.RankingConstantsConfig;
 
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -21,8 +22,8 @@ public class ModelsEvaluator extends AbstractComponent {
 
     private final ImmutableMap<String, Model> models;
 
-    public ModelsEvaluator(RankProfilesConfig config) {
-        models = ImmutableMap.copyOf(new RankProfilesConfigImporter().importFrom(config));
+    public ModelsEvaluator(RankProfilesConfig config, RankingConstantsConfig constantsConfig) {
+        models = ImmutableMap.copyOf(new RankProfilesConfigImporter().importFrom(config, constantsConfig));
     }
 
     /** Returns the models of this as an immutable map */
diff --git a/model-evaluation/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java b/model-evaluation/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java
index bfd6342218a..b9e7a27c013 100644
--- a/model-evaluation/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java
+++ b/model-evaluation/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java
@@ -1,33 +1,57 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.models.evaluation;
 
+import com.yahoo.filedistribution.fileacquirer.FileAcquirer;
+import com.yahoo.io.GrowableByteBuffer;
+import com.yahoo.io.IOUtils;
 import com.yahoo.searchlib.rankingexpression.ExpressionFunction;
 import com.yahoo.searchlib.rankingexpression.RankingExpression;
 import com.yahoo.searchlib.rankingexpression.parser.ParseException;
+import com.yahoo.searchlib.rankingexpression.rule.CompositeNode;
+import com.yahoo.searchlib.rankingexpression.rule.ExpressionNode;
+import com.yahoo.searchlib.rankingexpression.rule.ReferenceNode;
+import com.yahoo.tensor.Tensor;
+import com.yahoo.tensor.TensorType;
+import com.yahoo.tensor.serialization.TypedBinaryFormat;
 import com.yahoo.vespa.config.search.RankProfilesConfig;
+import com.yahoo.vespa.config.search.core.RankingConstantsConfig;
 
+import java.io.File;
+import java.io.IOException;
+import java.io.UncheckedIOException;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
 /**
- * Converts RankProfilesConfig instances to RankingExpressions for evaluation
+ * Converts RankProfilesConfig instances to RankingExpressions for evaluation.
+ * This class can be used by a single thread only.
  *
  * @author bratseth
  */
 class RankProfilesConfigImporter {
 
     /**
+     * Constants already imported in this while reading some expression.
+     * This is to avoid re-reading constants referenced
+     * multiple places, as that is potentially costly.
+     */
+    private Map<String, Constant> globalImportedConstants = new HashMap<>();
+
+    /**
      * Returns a map of the models contained in this config, indexed on name.
      * The map is modifiable and owned by the caller.
      */
-    Map<String, Model> importFrom(RankProfilesConfig config) {
+    Map<String, Model> importFrom(RankProfilesConfig config, RankingConstantsConfig constantsConfig) {
+        globalImportedConstants.clear();
         try {
             Map<String, Model> models = new HashMap<>();
             for (RankProfilesConfig.Rankprofile profile : config.rankprofile()) {
-                Model model = importProfile(profile);
+                Model model = importProfile(profile, constantsConfig);
                 models.put(model.name(), model);
             }
             return models;
@@ -37,11 +61,14 @@ class RankProfilesConfigImporter {
         }
     }
 
-    private Model importProfile(RankProfilesConfig.Rankprofile profile) throws ParseException {
+    private Model importProfile(RankProfilesConfig.Rankprofile profile, RankingConstantsConfig constantsConfig) throws ParseException {
         List<ExpressionFunction> functions = new ArrayList<>();
         Map<FunctionReference, ExpressionFunction> referencedFunctions = new HashMap<>();
         ExpressionFunction firstPhase = null;
         ExpressionFunction secondPhase = null;
+
+        List<Constant> constants = readConstants(constantsConfig);
+
         for (RankProfilesConfig.Rankprofile.Fef.Property property : profile.fef().property()) {
             Optional<FunctionReference> reference = FunctionReference.fromSerial(property.name());
             if ( reference.isPresent()) {
@@ -69,7 +96,7 @@ class RankProfilesConfigImporter {
             functions.add(secondPhase);
 
         try {
-            return new Model(profile.name(), functions, referencedFunctions);
+            return new Model(profile.name(), functions, referencedFunctions, constants);
         }
         catch (RuntimeException e) {
             throw new IllegalArgumentException("Could not load model '" + profile.name() + "'", e);
@@ -83,4 +110,30 @@ class RankProfilesConfigImporter {
         return null;
     }
 
+    private List<Constant> readConstants(RankingConstantsConfig constantsConfig) {
+        List<Constant> constants = new ArrayList<>();
+        for (RankingConstantsConfig.Constant constantConfig : constantsConfig.constant()) {
+            constants.add(new Constant(constantConfig.name(),
+                                       readTensorFromFile(TensorType.fromSpec(constantConfig.type()),
+                                                          constantConfig.fileref().value())));
+        }
+        return constants;
+    }
+
+    private Tensor readTensorFromFile(TensorType type, String fileName) {
+        try {
+            if (fileName.endsWith(".tbf"))
+                return TypedBinaryFormat.decode(Optional.of(type),
+                                                GrowableByteBuffer.wrap(IOUtils.readFileBytes(new File(fileName))));
+            // TODO: Support json and json.lz4
+
+            if (fileName.isEmpty()) // this is the case in unit tests
+                return Tensor.from(type, "{}");
+            throw new IllegalArgumentException("Unknown tensor file format (determined by file ending): " + fileName);
+        }
+        catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
 }
diff --git a/model-evaluation/src/test/java/ai/vespa/models/evaluation/ModelsEvaluatorTest.java b/model-evaluation/src/test/java/ai/vespa/models/evaluation/ModelsEvaluatorTest.java
index 23928c5b7e7..d94e5b2af1b 100644
--- a/model-evaluation/src/test/java/ai/vespa/models/evaluation/ModelsEvaluatorTest.java
+++ b/model-evaluation/src/test/java/ai/vespa/models/evaluation/ModelsEvaluatorTest.java
@@ -3,8 +3,10 @@ package ai.vespa.models.evaluation;
 
 import com.yahoo.config.subscription.ConfigGetter;
 import com.yahoo.config.subscription.FileSource;
+import com.yahoo.path.Path;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.vespa.config.search.RankProfilesConfig;
+import com.yahoo.vespa.config.search.core.RankingConstantsConfig;
 import org.junit.Test;
 
 import java.io.File;
@@ -18,15 +20,9 @@ public class ModelsEvaluatorTest {
 
     private static final double delta = 0.00000000001;
 
-    private ModelsEvaluator createModels() {
-        String configPath = "src/test/resources/config/rankexpression/rank-profiles.cfg";
-        RankProfilesConfig config = new ConfigGetter<>(new FileSource(new File(configPath)), RankProfilesConfig.class).getConfig("");
-        return new ModelsEvaluator(config);
-    }
-
     @Test
     public void testTensorEvaluation() {
-        ModelsEvaluator models = createModels();
+        ModelsEvaluator models = createModels("src/test/resources/config/rankexpression/");
         FunctionEvaluator function = models.evaluatorOf("macros", "fourtimessum");
         function.bind("var1", Tensor.from("{{x:0}:3,{x:1}:5}"));
         function.bind("var2", Tensor.from("{{x:0}:7,{x:1}:11}"));
@@ -35,7 +31,7 @@ public class ModelsEvaluatorTest {
 
     @Test
     public void testEvaluationDependingOnMacroTakingArguments() {
-        ModelsEvaluator models = createModels();
+        ModelsEvaluator models = createModels("src/test/resources/config/rankexpression/");
         FunctionEvaluator function = models.evaluatorOf("macros", "secondphase");
         function.bind("match", 3);
         function.bind("rankBoost", 5);
@@ -47,4 +43,13 @@ public class ModelsEvaluatorTest {
     // TODO: Test that rebinding doesn't work
     // TODO: Test with nested macros
 
+    private ModelsEvaluator createModels(String path) {
+        Path configDir = Path.fromString(path);
+        RankProfilesConfig config = new ConfigGetter<>(new FileSource(configDir.append("rank-profiles.cfg").toFile()),
+                                                       RankProfilesConfig.class).getConfig("");
+        RankingConstantsConfig constantsConfig = new ConfigGetter<>(new FileSource(configDir.append("ranking-constants.cfg").toFile()),
+                                                                    RankingConstantsConfig.class).getConfig("");
+        return new ModelsEvaluator(config, constantsConfig);
+    }
+
 }
diff --git a/model-evaluation/src/test/java/ai/vespa/models/evaluation/RankProfilesImporterTest.java b/model-evaluation/src/test/java/ai/vespa/models/evaluation/RankProfilesImporterTest.java
index e0f5674e016..84e01e58280 100644
--- a/model-evaluation/src/test/java/ai/vespa/models/evaluation/RankProfilesImporterTest.java
+++ b/model-evaluation/src/test/java/ai/vespa/models/evaluation/RankProfilesImporterTest.java
@@ -3,8 +3,10 @@ package ai.vespa.models.evaluation;
 
 import com.yahoo.config.subscription.ConfigGetter;
 import com.yahoo.config.subscription.FileSource;
+import com.yahoo.path.Path;
 import com.yahoo.searchlib.rankingexpression.ExpressionFunction;
 import com.yahoo.vespa.config.search.RankProfilesConfig;
+import com.yahoo.vespa.config.search.core.RankingConstantsConfig;
 import org.junit.Test;
 
 import java.io.File;
@@ -22,9 +24,8 @@ public class RankProfilesImporterTest {
 
     @Test
     public void testImportingModels() {
-        String configPath = "src/test/resources/config/models/rank-profiles.cfg";
-        RankProfilesConfig config = new ConfigGetter<>(new FileSource(new File(configPath)), RankProfilesConfig.class).getConfig("");
-        Map<String, Model> models = new RankProfilesConfigImporter().importFrom(config);
+        Map<String, Model> models = createModels("src/test/resources/config/models/");
+
         assertEquals(4, models.size());
 
         Model xgboost = models.get("xgboost_2_2");
@@ -36,6 +37,8 @@ public class RankProfilesImporterTest {
         assertFunction("default.add",
                        "join(reduce(join(rename(Placeholder, (d0, d1), (d0, d2)), constant(mnist_softmax_Variable), f(a,b)(a * b)), sum, d2), constant(mnist_softmax_Variable_1), f(a,b)(a + b))",
                        onnxMnistSoftmax);
+        assertEquals("tensor(d1[10],d2[784])",
+                     onnxMnistSoftmax.evaluatorOf("default.add").context().get("constant(mnist_softmax_Variable)").type().toString());
 
         Model tfMnistSoftmax = models.get("mnist_softmax_saved");
         assertFunction("serving_default.y",
@@ -50,9 +53,8 @@ public class RankProfilesImporterTest {
 
     @Test
     public void testImportingRankExpressions() {
-        String configPath = "src/test/resources/config/rankexpression/rank-profiles.cfg";
-        RankProfilesConfig config = new ConfigGetter<>(new FileSource(new File(configPath)), RankProfilesConfig.class).getConfig("");
-        Map<String, Model> models = new RankProfilesConfigImporter().importFrom(config);
+        Map<String, Model> models = createModels("src/test/resources/config/rankexpression/");
+
         assertEquals(18, models.size());
 
         Model macros = models.get("macros");
@@ -85,4 +87,13 @@ public class RankProfilesImporterTest {
         assertEquals(expression, function.getBody().getRoot().toString());
     }
 
+    private Map<String, Model> createModels(String path) {
+        Path configDir = Path.fromString(path);
+        RankProfilesConfig config = new ConfigGetter<>(new FileSource(configDir.append("rank-profiles.cfg").toFile()),
+                                                       RankProfilesConfig.class).getConfig("");
+        RankingConstantsConfig constantsConfig = new ConfigGetter<>(new FileSource(configDir.append("ranking-constants.cfg").toFile()),
+                                                                    RankingConstantsConfig.class).getConfig("");
+        return new RankProfilesConfigImporter().importFrom(config, constantsConfig);
+    }
+
 }
diff --git a/model-evaluation/src/test/resources/config/models/ranking-constants.cfg b/model-evaluation/src/test/resources/config/models/ranking-constants.cfg
new file mode 100644
index 00000000000..2b7495ace5e
--- /dev/null
+++ b/model-evaluation/src/test/resources/config/models/ranking-constants.cfg
@@ -0,0 +1,30 @@
+constant[0].name "mnist_saved_dnn_hidden1_weights_read"
+constant[0].fileref ""
+constant[0].type "tensor(d3[300],d4[784])"
+constant[1].name "mnist_saved_dnn_hidden2_weights_read"
+constant[1].fileref ""
+constant[1].type "tensor(d2[100],d3[300])"
+constant[2].name "mnist_softmax_saved_layer_Variable_1_read"
+constant[2].fileref ""
+constant[2].type "tensor(d1[10])"
+constant[3].name "mnist_saved_dnn_hidden1_bias_read"
+constant[3].fileref ""
+constant[3].type "tensor(d3[300])"
+constant[4].name "mnist_saved_dnn_hidden2_bias_read"
+constant[4].fileref ""
+constant[4].type "tensor(d2[100])"
+constant[5].name "mnist_softmax_Variable"
+constant[5].fileref ""
+constant[5].type "tensor(d1[10],d2[784])"
+constant[6].name "mnist_saved_dnn_outputs_weights_read"
+constant[6].fileref ""
+constant[6].type "tensor(d1[10],d2[100])"
+constant[7].name "mnist_softmax_saved_layer_Variable_read"
+constant[7].fileref ""
+constant[7].type "tensor(d1[10],d2[784])"
+constant[8].name "mnist_softmax_Variable_1"
+constant[8].fileref ""
+constant[8].type "tensor(d1[10])"
+constant[9].name "mnist_saved_dnn_outputs_bias_read"
+constant[9].fileref ""
+constant[9].type "tensor(d1[10])"
\ No newline at end of file
diff --git a/model-evaluation/src/test/resources/config/rankexpression/ranking-constants.cfg b/model-evaluation/src/test/resources/config/rankexpression/ranking-constants.cfg
new file mode 100644
index 00000000000..e69de29bb2d
--- /dev/null
+++ b/model-evaluation/src/test/resources/config/rankexpression/ranking-constants.cfg