Skip to content

Commit

Permalink
support model_task_type and qa_model_config in ml input
Browse files Browse the repository at this point in the history
Signed-off-by: Bhavana Ramaram <rbhavna@amazon.com>
  • Loading branch information
rbhavna committed Mar 19, 2024
1 parent acc0b8d commit 2943726
Showing 6 changed files with 412 additions and 43 deletions.
7 changes: 7 additions & 0 deletions common/src/main/java/org/opensearch/ml/common/MLModel.java
Original file line number Diff line number Diff line change
@@ -21,6 +21,7 @@
import org.opensearch.ml.common.controller.MLRateLimiter;
import org.opensearch.ml.common.model.MLModelFormat;
import org.opensearch.ml.common.model.MLModelState;
import org.opensearch.ml.common.model.QuestionAnsweringModelConfig;
import org.opensearch.ml.common.model.TextEmbeddingModelConfig;
import org.opensearch.ml.common.model.MetricsCorrelationModelConfig;

@@ -39,6 +40,7 @@ public class MLModel implements ToXContentObject {
@Deprecated
public static final String ALGORITHM_FIELD = "algorithm";
public static final String FUNCTION_NAME_FIELD = "function_name";
public static final String MODEL_TASK_TYPE_FIELD = "model_task_type";
public static final String MODEL_NAME_FIELD = "name";
public static final String MODEL_GROUP_ID_FIELD = "model_group_id";
// We use int type for version in first release 1.3. In 2.4, we changed to
@@ -219,6 +221,8 @@ public MLModel(StreamInput input) throws IOException {
if (input.readBoolean()) {
if (algorithm.equals(FunctionName.METRICS_CORRELATION)) {
modelConfig = new MetricsCorrelationModelConfig(input);
} else if (algorithm.equals(FunctionName.QUESTION_ANSWERING)) {
modelConfig = new QuestionAnsweringModelConfig(input);
} else {
modelConfig = new TextEmbeddingModelConfig(input);
}
@@ -499,6 +503,7 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
case USER:
user = User.parse(parser);
break;
case MODEL_TASK_TYPE_FIELD:
case ALGORITHM_FIELD:
case FUNCTION_NAME_FIELD:
algorithm = FunctionName.from(parser.text().toUpperCase(Locale.ROOT));
@@ -527,6 +532,8 @@ public static MLModel parse(XContentParser parser, String algorithmName) throws
case MODEL_CONFIG_FIELD:
if (FunctionName.METRICS_CORRELATION.name().equals(algorithmName)) {
modelConfig = MetricsCorrelationModelConfig.parse(parser);
} else if (FunctionName.QUESTION_ANSWERING.name().equals(algorithmName)) {
modelConfig = QuestionAnsweringModelConfig.parse(parser);
} else {
modelConfig = TextEmbeddingModelConfig.parse(parser);
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,191 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.ml.common.model;

import lombok.Builder;
import lombok.Getter;
import lombok.Setter;
import org.opensearch.core.ParseField;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.xcontent.NamedXContentRegistry;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.ml.common.FunctionName;

import java.io.IOException;
import java.util.Locale;

import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;

@Setter
@Getter
public class QuestionAnsweringModelConfig extends MLModelConfig {
public static final String PARSE_FIELD_NAME = FunctionName.QUESTION_ANSWERING.name();
public static final NamedXContentRegistry.Entry XCONTENT_REGISTRY = new NamedXContentRegistry.Entry(
QuestionAnsweringModelConfig.class,
new ParseField(PARSE_FIELD_NAME),
it -> parse(it)
);
public static final String FRAMEWORK_TYPE_FIELD = "framework_type";
public static final String POOLING_MODE_FIELD = "pooling_mode";
public static final String NORMALIZE_RESULT_FIELD = "normalize_result";
public static final String MODEL_MAX_LENGTH_FIELD = "model_max_length";

private final FrameworkType frameworkType;
private final PoolingMode poolingMode;
private final boolean normalizeResult;
private final Integer modelMaxLength;

@Builder(toBuilder = true)
public QuestionAnsweringModelConfig(String modelType, FrameworkType frameworkType, String allConfig,
PoolingMode poolingMode, boolean normalizeResult, Integer modelMaxLength) {
super(modelType, allConfig);
if (frameworkType == null) {
throw new IllegalArgumentException("framework type is null");
}
this.frameworkType = frameworkType;
this.poolingMode = poolingMode;
this.normalizeResult = normalizeResult;
this.modelMaxLength = modelMaxLength;
}

public static QuestionAnsweringModelConfig parse(XContentParser parser) throws IOException {
String modelType = null;
FrameworkType frameworkType = null;
String allConfig = null;
PoolingMode poolingMode = null;
boolean normalizeResult = false;
Integer modelMaxLength = null;

ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
String fieldName = parser.currentName();
parser.nextToken();

switch (fieldName) {
case MODEL_TYPE_FIELD:
modelType = parser.text();
break;
case FRAMEWORK_TYPE_FIELD:
frameworkType = FrameworkType.from(parser.text().toUpperCase(Locale.ROOT));
break;
case ALL_CONFIG_FIELD:
allConfig = parser.text();
break;
case POOLING_MODE_FIELD:
poolingMode = PoolingMode.from(parser.text().toUpperCase(Locale.ROOT));
break;
case NORMALIZE_RESULT_FIELD:
normalizeResult = parser.booleanValue();
break;
case MODEL_MAX_LENGTH_FIELD:
modelMaxLength = parser.intValue();
break;
default:
parser.skipChildren();
break;
}
}
return new QuestionAnsweringModelConfig(modelType, frameworkType, allConfig, poolingMode, normalizeResult, modelMaxLength);
}

@Override
public String getWriteableName() {
return PARSE_FIELD_NAME;
}

public QuestionAnsweringModelConfig(StreamInput in) throws IOException{
super(in);
frameworkType = in.readEnum(FrameworkType.class);
if (in.readBoolean()) {
poolingMode = in.readEnum(PoolingMode.class);
} else {
poolingMode = null;
}
normalizeResult = in.readBoolean();
modelMaxLength = in.readOptionalInt();
}

@Override
public void writeTo(StreamOutput out) throws IOException {
super.writeTo(out);
out.writeEnum(frameworkType);
if (poolingMode != null) {
out.writeBoolean(true);
out.writeEnum(poolingMode);
} else {
out.writeBoolean(false);
}
out.writeBoolean(normalizeResult);
out.writeOptionalInt(modelMaxLength);
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject();
if (modelType != null) {
builder.field(MODEL_TYPE_FIELD, modelType);
}
if (frameworkType != null) {
builder.field(FRAMEWORK_TYPE_FIELD, frameworkType);
}
if (allConfig != null) {
builder.field(ALL_CONFIG_FIELD, allConfig);
}
if (modelMaxLength != null) {
builder.field(MODEL_MAX_LENGTH_FIELD, modelMaxLength);
}
if (poolingMode != null) {
builder.field(POOLING_MODE_FIELD, poolingMode);
}
if (normalizeResult) {
builder.field(NORMALIZE_RESULT_FIELD, normalizeResult);
}
builder.endObject();
return builder;
}

public enum PoolingMode {
MEAN("mean"),
MEAN_SQRT_LEN("mean_sqrt_len"),
MAX("max"),
WEIGHTED_MEAN("weightedmean"),
CLS("cls"),
LAST_TOKEN("lasttoken");

private String name;

public String getName() {
return name;
}
PoolingMode(String name) {
this.name = name;
}

public static PoolingMode from(String value) {
try {
return PoolingMode.valueOf(value.toUpperCase(Locale.ROOT));
} catch (Exception e) {
throw new IllegalArgumentException("Wrong pooling method");
}
}
}
public enum FrameworkType {
HUGGINGFACE_TRANSFORMERS,
SENTENCE_TRANSFORMERS,
HUGGINGFACE_TRANSFORMERS_NEURON;

public static FrameworkType from(String value) {
try {
return FrameworkType.valueOf(value.toUpperCase(Locale.ROOT));
} catch (Exception e) {
throw new IllegalArgumentException("Wrong framework type");
}
}
}

}
Original file line number Diff line number Diff line change
@@ -23,6 +23,7 @@
import org.opensearch.ml.common.controller.MLRateLimiter;
import org.opensearch.ml.common.model.MLModelFormat;
import org.opensearch.ml.common.model.MetricsCorrelationModelConfig;
import org.opensearch.ml.common.model.QuestionAnsweringModelConfig;
import org.opensearch.ml.common.model.TextEmbeddingModelConfig;

import java.io.IOException;
@@ -41,6 +42,8 @@
public class MLRegisterModelInput implements ToXContentObject, Writeable {

public static final String FUNCTION_NAME_FIELD = "function_name";
public static final String MODEL_TASK_TYPE_FIELD = "model_task_type";

public static final String NAME_FIELD = "name";
public static final String MODEL_GROUP_ID_FIELD = "model_group_id";
public static final String DESCRIPTION_FIELD = "description";
@@ -166,6 +169,8 @@ public MLRegisterModelInput(StreamInput in) throws IOException {
if (in.readBoolean()) {
if (this.functionName.equals(FunctionName.METRICS_CORRELATION)) {
this.modelConfig = new MetricsCorrelationModelConfig(in);
} else if (this.functionName.equals(FunctionName.QUESTION_ANSWERING)) {
this.modelConfig = new QuestionAnsweringModelConfig(in);
} else {
this.modelConfig = new TextEmbeddingModelConfig(in);
}
@@ -357,6 +362,7 @@ public static MLRegisterModelInput parse(XContentParser parser, String modelName
String fieldName = parser.currentName();
parser.nextToken();
switch (fieldName) {
case MODEL_TASK_TYPE_FIELD:
case FUNCTION_NAME_FIELD:
functionName = FunctionName.from(parser.text().toUpperCase(Locale.ROOT));
break;
@@ -382,7 +388,11 @@ public static MLRegisterModelInput parse(XContentParser parser, String modelName
modelFormat = MLModelFormat.from(parser.text().toUpperCase(Locale.ROOT));
break;
case MODEL_CONFIG_FIELD:
modelConfig = TextEmbeddingModelConfig.parse(parser);
if (FunctionName.QUESTION_ANSWERING.equals(functionName)) {
modelConfig = QuestionAnsweringModelConfig.parse(parser);
} else {
modelConfig = TextEmbeddingModelConfig.parse(parser);
}
break;
case CONNECTOR_FIELD:
connector = createConnector(parser);
@@ -456,6 +466,7 @@ public static MLRegisterModelInput parse(XContentParser parser, boolean deployMo
parser.nextToken();

switch (fieldName) {
case MODEL_TASK_TYPE_FIELD:
case FUNCTION_NAME_FIELD:
functionName = FunctionName.from(parser.text().toUpperCase(Locale.ROOT));
break;
@@ -493,7 +504,11 @@ public static MLRegisterModelInput parse(XContentParser parser, boolean deployMo
modelFormat = MLModelFormat.from(parser.text().toUpperCase(Locale.ROOT));
break;
case MODEL_CONFIG_FIELD:
modelConfig = TextEmbeddingModelConfig.parse(parser);
if (FunctionName.QUESTION_ANSWERING.equals(functionName)) {
modelConfig = QuestionAnsweringModelConfig.parse(parser);
} else {
modelConfig = TextEmbeddingModelConfig.parse(parser);
}
break;
case MODEL_NODE_IDS_FIELD:
ensureExpectedToken(XContentParser.Token.START_ARRAY, parser.currentToken(), parser);
Original file line number Diff line number Diff line change
@@ -22,6 +22,8 @@
import org.opensearch.ml.common.controller.MLRateLimiter;
import org.opensearch.ml.common.model.MLModelFormat;
import org.opensearch.ml.common.model.MLModelState;
import org.opensearch.ml.common.model.MetricsCorrelationModelConfig;
import org.opensearch.ml.common.model.QuestionAnsweringModelConfig;
import org.opensearch.ml.common.model.TextEmbeddingModelConfig;

import java.io.IOException;
@@ -35,6 +37,7 @@
public class MLRegisterModelMetaInput implements ToXContentObject, Writeable {

public static final String FUNCTION_NAME_FIELD = "function_name";
public static final String MODEL_TASK_TYPE_FIELD = "model_task_type";
public static final String MODEL_NAME_FIELD = "name"; // mandatory
public static final String DESCRIPTION_FIELD = "description"; // optional
public static final String IS_ENABLED_FIELD = "is_enabled"; // optional
@@ -144,7 +147,11 @@ public MLRegisterModelMetaInput(StreamInput in) throws IOException {
this.modelContentSizeInBytes = in.readOptionalLong();
this.modelContentHashValue = in.readString();
if (in.readBoolean()) {
modelConfig = new TextEmbeddingModelConfig(in);
if (this.functionName.equals(FunctionName.QUESTION_ANSWERING)) {
this.modelConfig = new QuestionAnsweringModelConfig(in);
} else {
this.modelConfig = new TextEmbeddingModelConfig(in);
}
}
this.totalChunks = in.readInt();
this.backendRoles = in.readOptionalStringList();
@@ -298,6 +305,7 @@ public static MLRegisterModelMetaInput parse(XContentParser parser) throws IOExc
case MODEL_NAME_FIELD:
name = parser.text();
break;
case MODEL_TASK_TYPE_FIELD:
case FUNCTION_NAME_FIELD:
functionName = FunctionName.from(parser.text());
break;
@@ -329,7 +337,11 @@ public static MLRegisterModelMetaInput parse(XContentParser parser) throws IOExc
modelContentHashValue = parser.text();
break;
case MODEL_CONFIG_FIELD:
modelConfig = TextEmbeddingModelConfig.parse(parser);
if (FunctionName.QUESTION_ANSWERING.equals(functionName)) {
modelConfig = QuestionAnsweringModelConfig.parse(parser);
} else {
modelConfig = TextEmbeddingModelConfig.parse(parser);
}
break;
case TOTAL_CHUNKS_FIELD:
totalChunks = parser.intValue(false);
Loading

0 comments on commit 2943726

Please sign in to comment.