Add new and update existing command options

Author: oshando

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/cli/CliRunner.java

@@ -25,8 +25,8 @@ public class CliRunner implements Callable<Integer> {
     @CommandLine.Option(names = {"-f", "--feature"}, description = {"Select one or more feature sets: all, code, doc-auto or doc-manual"})
     private List<String> featureSet =  Collections.singletonList("code");
 
-    @CommandLine.Option(names = {"-l", "--learning"}, description = {"Learning modes: manual, auto"})
-    private String learningMode = "manual";
+    @CommandLine.Option(names = {"-l", "--learning"}, description = {"Learning modes: meka, weka, ml-plan"})
+    private String learningMode = "meka";
 
     @CommandLine.Option(names = {"-s", "--srm"}, description = {"SRM: all, source, sink, sanitizer, authentication, relevant"})
     private List<String> srmClasses = Collections.singletonList("all");
@@ -46,6 +46,13 @@ public class CliRunner implements Callable<Integer> {
     @CommandLine.Option(names = {"-sp", "--training-split"}, description = {"Percentage for training"})
     private double split = 0.7;
 
+    @CommandLine.Option(names = {"-p", "--phase"}, description = {"Phase: validate, predict"})
+    private String phase = "predict";
+
+    @CommandLine.Option(names = {"-pt", "--prediction-threshold"}, description = {"Threshold for predicting categories"})
+    private double predictionThreshold = 0.5;
+
+
     @Override
     public Integer call() throws Exception {
 
@@ -60,7 +67,9 @@ public Integer call() throws Exception {
                 exportArffData,
                 isDocumented,
                 iterations,
-                split);
+                split,
+                phase);
+        options.setPredictionThreshold(predictionThreshold);
 
         return new SwanCli().run(options);
     }

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/cli/SwanCli.java

@@ -31,7 +31,7 @@ public Integer run(SwanOptions options) throws Exception {
         }
 
         if(options.getSrmClasses().contains("all")){
-            options.setSrmClasses(Arrays.asList("source", "sink", "sanitizer", "authentication", "relevant"));
+            options.setSrmClasses(Arrays.asList("source", "sink", "sanitizer", "authentication"));
         }
 
         if(options.getCweClasses().contains("all")){

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/cli/SwanOptions.java

@@ -22,6 +22,27 @@ public class SwanOptions {
     private boolean isDocumented;
     private int iterations;
     private double trainTestSplit;
+    private String phase;
+    private double predictionThreshold;
+
+    public SwanOptions(String testDataDir, String trainDataDir, String datasetJson, String outputDir,
+                       List<String> featureSet, String learningMode, List<String> srmClasses,
+                       List<String> cweClasses, boolean exportArffData, boolean isDocumented,
+                       int iterations, double trainTestSplit, String phase) {
+        this.testDataDir = testDataDir;
+        this.trainDataDir = trainDataDir;
+        this.datasetJson = datasetJson;
+        this.outputDir = outputDir;
+        this.featureSet = featureSet;
+        this.learningMode = learningMode;
+        this.srmClasses = srmClasses;
+        this.cweClasses = cweClasses;
+        this.exportArffData = exportArffData;
+        this.isDocumented = isDocumented;
+        this.iterations = iterations;
+        this.trainTestSplit = trainTestSplit;
+        this.phase = phase;
+    }
 
     public SwanOptions(String testDataDir, String trainDataDir, String datasetJson, String outputDir,
                        List<String> featureSet, String learningMode, List<String> srmClasses, List<String> cweClasses,
@@ -140,6 +161,22 @@ public void setTrainTestSplit(double trainTestSplit) {
         this.trainTestSplit = trainTestSplit;
     }
 
+    public String getPhase() {
+        return phase;
+    }
+
+    public void setPhase(String phase) {
+        this.phase = phase;
+    }
+
+    public double getPredictionThreshold() {
+        return predictionThreshold;
+    }
+
+    public void setPredictionThreshold(double predictionThreshold) {
+        this.predictionThreshold = predictionThreshold;
+    }
+
     @Override
     public String toString() {
         return "SwanOptions{" +

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/model/ModelEvaluator.java

@@ -31,15 +31,21 @@
 public class ModelEvaluator {
 
     public enum Mode {
-        MANUAL,
-        AUTOMATIC
+        WEKA,
+        MEKA,
+        MLPLAN
+    }
+
+    public enum Phase {
+        VALIDATE,
+        PREDICT
     }
 
     private FeaturesHandler features;
     private SwanOptions options;
     private static final Logger logger = LoggerFactory.getLogger(ModelEvaluator.class);
 
-    public ModelEvaluator(FeaturesHandler features, SwanOptions options) {
+    public ModelEvaluator(FeaturesHandler features, SwanOptions options, Set<Method> methods) {
         this.features = features;
         this.options = options;
     }
@@ -49,8 +55,7 @@ public ModelEvaluator(FeaturesHandler features, SwanOptions options) {
      *
      * @return Hashmap containing the name of the classifier and it's F-Measure
      */
-    public HashMap<String, HashMap<String, String>> trainModel() {
-
+    public void trainModel() {
 
         switch (Mode.valueOf(options.getLearningMode().toUpperCase())) {