uqa-ml 0.1.9

Machine-learning model specs, training, and inference backends for UQA
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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

use thiserror::Error;
use uqa_core::DocId;
use uqa_operators::ExecutionContext;

use crate::model::{predict_cpu, predict_feature_batch_cpu, DeepModel, PredictResult};
use crate::training::{deep_learn, DeepLearnOutput, LearnOptions, TrainingSet};

#[derive(Debug, Error)]
pub enum MLError {
    #[error("{0}")]
    InvalidModel(String),
    #[error("{0}")]
    InvalidTrainingSet(String),
    #[error("{0}")]
    Backend(String),
}

pub type MLResult<T> = Result<T, MLError>;

pub(crate) fn try_vec_with_capacity<T>(capacity: usize, context: &str) -> MLResult<Vec<T>> {
    let mut values = Vec::new();
    values.try_reserve_exact(capacity).map_err(|error| {
        MLError::Backend(format!(
            "cannot allocate {capacity} elements for {context}: {error}"
        ))
    })?;
    Ok(values)
}

pub(crate) fn try_filled_vec<T: Clone>(length: usize, value: T, context: &str) -> MLResult<Vec<T>> {
    let mut values = try_vec_with_capacity(length, context)?;
    values.resize(length, value);
    Ok(values)
}

pub(crate) fn try_clone_slice<T: Clone>(values: &[T], context: &str) -> MLResult<Vec<T>> {
    let mut cloned = try_vec_with_capacity(values.len(), context)?;
    cloned.extend_from_slice(values);
    Ok(cloned)
}

pub trait MLBackend {
    fn name(&self) -> &'static str;

    fn predict(&self, model: &DeepModel, ctx: &ExecutionContext) -> MLResult<PredictResult>;

    fn predict_features(
        &self,
        model: &DeepModel,
        examples: &[(DocId, Vec<f64>)],
    ) -> MLResult<PredictResult>;

    fn deep_learn(
        &self,
        training_set: &TrainingSet,
        options: &LearnOptions,
    ) -> MLResult<DeepLearnOutput>;
}

#[derive(Debug, Default, Clone, Copy)]
pub struct CPUBackend;

impl MLBackend for CPUBackend {
    fn name(&self) -> &'static str {
        "cpu"
    }

    fn predict(&self, model: &DeepModel, ctx: &ExecutionContext) -> MLResult<PredictResult> {
        predict_cpu(model, ctx)
    }

    fn predict_features(
        &self,
        model: &DeepModel,
        examples: &[(DocId, Vec<f64>)],
    ) -> MLResult<PredictResult> {
        predict_feature_batch_cpu(model, examples)
    }

    fn deep_learn(
        &self,
        training_set: &TrainingSet,
        options: &LearnOptions,
    ) -> MLResult<DeepLearnOutput> {
        deep_learn(training_set, options)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn impossible_runtime_allocations_are_errors() {
        let error = try_filled_vec(usize::MAX, 0_u8, "test model channels")
            .expect_err("an impossible external dimension must not panic or abort");
        assert!(error.to_string().contains("cannot allocate"));
    }
}