Add paramete to ski dataset colelction

cqfn · May 15, 2020 · 8a00f49 · 8a00f49
1 parent b3829ea
commit 8a00f49
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/aibolit/__main__.py b/aibolit/__main__.py
@@ -105,8 +105,15 @@ def train():
         required=False,
         default=None
     )
+    parser.add_argument(
+        '--skip_collect_dataset',
+        required=False,
+        default=False,
+        action='store_true'
+    )
     args = parser.parse_args(sys.argv[2:])
-    collect_dataset(args.java_folder, args.max_classes)
+    if not args.skip_collect_dataset:
+        collect_dataset(args.java_folder, args.max_classes)
     train_process()