Fix update model config invalid error (opensearch-project#3994) (opensearch-project#3997)

* Fix update model config



* Add model config validation logic



* Simplified validateModelConfig, add version compatibility tests



---------


(cherry picked from commit c5229ca)

Signed-off-by: Nathalie Jonathan <nathhjo@amazon.com>
Co-authored-by: Nathalie Jonathan <143617992+nathaliellenaa@users.noreply.github.com>

Author: opensearch-trigger-bot[bot]

common/src/main/java/org/opensearch/ml/common/model/BaseModelConfig.java

@@ -9,6 +9,7 @@
 import static org.opensearch.ml.common.CommonValue.VERSION_3_1_0;
 
 import java.io.IOException;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -41,20 +42,60 @@ public class BaseModelConfig extends MLModelConfig {
         it -> parse(it)
     );
 
+    public static final String EMBEDDING_DIMENSION_FIELD = "embedding_dimension";
+    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";
+    public static final String QUERY_PREFIX = "query_prefix";
+    public static final String PASSAGE_PREFIX = "passage_prefix";
     public static final String ADDITIONAL_CONFIG_FIELD = "additional_config";
+
+    protected Integer embeddingDimension;
+    protected FrameworkType frameworkType;
+    protected PoolingMode poolingMode;
+    protected boolean normalizeResult;
+    protected Integer modelMaxLength;
+    protected String queryPrefix;
+    protected String passagePrefix;
     protected Map<String, Object> additionalConfig;
 
     @Builder(builderMethodName = "baseModelConfigBuilder")
-    public BaseModelConfig(String modelType, String allConfig, Map<String, Object> additionalConfig) {
+    public BaseModelConfig(
+        String modelType,
+        String allConfig,
+        Map<String, Object> additionalConfig,
+        Integer embeddingDimension,
+        FrameworkType frameworkType,
+        PoolingMode poolingMode,
+        boolean normalizeResult,
+        Integer modelMaxLength,
+        String queryPrefix,
+        String passagePrefix
+    ) {
         super(modelType, allConfig);
         this.additionalConfig = additionalConfig;
+        this.embeddingDimension = embeddingDimension;
+        this.frameworkType = frameworkType;
+        this.poolingMode = poolingMode;
+        this.normalizeResult = normalizeResult;
+        this.modelMaxLength = modelMaxLength;
+        this.queryPrefix = queryPrefix;
+        this.passagePrefix = passagePrefix;
         validateNoDuplicateKeys(allConfig, additionalConfig);
     }
 
     public static BaseModelConfig parse(XContentParser parser) throws IOException {
         String modelType = null;
         String allConfig = null;
         Map<String, Object> additionalConfig = null;
+        Integer embeddingDimension = null;
+        FrameworkType frameworkType = null;
+        PoolingMode poolingMode = null;
+        boolean normalizeResult = false;
+        Integer modelMaxLength = null;
+        String queryPrefix = null;
+        String passagePrefix = null;
 
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
         while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -71,12 +112,44 @@ public static BaseModelConfig parse(XContentParser parser) throws IOException {
                 case ADDITIONAL_CONFIG_FIELD:
                     additionalConfig = parser.map();
                     break;
+                case EMBEDDING_DIMENSION_FIELD:
+                    embeddingDimension = parser.intValue();
+                    break;
+                case FRAMEWORK_TYPE_FIELD:
+                    frameworkType = FrameworkType.from(parser.text().toUpperCase(Locale.ROOT));
+                    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;
+                case QUERY_PREFIX:
+                    queryPrefix = parser.text();
+                    break;
+                case PASSAGE_PREFIX:
+                    passagePrefix = parser.text();
+                    break;
                 default:
                     parser.skipChildren();
                     break;
             }
         }
-        return new BaseModelConfig(modelType, allConfig, additionalConfig);
+        return new BaseModelConfig(
+            modelType,
+            allConfig,
+            additionalConfig,
+            embeddingDimension,
+            frameworkType,
+            poolingMode,
+            normalizeResult,
+            modelMaxLength,
+            queryPrefix,
+            passagePrefix
+        );
     }
 
     @Override
@@ -89,6 +162,21 @@ public BaseModelConfig(StreamInput in) throws IOException {
         if (in.getVersion().onOrAfter(VERSION_3_1_0)) {
             this.additionalConfig = in.readMap();
         }
+        embeddingDimension = in.readOptionalInt();
+        if (in.readBoolean()) {
+            frameworkType = in.readEnum(FrameworkType.class);
+        } else {
+            frameworkType = null;
+        }
+        if (in.readBoolean()) {
+            poolingMode = in.readEnum(PoolingMode.class);
+        } else {
+            poolingMode = null;
+        }
+        normalizeResult = in.readBoolean();
+        modelMaxLength = in.readOptionalInt();
+        queryPrefix = in.readOptionalString();
+        passagePrefix = in.readOptionalString();
     }
 
     @Override
@@ -97,6 +185,23 @@ public void writeTo(StreamOutput out) throws IOException {
         if (out.getVersion().onOrAfter(VERSION_3_1_0)) {
             out.writeMap(additionalConfig);
         }
+        out.writeOptionalInt(embeddingDimension);
+        if (frameworkType != null) {
+            out.writeBoolean(true);
+            out.writeEnum(frameworkType);
+        } else {
+            out.writeBoolean(false);
+        }
+        if (poolingMode != null) {
+            out.writeBoolean(true);
+            out.writeEnum(poolingMode);
+        } else {
+            out.writeBoolean(false);
+        }
+        out.writeBoolean(normalizeResult);
+        out.writeOptionalInt(modelMaxLength);
+        out.writeOptionalString(queryPrefix);
+        out.writeOptionalString(passagePrefix);
     }
 
     @Override
@@ -111,10 +216,72 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params par
         if (additionalConfig != null) {
             builder.field(ADDITIONAL_CONFIG_FIELD, additionalConfig);
         }
+        if (embeddingDimension != null) {
+            builder.field(EMBEDDING_DIMENSION_FIELD, embeddingDimension);
+        }
+        if (frameworkType != null) {
+            builder.field(FRAMEWORK_TYPE_FIELD, frameworkType);
+        }
+        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);
+        }
+        if (queryPrefix != null) {
+            builder.field(QUERY_PREFIX, queryPrefix);
+        }
+        if (passagePrefix != null) {
+            builder.field(PASSAGE_PREFIX, passagePrefix);
+        }
         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");
+            }
+        }
+    }
+
     protected void validateNoDuplicateKeys(String allConfig, Map<String, Object> additionalConfig) {
         if (allConfig == null || additionalConfig == null || additionalConfig.isEmpty()) {
             return;