[ML] add new inference_config field to trained model config

client/rest-high-level/src/main/java/org/elasticsearch/client/ml/inference/MlInferenceNamedXContentProvider.java

@@ -19,6 +19,9 @@
 package org.elasticsearch.client.ml.inference;
 
 import org.elasticsearch.client.ml.inference.preprocessing.CustomWordEmbedding;
+import org.elasticsearch.client.ml.inference.trainedmodel.ClassificationConfig;
+import org.elasticsearch.client.ml.inference.trainedmodel.InferenceConfig;
+import org.elasticsearch.client.ml.inference.trainedmodel.RegressionConfig;
 import org.elasticsearch.client.ml.inference.trainedmodel.TrainedModel;
 import org.elasticsearch.client.ml.inference.trainedmodel.ensemble.Ensemble;
 import org.elasticsearch.client.ml.inference.trainedmodel.ensemble.LogisticRegression;
@@ -61,6 +64,14 @@ public List<NamedXContentRegistry.Entry> getNamedXContentParsers() {
             new ParseField(LangIdentNeuralNetwork.NAME),
             LangIdentNeuralNetwork::fromXContent));
 
+        // Inference Config
+        namedXContent.add(new NamedXContentRegistry.Entry(InferenceConfig.class,
+            ClassificationConfig.NAME,
+            ClassificationConfig::fromXContent));
+        namedXContent.add(new NamedXContentRegistry.Entry(InferenceConfig.class,
+            RegressionConfig.NAME,
+            RegressionConfig::fromXContent));
+
         // Aggregating output
         namedXContent.add(new NamedXContentRegistry.Entry(OutputAggregator.class,
             new ParseField(WeightedMode.NAME),

client/rest-high-level/src/main/java/org/elasticsearch/client/ml/inference/NamedXContentObjectHelper.java

@@ -54,4 +54,14 @@ public static XContentBuilder writeNamedObjects(XContentBuilder builder,
         }
         return builder;
     }
+
+    public static XContentBuilder writeNamedObject(XContentBuilder builder,
+                                                   ToXContent.Params params,
+                                                   String namedObjectName,
+                                                   NamedXContentObject namedObject) throws IOException {
+        builder.startObject(namedObjectName);
+        builder.field(namedObject.getName(), namedObject, params);
+        builder.endObject();
+        return builder;
+    }
 }

client/rest-high-level/src/main/java/org/elasticsearch/client/ml/inference/TrainedModelConfig.java

@@ -20,6 +20,7 @@
 
 import org.elasticsearch.Version;
 import org.elasticsearch.client.common.TimeUtil;
+import org.elasticsearch.client.ml.inference.trainedmodel.InferenceConfig;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.unit.ByteSizeValue;
@@ -36,6 +37,8 @@
 import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.client.ml.inference.NamedXContentObjectHelper.writeNamedObject;
