Add main and rework data pipeline

Signed-off-by: Pol Alvarez Vecino <[email protected]>

Author: kafkasl

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 = Counter(path_contexts)
         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()