[ML] Add new include flag to GET inference/<model_id> API for model training metadata

client/rest-high-level/src/main/java/org/elasticsearch/client/MLRequestConverters.java

@@ -779,9 +779,9 @@ static Request getTrainedModels(GetTrainedModelsRequest getTrainedModelsRequest)
  params.putParam(GetTrainedModelsRequest.DECOMPRESS_DEFINITION,
  Boolean.toString(getTrainedModelsRequest.getDecompressDefinition()));
  }
- if (getTrainedModelsRequest.getIncludeDefinition() != null) {
- params.putParam(GetTrainedModelsRequest.INCLUDE_MODEL_DEFINITION,
- Boolean.toString(getTrainedModelsRequest.getIncludeDefinition()));
+ if (getTrainedModelsRequest.getIncludes().isEmpty() == false) {
+ params.putParam(GetTrainedModelsRequest.INCLUDE,
+ Strings.collectionToCommaDelimitedString(getTrainedModelsRequest.getIncludes()));
  }
  if (getTrainedModelsRequest.getTags() != null) {
  params.putParam(GetTrainedModelsRequest.TAGS, Strings.collectionToCommaDelimitedString(getTrainedModelsRequest.getTags()));

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

@@ -26,21 +26,26 @@
 import org.elasticsearch.common.Nullable;
 
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 
 public class GetTrainedModelsRequest implements Validatable {
 
+ private static final String DEFINITION = "definition";
+ private static final String TOTAL_FEATURE_IMPORTANCE = "total_feature_importance";
  public static final String ALLOW_NO_MATCH = "allow_no_match";
- public static final String INCLUDE_MODEL_DEFINITION = "include_model_definition";
  public static final String FOR_EXPORT = "for_export";
  public static final String DECOMPRESS_DEFINITION = "decompress_definition";
  public static final String TAGS = "tags";
+ public static final String INCLUDE = "include";
 
  private final List<String> ids;
  private Boolean allowNoMatch;
- private Boolean includeDefinition;
+ private Set<String> includes = new HashSet<>();
  private Boolean decompressDefinition;
  private Boolean forExport;
  private PageParams pageParams;
@@ -86,19 +91,32 @@ public GetTrainedModelsRequest setPageParams(@Nullable PageParams pageParams) {
  return this;
  }
 
- public Boolean getIncludeDefinition() {
- return includeDefinition;
+ public Set<String> getIncludes() {
+ return Collections.unmodifiableSet(includes);
+ }
+
+ public GetTrainedModelsRequest includeDefinition() {
+ this.includes.add(DEFINITION);
+ return this;
+ }
+
+ public GetTrainedModelsRequest includeTotalFeatureImportance() {
+ this.includes.add(TOTAL_FEATURE_IMPORTANCE);
+ return this;
  }
 
  /**
  * Whether to include the full model definition.
  *
  * The full model definition can be very large.
- *
+ * @deprecated Use {@link GetTrainedModelsRequest#includeDefinition()}
  * @param includeDefinition If {@code true}, the definition is included.
  */
+ @Deprecated
  public GetTrainedModelsRequest setIncludeDefinition(Boolean includeDefinition) {
- this.includeDefinition = includeDefinition;
+ if (includeDefinition != null && includeDefinition) {
+ return this.includeDefinition();
+ }
  return this;
  }
 
@@ -173,13 +191,13 @@ public boolean equals(Object o) {
  return Objects.equals(ids, other.ids)
  && Objects.equals(allowNoMatch, other.allowNoMatch)
  && Objects.equals(decompressDefinition, other.decompressDefinition)
- && Objects.equals(includeDefinition, other.includeDefinition)
+ && Objects.equals(includes, other.includes)
  && Objects.equals(forExport, other.forExport)
  && Objects.equals(pageParams, other.pageParams);
  }
 
  @Override
  public int hashCode() {
- return Objects.hash(ids, allowNoMatch, pageParams, decompressDefinition, includeDefinition, forExport);
+ return Objects.hash(ids, allowNoMatch, pageParams, decompressDefinition, includes, forExport);
  }
 }

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

@@ -0,0 +1,208 @@
+/*
+ * 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.metadata;
+
+import org.elasticsearch.common.Nullable;
+import org.elasticsearch.common.ParseField;
+import org.elasticsearch.common.xcontent.ConstructingObjectParser;
+import org.elasticsearch.common.xcontent.ObjectParser;
+import org.elasticsearch.common.xcontent.ToXContentObject;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.common.xcontent.XContentParseException;
+import org.elasticsearch.common.xcontent.XContentParser;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
+
+public class TotalFeatureImportance implements ToXContentObject {
+
+ private static final String NAME = "total_feature_importance";
+ public static final ParseField FEATURE_NAME = new ParseField("feature_name");
+ public static final ParseField IMPORTANCE = new ParseField("importance");
+ public static final ParseField CLASSES = new ParseField("classes");
+ public static final ParseField MEAN_MAGNITUDE = new ParseField("mean_magnitude");
+ public static final ParseField MIN = new ParseField("min");
+ public static final ParseField MAX = new ParseField("max");
+
+ @SuppressWarnings("unchecked")
+ public static final ConstructingObjectParser<TotalFeatureImportance, Void> PARSER = new ConstructingObjectParser<>(NAME,
+ true,
+ a -> new TotalFeatureImportance((String)a[0], (Importance)a[1], (List<ClassImportance>)a[2]));
+
+ static {
+ PARSER.declareString(ConstructingObjectParser.constructorArg(), FEATURE_NAME);
+ PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), Importance.PARSER, IMPORTANCE);
+ PARSER.declareObjectArray(ConstructingObjectParser.optionalConstructorArg(), ClassImportance.PARSER, CLASSES);
+ }
+
+ public static TotalFeatureImportance fromXContent(XContentParser parser) {
+ return PARSER.apply(parser, null);
+ }
+
+ public final String featureName;
+ public final Importance importance;
+ public final List<ClassImportance> classImportances;
+
+ TotalFeatureImportance(String featureName, @Nullable Importance importance, @Nullable List<ClassImportance> classImportances) {
+ this.featureName = featureName;
+ this.importance = importance;
+ this.classImportances = classImportances == null ? Collections.emptyList() : classImportances;
+ }
+
+ @Override
+ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+ builder.startObject();
+ builder.field(FEATURE_NAME.getPreferredName(), featureName);
+ if (importance != null) {
+ builder.field(IMPORTANCE.getPreferredName(), importance);
+ }
+ if (classImportances.isEmpty() == false) {
+ builder.field(CLASSES.getPreferredName(), classImportances);
+ }
+ builder.endObject();
+ return builder;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ TotalFeatureImportance that = (TotalFeatureImportance) o;
+ return Objects.equals(that.importance, importance)
+ && Objects.equals(featureName, that.featureName)
+ && Objects.equals(classImportances, that.classImportances);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(featureName, importance, classImportances);
+ }
+
+ public static class Importance implements ToXContentObject {
+ private static final String NAME = "importance";
+
+ public static final ConstructingObjectParser<Importance, Void> PARSER = new ConstructingObjectParser<>(NAME,
+ true,
+ a -> new Importance((double)a[0], (double)a[1], (double)a[2]));
+
+ static {
+ PARSER.declareDouble(ConstructingObjectParser.constructorArg(), MEAN_MAGNITUDE);
+ PARSER.declareDouble(ConstructingObjectParser.constructorArg(), MIN);
+ PARSER.declareDouble(ConstructingObjectParser.constructorArg(), MAX);
+ }
+
+ private final double meanMagnitude;
+ private final double min;
+ private final double max;
+
+ public Importance(double meanMagnitude, double min, double max) {
+ this.meanMagnitude = meanMagnitude;
+ this.min = min;
+ this.max = max;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ Importance that = (Importance) o;
+ return Double.compare(that.meanMagnitude, meanMagnitude) == 0 &&
+ Double.compare(that.min, min) == 0 &&
+ Double.compare(that.max, max) == 0;
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(meanMagnitude, min, max);
+ }
+
+ @Override
+ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+ builder.startObject();
+ builder.field(MEAN_MAGNITUDE.getPreferredName(), meanMagnitude);
+ builder.field(MIN.getPreferredName(), min);
+ builder.field(MAX.getPreferredName(), max);
+ builder.endObject();
+ return builder;
+ }
+ }
+
+ public static class ClassImportance implements ToXContentObject {
+ private static final String NAME = "total_class_importance";
+
+ public static final ParseField CLASS_NAME = new ParseField("class_name");
+ public static final ParseField IMPORTANCE = new ParseField("importance");
+
+ public static final ConstructingObjectParser<ClassImportance, Void> PARSER = new ConstructingObjectParser<>(NAME,
+ true,
+ a -> new ClassImportance(a[0], (Importance)a[1]));
+
+ static {
+ PARSER.declareField(ConstructingObjectParser.constructorArg(), (p, c) -> {
+ if (p.currentToken() == XContentParser.Token.VALUE_STRING) {
+ return p.text();
+ } else if (p.currentToken() == XContentParser.Token.VALUE_NUMBER) {
+ return p.numberValue();
+ } else if (p.currentToken() == XContentParser.Token.VALUE_BOOLEAN) {
+ return p.booleanValue();
+ }
+ throw new XContentParseException("Unsupported token [" + p.currentToken() + "]");
+ }, CLASS_NAME, ObjectParser.ValueType.VALUE);
+ PARSER.declareObject(ConstructingObjectParser.constructorArg(), Importance.PARSER, IMPORTANCE);
+ }
+
+ public static ClassImportance fromXContent(XContentParser parser) {
+ return PARSER.apply(parser, null);
+ }
+
+ public final Object className;
+ public final Importance importance;
+
+ ClassImportance(Object className, Importance importance) {
+ this.className = className;
+ this.importance = importance;
+ }
+
+ @Override
+ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+ builder.startObject();
+ builder.field(CLASS_NAME.getPreferredName(), className);
+ builder.field(IMPORTANCE.getPreferredName(), importance);
+ builder.endObject();
+ return builder;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ ClassImportance that = (ClassImportance) o;
+ return Objects.equals(that.importance, importance) && Objects.equals(className, that.className);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(className, importance);
+ }
+
+ }
+}

client/rest-high-level/src/test/java/org/elasticsearch/client/MLRequestConvertersTests.java

@@ -894,7 +894,7 @@ public void testGetTrainedModels() {
  GetTrainedModelsRequest getRequest = new GetTrainedModelsRequest(modelId1, modelId2, modelId3)
  .setAllowNoMatch(false)
  .setDecompressDefinition(true)
- .setIncludeDefinition(false)
+ .includeDefinition()
  .setTags("tag1", "tag2")
  .setPageParams(new PageParams(100, 300));
 
@@ -908,7 +908,7 @@ public void testGetTrainedModels() {
  hasEntry("allow_no_match", "false"),
  hasEntry("decompress_definition", "true"),
  hasEntry("tags", "tag1,tag2"),
- hasEntry("include_model_definition", "false")
+ hasEntry("include", "definition")
  ));
  assertNull(request.getEntity());
  }

client/rest-high-level/src/test/java/org/elasticsearch/client/MachineLearningIT.java

@@ -2227,7 +2227,10 @@ public void testGetTrainedModels() throws Exception {
 
  {
  GetTrainedModelsResponse getTrainedModelsResponse = execute(
- new GetTrainedModelsRequest(modelIdPrefix + 0).setDecompressDefinition(true).setIncludeDefinition(true),
+ new GetTrainedModelsRequest(modelIdPrefix + 0)
+ .setDecompressDefinition(true)
+ .includeDefinition()
+ .includeTotalFeatureImportance(),
  machineLearningClient::getTrainedModels,
  machineLearningClient::getTrainedModelsAsync);
 
@@ -2238,7 +2241,10 @@ public void testGetTrainedModels() throws Exception {
  assertThat(getTrainedModelsResponse.getTrainedModels().get(0).getModelId(), equalTo(modelIdPrefix + 0));
 
  getTrainedModelsResponse = execute(
- new GetTrainedModelsRequest(modelIdPrefix + 0).setDecompressDefinition(false).setIncludeDefinition(true),
+ new GetTrainedModelsRequest(modelIdPrefix + 0)
+ .setDecompressDefinition(false)
+ .includeTotalFeatureImportance()
+ .includeDefinition(),
  machineLearningClient::getTrainedModels,
  machineLearningClient::getTrainedModelsAsync);
 
@@ -2249,7 +2255,8 @@ public void testGetTrainedModels() throws Exception {
  assertThat(getTrainedModelsResponse.getTrainedModels().get(0).getModelId(), equalTo(modelIdPrefix + 0));
 
  getTrainedModelsResponse = execute(
- new GetTrainedModelsRequest(modelIdPrefix + 0).setDecompressDefinition(false).setIncludeDefinition(false),
+ new GetTrainedModelsRequest(modelIdPrefix + 0)
+ .setDecompressDefinition(false),
  machineLearningClient::getTrainedModels,
  machineLearningClient::getTrainedModelsAsync);
  assertThat(getTrainedModelsResponse.getCount(), equalTo(1L));

client/rest-high-level/src/test/java/org/elasticsearch/client/documentation/MlClientDocumentationIT.java

@@ -3694,11 +3694,12 @@ public void testGetTrainedModels() throws Exception {
  // tag::get-trained-models-request
  GetTrainedModelsRequest request = new GetTrainedModelsRequest("my-trained-model") // <1>
  .setPageParams(new PageParams(0, 1)) // <2>
- .setIncludeDefinition(false) // <3>
- .setDecompressDefinition(false) // <4>
- .setAllowNoMatch(true) // <5>
- .setTags("regression") // <6>
- .setForExport(false); // <7>
+ .includeDefinition() // <3>
+ .includeTotalFeatureImportance() // <4>
+ .setDecompressDefinition(false) // <5>
+ .setAllowNoMatch(true) // <6>
+ .setTags("regression") // <7>
+ .setForExport(false); // <8>
  // end::get-trained-models-request
  request.setTags((List<String>)null);