uqa_ml/
model.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Serializable deep-fusion model specs plus CPU inference helpers.
8
9use std::collections::BTreeMap;
10
11use serde::{de::DeserializeOwned, Deserialize, Deserializer, Serialize};
12use uqa_core::{DocId, Value};
13use uqa_operators::{base::Direction, ExecutionContext, Operator};
14
15use crate::backend::{try_clone_slice, try_vec_with_capacity, MLBackend, MLError, MLResult};
16use crate::deep_fusion::{
17    AggregationKind, DeepFusionOperator, Gating, GlobalPoolMethod, Layer, PoolMethod,
18};
19
20/// Output of [`DeepModel`] inference: `(doc_id, score)` pairs plus,
21/// when the model ends in `Softmax`, per-doc class probability vectors.
22pub type PredictResult = (Vec<(DocId, f64)>, BTreeMap<DocId, Vec<f64>>);
23
24#[allow(clippy::upper_case_acronyms)]
25#[derive(Debug, Clone, Serialize, PartialEq)]
26#[serde(tag = "kind", rename_all = "snake_case")]
27pub enum DeepLayerSpec {
28    /// Runtime feature-vector input. Used by trained models and batched
29    /// feature inference.
30    Input {
31        dimensions: usize,
32    },
33    Embed {
34        embedding: Vec<f64>,
35    },
36    Dense {
37        weights: Vec<f64>,
38        bias: Vec<f64>,
39        output_channels: usize,
40        input_channels: usize,
41    },
42    Flatten,
43    GlobalPool {
44        method: PoolMethodSpec,
45    },
46    Softmax,
47    BatchNorm {
48        epsilon: f64,
49    },
50    Dropout {
51        p: f64,
52    },
53    #[serde(rename = "cnn_1d")]
54    CNN1D {
55        weights: Vec<f64>,
56        bias: Vec<f64>,
57        output_channels: usize,
58        input_channels: usize,
59        kernel_size: usize,
60        stride: usize,
61        padding: usize,
62    },
63    #[serde(rename = "cnn_2d")]
64    CNN2D {
65        weights: Vec<f64>,
66        bias: Vec<f64>,
67        output_channels: usize,
68        input_channels: usize,
69        input_height: usize,
70        input_width: usize,
71        kernel_height: usize,
72        kernel_width: usize,
73        stride_height: usize,
74        stride_width: usize,
75        padding_height: usize,
76        padding_width: usize,
77    },
78    Propagate {
79        edge_label: String,
80        aggregation: AggregationSpec,
81        direction: DirectionSpec,
82    },
83    Conv {
84        edge_label: String,
85        hop_weights: Vec<f64>,
86        direction: DirectionSpec,
87    },
88    Pool {
89        edge_label: String,
90        pool_size: usize,
91        method: PoolKindSpec,
92        direction: DirectionSpec,
93    },
94    Attention,
95    #[serde(rename = "rnn")]
96    RNN {
97        weights_input: Vec<f64>,
98        weights_hidden: Vec<f64>,
99        bias: Vec<f64>,
100        hidden_channels: usize,
101        input_channels: usize,
102        return_sequences: bool,
103    },
104    #[serde(rename = "lstm")]
105    LSTM {
106        weights_input: Vec<f64>,
107        weights_hidden: Vec<f64>,
108        bias: Vec<f64>,
109        hidden_channels: usize,
110        input_channels: usize,
111        return_sequences: bool,
112    },
113}
114
115impl<'de> Deserialize<'de> for DeepLayerSpec {
116    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
117    where
118        D: Deserializer<'de>,
119    {
120        let mut fields = BTreeMap::<String, serde_json::Value>::deserialize(deserializer)?;
121        let kind = take_layer_field::<String, D::Error>(&mut fields, "layer", "kind")?;
122        macro_rules! field {
123            ($name:literal) => {
124                take_layer_field::<_, D::Error>(&mut fields, &kind, $name)
125            };
126        }
127        Ok(match kind.as_str() {
128            "input" => Self::Input {
129                dimensions: field!("dimensions")?,
130            },
131            "embed" => Self::Embed {
132                embedding: field!("embedding")?,
133            },
134            "dense" => Self::Dense {
135                weights: field!("weights")?,
136                bias: field!("bias")?,
137                output_channels: field!("output_channels")?,
138                input_channels: field!("input_channels")?,
139            },
140            "flatten" => Self::Flatten,
141            "global_pool" => Self::GlobalPool {
142                method: field!("method")?,
143            },
144            "softmax" => Self::Softmax,
145            "batch_norm" => Self::BatchNorm {
146                epsilon: field!("epsilon")?,
147            },
148            "dropout" => Self::Dropout { p: field!("p")? },
149            "cnn_1d" => Self::CNN1D {
150                weights: field!("weights")?,
151                bias: field!("bias")?,
152                output_channels: field!("output_channels")?,
153                input_channels: field!("input_channels")?,
154                kernel_size: field!("kernel_size")?,
155                stride: field!("stride")?,
156                padding: field!("padding")?,
157            },
158            "cnn_2d" => Self::CNN2D {
159                weights: field!("weights")?,
160                bias: field!("bias")?,
161                output_channels: field!("output_channels")?,
162                input_channels: field!("input_channels")?,
163                input_height: field!("input_height")?,
164                input_width: field!("input_width")?,
165                kernel_height: field!("kernel_height")?,
166                kernel_width: field!("kernel_width")?,
167                stride_height: field!("stride_height")?,
168                stride_width: field!("stride_width")?,
169                padding_height: field!("padding_height")?,
170                padding_width: field!("padding_width")?,
171            },
172            "propagate" => Self::Propagate {
173                edge_label: field!("edge_label")?,
174                aggregation: field!("aggregation")?,
175                direction: field!("direction")?,
176            },
177            "conv" => Self::Conv {
178                edge_label: field!("edge_label")?,
179                hop_weights: field!("hop_weights")?,
180                direction: field!("direction")?,
181            },
182            "pool" => Self::Pool {
183                edge_label: field!("edge_label")?,
184                pool_size: field!("pool_size")?,
185                method: field!("method")?,
186                direction: field!("direction")?,
187            },
188            "attention" => Self::Attention,
189            "rnn" => Self::RNN {
190                weights_input: field!("weights_input")?,
191                weights_hidden: field!("weights_hidden")?,
192                bias: field!("bias")?,
193                hidden_channels: field!("hidden_channels")?,
194                input_channels: field!("input_channels")?,
195                return_sequences: field!("return_sequences")?,
196            },
197            "lstm" => Self::LSTM {
198                weights_input: field!("weights_input")?,
199                weights_hidden: field!("weights_hidden")?,
200                bias: field!("bias")?,
201                hidden_channels: field!("hidden_channels")?,
202                input_channels: field!("input_channels")?,
203                return_sequences: field!("return_sequences")?,
204            },
205            other => {
206                return Err(serde::de::Error::custom(format!(
207                    "unknown deep layer kind `{other}`"
208                )))
209            }
210        })
211    }
212}
213
214fn take_layer_field<T, E>(
215    fields: &mut BTreeMap<String, serde_json::Value>,
216    kind: &str,
217    name: &'static str,
218) -> Result<T, E>
219where
220    T: DeserializeOwned,
221    E: serde::de::Error,
222{
223    let value = fields
224        .remove(name)
225        .ok_or_else(|| E::custom(format!("{kind} layer is missing `{name}`")))?;
226    serde_json::from_value(value)
227        .map_err(|error| E::custom(format!("invalid `{name}` in {kind} layer: {error}")))
228}
229
230#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
231#[serde(rename_all = "snake_case")]
232pub enum PoolMethodSpec {
233    Avg,
234    Max,
235    AvgMax,
236}
237
238#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
239#[serde(rename_all = "snake_case")]
240pub enum AggregationSpec {
241    Mean,
242    Sum,
243    Max,
244}
245
246#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
247#[serde(rename_all = "snake_case")]
248pub enum DirectionSpec {
249    Out,
250    In,
251    Both,
252}
253
254#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
255#[serde(rename_all = "snake_case")]
256pub enum PoolKindSpec {
257    Avg,
258    Max,
259}
260
261#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
262#[serde(rename_all = "snake_case")]
263pub enum GatingSpec {
264    #[default]
265    None,
266    ReLU,
267    Swish,
268}
269
270#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
271pub struct DeepModel {
272    pub layers: Vec<DeepLayerSpec>,
273    /// Legacy Python catalogs omitted fusion tuning when the model only
274    /// carried a layer graph. Zero preserves the historical neutral value.
275    #[serde(default)]
276    pub alpha: f64,
277    #[serde(default)]
278    pub gating: GatingSpec,
279}
280
281impl DeepModel {
282    pub fn to_layers(&self) -> MLResult<Vec<Layer>> {
283        let mut layers = try_vec_with_capacity(self.layers.len(), "deep model runtime layers")?;
284        for spec in &self.layers {
285            layers.push(layer_from_spec(spec)?);
286        }
287        Ok(layers)
288    }
289
290    pub fn gating_runtime(&self) -> Gating {
291        match self.gating {
292            GatingSpec::None => Gating::None,
293            GatingSpec::ReLU => Gating::ReLU,
294            GatingSpec::Swish => Gating::Swish,
295        }
296    }
297
298    pub fn input_dimensions(&self) -> Option<usize> {
299        match self.layers.first() {
300            Some(DeepLayerSpec::Input { dimensions }) => Some(*dimensions),
301            _ => None,
302        }
303    }
304
305    /// Validate layer ordering, tensor shapes, dimensions, and numeric parameters.
306    pub fn validate(&self) -> MLResult<()> {
307        DeepFusionOperator::new(self.to_layers()?, self.alpha, self.gating_runtime()).map(|_| ())
308    }
309
310    /// Run CPU inference against an operator execution context.
311    pub fn predict(&self, ctx: &ExecutionContext) -> MLResult<PredictResult> {
312        predict_cpu(self, ctx)
313    }
314
315    /// Run inference through a specific backend.
316    pub fn predict_with_backend<B: MLBackend>(
317        &self,
318        backend: &B,
319        ctx: &ExecutionContext,
320    ) -> MLResult<PredictResult> {
321        self.validate()?;
322        backend.predict(self, ctx)
323    }
324
325    pub fn predict_features(&self, examples: &[(DocId, Vec<f64>)]) -> MLResult<PredictResult> {
326        predict_feature_batch_cpu(self, examples)
327    }
328
329    pub fn predict_features_with_backend<B: MLBackend>(
330        &self,
331        backend: &B,
332        examples: &[(DocId, Vec<f64>)],
333    ) -> MLResult<PredictResult> {
334        self.validate()?;
335        backend.predict_features(self, examples)
336    }
337}
338
339pub(crate) fn predict_cpu(model: &DeepModel, ctx: &ExecutionContext) -> MLResult<PredictResult> {
340    let layers = model.to_layers()?;
341    if layers.is_empty() {
342        return Err(MLError::InvalidModel(
343            "deep model requires at least one layer".into(),
344        ));
345    }
346    if model.input_dimensions().is_some() {
347        return Err(MLError::InvalidModel(
348            "Input-based models require predict_features rather than context-only predict".into(),
349        ));
350    }
351    let op = DeepFusionOperator::new(layers, model.alpha, model.gating_runtime())?;
352    let posting_list = op
353        .execute(ctx)
354        .map_err(|error| MLError::Backend(error.to_string()))?;
355    posting_list_to_prediction(&posting_list)
356}
357
358pub(crate) fn predict_feature_batch_cpu(
359    model: &DeepModel,
360    examples: &[(DocId, Vec<f64>)],
361) -> MLResult<PredictResult> {
362    let layers = model.to_layers()?;
363    if layers.is_empty() {
364        return Err(MLError::InvalidModel(
365            "deep model requires at least one layer".into(),
366        ));
367    }
368    let Some(expected_dims) = model.input_dimensions() else {
369        return Err(MLError::InvalidModel(
370            "feature prediction requires a model whose first layer is Input".into(),
371        ));
372    };
373    let op = DeepFusionOperator::new(layers, model.alpha, model.gating_runtime())?;
374    let mut scores = try_vec_with_capacity(examples.len(), "feature prediction scores")?;
375    let mut probs = BTreeMap::new();
376    for (doc_id, features) in examples {
377        if features.len() != expected_dims {
378            return Err(MLError::InvalidModel(format!(
379                "feature vector for doc {doc_id} has dimension {}, expected {expected_dims}",
380                features.len()
381            )));
382        }
383        let feature_copy = try_clone_slice(features, "feature prediction input")?;
384        let sample_prediction = op
385            .execute_features(*doc_id, feature_copy, &ExecutionContext::new())
386            .map_err(|error| MLError::Backend(error.to_string()))?;
387        let (mut sample_scores, sample_probs) = posting_list_to_prediction(&sample_prediction)?;
388        scores.append(&mut sample_scores);
389        probs.extend(sample_probs);
390    }
391    scores.sort_by_key(|(doc_id, _)| *doc_id);
392    Ok((scores, probs))
393}
394
395pub(crate) fn posting_list_to_prediction(pl: &uqa_core::PostingList) -> MLResult<PredictResult> {
396    let mut scores: Vec<(DocId, f64)> =
397        try_vec_with_capacity(pl.len(), "posting-list prediction scores")?;
398    let mut probs: BTreeMap<DocId, Vec<f64>> = BTreeMap::new();
399    for entry in pl.entries() {
400        if !entry.payload.score.is_finite() {
401            return Err(MLError::Backend(format!(
402                "prediction score for doc {} is not finite: {}",
403                entry.doc_id, entry.payload.score
404            )));
405        }
406        scores.push((entry.doc_id, entry.payload.score));
407        if let Some(Value::List(items)) = entry.payload.fields.get("class_probs") {
408            let v: Vec<f64> = items
409                .iter()
410                .enumerate()
411                .map(|(index, value)| match value {
412                    Value::Float(value) if value.is_finite() => Ok(*value),
413                    other => Err(MLError::Backend(format!(
414                        "class_probs[{index}] for doc {} is not a finite float: {other:?}",
415                        entry.doc_id
416                    ))),
417                })
418                .collect::<MLResult<_>>()?;
419            probs.insert(entry.doc_id, v);
420        }
421    }
422    Ok((scores, probs))
423}
424
425fn layer_from_spec(spec: &DeepLayerSpec) -> MLResult<Layer> {
426    Ok(match spec {
427        DeepLayerSpec::Input { dimensions } => Layer::Input {
428            dimensions: *dimensions,
429        },
430        DeepLayerSpec::Embed { embedding } => {
431            Layer::Embed(try_clone_slice(embedding, "embedding layer")?)
432        }
433        DeepLayerSpec::Dense {
434            weights,
435            bias,
436            output_channels,
437            input_channels,
438        } => Layer::Dense {
439            weights: try_clone_slice(weights, "dense weights")?,
440            bias: try_clone_slice(bias, "dense bias")?,
441            output_channels: *output_channels,
442            input_channels: *input_channels,
443        },
444        DeepLayerSpec::Flatten => Layer::Flatten,
445        DeepLayerSpec::GlobalPool { method } => Layer::GlobalPool(match method {
446            PoolMethodSpec::Avg => GlobalPoolMethod::Avg,
447            PoolMethodSpec::Max => GlobalPoolMethod::Max,
448            PoolMethodSpec::AvgMax => GlobalPoolMethod::AvgMax,
449        }),
450        DeepLayerSpec::Softmax => Layer::Softmax,
451        DeepLayerSpec::BatchNorm { epsilon } => Layer::BatchNorm { epsilon: *epsilon },
452        DeepLayerSpec::Dropout { p } => Layer::Dropout { p: *p },
453        DeepLayerSpec::CNN1D { .. } => cnn_1d_from_spec(spec)?,
454        DeepLayerSpec::CNN2D { .. } => cnn_2d_from_spec(spec)?,
455        DeepLayerSpec::Propagate {
456            edge_label,
457            aggregation,
458            direction,
459        } => Layer::Propagate {
460            edge_label: edge_label.clone(),
461            aggregation: match aggregation {
462                AggregationSpec::Mean => AggregationKind::Mean,
463                AggregationSpec::Sum => AggregationKind::Sum,
464                AggregationSpec::Max => AggregationKind::Max,
465            },
466            direction: direction_runtime(*direction),
467        },
468        DeepLayerSpec::Conv {
469            edge_label,
470            hop_weights,
471            direction,
472        } => Layer::Conv {
473            edge_label: edge_label.clone(),
474            hop_weights: try_clone_slice(hop_weights, "graph convolution hop weights")?,
475            direction: direction_runtime(*direction),
476        },
477        DeepLayerSpec::Pool {
478            edge_label,
479            pool_size,
480            method,
481            direction,
482        } => Layer::Pool {
483            edge_label: edge_label.clone(),
484            pool_size: *pool_size,
485            method: match method {
486                PoolKindSpec::Avg => PoolMethod::Avg,
487                PoolKindSpec::Max => PoolMethod::Max,
488            },
489            direction: direction_runtime(*direction),
490        },
491        DeepLayerSpec::Attention => Layer::Attention,
492        DeepLayerSpec::RNN { .. } => recurrent_from_spec(spec, false)?,
493        DeepLayerSpec::LSTM { .. } => recurrent_from_spec(spec, true)?,
494    })
495}
496
497fn cnn_1d_from_spec(spec: &DeepLayerSpec) -> MLResult<Layer> {
498    let DeepLayerSpec::CNN1D {
499        weights,
500        bias,
501        output_channels,
502        input_channels,
503        kernel_size,
504        stride,
505        padding,
506    } = spec
507    else {
508        return Err(MLError::InvalidModel(
509            "internal CNN1D conversion mismatch".into(),
510        ));
511    };
512    Ok(Layer::CNN1D {
513        weights: try_clone_slice(weights, "CNN1D weights")?,
514        bias: try_clone_slice(bias, "CNN1D bias")?,
515        output_channels: *output_channels,
516        input_channels: *input_channels,
517        kernel_size: *kernel_size,
518        stride: *stride,
519        padding: *padding,
520    })
521}
522
523fn cnn_2d_from_spec(spec: &DeepLayerSpec) -> MLResult<Layer> {
524    let DeepLayerSpec::CNN2D {
525        weights,
526        bias,
527        output_channels,
528        input_channels,
529        input_height,
530        input_width,
531        kernel_height,
532        kernel_width,
533        stride_height,
534        stride_width,
535        padding_height,
536        padding_width,
537    } = spec
538    else {
539        return Err(MLError::InvalidModel(
540            "internal CNN2D conversion mismatch".into(),
541        ));
542    };
543    Ok(Layer::CNN2D {
544        weights: try_clone_slice(weights, "CNN2D weights")?,
545        bias: try_clone_slice(bias, "CNN2D bias")?,
546        output_channels: *output_channels,
547        input_channels: *input_channels,
548        input_height: *input_height,
549        input_width: *input_width,
550        kernel_height: *kernel_height,
551        kernel_width: *kernel_width,
552        stride_height: *stride_height,
553        stride_width: *stride_width,
554        padding_height: *padding_height,
555        padding_width: *padding_width,
556    })
557}
558
559fn recurrent_from_spec(spec: &DeepLayerSpec, lstm: bool) -> MLResult<Layer> {
560    let (weights_input, weights_hidden, bias, hidden_channels, input_channels, return_sequences) =
561        match spec {
562            DeepLayerSpec::RNN {
563                weights_input,
564                weights_hidden,
565                bias,
566                hidden_channels,
567                input_channels,
568                return_sequences,
569            }
570            | DeepLayerSpec::LSTM {
571                weights_input,
572                weights_hidden,
573                bias,
574                hidden_channels,
575                input_channels,
576                return_sequences,
577            } => (
578                weights_input,
579                weights_hidden,
580                bias,
581                *hidden_channels,
582                *input_channels,
583                *return_sequences,
584            ),
585            _ => {
586                return Err(MLError::InvalidModel(
587                    "internal recurrent conversion mismatch".into(),
588                ));
589            }
590        };
591    let prefix = if lstm { "LSTM" } else { "RNN" };
592    let weights_input = try_clone_slice(weights_input, &format!("{prefix} input weights"))?;
593    let weights_hidden = try_clone_slice(weights_hidden, &format!("{prefix} hidden weights"))?;
594    let bias = try_clone_slice(bias, &format!("{prefix} bias"))?;
595    Ok(if lstm {
596        Layer::LSTM {
597            weights_input,
598            weights_hidden,
599            bias,
600            hidden_channels,
601            input_channels,
602            return_sequences,
603        }
604    } else {
605        Layer::RNN {
606            weights_input,
607            weights_hidden,
608            bias,
609            hidden_channels,
610            input_channels,
611            return_sequences,
612        }
613    })
614}
615
616fn direction_runtime(dir: DirectionSpec) -> Direction {
617    match dir {
618        DirectionSpec::Out => Direction::Out,
619        DirectionSpec::In => Direction::In,
620        DirectionSpec::Both => Direction::Both,
621    }
622}
623
624#[cfg(test)]
625mod tests {
626    use super::*;
627    use uqa_core::{Payload, PostingEntry, PostingList};
628
629    #[test]
630    fn legacy_layer_only_models_receive_neutral_fusion_defaults() {
631        let model: DeepModel = serde_json::from_str(r#"{"layers":[]}"#).unwrap();
632        assert_eq!(model.alpha, 0.0);
633        assert_eq!(model.gating, GatingSpec::None);
634    }
635
636    #[test]
637    fn recurrent_acronyms_serialize_as_plain_names() {
638        let rnn = DeepLayerSpec::RNN {
639            weights_input: vec![1.0],
640            weights_hidden: vec![0.0],
641            bias: vec![0.0],
642            hidden_channels: 1,
643            input_channels: 1,
644            return_sequences: true,
645        };
646        let lstm = DeepLayerSpec::LSTM {
647            weights_input: vec![0.0; 4],
648            weights_hidden: vec![0.0; 4],
649            bias: vec![0.0; 4],
650            hidden_channels: 1,
651            input_channels: 1,
652            return_sequences: false,
653        };
654
655        let rnn_json = serde_json::to_value(&rnn).unwrap();
656        let lstm_json = serde_json::to_value(&lstm).unwrap();
657        assert_eq!(rnn_json["kind"], "rnn");
658        assert_eq!(lstm_json["kind"], "lstm");
659
660        let cnn = DeepLayerSpec::CNN1D {
661            weights: vec![1.0],
662            bias: vec![0.0],
663            output_channels: 1,
664            input_channels: 1,
665            kernel_size: 1,
666            stride: 1,
667            padding: 0,
668        };
669        let cnn_json = serde_json::to_value(&cnn).unwrap();
670        assert_eq!(cnn_json["kind"], "cnn_1d");
671    }
672
673    #[test]
674    fn internally_tagged_layers_with_float_fields_round_trip_through_json() {
675        let model = DeepModel {
676            layers: vec![DeepLayerSpec::Dense {
677                weights: vec![1.0, 0.0],
678                bias: vec![0.5],
679                output_channels: 1,
680                input_channels: 2,
681            }],
682            alpha: 0.25,
683            gating: GatingSpec::None,
684        };
685        let json = serde_json::to_string(&model).unwrap();
686        let decoded = serde_json::from_str::<DeepModel>(&json)
687            .unwrap_or_else(|error| panic!("failed to decode {json}: {error}"));
688        assert_eq!(decoded, model);
689    }
690
691    #[test]
692    fn prediction_conversion_rejects_non_finite_scores() {
693        let postings =
694            PostingList::from_unsorted(vec![PostingEntry::new(7, Payload::with_score(f64::NAN))]);
695        let error = posting_list_to_prediction(&postings)
696            .expect_err("non-finite predictions must cross the API as errors");
697        assert!(error.to_string().contains("not finite"));
698    }
699}
uqa_ml/model.rs

uqa_ml/
model.rs
