Implements collect_logits and collect_labels in ray backend.

ludwig-ai · Jun 14, 2022 · fb4c31b · fb4c31b
1 parent 1a88d79
commit fb4c31b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/ludwig/backend/ray.py b/ludwig/backend/ray.py
@@ -653,6 +653,9 @@ def to_tensors(df: pd.DataFrame) -> pd.DataFrame:
         if collect_labels:
             columns = [f.proc_column for f in self.model.output_features.values()]
             labels = self.df_engine.from_ray_dataset(dataset.ds)[columns]
+            labels = labels.rename(
+                columns={f.proc_column: f"{f.feature_name}_labels" for f in self.model.output_features.values()}
+            )
             predictions = dask.dataframe.concat([predictions, labels])
 
         return predictions

diff --git a/ludwig/models/calibrator.py b/ludwig/models/calibrator.py
@@ -44,8 +44,10 @@ def train_calibration(self, dataset, dataset_name: str):
             if output_feature.calibration_module is not None:
                 feature_logits_key = f"{output_feature.feature_name}_logits"
                 if feature_logits_key in predictions:
-                    feature_logits = predictions[feature_logits_key]
-                    feature_labels = predictions[f"{output_feature.feature_name}_labels"]
+                    feature_logits = self.backend.df_engine.compute(predictions[feature_logits_key])
+                    feature_labels = self.backend.df_engine.compute(
+                        predictions[f"{output_feature.feature_name}_labels"]
+                    )
                     output_feature.calibration_module.train_calibration(
                         np.stack(feature_logits.values, axis=0), np.stack(feature_labels.values, axis=0)
                     )