Handle missing task names in return_df when data_params lacks tasks

Use .get() to safely access data_params["tasks"]["name"] so
return_df works on models that have not been fit (e.g. loaded from
checkpoint without training data).

Author: LiudengZhang

src/grelu/lightning/__init__.py

@@ -860,14 +860,17 @@ def predict_on_dataset(
 
             if return_df:
                 if (preds.ndim == 3) and (preds.shape[-1] == 1):
-                    task_names = self.data_params["tasks"]["name"]
                     n_tasks_pred = preds.shape[-2]
-                    if n_tasks_pred != len(task_names):
-                        warnings.warn(
-                            f"Prediction has {n_tasks_pred} task(s) but the model"
-                            f" has {len(task_names)} task name(s), likely due to a"
-                            " prediction transform. Using generic column names."
-                        )
+                    task_names = (
+                        self.data_params.get("tasks", {}).get("name", None)
+                    )
+                    if task_names is None or n_tasks_pred != len(task_names):
+                        if task_names is not None:
+                            warnings.warn(
+                                f"Prediction has {n_tasks_pred} task(s) but the model"
+                                f" has {len(task_names)} task name(s), likely due to a"
+                                " prediction transform. Using generic column names."
+                            )
                         task_names = [
                             f"task_{i}" for i in range(n_tasks_pred)
                         ]