com.datastax.insight.ml.spark.mllib.classification.SVM Maven / Gradle / Ivy

Go to download
package com.datastax.insight.ml.spark.mllib.classification;

import com.datastax.insight.spec.RDDOperator;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.mllib.classification.SVMModel;
import org.apache.spark.mllib.classification.SVMWithSGD;
import org.apache.spark.mllib.regression.LabeledPoint;
import scala.Tuple2;

public class SVM implements RDDOperator {
    public static SVMModel train(JavaRDD data,int numIterations){
        SVMModel model=SVMWithSGD.train(data.rdd(), numIterations);

        // Clear the default threshold.
        model.clearThreshold();

        return model;
    }

    public static JavaRDD> predict(JavaRDD data,SVMModel model){
        JavaRDD> scoreAndLabels = data.map(
                new Function>() {
                    public Tuple2 call(LabeledPoint p) {
                        Double score = model.predict(p.features());
                        return new Tuple2(score, p.label());
                    }
                }
        );
        return scoreAndLabels;
    }
}