Update value dtypes

peak_signal_to_noise_ratio.py

@@ -57,13 +57,32 @@ class StructuralSimilarityIndexMeasure(evaluate.Metric):
             description=_DESCRIPTION,
             citation=_CITATION,
             inputs_description=_KWARGS_DESCRIPTION,
-            features=datasets.Features({
-                "predictions": datasets.Sequence(datasets.Array2D(dtype="float32")),
-                "references": datasets.Sequence(datasets.Array2D(dtype="float32")),
-            }),
+            features=datasets.Features(self._get_feature_types()),
             reference_urls=["https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio"],
         )
 
+    def _get_feature_types(self):
+        if self.config_name == "multilist":
+            return {
+                # 1st Seq - num_samples, 2nd Seq - Height, 3rd Seq - Width
+                "predictions": datasets.Sequence(
+                    datasets.Sequence(datasets.Sequence(datasets.Value("float32")))
+                ),
+                "references": datasets.Sequence(
+                    datasets.Sequence(datasets.Sequence(datasets.Value("float32")))
+                ),
+            }
+        else:
+            return {
+                # 1st Seq - Height, 2rd Seq - Width
+                "predictions": datasets.Sequence(
+                    datasets.Sequence(datasets.Value("float32"))
+                ),
+                "references": datasets.Sequence(
+                    datasets.Sequence(datasets.Value("float32"))
+                ),
+            }
+
     def _compute(
         self, 
         predictions, 
@@ -72,12 +91,8 @@ class StructuralSimilarityIndexMeasure(evaluate.Metric):
         sample_weight=None,
     ) -> Dict[str, float]:
         samples = zip(predictions, references)
-        return {
-            "psnr": np.average(
-                list(map(
-                    lambda args: peak_signal_noise_ratio(*args, data_range), 
-                    samples
-                )), 
-                weights=sample_weight
-            )
-        }
+        return np.average(
+            list(map(lambda args: peak_signal_noise_ratio(*args, data_range), samples)), 
+            weights=sample_weight
+        )
+