Examples with KMeansTrainer org.tribuo.clustering.kmeans.KMeansTrainer used on opensource projects

Example 1 with KMeansTrainer

use of org.tribuo.clustering.kmeans.KMeansTrainer in project ml-commons by opensearch-project.

the class KMeans method trainAndPredict.

@Override
public MLOutput trainAndPredict(DataFrame dataFrame) {
    MutableDataset<ClusterID> trainDataset = TribuoUtil.generateDataset(dataFrame, new ClusteringFactory(), "KMeans training and predicting data from opensearch", TribuoOutputType.CLUSTERID);
    Integer centroids = Optional.ofNullable(parameters.getCentroids()).orElse(DEFAULT_CENTROIDS);
    Integer iterations = Optional.ofNullable(parameters.getIterations()).orElse(DEFAULT_ITERATIONS);
    KMeansTrainer trainer = new KMeansTrainer(centroids, iterations, distance, numThreads, seed);
    // won't store model in index
    KMeansModel kMeansModel = trainer.train(trainDataset);
    List<Prediction<ClusterID>> predictions = kMeansModel.predict(trainDataset);
    List<Map<String, Object>> listClusterID = new ArrayList<>();
    predictions.forEach(e -> listClusterID.add(Collections.singletonMap("ClusterID", e.getOutput().getID())));
    return MLPredictionOutput.builder().predictionResult(DataFrameBuilder.load(listClusterID)).build();
}

Example 2 with KMeansTrainer

use of org.tribuo.clustering.kmeans.KMeansTrainer in project ml-commons by opensearch-project.

the class KMeans method train.

@Override
public Model train(DataFrame dataFrame) {
    MutableDataset<ClusterID> trainDataset = TribuoUtil.generateDataset(dataFrame, new ClusteringFactory(), "KMeans training data from opensearch", TribuoOutputType.CLUSTERID);
    Integer centroids = Optional.ofNullable(parameters.getCentroids()).orElse(DEFAULT_CENTROIDS);
    Integer iterations = Optional.ofNullable(parameters.getIterations()).orElse(DEFAULT_ITERATIONS);
    KMeansTrainer trainer = new KMeansTrainer(centroids, iterations, distance, numThreads, seed);
    KMeansModel kMeansModel = trainer.train(trainDataset);
    Model model = new Model();
    model.setName(FunctionName.KMEANS.name());
    model.setVersion(1);
    model.setContent(ModelSerDeSer.serialize(kMeansModel));
    return model;
}