Examples with MLInput org.opensearch.ml.common.parameter.MLInput used on opensource projects

Example 46 with MLInput

use of org.opensearch.ml.common.parameter.MLInput in project ml-commons by opensearch-project.

the class IntegTestUtils method predictAndVerifyResult.

// Predict with the model generated, and verify the prediction result.
public static void predictAndVerifyResult(String taskId, MLInputDataset inputDataset) throws IOException {
    MLInput mlInput = MLInput.builder().algorithm(FunctionName.KMEANS).inputDataset(inputDataset).build();
    MLPredictionTaskRequest predictionRequest = new MLPredictionTaskRequest(taskId, mlInput);
    ActionFuture<MLTaskResponse> predictionFuture = client().execute(MLPredictionTaskAction.INSTANCE, predictionRequest);
    MLTaskResponse predictionResponse = predictionFuture.actionGet();
    XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
    builder.startObject();
    MLPredictionOutput mlPredictionOutput = (MLPredictionOutput) predictionResponse.getOutput();
    mlPredictionOutput.getPredictionResult().toXContent(builder, ToXContent.EMPTY_PARAMS);
    builder.endObject();
    String jsonStr = Strings.toString(builder);
    String expectedStr1 = "{\"column_metas\":[{\"name\":\"ClusterID\",\"column_type\":\"INTEGER\"}]," + "\"rows\":[{\"values\":[{\"column_type\":\"INTEGER\",\"value\":0}]}]}";
    String expectedStr2 = "{\"column_metas\":[{\"name\":\"ClusterID\",\"column_type\":\"INTEGER\"}]," + "\"rows\":[{\"values\":[{\"column_type\":\"INTEGER\",\"value\":1}]}]}";
    // The prediction result would not be a fixed value.
    assertTrue(expectedStr1.equals(jsonStr) || expectedStr2.equals(jsonStr));
}