great-ai/great_ai/models/save_model.py
2022-07-04 19:31:15 +02:00

24 lines
606 B
Python

from pathlib import Path
from typing import Optional, Union
from dill import dump
from ..context import get_context
def save_model(
model: Union[Path, str, object], key: str, *, keep_last_n: Optional[int] = None
) -> str:
file = get_context().large_file_implementation(
name=key, mode="wb", keep_last_n=keep_last_n
)
if isinstance(model, Path) or isinstance(model, str):
file.push(model)
else:
with file as f:
dump(model, f)
get_context().logger.info(f"Model {key} uploaded with version {file.version}")
return f"{key}:{file.version}"