Coverage for credoai/artifacts/model/regression_model.py: 64%

1"""Model artifact wrapping any regression model""" 

2from .base_model import Model 

3 

4 

5class RegressionModel(Model): 

6 """Class wrapper around classification model to be assessed 

7 

8 RegressionModel serves as an adapter between arbitrary regression models and the 

9 evaluations in Lens. Evaluations depend on 

10 RegressionModel instantiating `predict` 

11 

12 Parameters 

13 ---------- 

14 name : str 

15 Label of the model 

16 model_like : model_like 

17 A binary or multi-class classification model or pipeline. It must have a 

18 `predict` function that returns array containing the class labels for each sample. 

19 It can also optionally have a `predict_proba` function that returns array containing 

20 the class labels probabilities for each sample. 

21 """ 

22 

23 def __init__(self, name: str, model_like=None, tags=None): 

24 super().__init__("Regression", ["predict"], ["predict"], name, model_like, tags) 

25 

26 

27class DummyRegression: 

28 """Class wrapper around regression model predictions 

29 

30 This class can be used when a regression model is not available but its outputs are. 

31 The output include the array containing the predicted class labels and/or the array 

32 containing the class labels probabilities. 

33 Wrap the outputs with this class into a dummy classifier and pass it as 

34 the model to `RegressionModel`. 

35 

36 Parameters 

37 ---------- 

38 predict_output : array 

39 Array containing the output of a model's "predict" method 

40 """ 

41 

42 def __init__(self, name: str, predict_output=None, tags=None): 

43 self.predict_output = predict_output 

44 self.name = name 

45 self.tags = tags 

46 

47 def predict(self, X=None): 

48 return self.predict_output