modules/ml/src/main/java/org/apache/ignite/ml/knn/regression/KNNRegressionTrainer.java

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one or more
   3  * contributor license agreements.  See the NOTICE file distributed with
   4  * this work for additional information regarding copyright ownership.
   5  * The ASF licenses this file to You under the Apache License, Version 2.0
   6  * (the "License"); you may not use this file except in compliance with
   7  * the License.  You may obtain a copy of the License at
   8  *
   9  *      http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 package org.apache.ignite.ml.knn.regression;
  19
  20 import org.apache.ignite.ml.dataset.DatasetBuilder;
  21 import org.apache.ignite.ml.knn.KNNUtils;
  22 import org.apache.ignite.ml.math.Vector;
  23 import org.apache.ignite.ml.math.functions.IgniteBiFunction;
  24 import org.apache.ignite.ml.trainers.SingleLabelDatasetTrainer;
  25
  26 /**
  27  * kNN algorithm trainer to solve regression task.
  28  */
  29 public class KNNRegressionTrainer implements SingleLabelDatasetTrainer<KNNRegressionModel> {
  30     /**
  31      * Trains model based on the specified data.
  32      *
  33      * @param datasetBuilder Dataset builder.
  34      * @param featureExtractor Feature extractor.
  35      * @param lbExtractor Label extractor.
  36      * @return Model.
  37      */
  38     public <K, V> KNNRegressionModel fit(DatasetBuilder<K, V> datasetBuilder,
  39         IgniteBiFunction<K, V, Vector> featureExtractor, IgniteBiFunction<K, V, Double> lbExtractor) {
  40         return new KNNRegressionModel(KNNUtils.buildDataset(datasetBuilder, featureExtractor, lbExtractor));
  41     }
  42 }