Merge pull request #42 from secure-software-engineering/feature/train-dataset-arff

Improve training dataset loading

Author: oshando

swan-pipeline/pom.xml

@@ -106,7 +106,7 @@
             <groupId>edu.stanford.nlp</groupId>
             <artifactId>stanford-corenlp</artifactId>
             <version>4.3.0</version>
-            <classifier>models</classifier>
+            <classifier>models-english</classifier>
         </dependency>
         <dependency>
             <groupId>org.jsoup</groupId>

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/SwanPipeline.java

@@ -26,6 +26,7 @@ public class SwanPipeline {
 
     private static final Logger logger = LoggerFactory.getLogger(SwanPipeline.class);
     public static SwanOptions options;
+    private ModelEvaluator modelEvaluator;
 
     public SwanPipeline(SwanOptions options) {
         SwanPipeline.options = options;
@@ -47,7 +48,10 @@ public void run() throws IOException, InterruptedException {
         // Load methods in training dataset
         Dataset dataset = new Dataset();
         dataset.setTrain(SrmListUtils.importFile(options.getDatasetJson()));
-        soot.cleanupList(dataset.getTrain());
+
+        if (!options.getTrainDataDir().isEmpty())
+            soot.cleanupList(dataset.getTrain());
+
         logger.info("Loaded {} training methods, distribution={}", dataset.getTrainMethods().size(), Util.countCategories(dataset.getTrainMethods()));
 
         //Load methods from the test set
@@ -61,10 +65,14 @@ public void run() throws IOException, InterruptedException {
         IFeatureSet featureSet = featureSetSelector.select(dataset, options);
 
         //Train and evaluate model for SRM and CWE categories
-        ModelEvaluator modelEvaluator = new ModelEvaluator(featureSet, options, dataset.getTestMethods());
+        modelEvaluator = new ModelEvaluator(featureSet, options, dataset.getTestMethods());
         modelEvaluator.trainModel();
 
         long analysisTime = System.currentTimeMillis() - startAnalysisTime;
         logger.info("Total runtime {} minutes", analysisTime / 60000);
     }
+
+    public ModelEvaluator getModelEvaluator() {
+        return modelEvaluator;
+    }
 }

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

@@ -2,6 +2,7 @@
 
 import picocli.CommandLine;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.Callable;
@@ -11,16 +12,19 @@
 public class CliRunner implements Callable<Integer> {
 
     @CommandLine.Option(names = {"-test", "--test-data"}, description = {"Path of test JARs or class files"})
-    private String testDataDir = "/input/test-data";
+    private String testDataDir = "";
 
     @CommandLine.Option(names = {"-train", "--train-data"}, description = {"Path of training JARs or class files"})
-    private String trainDataDir = "/input/train-data";
+    private String trainDataDir = "";
 
     @CommandLine.Option(names = {"-d", "--dataset"}, description = {"Path to JSON dataset file"})
-    private String datasetJson = "/input/dataset/swan-dataset.json";
+    private String datasetJson = "/dataset/swan-dataset.json";
+
+    @CommandLine.Option(names = {"-in", "--train-instances"}, description = {"Path to ARFF files that contain training instances"})
+    private List<String> instancesArff = new ArrayList<>();
 
     @CommandLine.Option(names = {"-o", "--output"}, description = {"Directory to save output files"})
-    private String outputDir = "/swan-output";
+    private String outputDir = "";
 
     @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");
@@ -52,9 +56,7 @@ public class CliRunner implements Callable<Integer> {
     @CommandLine.Option(names = {"-pt", "--prediction-threshold"}, description = {"Threshold for predicting categories"})
     private double predictionThreshold = 0.5;
 
-
-    @Override
-    public Integer call() throws Exception {
+    public SwanOptions initializeOptions(){
 
         SwanOptions options = new SwanOptions(testDataDir,
                 trainDataDir,
@@ -70,7 +72,14 @@ public Integer call() throws Exception {
                 split,
                 phase);
         options.setPredictionThreshold(predictionThreshold);
+        options.setInstances(instancesArff);
+
+        return options;
+    }
+
+    @Override
+    public Integer call() throws Exception {
 
-        return new SwanCli().run(options);
+        return new SwanCli().run(initializeOptions());
     }
 }

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

@@ -4,7 +4,9 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 import java.util.concurrent.CancellationException;
 
 /**
@@ -13,39 +15,44 @@
 public class SwanCli {
 
     private static final Logger logger = LoggerFactory.getLogger(SwanCli.class);
+    private SwanPipeline swanPipeline;
 
     public Integer run(SwanOptions options) throws Exception {
 
         FileUtility fileUtility = new FileUtility();
 
-        if (options.getDatasetJson().contentEquals("/input/dataset/swan-dataset.json")) {
+        if (options.getDatasetJson().contentEquals("/dataset/swan-dataset.json")) {
             options.setDatasetJson(fileUtility.getResourceFile(options.getDatasetJson()).getAbsolutePath());
         }
 
-        if (options.getTrainDataDir().contentEquals("/input/train-data")) {
-            options.setTrainDataDir(fileUtility.getResourceDirectory("/input/train-data").getAbsolutePath());
-        }
-
-        if (options.getTestDataDir().contentEquals("/input/test-data")) {
-            options.setTestDataDir(fileUtility.getResourceDirectory("/input/test-data").getAbsolutePath());
-        }
-
-        if(options.getSrmClasses().contains("all")){
+        if (options.getSrmClasses().contains("all")) {
             options.setSrmClasses(Arrays.asList("source", "sink", "sanitizer", "authentication"));
         }
 
-        if(options.getCweClasses().contains("all")){
+        if (options.getCweClasses().contains("all")) {
             options.setCweClasses(Arrays.asList("cwe078", "cwe079", "cwe089", "cwe306", "cwe601", "cwe862", "cwe863"));
         }
 
-        if(options.getFeatureSet().contains("all")){
+        if (options.getFeatureSet().contains("all")) {
             options.setFeatureSet(Arrays.asList("code", "doc-manual", "doc-auto"));
         }
 
+        if (options.getInstances().isEmpty()) {
+
+            List<String> instances = new ArrayList<>();
+
+            for (String feature : options.getFeatureSet()){
+                String filepath = "/dataset/" + options.getToolkit() + "-" + feature + "-instances.arff";
+                instances.add(fileUtility.getResourceFile(filepath).getAbsolutePath());
+            }
+
+            options.setInstances(instances);
+        }
+
         logger.info("SWAN options: {}", options);
 
         try {
-            SwanPipeline swanPipeline = new SwanPipeline(options);
+            swanPipeline = new SwanPipeline(options);
             swanPipeline.run();
 
             return 0;
@@ -60,4 +67,8 @@ public Integer run(SwanOptions options) throws Exception {
             fileUtility.dispose();
         }
     }
+
+    public SwanPipeline getSwanPipeline() {
+        return swanPipeline;
+    }
 }

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

@@ -24,6 +24,7 @@ public class SwanOptions {
     private double trainTestSplit;
     private String phase;
     private double predictionThreshold;
+    private List<String> instancesArff;
 
     public SwanOptions(String testDataDir, String trainDataDir, String datasetJson, String outputDir,
                        List<String> featureSet, String toolkit, List<String> srmClasses,
@@ -177,12 +178,21 @@ public void setPredictionThreshold(double predictionThreshold) {
         this.predictionThreshold = predictionThreshold;
     }
 
+    public List<String> getInstances() {
+        return instancesArff;
+    }
+
+    public void setInstances(List<String> instancesArff) {
+        this.instancesArff = instancesArff;
+    }
+
     @Override
     public String toString() {
         return "SwanOptions{" +
                 "testData='" + testDataDir + '\'' +
                 ", trainData='" + trainDataDir + '\'' +
                 ", datasetJson='" + datasetJson + '\'' +
+                ", instances='" + instancesArff + '\'' +
                 ", outputDir='" + outputDir + '\'' +
                 ", featureSet='" + featureSet + '\'' +
                 ", learningMode='" + toolkit + '\'' +

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/features/FeatureSet.java

@@ -188,10 +188,9 @@ public ArrayList<Attribute> addDocAttributes(FeatureSet.Type instanceSet) {
         return attributes;
     }
 
-    public Instances createInstances(List<Type> featureSets, ArrayList<Attribute> attributes,
-                                     Set<Method> methods, Set<Category> categories, String name) {
 
-        Instances instances = new Instances(name, attributes, 0);
+    public Instances createInstances(Instances instances, List<Type> featureSets, ArrayList<Attribute> attributes,
+                                     Set<Method> methods, Set<Category> categories) {
 
         for (FeatureSet.Type featureSet : featureSets)
             switch (featureSet) {
@@ -207,6 +206,15 @@ public Instances createInstances(List<Type> featureSets, ArrayList<Attribute> at
     }
 
 
+    public Instances createInstances(List<Type> featureSets, ArrayList<Attribute> attributes,
+                                     Set<Method> methods, Set<Category> categories) {
+
+        Instances instances = new Instances("swan-srm", attributes, 0);
+
+        return createInstances(instances, featureSets, attributes, methods, categories);
+    }
+
+
     /**
      * Adds data for SWAN features to instance set.
      *
@@ -256,13 +264,13 @@ public ArrayList<Instance> getCodeInstances(Instances instances, Set<Method> met
 
                 switch (entry.getKey().applies(method)) {
                     case TRUE:
-                        inst.setValue(entry.getValue(), "true");
+                        inst.setValue(instances.attribute(String.valueOf(entry.getKey())), "true");
                         break;
                     case FALSE:
-                        inst.setValue(entry.getValue(), "false");
+                        inst.setValue(instances.attribute(String.valueOf(entry.getKey())), "false");
                         break;
                     default:
-                        inst.setMissing(entry.getValue());
+                        inst.setMissing(instances.attribute(String.valueOf(entry.getKey())));
                 }
             }
 

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/features/MekaFeatureSet.java

@@ -9,9 +9,13 @@
 import meka.filters.unsupervised.attribute.MekaClassAttributes;
 import weka.core.Attribute;
 import weka.core.Instances;
+import weka.core.converters.ArffLoader;
 import weka.filters.Filter;
 
+import java.io.File;
+import java.io.IOException;
 import java.util.*;
+import java.util.stream.Collectors;
 
 public class MekaFeatureSet extends FeatureSet implements IFeatureSet {
 
@@ -26,14 +30,34 @@ public void createFeatures() {
 
         List<FeatureSet.Type> featureSets = initializeFeatures();
 
+        Instances trainInstances = null;
+        Instances structure = null;
+
         //Create and set attributes for the train instances
-        ArrayList<Attribute> trainAttributes = createAttributes(getCategories(options.getAllClasses()), dataset.getTrainMethods(), featureSets);
-        Instances trainInstances = createInstances(featureSets, trainAttributes, dataset.getTrainMethods(), getCategories(options.getAllClasses()), "train-instances");
-        this.instances.put("train", convertToMekaInstances(trainInstances));
+        if (options.getInstances().isEmpty()) {
+            ArrayList<Attribute> trainAttributes = createAttributes(getCategories(options.getAllClasses()), dataset.getTrainMethods(), featureSets);
+            trainInstances = createInstances(featureSets, trainAttributes, dataset.getTrainMethods(), getCategories(options.getAllClasses()));
+        } else {
+            ArffLoader loader = new ArffLoader();
+
+            try {
+                loader.setSource(new File(options.getInstances().get(0)));
+                trainInstances = loader.getDataSet();
+                structure = loader.getStructure();
+            } catch (IOException e) {
+                e.printStackTrace();
+            }
+        }
 
         //Create and set attributes for the test instances.
+        Attribute idAttr = new Attribute("id", dataset.getTestMethods().stream().map(Method::getArffSafeSignature).collect(Collectors.toList()));
+        structure.replaceAttributeAt(idAttr, structure.attribute("id").index());
+        ArrayList<Attribute> aList = Collections.list(structure.enumerateAttributes());
+
         ArrayList<Attribute> testAttributes = createAttributes(getCategories(options.getAllClasses()), dataset.getTestMethods(), featureSets);
-        Instances testInstances = createInstances(featureSets, testAttributes, dataset.getTestMethods(), getCategories(options.getAllClasses()), "test-instances");
+        Instances testInstances = (createInstances(structure, featureSets, aList, dataset.getTestMethods(), getCategories(options.getAllClasses())));
+
+        this.instances.put("train", convertToMekaInstances(trainInstances));
         this.instances.put("test", convertToMekaInstances(testInstances));
     }
 
@@ -67,7 +91,7 @@ public Instances convertToMekaInstances(Instances instances) {
             filter.setAttributeIndices("1-11");
             filter.setInputFormat(instances);
             output = Filter.useFilter(instances, filter);
-            output.setRelationName(instances.relationName() + ":" + output.relationName());
+            output.setRelationName("swan-srm:" + output.relationName());
 
             Util.exportInstancesToArff(output);
         } catch (Exception e) {
@@ -85,5 +109,4 @@ public HashSet<Category> getCategories(List<String> cat) {
 
         return categories;
     }
-
 }

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/features/WekaFeatureSet.java

@@ -34,7 +34,7 @@ public void createFeatures() {
             ArrayList<Attribute> trainAttributes = createAttributes(category, dataset.getTrainMethods(), featureSets);
 
             String instanceName = category.getId().toLowerCase() + "-train-instances";
-            Instances trainInstances = createInstances(featureSets, trainAttributes, dataset.getTrainMethods(), Collections.singleton(category), instanceName);
+            Instances trainInstances = createInstances(featureSets, trainAttributes, dataset.getTrainMethods(), Collections.singleton(category));
             this.instances.put(category.getId().toLowerCase(), trainInstances);
             Util.exportInstancesToArff(trainInstances);
 

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/io/dataset/SrmList.java

@@ -31,7 +31,6 @@ public SrmList(String sourceFileDir) {
         methods = new HashSet<>();
     }
 
-
     public void load(final Set<Method> trainingSet) {
 
         Util.createSubclassAnnotations(methods, "classpath");

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/io/dataset/SrmListUtils.java

@@ -43,10 +43,9 @@ public static SrmList importFile(String file) throws IOException {
      */
     public static void exportFile(SrmList srmList, String file) throws IOException {
 
-        srmList.removeUnclassifiedMethods();
         ObjectMapper objectMapper = new ObjectMapper();
         objectMapper.writeValue(new File(file), srmList);
-        logger.info("{} methods exported to {}", srmList.getMethods().size(), file);
+        logger.info("{} SRMs exported to {}", srmList.getMethods().size(), file);
     }
 
     public static void removeUndocumentedMethods(SrmList list) {