-
Notifications
You must be signed in to change notification settings - Fork 138
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Signed-off-by: Yaliang Wu <ylwu@amazon.com>
- Loading branch information
Showing
7 changed files
with
171 additions
and
70 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,75 @@ | ||
/* | ||
* SPDX-License-Identifier: Apache-2.0 | ||
* | ||
* The OpenSearch Contributors require contributions made to | ||
* this file be licensed under the Apache-2.0 license or a | ||
* compatible open source license. | ||
* | ||
* Modifications Copyright OpenSearch Contributors. See | ||
* GitHub history for details. | ||
*/ | ||
|
||
package org.opensearch.ml.model; | ||
|
||
import java.io.IOException; | ||
import java.util.Base64; | ||
|
||
import lombok.Builder; | ||
import lombok.Getter; | ||
|
||
import org.opensearch.common.xcontent.ToXContentObject; | ||
import org.opensearch.common.xcontent.XContentBuilder; | ||
import org.opensearch.commons.authuser.User; | ||
import org.opensearch.ml.common.parameter.FunctionName; | ||
import org.opensearch.ml.engine.Model; | ||
|
||
@Getter | ||
public class MLModel implements ToXContentObject { | ||
public static final String ALGORITHM = "algorithm"; | ||
public static final String MODEL_NAME = "name"; | ||
public static final String MODEL_VERSION = "version"; | ||
public static final String MODEL_CONTENT = "content"; | ||
public static final String USER = "user"; | ||
|
||
private String name; | ||
private FunctionName algorithm; | ||
private Integer version; | ||
private String content; | ||
private User user; | ||
|
||
@Builder | ||
public MLModel(String name, FunctionName algorithm, Integer version, String content, User user) { | ||
this.name = name; | ||
this.algorithm = algorithm; | ||
this.version = version; | ||
this.content = content; | ||
this.user = user; | ||
} | ||
|
||
public MLModel(FunctionName algorithm, Model model) { | ||
this(model.getName(), algorithm, model.getVersion(), Base64.getEncoder().encodeToString(model.getContent()), null); | ||
} | ||
|
||
@Override | ||
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { | ||
builder.startObject(); | ||
if (name != null) { | ||
builder.field(MODEL_NAME, name); | ||
} | ||
if (algorithm != null) { | ||
builder.field(ALGORITHM, algorithm); | ||
} | ||
if (version != null) { | ||
builder.field(MODEL_VERSION, version); | ||
} | ||
if (content != null) { | ||
builder.field(MODEL_CONTENT, content); | ||
} | ||
if (user != null) { | ||
builder.field(USER, user); | ||
} | ||
builder.endObject(); | ||
return builder; | ||
} | ||
|
||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
44 changes: 44 additions & 0 deletions
44
plugin/src/test/java/org/opensearch/ml/model/MLModelTests.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,44 @@ | ||
/* | ||
* SPDX-License-Identifier: Apache-2.0 | ||
* | ||
* The OpenSearch Contributors require contributions made to | ||
* this file be licensed under the Apache-2.0 license or a | ||
* compatible open source license. | ||
* | ||
* Modifications Copyright OpenSearch Contributors. See | ||
* GitHub history for details. | ||
*/ | ||
|
||
package org.opensearch.ml.model; | ||
|
||
import static org.junit.Assert.assertEquals; | ||
import static org.opensearch.common.xcontent.ToXContent.EMPTY_PARAMS; | ||
|
||
import java.io.IOException; | ||
|
||
import org.junit.Test; | ||
import org.opensearch.common.xcontent.XContentBuilder; | ||
import org.opensearch.common.xcontent.XContentType; | ||
import org.opensearch.ml.common.parameter.FunctionName; | ||
import org.opensearch.ml.utils.TestHelper; | ||
|
||
public class MLModelTests { | ||
|
||
@Test | ||
public void toXContent() throws IOException { | ||
MLModel mlModel = MLModel.builder().algorithm(FunctionName.KMEANS).name("model_name").version(1).content("test_content").build(); | ||
XContentBuilder builder = XContentBuilder.builder(XContentType.JSON.xContent()); | ||
mlModel.toXContent(builder, EMPTY_PARAMS); | ||
String mlModelContent = TestHelper.xContentBuilderToString(builder); | ||
assertEquals("{\"name\":\"model_name\",\"algorithm\":\"KMEANS\",\"version\":1,\"content\":\"test_content\"}", mlModelContent); | ||
} | ||
|
||
@Test | ||
public void toXContent_NullValue() throws IOException { | ||
MLModel mlModel = MLModel.builder().build(); | ||
XContentBuilder builder = XContentBuilder.builder(XContentType.JSON.xContent()); | ||
mlModel.toXContent(builder, EMPTY_PARAMS); | ||
String mlModelContent = TestHelper.xContentBuilderToString(builder); | ||
assertEquals("{}", mlModelContent); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters