llm4ad 1.0__py2.py3-none-any.whl

llm4ad/__init__.py

@@ -0,0 +1,6 @@
+from . import base
+from . import method
+from . import task
+from .tools import profiler
+from .tools import llm
+

llm4ad/base/__init__.py

@@ -0,0 +1,9 @@
+from . import code, evaluate, sample, modify_code
+from .code import (
+    Function,
+    Program,
+    TextFunctionProgramConverter
+)
+from .evaluate import Evaluation, SecureEvaluator
+from .modify_code import ModifyCode
+from .sample import LLM, SampleTrimmer

llm4ad/base/code.py

@@ -0,0 +1,270 @@
+"""
+This file implements 2 classes representing unities of code:
+
+- Function, containing all the information we need about functions: name, args,
+  body and optionally a return type and a docstring.
+
+- Program, which contains a code preface (which could be imports, global
+  variables and classes, ...) and a list of Functions.
+
+- For example, a function is shown below,
+which is an un-executable program because 'np' is not defined, and 'WEIGHT' is not defined.
+--------------------------------------------------------------------------------------------
+def func(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    b = b + WEIGHT
+    return a + b
+--------------------------------------------------------------------------------------------
+
+- A program is an executable object as shown below.
+--------------------------------------------------------------------------------------------
+import numpy as np
+WEIGHT = 10
+
+def func(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    b = b + WEIGHT
+    return a + b
+--------------------------------------------------------------------------------------------
+"""
+from __future__ import annotations
+
+import ast
+import copy
+import dataclasses
+from typing import Any
+
+
+@dataclasses.dataclass
+class Function:
+    """A parsed Python function."""
+
+    name: str
+    args: str
+    body: str
+    return_type: str | None = None
+    docstring: str | None = None
+    score: Any | None = None
+    evaluate_time: float | None = None
+    sample_time: float | None = None
+
+    def __str__(self) -> str:
+        return_type = f' -> {self.return_type}' if self.return_type else ''
+
+        function = f'def {self.name}({self.args}){return_type}:\n'
+        if self.docstring:
+            # self.docstring is already indented on every line except the first one.
+            # Here, we assume the indentation is always four spaces.
+            new_line = '\n' if self.body else ''
+            function += f'    """{self.docstring}"""{new_line}'
+        # self.body is already indented.
+        function += self.body + '\n\n'
+        return function
+
+    def __setattr__(self, name: str, value: str) -> None:
+        # Ensure there aren't leading & trailing new lines in `body`.
+        if name == 'body':
+            value = value.strip('\n')
+        # Ensure there aren't leading & trailing quotes in `docstring``.
+        if name == 'docstring' and value is not None:
+            if '"""' in value:
+                value = value.strip()
+                value = value.replace('"""', '')
+        super().__setattr__(name, value)
+
+    def __eq__(self, other: Function):
+        assert isinstance(other, Function)
+        return (self.name == other.name and
+                self.args == other.args and
+                self.return_type == other.return_type and
+                self.body == other.body)
+
+
+@dataclasses.dataclass(frozen=True)
+class Program:
+    """A parsed Python program."""
+
+    # `preface` is everything from the beginning of the code till the first
+    # function is found.
+    preface: str
+    functions: list[Function]
+
+    def __str__(self) -> str:
+        program = f'{self.preface}\n' if self.preface else ''
+        program += '\n'.join([str(f) for f in self.functions])
+        return program
+
+    def find_function_index(self, function_name: str) -> int:
+        """Returns the index of input function name."""
+        function_names = [f.name for f in self.functions]
+        count = function_names.count(function_name)
+        if count == 0:
+            raise ValueError(
+                f'function {function_name} does not exist in program:\n{str(self)}'
+            )
+        if count > 1:
+            raise ValueError(
+                f'function {function_name} exists more than once in program:\n'
+                f'{str(self)}'
+            )
+        index = function_names.index(function_name)
+        return index
+
+    def get_function(self, function_name: str) -> Function:
+        index = self.find_function_index(function_name)
+        return self.functions[index]
+
+
+class _ProgramVisitor(ast.NodeVisitor):
+    """Parses code to collect all required information to produce a `Program`.
+    Note that we do not store function decorators.
+    """
+
+    def __init__(self, sourcecode: str):
+        self._codelines: list[str] = sourcecode.splitlines()
+        self._preface: str = ''
+        self._functions: list[Function] = []
+        self._current_function: str | None = None
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        """Collects all information about the function being parsed.
+        """
+        # We only care about first level functions.
+        if node.col_offset == 0:
+            self._current_function = node.name
+            # --------------------------------------------------------------------------------------------
+            # RZ: Here is an issue: some cases were not taken into consideration.
+            # For instance, if the first parsed function has decorators,
+            # these decorators will also be retained in the preceding code snippet.
+            # However, in the algorithm's flow, these decorators are invalid,
+            # leading to subsequent evaluations reporting errors.
+            # --------------------------------------------------------------------------------------------
+            # if not self._functions:
+            #     self._preface = '\n'.join(self._codelines[:node.lineno - 1])
+            # --------------------------------------------------------------------------------------------
+            # Fix bugs: if the first function is decorated, find the smallest line_no of the decorator,
+            # FIX bugs: and retrain the code above it as 'preface'.
+            # --------------------------------------------------------------------------------------------
+            if not self._functions:
+                has_decorators = bool(node.decorator_list)
+                if has_decorators:
+                    # find the minimum line number and retain the code above
+                    decorator_start_line = min(decorator.lineno for decorator in node.decorator_list)
+                    self._preface = '\n'.join(self._codelines[:decorator_start_line - 1])
+                else:
+                    self._preface = '\n'.join(self._codelines[:node.lineno - 1])
+            # --------------------------------------------------------------------------------------------
+            function_end_line = node.end_lineno
+            body_start_line = node.body[0].lineno - 1
+            # Extract the docstring.
+            docstring = None
+            if isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Str):
+                docstring = f'    """{ast.literal_eval(ast.unparse(node.body[0]))}"""'
+                if len(node.body) > 1:
+                    # ----------------------------------------------------------------------------
+                    # RZ: If the next line of the docstring is also comment that starts with "#"
+                    # We should preserve it.
+                    # ----------------------------------------------------------------------------
+                    # body_start_line = node.body[1].lineno - 1
+                    # ----------------------------------------------------------------------------
+                    # Fix bugs: If we find the docstring, the function body start with
+                    # Fix bugs: body[0].end_lineno, where body[0] refers to the docstring node.
+                    # ----------------------------------------------------------------------------
+                    body_start_line = node.body[0].end_lineno
+                    # ----------------------------------------------------------------------------
+                else:
+                    body_start_line = function_end_line
+
+            self._functions.append(Function(
+                name=node.name,
+                args=ast.unparse(node.args),
+                return_type=ast.unparse(node.returns) if node.returns else None,
+                docstring=docstring,
+                body='\n'.join(self._codelines[body_start_line:function_end_line]),
+            ))
+        self.generic_visit(node)
+
+    def return_program(self) -> Program:
+        return Program(preface=self._preface, functions=self._functions)
+
+
+class TextFunctionProgramConverter:
+    """Convert text to Program instance and Function instance,
+    Convert Program instance to Function instance, and Function instance to Program instance.
+    """
+
+    @classmethod
+    def text_to_program(cls, program_str: str) -> Program | None:
+        """Returns Program object by parsing input text using Python AST.
+        """
+        try:
+            # We assume that the program is composed of some preface (e.g. imports,
+            # classes, assignments, ...) followed by a sequence of functions.
+            tree = ast.parse(program_str)
+            visitor = _ProgramVisitor(program_str)
+            visitor.visit(tree)
+            return visitor.return_program()
+        except:
+            return None
+
+    @classmethod
+    def text_to_function(cls, program_str: str) -> Function | None:
+        """Returns Function object by parsing input text using Python AST.
+        """
+        try:
+            program = cls.text_to_program(program_str)
+            if len(program.functions) != 1:
+                raise ValueError(f'Only one function expected, got {len(program.functions)}'
+                                 f':\n{program.functions}')
+            return program.functions[0]
+        except ValueError as value_err:
+            raise value_err
+        except:
+            return None
+
+    @classmethod
+    def function_to_program(cls, function: str | Function, template_program: str | Program) -> Program | None:
+        try:
+            # convert function to Function instance
+            if isinstance(function, str):
+                function = cls.text_to_function(function)
+            else:
+                function = copy.deepcopy(function)
+
+            # convert template_program to Program instance
+            if isinstance(template_program, str):
+                template_program = cls.text_to_program(template_program)
+            else:
+                template_program = copy.deepcopy(template_program)
+
+            # assert that a program have one function
+            if len(template_program.functions) != 1:
+                raise ValueError(f'Only one function expected, got {len(template_program.functions)}'
+                                 f':\n{template_program.functions}')
+
+            # replace the function body with the new function body
+            template_program.functions[0].body = function.body
+            return template_program
+        except ValueError as value_err:
+            raise value_err
+        except:
+            return None
+
+    @classmethod
+    def program_to_function(cls, program: str | Program) -> Function | None:
+        try:
+            # convert program to Program instance
+            if isinstance(program, str):
+                program = cls.text_to_program(program)
+            else:
+                program = copy.deepcopy(program)
+
+            # assert that a program have one function
+            if len(program.functions) != 1:
+                raise ValueError(f'Only one function expected, got {len(program.functions)}'
+                                 f':\n{program.functions}')
+
+            # return the function
+            return program.functions[0]
+        except ValueError as value_err:
+            raise value_err
+        except:
+            return None

