Simplify API

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

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

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: