fix typing

Author: voorhs

autointent/metrics/retrieval.py

@@ -114,14 +114,14 @@ def ignore_oos(func: RetrievalMetricFn) -> RetrievalMetricFn:
     """Ignore OOS in metrics calculation (decorator)."""
 
     @wraps(func)
-    def wrapper(query_labels: list[Any | None], candidates_labels: list[Any]) -> float:
+    def wrapper(query_labels: LABELS_VALUE_TYPE, candidates_labels: CANDIDATE_TYPE) -> float:
         query_labels_filtered = [lab for lab in query_labels if lab is not None]
         candidates_labels_filtered = [
             cand for cand, lab in zip(candidates_labels, query_labels, strict=True) if lab is not None
         ]
-        return func(query_labels_filtered, candidates_labels_filtered)
+        return func(query_labels_filtered, candidates_labels_filtered)  # type: ignore[arg-type]
 
-    return wrapper
+    return wrapper  # type: ignore[return-value]
 
 
 @ignore_oos

autointent/metrics/scoring.py

@@ -2,7 +2,7 @@
 
 import logging
 from functools import wraps
-from typing import Any, Protocol
+from typing import Protocol
 
 import numpy as np
 from sklearn.metrics import coverage_error, label_ranking_average_precision_score, label_ranking_loss, roc_auc_score
@@ -34,10 +34,10 @@ def ignore_oos(func: ScoringMetricFn) -> ScoringMetricFn:
     """Ignore OOS in metrics calculation (decorator)."""
 
     @wraps(func)
-    def wrapper(labels: list[Any | None], scores: list[Any]) -> float:
+    def wrapper(labels: LABELS_VALUE_TYPE, scores: SCORES_VALUE_TYPE) -> float:
         labels_filtered = [lab for lab in labels if lab is not None]
         scores_filtered = [score for score, lab in zip(scores, labels, strict=True) if lab is not None]
-        return func(labels_filtered, scores_filtered)
+        return func(labels_filtered, scores_filtered)  # type: ignore[arg-type]
 
     return wrapper