Skip to main content

uqa_ml/
backend.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use thiserror::Error;
8use uqa_core::DocId;
9use uqa_operators::ExecutionContext;
10
11use crate::model::{predict_cpu, predict_feature_batch_cpu, DeepModel, PredictResult};
12use crate::training::{deep_learn, DeepLearnOutput, LearnOptions, TrainingSet};
13
14#[derive(Debug, Error)]
15pub enum MLError {
16    #[error("{0}")]
17    InvalidModel(String),
18    #[error("{0}")]
19    InvalidTrainingSet(String),
20    #[error("{0}")]
21    Backend(String),
22}
23
24pub type MLResult<T> = Result<T, MLError>;
25
26pub(crate) fn try_vec_with_capacity<T>(capacity: usize, context: &str) -> MLResult<Vec<T>> {
27    let mut values = Vec::new();
28    values.try_reserve_exact(capacity).map_err(|error| {
29        MLError::Backend(format!(
30            "cannot allocate {capacity} elements for {context}: {error}"
31        ))
32    })?;
33    Ok(values)
34}
35
36pub(crate) fn try_filled_vec<T: Clone>(length: usize, value: T, context: &str) -> MLResult<Vec<T>> {
37    let mut values = try_vec_with_capacity(length, context)?;
38    values.resize(length, value);
39    Ok(values)
40}
41
42pub(crate) fn try_clone_slice<T: Clone>(values: &[T], context: &str) -> MLResult<Vec<T>> {
43    let mut cloned = try_vec_with_capacity(values.len(), context)?;
44    cloned.extend_from_slice(values);
45    Ok(cloned)
46}
47
48pub trait MLBackend {
49    fn name(&self) -> &'static str;
50
51    fn predict(&self, model: &DeepModel, ctx: &ExecutionContext) -> MLResult<PredictResult>;
52
53    fn predict_features(
54        &self,
55        model: &DeepModel,
56        examples: &[(DocId, Vec<f64>)],
57    ) -> MLResult<PredictResult>;
58
59    fn deep_learn(
60        &self,
61        training_set: &TrainingSet,
62        options: &LearnOptions,
63    ) -> MLResult<DeepLearnOutput>;
64}
65
66#[derive(Debug, Default, Clone, Copy)]
67pub struct CPUBackend;
68
69impl MLBackend for CPUBackend {
70    fn name(&self) -> &'static str {
71        "cpu"
72    }
73
74    fn predict(&self, model: &DeepModel, ctx: &ExecutionContext) -> MLResult<PredictResult> {
75        predict_cpu(model, ctx)
76    }
77
78    fn predict_features(
79        &self,
80        model: &DeepModel,
81        examples: &[(DocId, Vec<f64>)],
82    ) -> MLResult<PredictResult> {
83        predict_feature_batch_cpu(model, examples)
84    }
85
86    fn deep_learn(
87        &self,
88        training_set: &TrainingSet,
89        options: &LearnOptions,
90    ) -> MLResult<DeepLearnOutput> {
91        deep_learn(training_set, options)
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn impossible_runtime_allocations_are_errors() {
101        let error = try_filled_vec(usize::MAX, 0_u8, "test model channels")
102            .expect_err("an impossible external dimension must not panic or abort");
103        assert!(error.to_string().contains("cannot allocate"));
104    }
105}