Examples with DataFrameInputDataset org.opensearch.ml.common.dataset.DataFrameInputDataset used on opensource projects

Example 11 with DataFrameInputDataset

use of org.opensearch.ml.common.dataset.DataFrameInputDataset in project ml-commons by opensearch-project.

the class MachineLearningClientTest method train.

@Test
public void train() {
    MLInput mlInput = MLInput.builder().algorithm(FunctionName.KMEANS).parameters(mlParameters).inputDataset(new DataFrameInputDataset(input)).build();
    assertEquals(modekId, ((MLTrainingOutput) machineLearningClient.train(mlInput, false).actionGet()).getModelId());
}

Example 12 with DataFrameInputDataset

use of org.opensearch.ml.common.dataset.DataFrameInputDataset in project ml-commons by opensearch-project.

the class MLPredictionTaskRequestTest method writeTo_Success.

@Test
public void writeTo_Success() throws IOException {
    MLPredictionTaskRequest request = MLPredictionTaskRequest.builder().mlInput(mlInput).build();
    BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
    request.writeTo(bytesStreamOutput);
    request = new MLPredictionTaskRequest(bytesStreamOutput.bytes().streamInput());
    assertEquals(FunctionName.KMEANS, request.getMlInput().getAlgorithm());
    KMeansParams params = (KMeansParams) request.getMlInput().getParameters();
    assertEquals(1, params.getCentroids().intValue());
    MLInputDataset inputDataset = request.getMlInput().getInputDataset();
    assertEquals(MLInputDataType.DATA_FRAME, inputDataset.getInputDataType());
    DataFrame dataFrame = ((DataFrameInputDataset) inputDataset).getDataFrame();
    assertEquals(1, dataFrame.size());
    assertEquals(1, dataFrame.columnMetas().length);
    assertEquals("key1", dataFrame.columnMetas()[0].getName());
    assertEquals(ColumnType.DOUBLE, dataFrame.columnMetas()[0].getColumnType());
    assertEquals(1, dataFrame.getRow(0).size());
    assertEquals(2.00, dataFrame.getRow(0).getValue(0).getValue());
    assertNull(request.getModelId());
}

Example 13 with DataFrameInputDataset

use of org.opensearch.ml.common.dataset.DataFrameInputDataset in project ml-commons by opensearch-project.

the class MLTrainingTaskRunner method startTrainingTask.

/**
 * Start training task
 * @param mlTask ML task
 * @param mlInput ML input
 * @param listener Action listener
 */
public void startTrainingTask(MLTask mlTask, MLInput mlInput, ActionListener<MLTaskResponse> listener) {
    ActionListener<MLTaskResponse> internalListener = wrappedCleanupListener(listener, mlTask.getTaskId());
    // track ML task count and add ML task into cache
    mlStats.getStat(ML_EXECUTING_TASK_COUNT).increment();
    mlStats.getStat(ML_TOTAL_REQUEST_COUNT).increment();
    mlStats.createCounterStatIfAbsent(requestCountStat(mlTask.getFunctionName(), ActionName.TRAIN)).increment();
    mlTaskManager.add(mlTask);
    try {
        if (mlInput.getInputDataset().getInputDataType().equals(MLInputDataType.SEARCH_QUERY)) {
            ActionListener<DataFrame> dataFrameActionListener = ActionListener.wrap(dataFrame -> {
                train(mlTask, mlInput.toBuilder().inputDataset(new DataFrameInputDataset(dataFrame)).build(), internalListener);
            }, e -> {
                log.error("Failed to generate DataFrame from search query", e);
                internalListener.onFailure(e);
            });
            mlInputDatasetHandler.parseSearchQueryInput(mlInput.getInputDataset(), new ThreadedActionListener<>(log, threadPool, TASK_THREAD_POOL, dataFrameActionListener, false));
        } else {
            threadPool.executor(TASK_THREAD_POOL).execute(() -> {
                train(mlTask, mlInput, internalListener);
            });
        }
    } catch (Exception e) {
        log.error("Failed to train " + mlInput.getAlgorithm(), e);
        internalListener.onFailure(e);
    }
}

Example 14 with DataFrameInputDataset

use of org.opensearch.ml.common.dataset.DataFrameInputDataset in project ml-commons by opensearch-project.

the class MLEngineTest method trainAndPredictWithKmeans.

@Test
public void trainAndPredictWithKmeans() {
    int dataSize = 100;
    MLAlgoParams parameters = KMeansParams.builder().build();
    DataFrame dataFrame = constructKMeansDataFrame(dataSize);
    MLInputDataset inputData = new DataFrameInputDataset(dataFrame);
    Input input = new MLInput(FunctionName.KMEANS, parameters, inputData);
    MLPredictionOutput output = (MLPredictionOutput) MLEngine.trainAndPredict(input);
    Assert.assertEquals(dataSize, output.getPredictionResult().size());
}