update

top5_error_rate.py

@@ -53,24 +53,13 @@ class Top5ErrorRate(evaluate.Metric):
     def _compute(
         self,
         *,
-        predictions: list[list[float]] = None,
-        references: list[float] = None,
+        predictions: list[list[int]] = None,
+        references: list[int] = None,
         **kwargs,
     ) -> Dict[str, Any]:
-        # 确保输入是numpy数组
-        predictions = np.array(predictions)
-        references = np.array(references)
 
-        # 获取每个样本的top-5预测类别
-        top5_pred = np.argsort(predictions, axis=1)[:, -5:]
-
-        # 计算top-5错误率
-        correct = 0
         total = len(references)
-
-        for i in range(total):
-            if references[i] in top5_pred[i]:
-                correct += 1
+        correct = sum(1 for pred, ref in zip(predictions, references) if ref in pred)
 
         error_rate = 1.0 - (correct / total)