omikuji 0.5.1

an efficient implementation of Partitioned Label Treesand its variations for extreme multi-label classification
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

use crate::util::create_progress_bar;
use crate::{DataSet, Index, IndexValueVec, Model};
use hashbrown::HashSet;
use itertools::izip;
use log::info;
use rayon::prelude::*;
use std::sync::Mutex;
use std::time;

fn precision_at_k(
    max_k: usize,
    true_labels: &[HashSet<Index>],
    predicted_labels: &[IndexValueVec],
) -> Vec<f32> {
    assert_eq!(true_labels.len(), predicted_labels.len());
    let mut ps = vec![0.; max_k];
    for (truth, predictions) in izip!(true_labels, predicted_labels) {
        let mut n_correct = 0;
        for k in 0..max_k.min(predictions.len()) {
            if truth.contains(&predictions[k].0) {
                n_correct += 1;
            }
            ps[k] += n_correct as f32 / (k + 1) as f32;
        }
    }
    for p in &mut ps {
        *p /= predicted_labels.len() as f32;
    }
    ps
}

pub fn test_all(
    model: &Model,
    test_dataset: &DataSet,
    beam_size: usize,
) -> (Vec<IndexValueVec>, Vec<f32>) {
    let n_examples = test_dataset.feature_lists.len();
    let pb = Mutex::new(create_progress_bar(n_examples as u64));
    let start_t = time::Instant::now();
    let predicted_labels = test_dataset
        .feature_lists
        .par_iter()
        .map(|feature_vec| {
            let predictions = model.predict(feature_vec, beam_size);
            pb.lock().expect("Failed to lock progress bar").add(1);
            predictions
        })
        .collect::<Vec<_>>();
    info!(
        "Done testing on {} examples; it took {:.2}s",
        n_examples,
        start_t.elapsed().as_secs_f32()
    );

    let precisions = precision_at_k(5, &test_dataset.label_sets, &predicted_labels);
    info!(
        "Precision@[1, 3, 5] = [{:.2}, {:.2}, {:.2}]",
        precisions[0] * 100.,
        precisions[2] * 100.,
        precisions[4] * 100.,
    );

    (predicted_labels, precisions)
}