Skip to content

Commit 4a16102

Browse files
benwtrentbenwtrent
authored
[7.x] [ML] add new inference_config field to trained model config (#54421) (#54647)
* [ML] add new inference_config field to trained model config (#54421) A new field called `inference_config` is now added to the trained model config object. This new field allows for default inference settings from analytics or some external model builder. The inference processor can still override whatever is set as the default in the trained model config. * fixing for backport
1 parent 2113c1f commit 4a16102

File tree

61 files changed

+1949
-302
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

61 files changed

+1949
-302
lines changed

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@
1919
package org.elasticsearch.client.ml.inference;
2020

2121
import org.elasticsearch.client.ml.inference.preprocessing.CustomWordEmbedding;
22+
import org.elasticsearch.client.ml.inference.trainedmodel.ClassificationConfig;
23+
import org.elasticsearch.client.ml.inference.trainedmodel.InferenceConfig;
24+
import org.elasticsearch.client.ml.inference.trainedmodel.RegressionConfig;
2225
import org.elasticsearch.client.ml.inference.trainedmodel.TrainedModel;
2326
import org.elasticsearch.client.ml.inference.trainedmodel.ensemble.Ensemble;
2427
import org.elasticsearch.client.ml.inference.trainedmodel.ensemble.LogisticRegression;
@@ -61,6 +64,14 @@ public List<NamedXContentRegistry.Entry> getNamedXContentParsers() {
6164
new ParseField(LangIdentNeuralNetwork.NAME),
6265
LangIdentNeuralNetwork::fromXContent));
6366

67+
// Inference Config
68+
namedXContent.add(new NamedXContentRegistry.Entry(InferenceConfig.class,
69+
ClassificationConfig.NAME,
70+
ClassificationConfig::fromXContent));
71+
namedXContent.add(new NamedXContentRegistry.Entry(InferenceConfig.class,
72+
RegressionConfig.NAME,
73+
RegressionConfig::fromXContent));
74+
6475
// Aggregating output
6576
namedXContent.add(new NamedXContentRegistry.Entry(OutputAggregator.class,
6677
new ParseField(WeightedMode.NAME),

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,4 +54,14 @@ public static XContentBuilder writeNamedObjects(XContentBuilder builder,
5454
}
5555
return builder;
5656
}
57+
58+
public static XContentBuilder writeNamedObject(XContentBuilder builder,
59+
ToXContent.Params params,
60+
String namedObjectName,
61+
NamedXContentObject namedObject) throws IOException {
62+
builder.startObject(namedObjectName);
63+
builder.field(namedObject.getName(), namedObject, params);
64+
builder.endObject();
65+
return builder;
66+
}
5767
}

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

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import org.elasticsearch.Version;
2222
import org.elasticsearch.client.common.TimeUtil;
23+
import org.elasticsearch.client.ml.inference.trainedmodel.InferenceConfig;
2324
import org.elasticsearch.common.ParseField;
2425
import org.elasticsearch.common.Strings;
2526
import org.elasticsearch.common.unit.ByteSizeValue;
@@ -36,6 +37,8 @@
3637
import java.util.Map;
3738
import java.util.Objects;
3839

40+
import static org.elasticsearch.client.ml.inference.NamedXContentObjectHelper.writeNamedObject;
41+
3942
public class TrainedModelConfig implements ToXContentObject {
4043

4144
public static final String NAME = "trained_model_config";
@@ -54,6 +57,7 @@ public class TrainedModelConfig implements ToXContentObject {
5457
public static final ParseField ESTIMATED_OPERATIONS = new ParseField("estimated_operations");
5558
public static final ParseField LICENSE_LEVEL = new ParseField("license_level");
5659
public static final ParseField DEFAULT_FIELD_MAP = new ParseField("default_field_map");
60+
public static final ParseField INFERENCE_CONFIG = new ParseField("inference_config");
5761

5862
public static final ObjectParser<Builder, Void> PARSER = new ObjectParser<>(NAME,
5963
true,
@@ -78,6 +82,9 @@ public class TrainedModelConfig implements ToXContentObject {
7882
PARSER.declareLong(TrainedModelConfig.Builder::setEstimatedOperations, ESTIMATED_OPERATIONS);
7983
PARSER.declareString(TrainedModelConfig.Builder::setLicenseLevel, LICENSE_LEVEL);
8084
PARSER.declareObject(TrainedModelConfig.Builder::setDefaultFieldMap, (p, c) -> p.mapStrings(), DEFAULT_FIELD_MAP);
85+
PARSER.declareNamedObject(TrainedModelConfig.Builder::setInferenceConfig,
86+
(p, c, n) -> p.namedObject(InferenceConfig.class, n, null),
87+
INFERENCE_CONFIG);
8188
}
8289

8390
public static TrainedModelConfig fromXContent(XContentParser parser) throws IOException {
@@ -98,6 +105,7 @@ public static TrainedModelConfig fromXContent(XContentParser parser) throws IOEx
98105
private final Long estimatedOperations;
99106
private final String licenseLevel;
100107
private final Map<String, String> defaultFieldMap;
108+
private final InferenceConfig inferenceConfig;
101109

102110
TrainedModelConfig(String modelId,
103111
String createdBy,
@@ -112,7 +120,8 @@ public static TrainedModelConfig fromXContent(XContentParser parser) throws IOEx
112120
Long estimatedHeapMemory,
113121
Long estimatedOperations,
114122
String licenseLevel,
115-
Map<String, String> defaultFieldMap) {
123+
Map<String, String> defaultFieldMap,
124+
InferenceConfig inferenceConfig) {
116125
this.modelId = modelId;
117126
this.createdBy = createdBy;
118127
this.version = version;
@@ -127,6 +136,7 @@ public static TrainedModelConfig fromXContent(XContentParser parser) throws IOEx
127136
this.estimatedOperations = estimatedOperations;
128137
this.licenseLevel = licenseLevel;
129138
this.defaultFieldMap = defaultFieldMap == null ? null : Collections.unmodifiableMap(defaultFieldMap);
139+
this.inferenceConfig = inferenceConfig;
130140
}
131141

132142
public String getModelId() {
@@ -189,6 +199,10 @@ public Map<String, String> getDefaultFieldMap() {
189199
return defaultFieldMap;
190200
}
191201

202+
public InferenceConfig getInferenceConfig() {
203+
return inferenceConfig;
204+
}
205+
192206
public static Builder builder() {
193207
return new Builder();
194208
}
@@ -238,6 +252,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
238252
if (defaultFieldMap != null) {
239253
builder.field(DEFAULT_FIELD_MAP.getPreferredName(), defaultFieldMap);
240254
}
255+
if (inferenceConfig != null) {
256+
writeNamedObject(builder, params, INFERENCE_CONFIG.getPreferredName(), inferenceConfig);
257+
}
241258
builder.endObject();
242259
return builder;
243260
}
@@ -265,6 +282,7 @@ public boolean equals(Object o) {
265282
Objects.equals(estimatedOperations, that.estimatedOperations) &&
266283
Objects.equals(licenseLevel, that.licenseLevel) &&
267284
Objects.equals(defaultFieldMap, that.defaultFieldMap) &&
285+
Objects.equals(inferenceConfig, that.inferenceConfig) &&
268286
Objects.equals(metadata, that.metadata);
269287
}
270288

@@ -283,6 +301,7 @@ public int hashCode() {
283301
metadata,
284302
licenseLevel,
285303
input,
304+
inferenceConfig,
286305
defaultFieldMap);
287306
}
288307

@@ -303,6 +322,7 @@ public static class Builder {
303322
private Long estimatedOperations;
304323
private String licenseLevel;
305324
private Map<String, String> defaultFieldMap;
325+
private InferenceConfig inferenceConfig;
306326

307327
public Builder setModelId(String modelId) {
308328
this.modelId = modelId;
@@ -387,6 +407,11 @@ public Builder setDefaultFieldMap(Map<String, String> defaultFieldMap) {
387407
return this;
388408
}
389409

410+
public Builder setInferenceConfig(InferenceConfig inferenceConfig) {
411+
this.inferenceConfig = inferenceConfig;
412+
return this;
413+
}
414+
390415
public TrainedModelConfig build() {
391416
return new TrainedModelConfig(
392417
modelId,
@@ -402,7 +427,8 @@ public TrainedModelConfig build() {
402427
estimatedHeapMemory,
403428
estimatedOperations,
404429
licenseLevel,
405-
defaultFieldMap);
430+
defaultFieldMap,
431+
inferenceConfig);
406432
}
407433
}
408434

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

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.elasticsearch.client.ml.inference.trainedmodel;
20+
21+
import org.elasticsearch.common.ParseField;
22+
import org.elasticsearch.common.xcontent.ConstructingObjectParser;
23+
import org.elasticsearch.common.xcontent.ToXContent;
24+
import org.elasticsearch.common.xcontent.XContentBuilder;
25+
import org.elasticsearch.common.xcontent.XContentParser;
26+
27+
import java.io.IOException;
28+
import java.util.Objects;
29+
30+
import static org.elasticsearch.common.xcontent.ConstructingObjectParser.optionalConstructorArg;
31+
32+
public class ClassificationConfig implements InferenceConfig {
33+
34+
public static final ParseField NAME = new ParseField("classification");
35+
36+
public static final ParseField RESULTS_FIELD = new ParseField("results_field");
37+
public static final ParseField NUM_TOP_CLASSES = new ParseField("num_top_classes");
38+
public static final ParseField TOP_CLASSES_RESULTS_FIELD = new ParseField("top_classes_results_field");
39+
public static final ParseField NUM_TOP_FEATURE_IMPORTANCE_VALUES = new ParseField("num_top_feature_importance_values");
40+
41+
42+
private final Integer numTopClasses;
43+
private final String topClassesResultsField;
44+
private final String resultsField;
45+
private final Integer numTopFeatureImportanceValues;
46+
47+
private static final ConstructingObjectParser<ClassificationConfig, Void> PARSER =
48+
new ConstructingObjectParser<>(NAME.getPreferredName(), true, args -> new ClassificationConfig(
49+
(Integer) args[0], (String) args[1], (String) args[2], (Integer) args[3]));
50+
51+
static {
52+
PARSER.declareInt(optionalConstructorArg(), NUM_TOP_CLASSES);
53+
PARSER.declareString(optionalConstructorArg(), RESULTS_FIELD);
54+
PARSER.declareString(optionalConstructorArg(), TOP_CLASSES_RESULTS_FIELD);
55+
PARSER.declareInt(optionalConstructorArg(), NUM_TOP_FEATURE_IMPORTANCE_VALUES);
56+
}
57+
58+
public static ClassificationConfig fromXContent(XContentParser parser) {
59+
return PARSER.apply(parser, null);
60+
}
61+
62+
public ClassificationConfig() {
63+
this(null, null, null, null);
64+
}
65+
66+
public ClassificationConfig(Integer numTopClasses, String resultsField, String topClassesResultsField, Integer featureImportance) {
67+
this.numTopClasses = numTopClasses;
68+
this.topClassesResultsField = topClassesResultsField;
69+
this.resultsField = resultsField;
70+
this.numTopFeatureImportanceValues = featureImportance;
71+
}
72+
73+
public Integer getNumTopClasses() {
74+
return numTopClasses;
75+
}
76+
77+
public String getTopClassesResultsField() {
78+
return topClassesResultsField;
79+
}
80+
81+
public String getResultsField() {
82+
return resultsField;
83+
}
84+
85+
public Integer getNumTopFeatureImportanceValues() {
86+
return numTopFeatureImportanceValues;
87+
}
88+
89+
@Override
90+
public boolean equals(Object o) {
91+
if (this == o) return true;
92+
if (o == null || getClass() != o.getClass()) return false;
93+
ClassificationConfig that = (ClassificationConfig) o;
94+
return Objects.equals(numTopClasses, that.numTopClasses)
95+
&& Objects.equals(topClassesResultsField, that.topClassesResultsField)
96+
&& Objects.equals(resultsField, that.resultsField)
97+
&& Objects.equals(numTopFeatureImportanceValues, that.numTopFeatureImportanceValues);
98+
}
99+
100+
@Override
101+
public int hashCode() {
102+
return Objects.hash(numTopClasses, topClassesResultsField, resultsField, numTopFeatureImportanceValues);
103+
}
104+
105+
@Override
106+
public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
107+
builder.startObject();
108+
if (numTopClasses != null) {
109+
builder.field(NUM_TOP_CLASSES.getPreferredName(), numTopClasses);
110+
}
111+
if (topClassesResultsField != null) {
112+
builder.field(TOP_CLASSES_RESULTS_FIELD.getPreferredName(), topClassesResultsField);
113+
}
114+
if (resultsField != null) {
115+
builder.field(RESULTS_FIELD.getPreferredName(), resultsField);
116+
}
117+
if (numTopFeatureImportanceValues != null) {
118+
builder.field(NUM_TOP_FEATURE_IMPORTANCE_VALUES.getPreferredName(), numTopFeatureImportanceValues);
119+
}
120+
builder.endObject();
121+
return builder;
122+
}
123+
124+
@Override
125+
public String getName() {
126+
return NAME.getPreferredName();
127+
}
128+
129+
}

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
/*
2+
* Licensed to Elasticsearch under one or more contributor
3+
* license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright
5+
* ownership. Elasticsearch licenses this file to you under
6+
* the Apache License, Version 2.0 (the "License"); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.elasticsearch.client.ml.inference.trainedmodel;
20+
21+
import org.elasticsearch.client.ml.inference.NamedXContentObject;
22+
23+
24+
public interface InferenceConfig extends NamedXContentObject {
25+
26+
}

0 commit comments

Comments
 (0)