Feature/cli (#174)

AHReccese · web-flow · commit 53d39acd4fdc · 2025-01-11T02:50:44.000+03:30
* add art

* add __main__ (for CLI) to `setup.py`

* add overview to `pymilo_param.py`

* add url regex

* add sklearn supported categories dictionary

* add `print_supported_ml_models` for CLI

* add `pymilo_help` for CLI

* develop `get_sklearn_class` function

* implement CLI for pymilo(+pymilo server)

* apply `autopep8.sh`

* `CHANGELOG.md` updated

* add docstring to `main` function

* update docstring in `get_sklearn_class` function

* apply minor refactoring

* refactor to run `communicator.run()` once

* `CHANGELOG.md` updated

* apply refactoring to `if-else`'s
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,12 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 ## [Unreleased]
 ### Added
+- `CLI` handler
+- `print_supported_ml_models` function in `pymilo_func.py`
+- `pymilo_help` function in `pymilo_func.py`
+- `SKLEARN_SUPPORTED_CATEGORIES` in `pymilo_param.py`
+- `OVERVIEW` in `pymilo_param.py`
+- `get_sklearn_class` in `utils.util.py`
 ### Changed
 - `to_pymilo_issue` function in `PymiloException`
 - `valid_url_valid_file` testcase added in `test_exceptions.py`
diff --git a/pymilo/__main__.py b/pymilo/__main__.py
@@ -0,0 +1,118 @@
+# -*- coding: utf-8 -*-
+"""PyMilo main."""
+import re
+import argparse
+from art import tprint
+from .pymilo_param import PYMILO_VERSION, URL_REGEX
+from .pymilo_func import print_supported_ml_models, pymilo_help
+from .pymilo_obj import Import
+from .utils.util import get_sklearn_class
+
+ml_streaming_support = True
+try:
+    from .streaming import PymiloServer, Compression, CommunicationProtocol
+except BaseException:
+    ml_streaming_support = False
+
+
+def main():
+    """
+    CLI main function.
+
+    :return: None
+    """
+    parser = argparse.ArgumentParser(description='Run the Pymilo server with a specified compression method.')
+    parser.add_argument(
+        '--compression',
+        type=str,
+        choices=['NULL', 'GZIP', 'ZLIB', 'LZMA', 'BZ2'],
+        default='NULL',
+        help='Specify the compression method (NULL, GZIP, ZLIB, LZMA, or BZ2). Default is NULL.'
+    )
+    parser.add_argument(
+        '--port',
+        type=int,
+        default=8000,
+        help='Specify PyMiloServer port number',
+        metavar="",
+    )
+    parser.add_argument(
+        '--protocol',
+        type=str,
+        choices=['REST', 'WEBSOCKET'],
+        default='REST',
+        help='Specify the communication protocol (REST or WEBSOCKET). Default is REST.'
+    )
+    parser.add_argument(
+        '--load',
+        type=str,
+        default=None,
+        help='the `load` command specifies the path to the JSON file of the previously exported ML model by PyMilo.',
+        metavar="",
+    )
+    parser.add_argument(
+        '--init',
+        type=str,
+        default=None,
+        help='the `init` command specifies the ML model to initialize the PyMilo Server with.',
+        metavar="",
+    )
+    parser.add_argument(
+        '--bare',
+        default=False,
+        action='store_true',
+        help='The `bare` command starts the PyMilo Server without an internal ML model.',
+    )
+    parser.add_argument('--version', action='store_true', default=False, help='PyMilo version')
+    parser.add_argument('-v', action='store_true', default=False, help='PyMilo version')
+    args = parser.parse_args()
+    if args.version or args.v:
+        print(PYMILO_VERSION)
+        return
+    if not ml_streaming_support:
+        print("Error: ML Streaming is not installed.")
+        print("To install ML Streaming, run the following command:")
+        print("pip install pymilo[streaming]")
+        print("For more information, visit the PyMilo README at https://github.com/openscilab/pymilo")
+        tprint("PyMilo")
+        tprint("V:" + PYMILO_VERSION)
+        pymilo_help()
+        parser.print_help()
+        return
+    run_ps = False
+    _model = None
+    _port = args.port
+    _compressor = Compression[args.compression]
+    _communication_protocol = CommunicationProtocol[args.protocol]
+    if args.load:
+        path = args.load
+        run_ps = True
+        _model = Import(url=path) if re.match(URL_REGEX, path) else Import(file_adr=path)
+    elif args.init:
+        model_name = args.init
+        model_class = get_sklearn_class(model_name)
+        if model_class is None:
+            print(
+                "The given ML model name is neither valid nor supported, use the list below: \n{print_supported_ml_models}")
+            print_supported_ml_models()
+            return
+        run_ps = True
+        _model = model_class()
+    elif args.bare:
+        run_ps = True
+        _model = model_class()
+    if not run_ps:
+        tprint("PyMilo")
+        tprint("V:" + PYMILO_VERSION)
+        pymilo_help()
+        parser.print_help()
+    else:
+        PymiloServer(
+            model=_model,
+            port=_port,
+            compressor=_compressor,
+            communication_protocol=_communication_protocol,
+        ).communicator.run()
+
+if __name__ == '__main__':
+    main()
diff --git a/pymilo/pymilo_func.py b/pymilo/pymilo_func.py
@@ -5,6 +5,7 @@
 
 from .chains.ensemble_chain import get_transporter
 from .transporters.transporter import Command
+from .pymilo_param import SKLEARN_SUPPORTED_CATEGORIES, NOT_SUPPORTED, OVERVIEW
 
 
 def get_sklearn_version():
@@ -62,3 +63,28 @@ def compare_model_outputs(exported_output,
             return False  # TODO: throw exception
         total_error += np.abs(imported_output[key] - exported_output[key])
     return np.abs(total_error) < epsilon_error
+
+
+def print_supported_ml_models():
+    """
+    Print the supported sklearn ML models categorized by type.
+
+    :return: None
+    """
+    print("Supported Machine Learning Models:")
+    for category, table in SKLEARN_SUPPORTED_CATEGORIES.items():
+        print(f"**{category}**:")
+        for model_name in table:
+            if table[model_name] != NOT_SUPPORTED:
+                print(f"- {model_name}")
+
+
+def pymilo_help():
+    """
+    Print PyMilo details.
+
+    :return: None
+    """
+    print(OVERVIEW)
+    print("Repo : https://github.com/openscilab/pymilo")
+    print("Webpage : https://openscilab.com/\n")
diff --git a/pymilo/pymilo_param.py b/pymilo/pymilo_param.py
@@ -82,13 +82,16 @@
 except BaseException:
     pass
 
-
+OVERVIEW = """
+PyMilo is an open source Python package that provides a simple, efficient, and safe way for users to export pre-trained machine learning models in a transparent way.
+"""
 PYMILO_VERSION = "1.1"
 NOT_SUPPORTED = "NOT_SUPPORTED"
 PYMILO_VERSION_DOES_NOT_EXIST = "Corrupted JSON file, `pymilo_version` doesn't exist in this file."
 UNEQUAL_PYMILO_VERSIONS = "warning: Installed PyMilo version differs from the PyMilo version used to create the JSON file."
 UNEQUAL_SKLEARN_VERSIONS = "warning: Installed Scikit version differs from the Scikit version used to create the JSON file and it may prevent PyMilo from transporting seamlessly."
 INVALID_IMPORT_INIT_PARAMS = "Invalid input parameters, you should either pass a valid file_adr or a json_dump or a url to initiate Import class."
+URL_REGEX = r'^(http|https)://[a-zA-Z0-9.-_]+\.[a-zA-Z]{2,}(/\S*)?$'
 DOWNLOAD_MODEL_FAILED = "Failed to download the JSON file, Server didn't respond."
 INVALID_DOWNLOADED_MODEL = "The downloaded content is not a valid JSON file."
 BATCH_IMPORT_INVALID_DIRECTORY = "The given directory does not exist."
@@ -276,3 +279,15 @@
     "ENSEMBLE": "exported_ensembles",
     "CROSS_DECOMPOSITION": "exported_cross_decomposition",
 }
+
+SKLEARN_SUPPORTED_CATEGORIES = {
+    "LINEAR_MODEL": SKLEARN_LINEAR_MODEL_TABLE,
+    "NEURAL_NETWORK": SKLEARN_NEURAL_NETWORK_TABLE,
+    "DECISION_TREE": SKLEARN_DECISION_TREE_TABLE,
+    "CLUSTERING": SKLEARN_CLUSTERING_TABLE,
+    "NAIVE_BAYES": SKLEARN_NAIVE_BAYES_TABLE,
+    "SVM": SKLEARN_SVM_TABLE,
+    "NEIGHBORS": SKLEARN_NEIGHBORS_TABLE,
+    "ENSEMBLE": SKLEARN_ENSEMBLE_TABLE,
+    "CROSS_DECOMPOSITION": SKLEARN_CROSS_DECOMPOSITION_TABLE,
+}
diff --git a/pymilo/utils/util.py b/pymilo/utils/util.py
@@ -3,7 +3,7 @@
 import requests
 import importlib
 from inspect import signature
-from ..pymilo_param import DOWNLOAD_MODEL_FAILED, INVALID_DOWNLOADED_MODEL
+from ..pymilo_param import DOWNLOAD_MODEL_FAILED, INVALID_DOWNLOADED_MODEL, SKLEARN_SUPPORTED_CATEGORIES
 
 
 def get_sklearn_type(model):
@@ -165,3 +165,19 @@ def download_model(url):
             return response.json()
     except ValueError:
         raise Exception(INVALID_DOWNLOADED_MODEL)
+
+
+def get_sklearn_class(model_name):
+    """
+    Return the sklearn class of the requested model name.
+
+    :param model_name: model name
+    :type model_name: str
+
+    :return: sklearn ML model class
+    """
+    for _, category_models in SKLEARN_SUPPORTED_CATEGORIES.items():
+        if model_name in category_models:
+            return category_models[model_name]
+    # todo raise exception
+    return None
diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,5 @@
+art>=1.8
 numpy>=1.9.0
+requests>=2.0.0
 scikit-learn>=0.22.2
 scipy>=0.19.1
-requests>=2.0.0
diff --git a/setup.py b/setup.py
@@ -79,4 +79,9 @@ def read_description():
         'Topic :: Scientific/Engineering :: Physics',
     ],
     license='MIT',
+    entry_points={
+            'console_scripts': [
+                'pymilo = pymilo.__main__:main',
+            ]
+    }
 )
