Merge pull request #5 from kafkasl/master

Add rest of data pipeline

Author: r0mainK

src/__main__.py

@@ -0,0 +1,54 @@
+import argparse
+import sys
+
+from sourced.ml.cmd.args import add_repo2_args
+from sourced.ml.cmd import ArgumentDefaultsHelpFormatterNoNone
+from cmd.code2vec_extract_features import code2vec_extract_features
+
+
+def get_parser() -> argparse.ArgumentParser:
+    """
+    Creates the cmdline argument parser.
+    """
+    parser = argparse.ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatterNoNone)
+
+    # sourced.engine args
+
+    subparsers = parser.add_subparsers(help="Commands", dest="command")
+
+    extract_parser = subparsers.add_parser("extract",
+                                           help="Extract features from input repositories",
+                                           formatter_class=ArgumentDefaultsHelpFormatterNoNone)
+
+    extract_parser.set_defaults(handler=code2vec_extract_features)
+
+    add_repo2_args(extract_parser)
+
+    # code2vec specific args
+    extract_parser.add_argument('--max-length', type=int, default=5, help="Max path length.",
+                                required=False)
+    extract_parser.add_argument('--max-width', type=int, default=2, help="Max path width.",
+                                required=False)
+    extract_parser.add_argument('-o', '--output', type=str,
+                                help="Output path for the Code2VecFeatures model", required=True)
+    return parser
+
+
+def main():
+    parser = get_parser()
+
+    args = parser.parse_args()
+
+    try:
+        handler = args.handler
+    except AttributeError:
+        def print_usage(_):
+            parser.print_usage()
+
+        handler = print_usage
+
+    return handler(args)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/algorithms/uast_to_bag_paths.py

@@ -30,9 +30,6 @@ def __call__(self, uast):
         dict_of_paths = {str(path): val for path, val in Counter(path_contexts).items()}
         self._log.info("Extracted paths successfully")
 
-        from pprint import pprint
-        pprint(dict_of_paths)
-
         return dict_of_paths
 
     def _get_log_name(self):

src/cmd/__init__.py

@@ -1,49 +1,29 @@
 import logging
-import argparse
 from uuid import uuid4
 
 from extractors.paths import UastPathsBagExtractor
+from transformers.vocabulary2id import Vocabulary2Id
 from sourced.ml.transformers import UastDeserializer, Uast2BagFeatures, create_uast_source, \
-    UastRow2Document, Collector
+    UastRow2Document, Moder
 from sourced.ml.utils.engine import pipeline_graph, pause
-from sourced.ml.cmd.args import add_repo2_args
+
 
 @pause
-def code2vec(args):
+def code2vec_extract_features(args):
     log = logging.getLogger("code2vec")
     session_name = "code2vec-%s" % uuid4()
     root, start_point = create_uast_source(args, session_name)
 
     res = start_point \
+        .link(Moder("func")) \
         .link(UastRow2Document()) \
         .link(UastDeserializer()) \
         .link(Uast2BagFeatures([UastPathsBagExtractor(args.max_length, args.max_width)])) \
-        .link(Collector()) \
+        .link(Vocabulary2Id(root.session.sparkContext, args.output)) \
         .execute()
 
     # TODO: Add rest of data pipeline: extract distinct paths and terminal nodes for embedding mapping
     # TODO: Add transformer to write bags and vocabs to a model
     # TODO: Add ML pipeline
 
     pipeline_graph(args, log, root)
-
-
-def main():
-    parser = argparse.ArgumentParser()
-
-    # sourced.engine args
-    add_repo2_args(parser)
-
-    # code2vec specific args
-    parser.add_argument('-g', '--max_length', type=int, default=5, help="Max path length.",
-                        required=False)
-    parser.add_argument('-w', '--max_width', type=int, default=2, help="Max path width.",
-                        required=False)
-
-    args = parser.parse_args()
-
-    code2vec(args)
-
-
-if __name__ == '__main__':
-    main()

src/code2vec.py renamed to src/cmd/code2vec_extract_features.py

@@ -0,0 +1,106 @@
+from modelforge import register_model, Model
+from itertools import islice
+
+
+@register_model
+class Code2VecFeatures(Model):
+    """
+    Code2VecFeatures model - path contexts from source code.
+    """
+    NAME = "code2vec_features"
+
+    def construct(self, value2index, path2index, value2freq, path2freq, path_contexts):
+        self._value2index = value2index
+        self._path2index = path2index
+        self._value2freq = value2freq
+        self._path2freq = path2freq
+
+        self._path_contexts = path_contexts
+        return self
+
+    def _load_tree(self, tree):
+        self.construct(value2index=tree["value2index"],
+                       path2index=tree["path2index"],
+                       value2freq=tree["value2freq"],
+                       path2freq=tree["path2freq"],
+                       path_contexts=tree["path_contexts"])
+
+    @property
+    def value2index(self):
+        """
+        Dict mapping value -> ID.
+        """
+        return self._value2index
+
+    @property
+    def path2index(self):
+        """
+        Dict mapping path -> ID.
+        """
+        return self._path2index
+
+    @property
+    def value2freq(self):
+        """
+         Dict mapping value -> frequency.
+        """
+        return self._value2freq
+
+    @property
+    def path2freq(self):
+        """
+         Dict mapping path -> frequency.
+        """
+        return self._path2freq
+
+    @property
+    def path_contexts(self):
+        """
+        List with the processed source code identifiers.
+        """
+        return self._path_contexts
+
+    def value2index_items(self):
+        """
+        Returns the tuples belonging to value -> index mapping.
+        """
+        return self._value2index.items()
+
+    def path2index_items(self):
+        """
+        Returns the tuples belonging to path -> index mapping.
+        """
+        return self._path2index.items()
+
+    def value2freq_items(self):
+        """
+        Returns the tuples belonging to value -> freq mapping.
+        """
+        return self._value2freq.items()
+
+    def path2freq_items(self):
+        """
+        Returns the tuples belonging to path -> freq mapping.
+        """
+        return self._path2freq.items()
+
+    def _generate_tree(self):
+        return {"value2index": self._value2index,
+                "path2index": self._path2index,
+                "value2freq": self._value2freq,
+                "path2freq": self._path2freq,
+                "path_contexts": self._path_contexts}
+
+    def dump(self):
+        return "Number of values: %s\n" \
+               "Number of paths: %s\n" \
+               "First 10 value -> ID: %s\n" \
+               "First 10 path -> ID: %s\n" \
+               "First 10 value -> frequency: %s\n" \
+               "First 10 path -> frequency: %s" % \
+               (len(self._value2index_freq),
+                len(self.path2index_freq),
+                list(islice(self._value2index, 10)),
+                list(islice(self._path2index, 10)),
+                list(islice(self._value2freq, 10)),
+                list(islice(self._path2freq, 10)))

src/models/__init__.py

@@ -0,0 +1,95 @@
+import operator
+
+from pyspark import RDD, Row
+from models.code2vec_features import Code2VecFeatures
+
+from ast import literal_eval as make_tuple
+from sourced.ml.transformers import Transformer
+
+
+class Vocabulary2Id(Transformer):
+    def __init__(self, sc, output: str, **kwargs):
+        super().__init__(**kwargs)
+        self.output = output
+        self.sc = sc
+
+    def __call__(self, rows: RDD):
+        value2index, path2index, value2freq, path2freq = self.build_vocabularies(rows)
+
+        doc2path_contexts = self.build_doc2pc(value2index, path2index, rows)
+
+        doc2path_contexts = doc2path_contexts.collect()
+
+        Code2VecFeatures().construct(value2index=value2index,
+                                     path2index=path2index,
+                                     value2freq=value2freq,
+                                     path2freq=path2freq,
+                                     path_contexts=doc2path_contexts).save(
+            self.output)
+
+    @staticmethod
+    def _unstringify_path_context(row):
+        """
+        Takes a row containing ((pc, doc), freq) and returns a tuple (u, path, v)
+        (removes namespace prefix v.)
+        """
+        return make_tuple(row[0][0][2:])
+
+    def build_vocabularies(self, rows: RDD):
+        """
+        Process rows to gather values and paths with their frequencies.
+        :param rows: row structure is ((key, doc), val) where:
+            * key: str with the path context
+            * doc: file name
+            * val: number of occurrences of key in doc
+        """
+
+        def _flatten_row(row: Row):
+            # 2: removes the namespace v. from the string to parse it as tuple
+            k = Vocabulary2Id._unstringify_path_context(row)
+            return [(k[0], 1), (k[1], 1), (k[2], 1)]
+
+        rows = rows \
+            .flatMap(_flatten_row) \
+            .reduceByKey(operator.add) \
+            .persist()
+
+        values = rows.filter(lambda x: type(x[0]) == str).collect()
+        paths = rows.filter(lambda x: type(x[0]) == tuple).collect()
+
+        value2index = {w: id for id, (w, _) in enumerate(values)}
+        path2index = {w: id for id, (w, _) in enumerate(paths)}
+        value2freq = {w: freq for _, (w, freq) in enumerate(values)}
+        path2freq = {w: freq for _, (w, freq) in enumerate(paths)}
+
+        rows.unpersist()
+
+        return value2index, path2index, value2freq, path2freq
+
+    def build_doc2pc(self, value2index: dict, path2index: dict, rows: RDD):
+        """
+        Process rows and build elements (doc, [path_context_1, path_context_2, ...])
+        :param value2index_freq: value -> id
+        :param path2index_freq: path -> id
+        """
+
+        bc_value2index = self.sc.broadcast(value2index)
+        bc_path2index = self.sc.broadcast(path2index)
+
+        def _doc2pc(row: Row):
+            (u, path, v), doc = Vocabulary2Id._unstringify_path_context(row), row[0][1]
+
+            return doc, (bc_value2index.value[u], bc_path2index.value[path],
+                         bc_value2index.value[v])
+
+        rows = rows \
+            .map(_doc2pc) \
+            .distinct() \
+            .combineByKey(lambda value: [value],
+                          lambda x, value: x + [value],
+                          lambda x, y: x + y)
+
+        bc_value2index.unpersist(blocking=True)
+        bc_path2index.unpersist(blocking=True)
+
+        return rows