🌐 AI搜索 & 代理 主页
Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions pgml-extension/src/metrics.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
/// Module providing various metrics used to rank the algorithms.
use pgrx::*;
use std::collections::{BTreeSet, HashMap};

use ndarray::{Array2, ArrayView1};
Expand Down Expand Up @@ -51,13 +52,17 @@ impl ConfusionMatrix {
y_hat: &ArrayView1<usize>,
num_classes: usize,
) -> ConfusionMatrix {
assert_eq!(ground_truth.len(), y_hat.len());
if ground_truth.len() != y_hat.len() {
error!("Can't compute metrics when the ground truth labels are a different size than the predicted labels. {} != {}", ground_truth.len(), y_hat.len())
};

// Distinct classes.
let mut classes = ground_truth.iter().collect::<BTreeSet<_>>();
classes.extend(&mut y_hat.iter().collect::<BTreeSet<_>>().into_iter());

assert_eq!(num_classes, classes.len());
if num_classes != classes.len() {
error!("Can't compute metrics when the number of classes in the test set is different than the number of classes in the training set. {} != {}", num_classes, classes.len())
};

// Class value = index in the confusion matrix
// e.g. class value 5 will be index 4 if there are classes 1, 2, 3 and 4 present.
Expand Down