Examples with MLTaskResponse org.opensearch.ml.common.transport.MLTaskResponse used on opensource projects

Example 21 with MLTaskResponse

use of org.opensearch.ml.common.transport.MLTaskResponse in project ml-commons by opensearch-project.

the class IntegTestUtils method trainModel.

// Train a model.
public static String trainModel(MLInputDataset inputDataset) throws ExecutionException, InterruptedException {
    MLInput mlInput = MLInput.builder().algorithm(FunctionName.KMEANS).inputDataset(inputDataset).build();
    // TODO: support train test in sync way
    MLTrainingTaskRequest trainingRequest = new MLTrainingTaskRequest(mlInput, true);
    ActionFuture<MLTaskResponse> trainingFuture = client().execute(MLTrainingTaskAction.INSTANCE, trainingRequest);
    MLTaskResponse trainingResponse = trainingFuture.actionGet();
    assertNotNull(trainingResponse);
    MLTrainingOutput modelTrainingOutput = (MLTrainingOutput) trainingResponse.getOutput();
    String modelId = modelTrainingOutput.getModelId();
    String status = modelTrainingOutput.getStatus();
    assertNotNull(modelId);
    assertFalse(modelId.isEmpty());
    assertEquals("CREATED", status);
    return modelId;
}

Example 22 with MLTaskResponse

use of org.opensearch.ml.common.transport.MLTaskResponse 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));
}