1use super::{
10 apply_attention, apply_batch_norm, apply_cnn_1d_layer, apply_cnn_2d_layer, apply_conv,
11 apply_dense, apply_dropout, apply_embed, apply_flatten, apply_global_pool, apply_lstm_layer,
12 apply_pool, apply_propagate, apply_rnn_layer, apply_signal, apply_softmax, build_result,
13 runtime_model_error, validate_state, BTreeMap, DeepFusionOperator, ExecutionContext,
14 ForwardState, IndexStats, Layer, Operator, OperatorResult, StorageBackendError,
15 StorageBackendResult,
16};
17
18impl DeepFusionOperator {
19 fn validate_feature_input(&self, features: &[f64]) -> StorageBackendResult<()> {
20 let Some(Layer::Input { dimensions }) = self.layers.first() else {
21 return Err(runtime_model_error(
22 "execute_features requires an Input-based deep-fusion model",
23 ));
24 };
25 if features.len() != *dimensions {
26 return Err(StorageBackendError::Other(format!(
27 "deep-fusion feature vector has dimension {}, expected {dimensions}",
28 features.len()
29 )));
30 }
31 if let Some((index, value)) = features
32 .iter()
33 .enumerate()
34 .find(|(_, value)| !value.is_finite())
35 {
36 return Err(StorageBackendError::Other(format!(
37 "deep-fusion feature {index} must be finite, got {value}"
38 )));
39 }
40 Ok(())
41 }
42
43 pub(super) fn coverage_default(coverage: usize, total: usize) -> f64 {
44 uqa_operators::hybrid::coverage_based_default(coverage, total, 0.01)
45 }
46
47 pub fn execute_features(
48 &self,
49 doc_id: u64,
50 features: Vec<f64>,
51 ctx: &ExecutionContext,
52 ) -> OperatorResult {
53 self.validate_feature_input(&features)?;
54 let mut state = ForwardState {
55 num_channels: features.len(),
56 channel_map: BTreeMap::from([(doc_id, features)]),
57 softmax_applied: false,
58 };
59 self.apply_layers(ctx, &mut state)?;
60 build_result(
61 &state.channel_map,
62 state.num_channels,
63 state.softmax_applied,
64 )
65 }
66
67 fn apply_layers(
68 &self,
69 ctx: &ExecutionContext,
70 state: &mut ForwardState,
71 ) -> StorageBackendResult<()> {
72 for layer in &self.layers {
73 self.apply_layer(ctx, state, layer)?;
74 validate_state(state)?;
75 }
76 Ok(())
77 }
78
79 fn apply_layer(
80 &self,
81 ctx: &ExecutionContext,
82 state: &mut ForwardState,
83 layer: &Layer,
84 ) -> StorageBackendResult<()> {
85 match layer {
86 Layer::Input { dimensions } => {
87 state.num_channels = *dimensions;
88 }
89 Layer::Embed(embedding) => apply_embed(embedding, state)?,
90 Layer::Signal(signals) => {
91 apply_signal(signals, ctx, self.alpha, self.gating, state)?;
92 }
93 Layer::Dense {
94 weights,
95 bias,
96 output_channels,
97 input_channels,
98 } => apply_dense(
99 weights,
100 bias,
101 *output_channels,
102 *input_channels,
103 self.gating,
104 state,
105 )?,
106 Layer::Flatten => apply_flatten(state)?,
107 Layer::GlobalPool(method) => apply_global_pool(*method, state)?,
108 Layer::Softmax => apply_softmax(state)?,
109 Layer::BatchNorm { epsilon } => apply_batch_norm(*epsilon, state)?,
110 Layer::Dropout { p } => apply_dropout(*p, state),
111 Layer::CNN1D { .. } => apply_cnn_1d_layer(layer, self.gating, state)?,
112 Layer::CNN2D { .. } => apply_cnn_2d_layer(layer, self.gating, state)?,
113 Layer::Propagate {
114 edge_label,
115 aggregation,
116 direction,
117 } => apply_propagate(
118 edge_label,
119 *aggregation,
120 *direction,
121 ctx,
122 self.gating,
123 state,
124 )?,
125 Layer::Conv {
126 edge_label,
127 hop_weights,
128 direction,
129 } => apply_conv(edge_label, hop_weights, *direction, ctx, self.gating, state)?,
130 Layer::Pool {
131 edge_label,
132 pool_size,
133 method,
134 direction,
135 } => apply_pool(edge_label, *pool_size, *method, *direction, ctx, state)?,
136 Layer::Attention => apply_attention(state)?,
137 Layer::RNN { .. } => apply_rnn_layer(layer, state)?,
138 Layer::LSTM { .. } => apply_lstm_layer(layer, state)?,
139 }
140 Ok(())
141 }
142}
143
144impl Operator for DeepFusionOperator {
145 fn execute(&self, ctx: &ExecutionContext) -> OperatorResult {
146 if matches!(self.layers.first(), Some(Layer::Input { .. })) {
147 return Err(runtime_model_error(
148 "Input-based deep-fusion models require execute_features",
149 ));
150 }
151 let mut state = ForwardState {
152 channel_map: BTreeMap::new(),
153 num_channels: 1,
154 softmax_applied: false,
155 };
156 self.apply_layers(ctx, &mut state)?;
157 build_result(
158 &state.channel_map,
159 state.num_channels,
160 state.softmax_applied,
161 )
162 }
163
164 fn cost_estimate(&self, stats: &IndexStats) -> f64 {
165 let mut total = 0.0f64;
166 for layer in &self.layers {
167 match layer {
168 Layer::Signal(signals) => {
169 for s in signals {
170 total += s.cost_estimate(stats);
171 }
172 }
173 Layer::Input { dimensions } => total += *dimensions as f64,
174 Layer::Embed(emb) => total += emb.len() as f64,
175 Layer::Dense {
176 output_channels,
177 input_channels,
178 ..
179 } => total += (*output_channels as f64) * (*input_channels as f64),
180 Layer::Flatten
181 | Layer::GlobalPool(_)
182 | Layer::Softmax
183 | Layer::BatchNorm { .. }
184 | Layer::Dropout { .. }
185 | Layer::CNN1D { .. }
186 | Layer::CNN2D { .. }
187 | Layer::RNN { .. }
188 | Layer::LSTM { .. }
189 | Layer::Propagate { .. }
190 | Layer::Conv { .. }
191 | Layer::Pool { .. } => total += stats.total_docs as f64,
192 Layer::Attention => {
193 let n = stats.total_docs as f64;
194 total += n * n;
195 }
196 }
197 }
198 total
199 }
200}