+
 public class TrainedModelConfig implements ToXContentObject {
 
     public static final String NAME = "trained_model_config";
@@ -54,6 +57,7 @@ public class TrainedModelConfig implements ToXContentObject {
     public static final ParseField ESTIMATED_OPERATIONS = new ParseField("estimated_operations");
     public static final ParseField LICENSE_LEVEL = new ParseField("license_level");
     public static final ParseField DEFAULT_FIELD_MAP = new ParseField("default_field_map");
+    public static final ParseField INFERENCE_CONFIG = new ParseField("inference_config");
 
     public static final ObjectParser<Builder, Void> PARSER = new ObjectParser<>(NAME,
             true,
@@ -78,6 +82,9 @@ public class TrainedModelConfig implements ToXContentObject {
         PARSER.declareLong(TrainedModelConfig.Builder::setEstimatedOperations, ESTIMATED_OPERATIONS);
         PARSER.declareString(TrainedModelConfig.Builder::setLicenseLevel, LICENSE_LEVEL);
         PARSER.declareObject(TrainedModelConfig.Builder::setDefaultFieldMap, (p, c) -> p.mapStrings(), DEFAULT_FIELD_MAP);
+        PARSER.declareNamedObject(TrainedModelConfig.Builder::setInferenceConfig,
+            (p, c, n) -> p.namedObject(InferenceConfig.class, n, null),
+            INFERENCE_CONFIG);
     }
 
     public static TrainedModelConfig fromXContent(XContentParser parser) throws IOException {
@@ -98,6 +105,7 @@ public static TrainedModelConfig fromXContent(XContentParser parser) throws IOEx
     private final Long estimatedOperations;
     private final String licenseLevel;
     private final Map<String, String> defaultFieldMap;
+    private final InferenceConfig inferenceConfig;
 
     TrainedModelConfig(String modelId,
                        String createdBy,
@@ -112,7 +120,8 @@ public static TrainedModelConfig fromXContent(XContentParser parser) throws IOEx
                        Long estimatedHeapMemory,
                        Long estimatedOperations,
                        String licenseLevel,
-                       Map<String, String> defaultFieldMap) {
+                       Map<String, String> defaultFieldMap,
+                       InferenceConfig inferenceConfig) {
         this.modelId = modelId;
         this.createdBy = createdBy;
         this.version = version;
@@ -127,6 +136,7 @@ public static TrainedModelConfig fromXContent(XContentParser parser) throws IOEx
         this.estimatedOperations = estimatedOperations;
         this.licenseLevel = licenseLevel;
         this.defaultFieldMap = defaultFieldMap == null ? null : Collections.unmodifiableMap(defaultFieldMap);
+        this.inferenceConfig = inferenceConfig;
     }
 
     public String getModelId() {
@@ -189,6 +199,10 @@ public Map<String, String> getDefaultFieldMap() {
         return defaultFieldMap;
     }
 
+    public InferenceConfig getInferenceConfig() {
+        return inferenceConfig;
+    }
+
     public static Builder builder() {
         return new Builder();
     }
@@ -238,6 +252,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         if (defaultFieldMap != null) {
             builder.field(DEFAULT_FIELD_MAP.getPreferredName(), defaultFieldMap);
         }
+        if (inferenceConfig != null) {
+            writeNamedObject(builder, params, INFERENCE_CONFIG.getPreferredName(), inferenceConfig);
+        }
         builder.endObject();
         return builder;
     }
@@ -265,6 +282,7 @@ public boolean equals(Object o) {
             Objects.equals(estimatedOperations, that.estimatedOperations) &&
             Objects.equals(licenseLevel, that.licenseLevel) &&
             Objects.equals(defaultFieldMap, that.defaultFieldMap) &&
+            Objects.equals(inferenceConfig, that.inferenceConfig) &&
             Objects.equals(metadata, that.metadata);
     }
 
@@ -283,6 +301,7 @@ public int hashCode() {
             metadata,
             licenseLevel,
             input,
+            inferenceConfig,
             defaultFieldMap);
     }
 
@@ -303,6 +322,7 @@ public static class Builder {
         private Long estimatedOperations;
         private String licenseLevel;
         private Map<String, String> defaultFieldMap;
+        private InferenceConfig inferenceConfig;
 
         public Builder setModelId(String modelId) {
             this.modelId = modelId;
@@ -387,6 +407,11 @@ public Builder setDefaultFieldMap(Map<String, String> defaultFieldMap) {
             return this;
         }
 
+        public Builder setInferenceConfig(InferenceConfig inferenceConfig) {
+            this.inferenceConfig = inferenceConfig;
+            return this;
+        }
+
         public TrainedModelConfig build() {
             return new TrainedModelConfig(
                 modelId,
@@ -402,7 +427,8 @@ public TrainedModelConfig build() {
                 estimatedHeapMemory,
                 estimatedOperations,
                 licenseLevel,
-                defaultFieldMap);
+                defaultFieldMap,
+                inferenceConfig);
         }
     }
 

client/rest-high-level/src/main/java/org/elasticsearch/client/ml/inference/trainedmodel/ClassificationConfig.java

@@ -0,0 +1,129 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.client.ml.inference.trainedmodel;
+
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.xcontent.ConstructingObjectParser;
+import org.elasticsearch.common.xcontent.ToXContent;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentParser;
+
+import java.io.IOException;
+import java.util.Objects;
+
+import static org.elasticsearch.common.xcontent.ConstructingObjectParser.optionalConstructorArg;
+
+public class ClassificationConfig implements InferenceConfig {
+
+    public static final ParseField NAME = new ParseField("classification");
+
+    public static final ParseField RESULTS_FIELD = new ParseField("results_field");
+    public static final ParseField NUM_TOP_CLASSES = new ParseField("num_top_classes");
+    public static final ParseField TOP_CLASSES_RESULTS_FIELD = new ParseField("top_classes_results_field");
+    public static final ParseField NUM_TOP_FEATURE_IMPORTANCE_VALUES = new ParseField("num_top_feature_importance_values");
+
+
+    private final Integer numTopClasses;
+    private final String topClassesResultsField;
+    private final String resultsField;
+    private final Integer numTopFeatureImportanceValues;
+
+    private static final ConstructingObjectParser<ClassificationConfig, Void> PARSER =
+            new ConstructingObjectParser<>(NAME.getPreferredName(), true, args -> new ClassificationConfig(
+                    (Integer) args[0], (String) args[1], (String) args[2], (Integer) args[3]));
+
+    static {
+        PARSER.declareInt(optionalConstructorArg(), NUM_TOP_CLASSES);
+        PARSER.declareString(optionalConstructorArg(), RESULTS_FIELD);
+        PARSER.declareString(optionalConstructorArg(), TOP_CLASSES_RESULTS_FIELD);
+        PARSER.declareInt(optionalConstructorArg(), NUM_TOP_FEATURE_IMPORTANCE_VALUES);
+    }
+
+    public static ClassificationConfig fromXContent(XContentParser parser) {
+        return PARSER.apply(parser, null);
+    }
+
+    public ClassificationConfig() {
+        this(null, null, null, null);
+    }
+
+    public ClassificationConfig(Integer numTopClasses, String resultsField, String topClassesResultsField, Integer featureImportance) {
+        this.numTopClasses = numTopClasses;
+        this.topClassesResultsField = topClassesResultsField;
+        this.resultsField = resultsField;
+        this.numTopFeatureImportanceValues = featureImportance;
+    }
+
+    public Integer getNumTopClasses() {
+        return numTopClasses;
+    }
+
+    public String getTopClassesResultsField() {
+        return topClassesResultsField;
+    }
+
+    public String getResultsField() {
+        return resultsField;
+    }
+
+    public Integer getNumTopFeatureImportanceValues() {
+        return numTopFeatureImportanceValues;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        ClassificationConfig that = (ClassificationConfig) o;
+        return Objects.equals(numTopClasses, that.numTopClasses)
+            && Objects.equals(topClassesResultsField, that.topClassesResultsField)
+            && Objects.equals(resultsField, that.resultsField)
+            && Objects.equals(numTopFeatureImportanceValues, that.numTopFeatureImportanceValues);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(numTopClasses, topClassesResultsField, resultsField, numTopFeatureImportanceValues);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
+        builder.startObject();
+        if (numTopClasses != null) {
+            builder.field(NUM_TOP_CLASSES.getPreferredName(), numTopClasses);
+        }
+        if (topClassesResultsField != null) {
+            builder.field(TOP_CLASSES_RESULTS_FIELD.getPreferredName(), topClassesResultsField);
+        }
+        if (resultsField != null) {
+            builder.field(RESULTS_FIELD.getPreferredName(), resultsField);
+        }
+        if (numTopFeatureImportanceValues != null) {
+            builder.field(NUM_TOP_FEATURE_IMPORTANCE_VALUES.getPreferredName(), numTopFeatureImportanceValues);
+        }
+        builder.endObject();
+        return builder;
+    }
+
+    @Override
+    public String getName() {
+        return NAME.getPreferredName();
+    }
+
+}

client/rest-high-level/src/main/java/org/elasticsearch/client/ml/inference/trainedmodel/InferenceConfig.java

@@ -0,0 +1,26 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.client.ml.inference.trainedmodel;
+
+import org.elasticsearch.client.ml.inference.NamedXContentObject;
+
+
+public interface InferenceConfig extends NamedXContentObject {
+
+}