1use std::path::Path;
7
8use anyhow::Result;
9use burn::tensor::{activation, Device, Tensor, TensorData};
10use zsfm_burn::{open_gguf_weights, GgufWeights};
11
12use crate::config::TtmConfig;
13
14fn device() -> Device {
15 Device::flex()
16}
17
18fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
19 weights.load_2d(name)
20}
21
22fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
23 weights.load_1d(name)
24}
25
26struct BurnMixerLayer {
27 patch_norm_w: Tensor<1>,
28 patch_norm_b: Tensor<1>,
29 patch_fc1_w: Tensor<2>,
30 patch_fc1_b: Tensor<1>,
31 patch_fc2_w: Tensor<2>,
32 patch_fc2_b: Tensor<1>,
33 patch_gate_w: Tensor<2>,
34 patch_gate_b: Tensor<1>,
35 feat_norm_w: Tensor<1>,
36 feat_norm_b: Tensor<1>,
37 feat_fc1_w: Tensor<2>,
38 feat_fc1_b: Tensor<1>,
39 feat_fc2_w: Tensor<2>,
40 feat_fc2_b: Tensor<1>,
41 feat_gate_w: Tensor<2>,
42 feat_gate_b: Tensor<1>,
43}
44
45struct BurnAdaptiveLevel {
46 factor: usize,
47 layers: Vec<BurnMixerLayer>,
48}
49
50pub struct BurnTtmModel {
51 config: TtmConfig,
52 patcher_w: Tensor<2>,
53 patcher_b: Tensor<1>,
54 enc_levels: Vec<BurnAdaptiveLevel>,
55 dec_adapter_w: Tensor<2>,
56 dec_adapter_b: Tensor<1>,
57 dec_layers: Vec<BurnMixerLayer>,
58 head_w: Tensor<2>,
59 head_b: Tensor<1>,
60}
61
62fn load_mixer_layer(weights: &GgufWeights, prefix: &str) -> Result<BurnMixerLayer> {
63 Ok(BurnMixerLayer {
64 patch_norm_w: load_v(weights, &format!("{prefix}.patch_norm.weight"))?,
65 patch_norm_b: load_v(weights, &format!("{prefix}.patch_norm.bias"))?,
66 patch_fc1_w: load_t(weights, &format!("{prefix}.patch_fc1.weight"))?,
67 patch_fc1_b: load_v(weights, &format!("{prefix}.patch_fc1.bias"))?,
68 patch_fc2_w: load_t(weights, &format!("{prefix}.patch_fc2.weight"))?,
69 patch_fc2_b: load_v(weights, &format!("{prefix}.patch_fc2.bias"))?,
70 patch_gate_w: load_t(weights, &format!("{prefix}.patch_gate.weight"))?,
71 patch_gate_b: load_v(weights, &format!("{prefix}.patch_gate.bias"))?,
72 feat_norm_w: load_v(weights, &format!("{prefix}.feat_norm.weight"))?,
73 feat_norm_b: load_v(weights, &format!("{prefix}.feat_norm.bias"))?,
74 feat_fc1_w: load_t(weights, &format!("{prefix}.feat_fc1.weight"))?,
75 feat_fc1_b: load_v(weights, &format!("{prefix}.feat_fc1.bias"))?,
76 feat_fc2_w: load_t(weights, &format!("{prefix}.feat_fc2.weight"))?,
77 feat_fc2_b: load_v(weights, &format!("{prefix}.feat_fc2.bias"))?,
78 feat_gate_w: load_t(weights, &format!("{prefix}.feat_gate.weight"))?,
79 feat_gate_b: load_v(weights, &format!("{prefix}.feat_gate.bias"))?,
80 })
81}
82
83fn linear(x: Tensor<2>, w: Tensor<2>, b: Tensor<1>) -> Tensor<2> {
84 zsfm_burn::linear::burn_linear_bias(x, w, b)
85}
86
87fn layer_norm(x: Tensor<2>, w: Tensor<1>, b: Tensor<1>, eps: f32) -> Tensor<2> {
88 zsfm_burn::norm::burn_layer_norm(x, w, b, eps)
89}
90
91fn gelu(x: Tensor<2>) -> Tensor<2> {
92 activation::gelu(x)
93}
94
95fn softmax_last_dim(x: Tensor<2>) -> Tensor<2> {
96 activation::softmax(x, 1)
97}
98
99impl BurnTtmModel {
100 pub fn load(gguf_path: &Path, config: TtmConfig) -> Result<Self> {
101 let weights = open_gguf_weights(gguf_path)?;
102
103 let patcher_w = load_t(&weights, "enc.patcher.weight")?;
104 let patcher_b = load_v(&weights, "enc.patcher.bias")?;
105
106 let n_levels = config.adaptive_patching_levels;
107 let mut enc_levels = Vec::with_capacity(n_levels);
108 for l in 0..n_levels {
109 let factor = 1usize << (n_levels - 1 - l);
110 let mut layers = Vec::with_capacity(config.num_layers);
111 for n in 0..config.num_layers {
112 layers.push(load_mixer_layer(
113 &weights,
114 &format!("enc.blk.{l}.layer.{n}"),
115 )?);
116 }
117 enc_levels.push(BurnAdaptiveLevel { factor, layers });
118 }
119
120 let dec_adapter_w = load_t(&weights, "dec.adapter.weight")?;
121 let dec_adapter_b = load_v(&weights, "dec.adapter.bias")?;
122 let mut dec_layers = Vec::with_capacity(config.decoder_num_layers);
123 for n in 0..config.decoder_num_layers {
124 dec_layers.push(load_mixer_layer(&weights, &format!("dec.blk.{n}"))?);
125 }
126 let head_w = load_t(&weights, "head.weight")?;
127 let head_b = load_v(&weights, "head.bias")?;
128
129 Ok(Self {
130 config,
131 patcher_w,
132 patcher_b,
133 enc_levels,
134 dec_adapter_w,
135 dec_adapter_b,
136 dec_layers,
137 head_w,
138 head_b,
139 })
140 }
141
142 pub fn forecast(&self, context: &[f32]) -> Result<Vec<f32>> {
143 let cfg = &self.config;
144 let (scaled, mean, std) = std_scale(context);
145 let patches = patchify(&scaled, cfg.patch_length, cfg.patch_stride, cfg.num_patches);
146 let flat: Vec<f32> = patches.into_iter().flatten().collect();
147 let dev = device();
148 let mut h = Tensor::<2>::from_data(
149 TensorData::new(flat, [cfg.num_patches, cfg.patch_length]),
150 &dev,
151 );
152
153 h = linear(h, self.patcher_w.clone(), self.patcher_b.clone());
154 for level in &self.enc_levels {
155 h = self.forward_adaptive_level(h, level);
156 }
157 h = linear(h, self.dec_adapter_w.clone(), self.dec_adapter_b.clone());
158 for layer in &self.dec_layers {
159 h = forward_mixer_layer(h, layer, cfg.norm_eps as f32);
160 }
161 let flat_dim = cfg.num_patches * cfg.decoder_d_model;
162 let h = h.reshape([flat_dim]);
163 let h = h.unsqueeze_dim::<2>(0);
164 let h = linear(h, self.head_w.clone(), self.head_b.clone());
165 let h: Tensor<1> = h.squeeze_dim(0);
166 let forecast: Vec<f32> = h.to_data().try_to_vec()?;
167 Ok(forecast.iter().map(|&v| v * std + mean).collect())
168 }
169
170 fn forward_adaptive_level(&self, hidden: Tensor<2>, level: &BurnAdaptiveLevel) -> Tensor<2> {
171 let factor = level.factor;
172 let [p, f] = hidden.dims();
173 let mut h = if factor > 1 {
174 hidden.reshape([p * factor, f / factor])
175 } else {
176 hidden
177 };
178 for layer in &level.layers {
179 h = forward_mixer_layer(h, layer, self.config.norm_eps as f32);
180 }
181 if factor > 1 {
182 h.reshape([p, f])
183 } else {
184 h
185 }
186 }
187}
188
189fn forward_mixer_layer(h: Tensor<2>, layer: &BurnMixerLayer, eps: f32) -> Tensor<2> {
190 forward_feat_mixer(forward_patch_mixer(h.clone(), layer, eps), layer, eps)
191}
192
193fn forward_patch_mixer(h: Tensor<2>, w: &BurnMixerLayer, eps: f32) -> Tensor<2> {
194 let normed = layer_norm(
195 h.clone(),
196 w.patch_norm_w.clone(),
197 w.patch_norm_b.clone(),
198 eps,
199 );
200 let t = normed.transpose();
201 let m = gelu(linear(t, w.patch_fc1_w.clone(), w.patch_fc1_b.clone()));
202 let m = linear(m, w.patch_fc2_w.clone(), w.patch_fc2_b.clone());
203 let gate = softmax_last_dim(linear(
204 m.clone(),
205 w.patch_gate_w.clone(),
206 w.patch_gate_b.clone(),
207 ));
208 let gated = m * gate;
209 gated.transpose() + h
210}
211
212fn forward_feat_mixer(h: Tensor<2>, w: &BurnMixerLayer, eps: f32) -> Tensor<2> {
213 let normed = layer_norm(h.clone(), w.feat_norm_w.clone(), w.feat_norm_b.clone(), eps);
214 let m = gelu(linear(normed, w.feat_fc1_w.clone(), w.feat_fc1_b.clone()));
215 let m = linear(m, w.feat_fc2_w.clone(), w.feat_fc2_b.clone());
216 let gate = softmax_last_dim(linear(
217 m.clone(),
218 w.feat_gate_w.clone(),
219 w.feat_gate_b.clone(),
220 ));
221 let gated = m * gate;
222 gated + h
223}
224
225fn std_scale(x: &[f32]) -> (Vec<f32>, f32, f32) {
226 let n = x.len() as f64;
227 let mean = x.iter().map(|&v| v as f64).sum::<f64>() / n;
228 let var = x.iter().map(|&v| (v as f64 - mean).powi(2)).sum::<f64>() / n;
229 let std = (var + 1e-5).sqrt() as f32;
230 let scaled: Vec<f32> = x.iter().map(|&v| (v as f32 - mean as f32) / std).collect();
231 (scaled, mean as f32, std)
232}
233
234fn patchify(
235 x: &[f32],
236 patch_length: usize,
237 patch_stride: usize,
238 num_patches: usize,
239) -> Vec<Vec<f32>> {
240 let new_seq_len = patch_length + patch_stride * (num_patches - 1);
241 let padded: Vec<f32> = if x.len() < new_seq_len {
242 let pad_val = x.first().copied().unwrap_or(0.0);
243 let mut v = vec![pad_val; new_seq_len - x.len()];
244 v.extend_from_slice(x);
245 v
246 } else {
247 x[x.len() - new_seq_len..].to_vec()
248 };
249 (0..num_patches)
250 .map(|i| padded[i * patch_stride..i * patch_stride + patch_length].to_vec())
251 .collect()
252}
253
254#[cfg(test)]
255mod tests {
256 use super::*;
257 use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
258
259 fn pseudo(n: usize, seed: u64) -> Vec<f32> {
260 let mut x = seed;
261 (0..n)
262 .map(|_| {
263 x = x
264 .wrapping_mul(6364136223846793005)
265 .wrapping_add(1442695040888963407);
266 ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
267 })
268 .collect()
269 }
270
271 fn test_config() -> TtmConfig {
272 TtmConfig {
273 context_length: 64,
274 prediction_length: 4,
275 patch_length: 4,
276 patch_stride: 2,
277 d_model: 8,
278 num_layers: 1,
279 decoder_d_model: 8,
280 decoder_num_layers: 1,
281 expansion_factor: 2,
282 adaptive_patching_levels: 1,
283 num_patches: 4,
284 scaling: "std".into(),
285 gated_attn: true,
286 norm_eps: 1e-5,
287 }
288 }
289
290 fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64) {
292 *seed += 1;
293 let data = pseudo(rows * cols, *seed);
294 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
295 w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
296 }
297
298 fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
299 *seed += 1;
300 let data = pseudo(n, *seed);
301 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
302 w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
303 }
304
305 fn put_layer(w: &mut GGUFWriter, prefix: &str, p: usize, f: usize, seed: &mut u64) {
306 put1(w, &format!("{prefix}.patch_norm.weight"), f, seed);
307 put1(w, &format!("{prefix}.patch_norm.bias"), f, seed);
308 put2(w, &format!("{prefix}.patch_fc1.weight"), 2 * p, p, seed);
309 put1(w, &format!("{prefix}.patch_fc1.bias"), 2 * p, seed);
310 put2(w, &format!("{prefix}.patch_fc2.weight"), p, 2 * p, seed);
311 put1(w, &format!("{prefix}.patch_fc2.bias"), p, seed);
312 put2(w, &format!("{prefix}.patch_gate.weight"), p, p, seed);
313 put1(w, &format!("{prefix}.patch_gate.bias"), p, seed);
314 put1(w, &format!("{prefix}.feat_norm.weight"), f, seed);
315 put1(w, &format!("{prefix}.feat_norm.bias"), f, seed);
316 put2(w, &format!("{prefix}.feat_fc1.weight"), 2 * f, f, seed);
317 put1(w, &format!("{prefix}.feat_fc1.bias"), 2 * f, seed);
318 put2(w, &format!("{prefix}.feat_fc2.weight"), f, 2 * f, seed);
319 put1(w, &format!("{prefix}.feat_fc2.bias"), f, seed);
320 put2(w, &format!("{prefix}.feat_gate.weight"), f, f, seed);
321 put1(w, &format!("{prefix}.feat_gate.bias"), f, seed);
322 }
323
324 #[test]
325 fn burn_forecast_golden() {
326 let dir = std::env::temp_dir().join(format!("ttm-burn-{}", std::process::id()));
327 let _ = std::fs::remove_dir_all(&dir);
328 std::fs::create_dir_all(&dir).unwrap();
329 let path = dir.join("model.gguf");
330
331 let mut w = GGUFWriter::new();
332 w.add_metadata("general.architecture", GGUFMetaValue::String("ttm".into()));
333 let mut seed = 100u64;
334 put2(&mut w, "enc.patcher.weight", 8, 4, &mut seed);
336 put1(&mut w, "enc.patcher.bias", 8, &mut seed);
337 put_layer(&mut w, "enc.blk.0.layer.0", 4, 8, &mut seed);
338 put2(&mut w, "dec.adapter.weight", 8, 8, &mut seed);
339 put1(&mut w, "dec.adapter.bias", 8, &mut seed);
340 put_layer(&mut w, "dec.blk.0", 4, 8, &mut seed);
341 put2(&mut w, "head.weight", 4, 32, &mut seed);
342 put1(&mut w, "head.bias", 4, &mut seed);
343
344 let f = std::fs::File::create(&path).unwrap();
345 w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
346
347 let cfg = test_config();
348 let burn = BurnTtmModel::load(&path, cfg).unwrap();
349
350 let ctx: Vec<f32> = (0..16).map(|i| 10.0 + 0.5 * i as f32).collect();
351 let t0 = std::time::Instant::now();
352 let a = burn.forecast(&ctx).unwrap();
353 let ms_a = t0.elapsed().as_secs_f64() * 1000.0;
354 let b = burn.forecast(&ctx).unwrap();
355 assert_eq!(a.len(), 4);
356 assert!(a.iter().all(|v| v.is_finite()));
357 assert!(a.iter().any(|&v| v != 0.0));
358 assert_eq!(a, b, "forecast must be deterministic");
359 println!("ttm synthetic: burn {ms_a:.3}ms out={a:?}");
360 let golden: [f32; 4] = [-414.10815, -425.1715, -473.70477, -365.27527];
362 assert_eq!(a.len(), golden.len());
363 for (x, y) in a.iter().zip(golden.iter()) {
364 assert!((x - y).abs() < 1e-2, "golden drift: {x} vs {y}");
365 }
366
367 let _ = std::fs::remove_dir_all(&dir);
368 }
369}