Add '--class_name' and '--method_name' as new arguments

nok · Oct 8, 2017 · 6f2a1d9 · 6f2a1d9
1 parent b42bad2
commit 6f2a1d9
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/readme.md b/readme.md
@@ -239,8 +239,8 @@ joblib.dump(clf, 'estimator.pkl')
 After that the model can be transpiled by using the following command:
 
 ```
-python -m sklearn_porter --input <PICKLE_FILE> [--output <DEST_DIR>] [--pipe] [--c] [--java] [--js] [--go] [--php] [--ruby]
-python -m sklearn_porter -i <PICKLE_FILE> [-o <DEST_DIR>] [-p] [--c] [--java] [--js] [--go] [--php] [--ruby]
+python -m sklearn_porter --input <PICKLE_FILE> [--output <DEST_DIR>] [--class_name <NAME>] [--method_name <NAME>] [--pipe] [--c] [--java] [--js] [--go] [--php] [--ruby]
+python -m sklearn_porter -i <PICKLE_FILE> [-o <DEST_DIR>] [--class_name <NAME>] [--method_name <NAME>] [-p] [--c] [--java] [--js] [--go] [--php] [--ruby]
 ```
 
 For instance the following command transpiles the estimator to the target programming language JavaScript:

diff --git a/sklearn_porter/__main__.py b/sklearn_porter/__main__.py
@@ -19,15 +19,25 @@ def parse_args(args):
         '--input', '-i',
         required=True,
         help=(
-            'Set the path of an exported estimator '
+            'Path to an exported estimator '
             'in pickle (.pkl) format.'))
     parser.add_argument(
         '--output', '-o',
         required=False,
         help=(
-            'Set the destination directory, '
+            'Path to the destination directory '
             'where the transpiled estimator will be '
             'stored.'))
+    parser.add_argument(
+        '--class_name',
+        default='Brain',
+        required=False,
+        help='Define the class name in the final output.')
+    parser.add_argument(
+        '--method_name',
+        default='predict',
+        required=False,
+        help='Define the method name in the final output.')
     parser.add_argument(
         '--pipe', '-p',
         required=False,
@@ -81,7 +91,11 @@ def main():
     # Port estimator:
     try:
         porter = Porter(estimator, language=language)
-        output = porter.export(details=True)
+        class_name = str(args.get('class_name'))
+        method_name = str(args.get('method_name'))
+        output = porter.export(class_name=class_name,
+                               method_name=method_name,
+                               details=True)
     except Exception as e:
         sys.exit('Error: {}'.format(str(e)))
     else: