wehub-resource-sync
81 lines
2.9 KiB
Python
81 lines
2.9 KiB
Python
import warnings
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.utils.import_hooks import register_post_import_hook
|
|
from mlflow.utils.plugins import get_entry_points
|
|
|
|
|
|
class ModelEvaluatorRegistry:
|
|
"""
|
|
Scheme-based registry for model evaluator implementations
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._registry = {}
|
|
self._builtin_evaluators = {}
|
|
|
|
def register(self, scheme, evaluator):
|
|
"""Register model evaluator provided by other packages"""
|
|
self._registry[scheme] = evaluator
|
|
|
|
def register_builtin(self, scheme, evaluator):
|
|
"""Register built-in model evaluator"""
|
|
self._registry[scheme] = evaluator
|
|
self._builtin_evaluators[scheme] = evaluator
|
|
|
|
def register_entrypoints(self):
|
|
# Register ModelEvaluator implementation provided by other packages
|
|
for entrypoint in get_entry_points("mlflow.model_evaluator"):
|
|
try:
|
|
self.register(entrypoint.name, entrypoint.load())
|
|
except (AttributeError, ImportError) as exc:
|
|
warnings.warn(
|
|
'Failure attempting to register model evaluator for scheme "{}": {}'.format(
|
|
entrypoint.name, str(exc)
|
|
),
|
|
stacklevel=2,
|
|
)
|
|
|
|
def get_evaluator(self, evaluator_name):
|
|
"""
|
|
Get an evaluator instance from the registry based on the name of evaluator
|
|
"""
|
|
evaluator_cls = self._registry.get(evaluator_name)
|
|
if evaluator_cls is None:
|
|
raise MlflowException(
|
|
f"Could not find a registered model evaluator for: {evaluator_name}. "
|
|
f"Currently registered evaluator names are: {list(self._registry.keys())}"
|
|
)
|
|
return evaluator_cls()
|
|
|
|
def is_builtin(self, name):
|
|
return name in self._builtin_evaluators
|
|
|
|
def is_registered(self, name):
|
|
return name in self._registry
|
|
|
|
|
|
_model_evaluation_registry = ModelEvaluatorRegistry()
|
|
|
|
|
|
def register_evaluators(module):
|
|
from mlflow.models.evaluation.evaluators.classifier import ClassifierEvaluator
|
|
from mlflow.models.evaluation.evaluators.default import DefaultEvaluator
|
|
from mlflow.models.evaluation.evaluators.regressor import RegressorEvaluator
|
|
from mlflow.models.evaluation.evaluators.shap import ShapEvaluator
|
|
|
|
# Built-in evaluators
|
|
module._model_evaluation_registry.register_builtin(DefaultEvaluator.name, DefaultEvaluator)
|
|
module._model_evaluation_registry.register_builtin(
|
|
ClassifierEvaluator.name, ClassifierEvaluator
|
|
)
|
|
module._model_evaluation_registry.register_builtin(RegressorEvaluator.name, RegressorEvaluator)
|
|
module._model_evaluation_registry.register_builtin(ShapEvaluator.name, ShapEvaluator)
|
|
|
|
# Plugin evaluators
|
|
module._model_evaluation_registry.register_entrypoints()
|
|
|
|
|
|
# Put it in post-importing hook to avoid circuit importing
|
|
register_post_import_hook(register_evaluators, __name__, overwrite=True)
|