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 d043a40c29
commit dac4ce7b7c
No known key found for this signature in database
GPG key ID: 0EA1BC97D0AB076E
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
def load_model(
key: str, version: Optional[int] = None, return_path: bool = False
) -> Tuple[Any, int]:
def load_model(key: str, version: Optional[int] = None) -> Tuple[Any, int]:
file = get_context().large_file_implementation(name=key, mode="rb", version=version)
if return_path:
return file.get(), file.version
path = file.get()
if path.is_dir():
return path, file.version
with file as f:
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,
*,
version: Union[int, Literal["latest"]],
return_path: bool = False,
model_kwarg_name: str = "model",
) -> Callable[[F], F]:
assert (
@ -24,7 +23,6 @@ def use_model(
model, actual_version = load_model(
key=key,
version=None if version == "latest" else version,
return_path=return_path,
)
def decorator(func: F) -> F: