breaking file input and function input into different class methods

src/sasctl/pzmm/write_score_wrapper.py

@@ -4,83 +4,93 @@
 from pathlib import Path
 import re
 
-"""
-- Collect inputs and outputs from inputVar.json and outputVar.json
-   - Throw error or include optional arguments
-- How to handle imports?
-   - pass a list of imports, etc
-- How do we handle model loading?
-   - re.sub() with "Path(settings.pickle_path) / " and file name
-   - Document replacement of path with inclusion of settings.pickle_path
-- What form do we want to take the score snippet in? add boolean to check whether user is passing in a function or passing in the score snipped itslef if so adjust the code accordingly
-   - function...we'll parse and drop in the contents
-   - file, but tell us the score function...include other functions and pull out score function info
-"""
-
-
 class ScoreWrapper:
     score_wrapper: str = ""
 
     @classmethod
-    def write_score_wrapper(cls,
-                            imports: List[str],
-                            model_load: str,
-                            function_definition: str,
-                            model_scoring_code: str,
-                            is_function: bool,
-                            is_file: bool) -> str:
+    def write_score_wrapper_function_input(cls, 
+                                           imports: List[str], 
+                                           function_definition: str, 
+                                           function_body: str, 
+                                           model_load: str, 
+                                           model_name_with_file_extension: str):
         """
-        Method to generate scoring code and add it to cls.score_wrapper.
-
+        Method to generate scoring code from a function and add it to cls.score_wrapper.
 
         Parameters:
         imports (List[str]): List of modules to import.
-        model_load (str): Name of the model to load.
         function_definition (str): Function definition code.
-        model_scoring_code (str): Model scoring code. how the function is called.
+        function_body (str): Function body code.
+        model_load (str): Name of the model to load.
+        model_name_with_file_extension (str): Name of the model file with extension.
 
-        Returns: cls.score_wrapper (str): Score wrapper code.
+        Returns:
+        cls.score_wrapper (str): The scoring code.
         """
-        # Adding imports to score_wrapper
         for module in imports:
             cls.score_wrapper += f"import {module}\n"
-            cls.score_wrapper += "from pathlib import Path\n"
-            cls.score_wrapper += "import settings\n"
-
-        # Adding model load to score_wrapper
-        cls.score_wrapper += f"{model_load} = Path(settings.pickle_path) / '{model_load}'\n"
+
+        cls.score_wrapper += "from pathlib import Path\n"
+        cls.score_wrapper += "import settings\n\n\n"
+
+        cls.score_wrapper += f"model_path = Path(settings.pickle_path) / '{model_name_with_file_extension}'\n\n"
+        cls.score_wrapper += f"with open(Path(settings.pickle_path) / '{model_name_with_file_extension}', 'rb') as f:\n\tmodel = {model_load}\n\n"
+
+        cls.score_wrapper += f"{function_definition}:\n"
+        cls.score_wrapper += "\tglobal model\n"
+        cls.score_wrapper += "\ttry:\n"
+        cls.score_wrapper += f"\t\t{function_body}\n"
+        cls.score_wrapper += "\texcept Exception as e:\n"
+        cls.score_wrapper += "\t\tprint(f'Error: {e}')\n"
+        return cls.score_wrapper
 
-        # Adding function definition to score_wrapper
-        if is_function:
-            # if the function definition is given.
-            cls.score_wrapper += f"{function_definition}\n"
-        elif is_file:
-            # file, but tell us the score function...include other functions and pull out score function info use re to find all the functions specified maybe add parameters for if is_file is true to know how many functions user wants to include
-            # Extract all function definitions from the file
-            with open(function_definition, 'r') as f:
-                file_content = f.read()
-            # This only works for function definitions on one line like this def score(input_data):
-            function_defs = re.findall(r"(def .+?:)", file_content)
-            for function in function_defs:
-                cls.score_wrapper += f"{function}\n"
+    @classmethod
+    def write_score_wrapper_file_input(cls, 
+                                       imports: List[str], 
+                                       file_path: str, 
+                                       model_load: str, 
+                                       model_scoring_code: str, 
+                                       model_name_with_file_extension: str):
+        """
+        Method to generate scoring code from a file and add it to cls.score_wrapper.
 
-        # Adding model scoring code inside a try/except block to score_wrapper
-        cls.score_wrapper += "try:\n"
-        cls.score_wrapper += f"    global {model_load}\n"
-        cls.score_wrapper += f"    {model_scoring_code}\n"
-        cls.score_wrapper += "except Exception as e:\n"
-        cls.score_wrapper += "    print(f'Error: {e}')\n"
+        Parameters:
+        imports (List[str]): List of modules to import.
+        file_path (str): Path to the file containing the scoring code.
+        model_load (str): Name of the model to load.
+        model_scoring_code (str): Model scoring code. How the function is called.
+        model_name_with_file_extension (str): Name of the model file with extension.
 
+        Returns:
+        cls.score_wrapper (str): The scoring code.
+        """
+        for module in imports:
+            cls.score_wrapper += f"import {module}\n"
+
+        cls.score_wrapper += "from pathlib import Path\n"
+        cls.score_wrapper += "import settings\n\n\n"
+
+        cls.score_wrapper += f"model_path = Path(settings.pickle_path) / '{model_name_with_file_extension}'\n\n"
+        cls.score_wrapper += f"with open(Path(settings.pickle_path) / '{model_name_with_file_extension}', 'rb') as f:\n\tmodel = {model_load}\n\n"
+
+        with open(file_path, 'r') as f:
+            file_content = f.read()
+        function_defs = re.findall(r"(def .+?:.*?)(?=def |\Z)", file_content, re.DOTALL)
+        for function in function_defs:
+            cls.score_wrapper += f"{function}\n"
+
         return cls.score_wrapper
-
+    @classmethod
     def write_wrapper_to_file(cls,
                               path: str,
                               model_prefix: str,
                               file_name: str) -> None:
         """
         Method to write the score wrapper to a file.
+
         Parameters:
         file_name (str): Name of the file to write the score wrapper to.
+
         Returns:
         None
         """
@@ -89,30 +99,3 @@ def write_wrapper_to_file(cls,
             score_wrapper_file.write(cls.score_wrapper)
 
 
-"""
-<imports provided by user>
-import settings
-<model load>
-   add Path(settings.pickle_path
-function definition
-docstring: output string
-globalize model and try/except block
-<model scoring code>
-clean up output
-"""
-
-# example of definition and model scoring code
-function_definition = """
-def score(input_data):
-   # Preprocess data
-   processed_data = preprocess(input_data)
-   # Score the processed data
-   return my_model.predict(processed_data)
-"""
-
-model_scoring_code = """
-input_data = get_input_data()
-score(input_data)
-"""
-
-