fix reimports from ml-wrappers package for WrappedClassificationModel…

…, WrappedRegressionModel (#485)
interpretml · Jan 13, 2022 · c07b34f · c07b34f
1 parent 8bb715c
commit c07b34f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/python/interpret_community/common/model_wrapper/__init__.py b/python/interpret_community/common/model_wrapper/__init__.py
@@ -4,6 +4,8 @@
 
 """Reimports helpful model wrapper and utils for implicitly rewrapping the model to conform to explainer contracts."""
 
-from ml_wrappers.model import WrappedPytorchModel, _wrap_model, wrap_model
+from ml_wrappers.model import (WrappedClassificationModel, WrappedPytorchModel,
+                               WrappedRegressionModel, _wrap_model, wrap_model)
 
-__all__ = ['WrappedPytorchModel', '_wrap_model', 'wrap_model']
+__all__ = ['WrappedClassificationModel', 'WrappedPytorchModel',
+           'WrappedRegressionModel', '_wrap_model', 'wrap_model']
diff --git a/python/setup.py b/python/setup.py
@@ -39,7 +39,7 @@
     'numpy',
     'pandas',
     'scipy',
-    'ml-wrappers==0.0.5',
+    'ml-wrappers==0.0.6',
     'scikit-learn',
     'packaging',
     'interpret-core[required]>=0.1.20, <=0.2.7',