From 9fed69eb3ca676e607f656ec6831decfe0d816f0 Mon Sep 17 00:00:00 2001 From: Tim Head Date: Thu, 9 Nov 2023 16:27:05 +0100 Subject: [PATCH] Add rich HTML representation to estimators (#5630) This adds a Jupyter (and other notebook) rich display hook that produces a HTML widget to represent an estimator in notebooks. This adds the basics of having estimators displayed as HTML widgets in notebooks and other editors that use the Jupyter notebook "rich display" system. Screenshot 2023-10-26 at 15 33 37 This doesn't yet contain the cool feature of changing colour depending on fit status or the link to the documentation. For that we'd have to depend on a newer version of scikit-learn (or vendor the logic). In this case "newer" actually means "the next version to be released". WDYT? Authors: - Tim Head (https://github.com/betatim) - Simon Adorf (https://github.com/csadorf) - Dante Gama Dessavre (https://github.com/dantegd) Approvers: - Simon Adorf (https://github.com/csadorf) URL: https://github.com/rapidsai/cuml/pull/5630 --- python/cuml/internals/base.pyx | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/cuml/internals/base.pyx b/python/cuml/internals/base.pyx index 4fb03fdac9..09dda66064 100644 --- a/python/cuml/internals/base.pyx +++ b/python/cuml/internals/base.pyx @@ -28,6 +28,8 @@ from cuml.internals.safe_imports import ( np = cpu_only_import('numpy') nvtx_annotate = gpu_only_import_from("nvtx", "annotate", alt=null_decorator) +from sklearn.utils import estimator_html_repr + import cuml import cuml.common import cuml.internals.logger as logger @@ -443,6 +445,12 @@ class Base(TagsMixin, return {'preserves_dtype': [self.dtype]} return {} + def _repr_mimebundle_(self, **kwargs): + """Prepare representations used by jupyter kernels to display estimator""" + output = {"text/plain": repr(self)} + output["text/html"] = estimator_html_repr(self) + return output + def set_nvtx_annotations(self): for func_name in ['fit', 'transform', 'predict', 'fit_transform', 'fit_predict']: