Handle DatasetDict in clean_and_validate_predictions (#22)

Co-authored-by: Oleksandr Shchur <[email protected]>

Author: abdulfatir

src/fev/task.py

@@ -599,17 +599,16 @@ def _to_dataset(preds: datasets.Dataset | list[dict]) -> datasets.Dataset:
                 raise ValueError(f"predictions must be of type `datasets.Dataset` (received {type(preds)})")
             return preds
 
-        if self.is_multivariate:
-            if isinstance(predictions, datasets.DatasetDict):
-                pass
-            elif isinstance(predictions, dict):
-                predictions = datasets.DatasetDict({col: _to_dataset(preds) for col, preds in predictions.items()})
+        if not isinstance(predictions, datasets.DatasetDict):
+            if self.is_multivariate:
+                if isinstance(predictions, dict):
+                    predictions = datasets.DatasetDict({col: _to_dataset(preds) for col, preds in predictions.items()})
+                else:
+                    raise ValueError(
+                        f"predictions for multivariate tasks must be of type `datasets.DatasetDict` or `dict` (received {type(predictions)})"
+                    )
             else:
-                raise ValueError(
-                    f"predictions for multivariate tasks must be of type `datasets.DatasetDict` or `dict` (received {type(predictions)})"
-                )
-        else:
-            predictions = datasets.DatasetDict({self.target_column: _to_dataset(predictions)})
+                predictions = datasets.DatasetDict({self.target_column: _to_dataset(predictions)})
 
         predictions = predictions.cast(self.predictions_schema).with_format("numpy")
         for target_column, predictions_for_column in predictions.items():

test/test_task.py

@@ -145,6 +145,31 @@ def test_when_multivariate_task_is_created_then_data_contains_correct_columns(ta
     assert set(future_data.column_names) == set(all_column_names) - set(target_column)
 
 
+@pytest.mark.parametrize("return_list", [True, False])
+def test_when_predictions_provided_as_dataset_dict_for_univariate_task_then_predictions_can_be_scores(return_list):
+    def naive_forecast_univariate(task: fev.Task) -> list[dict]:
+        past_data, future_data = task.get_input_data()
+        predictions = []
+        for ts in past_data:
+            predictions.append({"predictions": [ts[task.target_column][-1] for _ in range(task.horizon)]})
+        if return_list:
+            return predictions
+        else:
+            return datasets.DatasetDict({task.target_column: datasets.Dataset.from_list(predictions)})
+
+    task = fev.Task(
+        dataset_path="autogluon/chronos_datasets",
+        dataset_config="monash_m1_yearly",
+        eval_metric="MASE",
+        extra_metrics=["WAPE"],
+        horizon=4,
+    )
+    predictions = naive_forecast_univariate(task)
+    summary = task.evaluation_summary(predictions, model_name="naive")
+    for metric in ["MASE", "WAPE"]:
+        assert np.isfinite(summary[metric])
+
+
 @pytest.mark.parametrize("target_column", [["OT"], ["OT", "LULL", "HULL"]])
 @pytest.mark.parametrize("return_dict", [True, False])
 def test_when_multivariate_task_is_used_then_predictions_can_be_scored(target_column, return_dict):