1use 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}