fix CPE precision evaluation error

1 files changed, 3 insertions, 5 deletions

diff --git a/sec_certs/model/evaluation.py b/sec_certs/model/evaluation.py
index 77cc2431..8885ddf5 100644
--- a/sec_certs/model/evaluation.py
+++ b/sec_certs/model/evaluation.py
@@ -33,18 +33,16 @@ def compute_precision(y: np.array, y_pred: np.array, **kwargs):
     return np.mean(prec)
 
 
-def evaluate(x_valid: List[Union[CommonCriteriaCert, FIPSCertificate]], y_valid: List[Optional[List[CPE]]], outpath: Optional[Union[Path, str]]):
+def evaluate(x_valid: List[Union[CommonCriteriaCert, FIPSCertificate]], y_valid: List[Optional[List[str]]], outpath: Optional[Union[Path, str]]):
     y_pred = [x.heuristics.cpe_matches for x in x_valid]
-    y_valid_uris = [[x.uri for x in y] if y else None for y in y_valid]
-
-    precision = compute_precision(y_valid_uris, y_pred)
+    precision = compute_precision(y_valid, y_pred)
 
     correctly_classified = []
     badly_classified = []
     n_new_certs_with_match = 0
     n_newly_identified = 0
 
-    for cert, predicted_cpes, verified_cpes in zip(x_valid, y_pred, y_valid_uris):
+    for cert, predicted_cpes, verified_cpes in zip(x_valid, y_pred, y_valid):
         verified_cpes_set = set(verified_cpes) if verified_cpes else set()
         predicted_cpes_set = set(predicted_cpes) if predicted_cpes else set()