diff --git a/great_ai/deploy/great_ai.py b/great_ai/deploy/great_ai.py index 037c72b..48df0e5 100644 --- a/great_ai/deploy/great_ai.py +++ b/great_ai/deploy/great_ai.py @@ -49,6 +49,8 @@ class GreatAI(Generic[T, V]): self, func: Callable[..., Union[V, Awaitable[V]]], ): + """Do not call this function directly, use GreatAI.create instead.""" + func = automatically_decorate_parameters(func) get_function_metadata_store(func).is_finalised = True @@ -94,6 +96,43 @@ class GreatAI(Generic[T, V]): def create( func: Union[Callable[..., Awaitable[V]], Callable[..., V]], ) -> Union["GreatAI[Awaitable[Trace[V]], V]", "GreatAI[Trace[V], V]"]: + """Decorate a function by wrapping it in a GreatAI instance. + + The function can be typed, synchronous or async. If it has + unwrapped parameters (parameters not affected by a @parameter + or @use_model decorator), those will be automatically wrapped. + + The return value is replaced by a Trace (or Awaitable[Trace]), + while the original return value is available under the `.output` + property. + + For configuration options, see great_ai.configure. + + Examples: + >>> @GreatAI.create + ... def my_function(a): + ... return a + 2 + >>> my_function(3).output + 5 + + >>> @GreatAI.create + ... def my_function(a: int) -> int: + ... return a + 2 + >>> my_function(3) + Trace[int]... + + >>> my_function('3').output + Traceback (most recent call last): + ... + TypeError: type of a must be int; got str instead + + Args: + func: The prediction function that needs to be decorated. + + Returns: + A GreatAI instance wrapping `func`. + + """ return GreatAI[Trace[V], V]( func, ) diff --git a/great_ai/models/save_model.py b/great_ai/models/save_model.py index df0939f..2b7049b 100644 --- a/great_ai/models/save_model.py +++ b/great_ai/models/save_model.py @@ -9,6 +9,30 @@ from ..context import get_context def save_model( model: Union[Path, str, object], key: str, *, keep_last_n: Optional[int] = None ) -> str: + """Save (and optionally serialise) a model in order to use by `use_model`. + + The `model` can be a Path or string representing a path in which case the + local file/folder is read and saved using the current LargeFile implementation. + In case `model` is an object, it is serialised using `dill` before uploading it. + + Examples: + >>> from great_ai import use_model + >>> save_model(3, 'my_number') + 'my_number:...' + + >>> @use_model('my_number') + ... def my_function(a, model): + ... return a + model + >>> my_function(4) + 7 + + Args: + model: The object or path to be uploaded. + key: The model's name. + keep_last_n: If specified, remove old models and only keep the latest n. Directly passed to LargeFile. + Returns: + The key and version of the saved model separated by a colon. Example: "key:version" + """ file = get_context().large_file_implementation( name=key, mode="wb", keep_last_n=keep_last_n ) diff --git a/great_ai/models/use_model.py b/great_ai/models/use_model.py index f2d1c14..685fb52 100644 --- a/great_ai/models/use_model.py +++ b/great_ai/models/use_model.py @@ -30,6 +30,38 @@ def use_model( version: Union[int, Literal["latest"]] = "latest", model_kwarg_name: str = "model", ) -> Callable[[F], F]: + """Inject a model into a function. + + Load a model specified by `key` and `version` using the + currently active `LargeFile` implementation. If it's a + single object, it is deserialised using `dill`. If it's + a directory of files, a `pathlib.Path` instance is given. + + By default, the function's `model` parameter is replaced + by the loaded model. This can be customised by changing + `model_kwarg_name`. + + Multiple models can be loaded by decorating the same + function with `use_model` multiple times. + + Examples: + >>> from great_ai import save_model + >>> save_model(3, 'my_number') + 'my_number:...' + >>> @use_model('my_number') + ... def my_function(a, model): + ... return a + model + >>> my_function(4) + 7 + + Args: + key: The model's name as stored by the LargeFile implementation. + version: The model's version as stored by the LargeFile implementation. + model_kwarg_name: the parameter to use for injecting the loaded model + Returns: + A decorator for model injection. + """ + assert ( isinstance(version, int) or version == "latest" ), "Only integers or the string literal `latest` is allowed as a version" diff --git a/great_ai/parameters/parameter.py b/great_ai/parameters/parameter.py index dab0650..1410be2 100644 --- a/great_ai/parameters/parameter.py +++ b/great_ai/parameters/parameter.py @@ -17,6 +17,35 @@ def parameter( validator: Callable[[Any], bool] = lambda _: True, disable_logging: bool = False, ) -> Callable[[F], F]: + """Control the validation and logging of function parameters. + + Examples: + >>> @parameter('a') + ... def my_function(a: int): + ... return a + 2 + >>> my_function(4) + 6 + >>> my_function('3') + Traceback (most recent call last): + ... + TypeError: type of a must be int; got str instead + + >>> @parameter('positive_a', validator=lambda v: v > 0) + ... def my_function(positive_a: int): + ... return a + 2 + >>> my_function(-1) + Traceback (most recent call last): + ... + great_ai.errors.argument_validation_error.ArgumentValidationError: ... + + Args: + parameter_name: Name of parameter to consider + validator: Optional validator to run against the concrete argument. ArgumentValidationError is thrown when the return value is False. + disable_logging: Do not save the value in any active TracingContext. + Returns: + A decorator for argument validation. + """ + def decorator(func: F) -> F: get_function_metadata_store(func).input_parameter_names.append(parameter_name) assert_function_is_not_finalised(func) diff --git a/great_ai/utilities/unique.py b/great_ai/utilities/unique.py index 534133a..bd8e8d5 100644 --- a/great_ai/utilities/unique.py +++ b/great_ai/utilities/unique.py @@ -11,8 +11,10 @@ def unique(values: Iterable[T], *, key: Callable[[T], Any] = lambda v: v) -> Lis Examples: >>> unique([1, 1, 5, 3, 3]) [1, 5, 3] + >>> unique([{'a': 1, 'b': 2}, {'a': 1, 'b': 3}], key=lambda v: v['a']) [{'a': 1, 'b': 2}] + >>> unique([{'a': 1, 'b': 2}, {'a': 1, 'b': 3}], key=lambda v: v['b']) [{'a': 1, 'b': 2}, {'a': 1, 'b': 3}]