add Tuple as possible type hint for EvalPredictions label_ids (#14473)

* Update trainer_utils.py * add Tuple type hints to all label_ids outputs affects EvalLoopOutput and PredicctionOutput
huggingface · Nov 21, 2021 · c9d2cf8 · c9d2cf8
1 parent a59e7c1
commit c9d2cf8
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/transformers/trainer_utils.py b/src/transformers/trainer_utils.py
@@ -75,19 +75,19 @@ class EvalPrediction(NamedTuple):
     """
 
     predictions: Union[np.ndarray, Tuple[np.ndarray]]
-    label_ids: np.ndarray
+    label_ids: Union[np.ndarray, Tuple[np.ndarray]]
 
 
 class EvalLoopOutput(NamedTuple):
     predictions: Union[np.ndarray, Tuple[np.ndarray]]
-    label_ids: Optional[np.ndarray]
+    label_ids: Optional[Union[np.ndarray, Tuple[np.ndarray]]]
     metrics: Optional[Dict[str, float]]
     num_samples: Optional[int]
 
 
 class PredictionOutput(NamedTuple):
     predictions: Union[np.ndarray, Tuple[np.ndarray]]
-    label_ids: Optional[np.ndarray]
+    label_ids: Optional[Union[np.ndarray, Tuple[np.ndarray]]]
     metrics: Optional[Dict[str, float]]