andras/great-ai
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases 4 Packages Wiki Activity Actions

Simplify API

Browse source
This commit is contained in:
Andras Schmelczer 2022-07-01 19:50:43 +02:00
parent 8f1418c346
commit b56273e267
2 changed files with 5 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

10
src/great_ai/great_ai/models/load_model.py
View file

@ -5,13 +5,13 @@ from joblib import load
from ..context import get_context from ..context import get_context
def load_model( def load_model(key: str, version: Optional[int] = None) -> Tuple[Any, int]:
key: str, version: Optional[int] = None, return_path: bool = False
) -> Tuple[Any, int]:
file = get_context().large_file_implementation(name=key, mode="rb", version=version) file = get_context().large_file_implementation(name=key, mode="rb", version=version)
if return_path: path = file.get()
return file.get(), file.version
if path.is_dir():
return path, file.version
with file as f: with file as f:
return load(f), file.version return load(f), file.version

2
src/great_ai/great_ai/models/use_model.py
View file

@ -14,7 +14,6 @@ def use_model(
key: str, key: str,
*, *,
version: Union[int, Literal["latest"]], version: Union[int, Literal["latest"]],
return_path: bool = False,
model_kwarg_name: str = "model", model_kwarg_name: str = "model",
) -> Callable[[F], F]: ) -> Callable[[F], F]:
assert ( assert (
@ -24,7 +23,6 @@ def use_model(
model, actual_version = load_model( model, actual_version = load_model(
key=key, key=key,
version=None if version == "latest" else version, version=None if version == "latest" else version,
return_path=return_path,
) )
def decorator(func: F) -> F: def decorator(func: F) -> F: