def test_get_model(mocker, raw_model, X, y, needs_proba):
model = TravaModel(raw_model=raw_model, model_id=model_id)
assert model.get_model(for_train=True) == raw_model
assert model.get_model(for_train=False) == raw_model
y_predict_proba = mocker.Mock()
if needs_proba:
raw_model.predict_proba.return_value = y_predict_proba
y_pred = mocker.Mock()
raw_model.predict.return_value = y_pred
model.fit(X=X, y=y)
model.predict(X=X, y=y)
model.unload_model()
train_cached_model = model.get_model(for_train=True)
test_cached_model = model.get_model(for_train=False)
assert train_cached_model != raw_model
assert test_cached_model != raw_model
assert train_cached_model.predict(X) == y_pred
if needs_proba:
assert train_cached_model.predict_proba(X) == y_predict_proba
def test_get_model_unload(mocker, raw_model, for_train):
trava_model = TravaModel(raw_model=raw_model, model_id=model_id)
trava_model.unload_model()
with pytest.raises(ValueError):
trava_model.get_model(for_train=for_train)
if for_train:
y_pred_key = "_y_train_pred"
else:
y_pred_key = "_y_test_pred"
y_pred_mock = mocker.Mock()
mocker.patch.object(trava_model, y_pred_key, y_pred_mock)
assert trava_model.get_model(for_train=for_train).predict(X=None) == y_pred_mock
def __call__(self, trava_model: TravaModel, for_train: bool, X, X_raw, y, **kwargs):
if self._requires_raw_model and not trava_model.raw_model():
raise Exception("Cannot perform eval on model {} " "because it was unloaded.".format(trava_model.model_id))
if self._requires_X_y and (X is None or X_raw is None or y is None):
raise Exception(
"Cannot perform eval on model {} "
"because data is required and was unloaded.".format(trava_model.model_id)
)
return self._scorer(
model=trava_model.get_model(for_train=for_train),
model_info=trava_model,
for_train=for_train,
X=X,
X_raw=X_raw,
y=y,
**kwargs
)