_

Author: jeromedockes

skrub/_data_ops/_estimator.py

@@ -289,14 +289,17 @@ def report(self, *, environment, mode, **full_report_kwargs):
         self._set_is_fitted(mode)
         return result
 
-    def _score(self, environment):
+    def _score(self, environment, return_predictions=False):
         score_node = find_scoring_node(self.data_op)
         if score_node is None:
-            return self._eval_in_mode("score", environment)
+            result = self._eval_in_mode("score", environment)
+            return (result, {}) if return_predictions else result
         estimator = self.__skrub_to_Xy_pipeline__(environment)
         cv_data = _compute_X_y_and_cv(self.data_op, environment)
         X, y = cv_data["X"], cv_data.get("y")
-        return estimator._score(X, y, cast_to_float=False)
+        return estimator._score(
+            X, y, cast_to_float=False, return_predictions=return_predictions
+        )
 
     def __getattr__(self, name):
         if name not in supported_modes(self.data_op):
@@ -899,38 +902,51 @@ def _process_scores(scorer_info, scorer_output):
                     name = "score"
         return [(name, scorer_output)]
 
-    def _score(self, X, y=None, cast_to_float=True):
+    def _score(self, X, y=None, cast_to_float=True, return_predictions=False):
         score_node = find_scoring_node(self.data_op)
         if score_node is None:
-            return self._eval_in_mode("score", X, y)
+            result = self._eval_in_mode("score", X, y)
+            return (result, {}) if return_predictions else result
         env = self._get_env(X, y)
         scorers = evaluate(
             score_node._skrub_impl.scorers, mode="fit_transform", environment=env
         )
         all_scores = []
-        caching_estimator = _CachingXyPipeline(self.data_op, self.environment)
+        predictions = env.get("_skrub_predictions", {})
+        cache = {(k, id(X)): v for k, v in predictions.items()}
+        caching_estimator = _CachingXyPipeline(
+            self.data_op, self.environment, cache=cache
+        )
         _copy_attr(self, caching_estimator, ["_is_fitted"])
         for scorer_info in scorers:
             scorer = self._prepare_scorer(scorer_info["scoring"], scorer_info["kwargs"])
             scorer_output = scorer(caching_estimator, X, y)
             all_scores.extend(self._process_scores(scorer_info, scorer_output))
+        predictions.update({k: v for ((k, X_id), v) in cache.items() if X_id == id(X)})
         rename = unique_renaming()
         result = {rename(name): score for name, score in all_scores}
-        if cast_to_float and len(result) == 1:
+        if cast_to_float and len(result) == 1 and not return_predictions:
             # If there is a single score stick to scikit-learn interface which
             # returns a number.
             return next(iter(result.values()))
-        return result
+        return (result, predictions) if return_predictions else result
 
 
 class _CachingXyPipeline(_XyPipeline):
+    def __init__(self, data_op, environment, cache):
+        super().__init__(data_op, environment)
+        self.cache = cache
+
     def _eval_in_mode(self, mode, X, y=None):
-        if not hasattr(self, "_cache"):
-            self._cache = {}
-        key = (mode, id(X), id(y))
-        if key not in self._cache:
-            self._cache[key] = super()._eval_in_mode(mode, X, y=y)
-        return self._cache[key]
+        if y is not None:
+            return super()._eval_in_mode(mode, X, y=y)
+        key = (mode, id(X))
+        if key not in self.cache:
+            self.cache[key] = super()._eval_in_mode(mode, X, y=y)
+        return self.cache[key]
+
+    def _score(self, X, y=None):
+        return super()._eval_in_mode("score", X, y=y)
 
 
 def _compute_X_y_and_cv(data_op, environment):