18 lines
605 B
Python

from collections import namedtuple
from image_prediction.locations import MLRUNS_DIR
from image_prediction.model_loader.loader import ModelLoader
from image_prediction.model_loader.loaders.mlflow import MlflowLoader
ModelClassesPair = namedtuple("ModelClassesPair", ["model", "classes"])
def load_model_and_classes(identifier, model_loader: ModelLoader = None) -> ModelClassesPair:
if not model_loader:
model_loader = MlflowLoader(MLRUNS_DIR)
model = model_loader.load_model(identifier)
classes = model_loader.load_classes(identifier)
return ModelClassesPair(model, classes)