use of org.opensearch.ml.common.dataset.MLInputDataset in project ml-commons by opensearch-project.
the class MLEngineTest method trainLinearRegressionModel.
private Model trainLinearRegressionModel() {
LinearRegressionParams parameters = LinearRegressionParams.builder().objectiveType(LinearRegressionParams.ObjectiveType.SQUARED_LOSS).optimizerType(LinearRegressionParams.OptimizerType.ADAM).learningRate(0.01).epsilon(1e-6).beta1(0.9).beta2(0.99).target("price").build();
DataFrame trainDataFrame = constructLinearRegressionTrainDataFrame();
MLInputDataset inputDataset = DataFrameInputDataset.builder().dataFrame(trainDataFrame).build();
Input mlInput = MLInput.builder().algorithm(FunctionName.LINEAR_REGRESSION).parameters(parameters).inputDataset(inputDataset).build();
return MLEngine.train(mlInput);
}
use of org.opensearch.ml.common.dataset.MLInputDataset in project ml-commons by opensearch-project.
the class MLEngineTest method predictKMeans.
@Test
public void predictKMeans() {
Model model = trainKMeansModel();
DataFrame predictionDataFrame = constructKMeansDataFrame(10);
MLInputDataset inputDataset = DataFrameInputDataset.builder().dataFrame(predictionDataFrame).build();
Input mlInput = MLInput.builder().algorithm(FunctionName.KMEANS).inputDataset(inputDataset).build();
MLPredictionOutput output = (MLPredictionOutput) MLEngine.predict(mlInput, model);
DataFrame predictions = output.getPredictionResult();
Assert.assertEquals(10, predictions.size());
predictions.forEach(row -> Assert.assertTrue(row.getValue(0).intValue() == 0 || row.getValue(0).intValue() == 1));
}
use of org.opensearch.ml.common.dataset.MLInputDataset in project ml-commons by opensearch-project.
the class RestMLTrainAndPredictIT method trainAndPredictKmeansWithParmas.
private void trainAndPredictKmeansWithParmas(KMeansParams params, Consumer<Map<Double, Integer>> function) throws IOException {
SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
sourceBuilder.query(new MatchAllQueryBuilder());
sourceBuilder.size(1000);
sourceBuilder.fetchSource(new String[] { "petal_length_in_cm", "petal_width_in_cm" }, null);
MLInputDataset inputData = SearchQueryInputDataset.builder().indices(ImmutableList.of(irisIndex)).searchSourceBuilder(sourceBuilder).build();
trainAndPredictKmeansWithIrisData(params, inputData, function);
}
use of org.opensearch.ml.common.dataset.MLInputDataset in project ml-commons by opensearch-project.
the class PredictionITTests method initTestingData.
@Before
public void initTestingData() throws ExecutionException, InterruptedException {
generateMLTestingData();
SearchSourceBuilder searchSourceBuilder = generateSearchSourceBuilder();
MLInputDataset inputDataset = new SearchQueryInputDataset(Collections.singletonList(TESTING_INDEX_NAME), searchSourceBuilder);
taskId = trainModel(inputDataset);
waitModelAvailable(taskId);
}
use of org.opensearch.ml.common.dataset.MLInputDataset in project ml-commons by opensearch-project.
the class MLCommonsRestTestCase method predict.
public void predict(RestClient client, FunctionName functionName, String modelId, String indexName, MLAlgoParams params, SearchSourceBuilder searchSourceBuilder, Consumer<Map<String, Object>> function) throws IOException {
MLInputDataset inputData = SearchQueryInputDataset.builder().indices(ImmutableList.of(indexName)).searchSourceBuilder(searchSourceBuilder).build();
MLInput kmeansInput = MLInput.builder().algorithm(functionName).parameters(params).inputDataset(inputData).build();
String endpoint = "/_plugins/_ml/_predict/" + functionName.name().toLowerCase(Locale.ROOT) + "/" + modelId;
Response response = TestHelper.makeRequest(client, "POST", endpoint, ImmutableMap.of(), TestHelper.toHttpEntity(kmeansInput), null);
verifyResponse(function, response);
}
Aggregations