fix

jer.py

@@ -14,7 +14,7 @@
 """TODO: Add a description here."""
 
 from operator import eq
-from typing import Iterable
+from typing import Callable, Iterable, Union
 
 import evaluate
 import datasets
@@ -56,6 +56,8 @@ Examples:
     {'recall': 0.5, 'precision': 1.0, 'f1': 0.6666666666666666}
 """
 
+Triplet = Union[str, tuple, int]
+
 @evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
 class jer(evaluate.Metric):
     """TODO: Short description of my evaluation module."""
@@ -82,18 +84,23 @@ class jer(evaluate.Metric):
 
     def _download_and_prepare(self, dl_manager):
         """Optional: download external resources useful to compute the scores"""
-        # TODO: Download external resources if needed
         pass
 
     def _compute(self, predictions, references, eq_fn=eq):
         """Returns the scores"""
         score_dicts = [
-            self._compute_single(prediction=prediction, reference=reference) 
+            self._compute_single(prediction=prediction, reference=reference, eq_fn=eq_fn) 
             for prediction, reference in zip(predictions, references)
         ]
         return {('mean_' + key): np.mean([scores[key] for scores in score_dicts]) for key in score_dicts[0].keys()}
     
-    def _compute_single(self, *, prediction: Iterable[str | tuple | int], reference: Iterable[str | tuple | int], eq_fn=eq):
+    def _compute_single(
+            self, 
+            *, 
+            prediction: Iterable[Triplet], 
+            reference: Iterable[Triplet], 
+            eq_fn: Callable[[Triplet, Triplet], bool],
+        ):
         reference_set = set(reference)
         assert len(reference) == len(reference_set), f"Duplicates found in the reference list {reference}"
         prediction_set = set(prediction)