llm4ad/base/evaluate.py

@@ -0,0 +1,285 @@
+from __future__ import annotations
+
+import multiprocessing
+import sys
+import time
+from abc import ABC, abstractmethod
+from typing import Any
+
+from .code import TextFunctionProgramConverter, Program
+from .modify_code import ModifyCode
+
+
+class Evaluation(ABC):
+    def __init__(
+            self,
+            template_program: str,
+            task_description: str = '',
+            use_numba_accelerate: bool = False,
+            use_protected_div: bool = False,
+            protected_div_delta: float = 1e-5,
+            random_seed: int | None = None,
+            timeout_seconds: int | float = None,
+            *,
+            exec_code: bool = True,
+            safe_evaluate: bool = True,
+            daemon_eval_process: bool = False,
+            **kwargs
+    ):
+        """Evaluator for executing generated code.
+        Args:
+            task_description: a brief description of the algorithm design task and specification.
+            template_program: the seed program (in str) as the initial function of the run.
+                the template_program should be executable, i.e., incorporating package import, and function definition, and function body.
+                The function body can be 'pass' in EoH as EoH does not request the template_program tobe valid.
+            use_numba_accelerate: Wrap the function with '@numba.jit(nopython=True)'.
+            use_protected_div   : Modify 'a / b' => 'a / (b + delta)'.
+            protected_div_delta : Delta value in protected div.
+            random_seed         : If is not None, set random seed in the first line of the function body.
+            timeout_seconds     : Terminate the evaluation after timeout seconds.
+            exec_code           : Using 'exec()' to compile the code and provide the callable function.
+                If is set to 'False', the 'callable_func' argument in 'self.evaluate_program' is always 'None'.
+                If is set to 'True', the user should provide the score of the program based on 'program_str' argument in 'self.evaluate_program'.
+            safe_evaluate       : Evaluate in safe mode using a new process. If is set to False,
+                the evaluation will not be terminated after timeout seconds. The user should consider how to
+                terminate evaluating in time.
+            daemon_eval_process : Set the evaluate process as a daemon process. If set to True,
+                you can not set new processes in the evaluator. Which means in self.evaluate_program(),
+                you can not create new processes.
+
+        -Assume that: use_numba_accelerate=True, self.use_protected_div=True, and self.random_seed=2024.
+        -The original function:
+        --------------------------------------------------------------------------------
+        import numpy as np
+
+        def f(a, b):
+            a = np.random.random()
+            return a / b
+        --------------------------------------------------------------------------------
+        -The modified function will be:
+        --------------------------------------------------------------------------------
+        import numpy as np
+        import numba
+
+        @numba.jit(nopython=True)
+        def f():
+            np.random.seed(2024)
+            a = np.random.random()
+            return _protected_div(a, b)
+
+        def _protected_div(a, b, delta=1e-5):
+            return a / (b + delta)
+        --------------------------------------------------------------------------------
+        As shown above, the 'import numba', 'numba.jit()' decorator, and '_protected_dev' will be added by this function.
+        """
+
+        self.template_program = template_program
+        self.task_description = task_description
+        self.use_numba_accelerate = use_numba_accelerate
+        self.use_protected_div = use_protected_div
+        self.protected_div_delta = protected_div_delta
+        self.random_seed = random_seed
+        self.timeout_seconds = timeout_seconds
+        self.exec_code = exec_code
+        self.safe_evaluate = safe_evaluate
+        self.daemon_eval_process = daemon_eval_process
+
+    @abstractmethod
+    def evaluate_program(self, program_str: str, callable_func: callable) -> Any | None:
+        """Evaluate a given function. You can use compiled function (function_callable),
+        as well as the original function strings for evaluation.
+        Args:
+            program_str: The function in string. You can ignore this argument when implementation. (See below).
+            callable_func: The callable heuristic function. You can call it using `callable_func(args, kwargs)`.
+        Return:
+            Returns the fitness value.
+
+        Assume that: self.use_numba_accelerate = True, self.use_protected_div = True,
+        and self.random_seed = 2024, the argument 'function_str' will be something like below:
+        --------------------------------------------------------------------------------
+        import numpy as np
+        import numba
+
+        @numba.jit(nopython=True)
+        def f(a, b):
+            np.random.seed(2024)
+            a = a + np.random.random()
+            return _protected_div(a, b)
+
+        def _protected_div(a, b, delta=1e-5):
+            return a / (b + delta)
+        --------------------------------------------------------------------------------
+        As shown above, the 'import numba', 'numba.jit()' decorator,
+        and '_protected_dev' will be added by this function.
+        """
+        raise NotImplementedError('Must provide a evaluator for a function.')
+
+
+class SecureEvaluator:
+    def __init__(self, evaluator: Evaluation, debug_mode=False, *, fork_proc: str | bool = 'auto', **kwargs):
+        assert fork_proc in [True, False, 'auto', 'default']
+        self._evaluator = evaluator
+        self._debug_mode = debug_mode
+
+        if self._evaluator.safe_evaluate:
+            if fork_proc == 'auto':
+                # force MacOS and Linux use 'fork' to generate new process
+                if sys.platform.startswith('darwin') or sys.platform.startswith('linux'):
+                    multiprocessing.set_start_method('fork', force=True)
+            elif fork_proc is True:
+                multiprocessing.set_start_method('fork', force=True)
+            elif fork_proc is False:
+                multiprocessing.set_start_method('spawn', force=True)
+
+    def _modify_program_code(self, program_str: str) -> str:
+        function_name = TextFunctionProgramConverter.text_to_function(program_str).name
+        if self._evaluator.use_numba_accelerate:
+            program_str = ModifyCode.add_numba_decorator(
+                program_str, function_name=function_name
+            )
+        if self._evaluator.use_protected_div:
+            program_str = ModifyCode.replace_div_with_protected_div(
+                program_str, self._evaluator.protected_div_delta, self._evaluator.use_numba_accelerate
+            )
+        if self._evaluator.random_seed is not None:
+            program_str = ModifyCode.add_numpy_random_seed_to_func(
+                program_str, function_name, self._evaluator.random_seed
+            )
+        return program_str
+
+    def evaluate_program(self, program: str | Program):
+        try:
+            program_str = str(program)
+            # record function name BEFORE modifying program code
+            function_name = TextFunctionProgramConverter.text_to_function(program_str).name
+
+            program_str = self._modify_program_code(program_str)
+            if self._debug_mode:
+                print(f'DEBUG: evaluated program:\n{program_str}\n')
+
+            # safe evaluate
+            if self._evaluator.safe_evaluate:
+                result_queue = multiprocessing.Queue()
+                process = multiprocessing.Process(
+                    target=self._evaluate_in_safe_process,
+                    args=(program_str, function_name, result_queue),
+                    daemon=self._evaluator.daemon_eval_process
+                )
+                process.start()
+
+                if self._evaluator.timeout_seconds is not None:
+                    try:
+                        # get the result in timeout seconds
+                        result = result_queue.get(timeout=self._evaluator.timeout_seconds)
+                        # after getting the result, terminate/kill the process
+                        process.terminate()
+                        process.join()
+                    except:
+                        # timeout
+                        if self._debug_mode:
+                            print(f'DEBUG: the evaluation time exceeds {self._evaluator.timeout_seconds}s.')
+                        process.terminate()
+                        process.join()
+                        result = None
+                else:
+                    result = result_queue.get()
+                    process.terminate()
+                    process.join()
+                return result
+            else:
+                return self._evaluate(program_str, function_name)
+        except Exception as e:
+            if self._debug_mode:
+                print(e)
+            return None
+
+    def evaluate_program_record_time(self, program: str | Program):
+        evaluate_start = time.time()
+        try:
+            program_str = str(program)
+            # record function name BEFORE modifying program code
+            function_name = TextFunctionProgramConverter.text_to_function(program_str).name
+
+            program_str = self._modify_program_code(program_str)
+            if self._debug_mode:
+                print(f'DEBUG: evaluated program:\n{program_str}\n')
+
+            # safe evaluate
+            if self._evaluator.safe_evaluate:
+                result_queue = multiprocessing.Queue()
+                process = multiprocessing.Process(
+                    target=self._evaluate_in_safe_process,
+                    args=(program_str, function_name, result_queue),
+                    daemon=self._evaluator.daemon_eval_process
+                )
+                process.start()
+
+                if self._evaluator.timeout_seconds is not None:
+                    try:
+                        # get the result in timeout seconds
+                        result = result_queue.get(timeout=self._evaluator.timeout_seconds)
+                        # after getting the result, terminate/kill the process
+                        process.terminate()
+                        process.join()
+                    except:
+                        # timeout
+                        if self._debug_mode:
+                            print(f'DEBUG: the evaluation time exceeds {self._evaluator.timeout_seconds}s.')
+                        process.terminate()
+                        process.join()
+                        result = None
+                else:
+                    result = result_queue.get()
+                    process.terminate()
+                    process.join()
+                # return evaluate result and evaluate time
+                return result, time.time() - evaluate_start
+            else:
+                start = time.time()
+                score = self._evaluate(program_str, function_name)
+                return score, time.time() - start
+        except Exception as e:
+            if self._debug_mode:
+                print(e)
+            # return None and evaluate time
+            return None, time.time() - evaluate_start
+
+    def _evaluate_in_safe_process(self, program_str: str, function_name, result_queue: multiprocessing.Queue):
+        try:
+            if self._evaluator.exec_code:
+                # compile the program, and maps the global func/var/class name to its address
+                all_globals_namespace = {}
+                # execute the program, map func/var/class to global namespace
+                exec(program_str, all_globals_namespace)
+                # get the pointer of 'function_to_run'
+                program_callable = all_globals_namespace[function_name]
+            else:
+                program_callable = None
+
+            # get evaluate result
+            res = self._evaluator.evaluate_program(program_str, program_callable)
+            result_queue.put(res)
+        except Exception as e:
+            if self._debug_mode:
+                print(e)
+            result_queue.put(None)
+
+    def _evaluate(self, program_str: str, function_name):
+        try:
+            if self._evaluator.exec_code:
+                # compile the program, and maps the global func/var/class name to its address
+                all_globals_namespace = {}
+                # execute the program, map func/var/class to global namespace
+                exec(program_str, all_globals_namespace)
+                # get the pointer of 'function_to_run'
+                program_callable = all_globals_namespace[function_name]
+            else:
+                program_callable = None
+
+            # get evaluate result
+            res = self._evaluator.evaluate_program(program_str, program_callable)
+            return res
+        except Exception as e:
+            if self._debug_mode:
+                print(e)
+            return None