1use std::path::Path;
8
9use anyhow::{Context, Result};
10use burn::tensor::{activation, Device, Tensor, TensorData};
11use zsfm_burn::linear::burn_linear_nd;
12use zsfm_burn::norm::burn_layer_norm_nd;
13use zsfm_burn::{open_gguf_weights, GgufWeights};
14
15use crate::config::{MitraConfig, Task};
16
17const LN_EPS: f32 = 1e-5;
18
19fn device() -> Device {
20 Device::flex()
21}
22
23fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
24 weights.load_2d(name)
25}
26
27fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
28 weights.load_1d(name)
29}
30
31fn linear<const D: usize>(x: Tensor<D>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<D> {
32 burn_linear_nd(x, w, b)
33}
34
35fn layer_norm<const D: usize>(x: Tensor<D>, w: Tensor<1>, b: Tensor<1>) -> Tensor<D> {
36 burn_layer_norm_nd(x, w, b, LN_EPS)
37}
38
39fn gelu<const D: usize>(x: Tensor<D>) -> Tensor<D> {
40 activation::gelu(x)
41}
42
43fn softmax_last<const D: usize>(x: Tensor<D>) -> Tensor<D> {
44 activation::softmax(x, D - 1)
45}
46
47struct BurnAttnW {
48 q_w: Tensor<2>,
49 q_b: Tensor<1>,
50 k_w: Tensor<2>,
51 k_b: Tensor<1>,
52 v_w: Tensor<2>,
53 v_b: Tensor<1>,
54 o_w: Tensor<2>,
55 o_b: Tensor<1>,
56}
57
58struct BurnBlockW {
59 ln1_w: Tensor<1>,
60 ln1_b: Tensor<1>,
61 ln2_w: Tensor<1>,
62 ln2_b: Tensor<1>,
63 ln3_w: Tensor<1>,
64 ln3_b: Tensor<1>,
65 ln4_w: Tensor<1>,
66 ln4_b: Tensor<1>,
67 attn_row: BurnAttnW,
68 attn_feat: BurnAttnW,
69 mlp1_fc1_w: Tensor<2>,
70 mlp1_fc1_b: Tensor<1>,
71 mlp1_fc2_w: Tensor<2>,
72 mlp1_fc2_b: Tensor<1>,
73 mlp2_fc1_w: Tensor<2>,
74 mlp2_fc1_b: Tensor<1>,
75 mlp2_fc2_w: Tensor<2>,
76 mlp2_fc2_b: Tensor<1>,
77}
78
79pub struct BurnMitraModel {
80 config: MitraConfig,
81 x_embed_w: Tensor<2>,
82 x_embed_b: Tensor<1>,
83 y_embed_w: Tensor<2>,
84 y_embed_b: Option<Tensor<1>>,
85 y_mask_w: Tensor<1>,
86 blocks: Vec<BurnBlockW>,
87 norm_f_w: Tensor<1>,
88 norm_f_b: Tensor<1>,
89 head_w: Tensor<2>,
90 head_b: Tensor<1>,
91}
92
93fn load_attn(weights: &GgufWeights, prefix: &str) -> Result<BurnAttnW> {
94 Ok(BurnAttnW {
95 q_w: load_t(weights, &format!("{prefix}.q.weight"))?,
96 q_b: load_v(weights, &format!("{prefix}.q.bias"))?,
97 k_w: load_t(weights, &format!("{prefix}.k.weight"))?,
98 k_b: load_v(weights, &format!("{prefix}.k.bias"))?,
99 v_w: load_t(weights, &format!("{prefix}.v.weight"))?,
100 v_b: load_v(weights, &format!("{prefix}.v.bias"))?,
101 o_w: load_t(weights, &format!("{prefix}.o.weight"))?,
102 o_b: load_v(weights, &format!("{prefix}.o.bias"))?,
103 })
104}
105
106impl BurnMitraModel {
107 pub fn load(gguf_path: &Path, config: MitraConfig) -> Result<Self> {
108 let weights = open_gguf_weights(gguf_path)
109 .with_context(|| format!("open {}", gguf_path.display()))?;
110
111 let x_embed_w = load_t(&weights, "x_embed.weight")?;
112 let x_embed_b = load_v(&weights, "x_embed.bias")?;
113 let y_embed_w = load_t(&weights, "y_embed.weight")?;
114 let y_embed_b = match config.task {
115 Task::Regression => Some(load_v(&weights, "y_embed.bias")?),
116 Task::Classification => None,
117 };
118 let y_mask_w = load_v(&weights, "y_mask.weight")?;
119
120 let mut blocks = Vec::with_capacity(config.n_layers);
121 for n in 0..config.n_layers {
122 let p = |s: &str| format!("blk.{n}.{s}");
123 blocks.push(BurnBlockW {
124 ln1_w: load_v(&weights, &p("ln1.weight"))?,
125 ln1_b: load_v(&weights, &p("ln1.bias"))?,
126 ln2_w: load_v(&weights, &p("ln2.weight"))?,
127 ln2_b: load_v(&weights, &p("ln2.bias"))?,
128 ln3_w: load_v(&weights, &p("ln3.weight"))?,
129 ln3_b: load_v(&weights, &p("ln3.bias"))?,
130 ln4_w: load_v(&weights, &p("ln4.weight"))?,
131 ln4_b: load_v(&weights, &p("ln4.bias"))?,
132 attn_row: load_attn(&weights, &p("attn_row"))?,
133 attn_feat: load_attn(&weights, &p("attn_feat"))?,
134 mlp1_fc1_w: load_t(&weights, &p("mlp1_fc1.weight"))?,
135 mlp1_fc1_b: load_v(&weights, &p("mlp1_fc1.bias"))?,
136 mlp1_fc2_w: load_t(&weights, &p("mlp1_fc2.weight"))?,
137 mlp1_fc2_b: load_v(&weights, &p("mlp1_fc2.bias"))?,
138 mlp2_fc1_w: load_t(&weights, &p("mlp2_fc1.weight"))?,
139 mlp2_fc1_b: load_v(&weights, &p("mlp2_fc1.bias"))?,
140 mlp2_fc2_w: load_t(&weights, &p("mlp2_fc2.weight"))?,
141 mlp2_fc2_b: load_v(&weights, &p("mlp2_fc2.bias"))?,
142 });
143 }
144
145 Ok(Self {
146 config,
147 x_embed_w,
148 x_embed_b,
149 y_embed_w,
150 y_embed_b,
151 y_mask_w,
152 blocks,
153 norm_f_w: load_v(&weights, "norm_f.weight")?,
154 norm_f_b: load_v(&weights, "norm_f.bias")?,
155 head_w: load_t(&weights, "head.weight")?,
156 head_b: load_v(&weights, "head.bias")?,
157 })
158 }
159
160 pub fn predict_classification(
161 &self,
162 x_support: &[Vec<f32>],
163 y_support: &[usize],
164 x_query: &[Vec<f32>],
165 n_classes: usize,
166 ) -> Result<Vec<Vec<f32>>> {
167 anyhow::ensure!(
168 self.config.task == Task::Classification,
169 "model was loaded as a regressor"
170 );
171 let (x_s, x_q, kept) = preprocess_x(x_support, x_query);
172 let n_s = x_s.len();
173 let n_q = x_q.len();
174 let n_feat = kept.len();
175 let x_emb = self.embed_x(&x_s, &x_q, n_s, n_q, n_feat);
176 let y_support_emb = self.embed_y_classes_support(y_support, n_s)?;
177 let y_query_mask = self.embed_y_mask(n_q);
178 let logits = self.forward(x_emb, y_support_emb, y_query_mask, n_s, n_q, n_feat);
179 let logits: Vec<f32> = to_host(&logits)?;
180 let dim_out = self.config.dim_output;
181 Ok((0..n_q)
182 .map(|i| logits[i * dim_out..i * dim_out + n_classes].to_vec())
183 .collect())
184 }
185
186 pub fn predict_regression(
187 &self,
188 x_support: &[Vec<f32>],
189 y_support: &[f32],
190 x_query: &[Vec<f32>],
191 ) -> Result<Vec<f32>> {
192 anyhow::ensure!(
193 self.config.task == Task::Regression,
194 "model was loaded as a classifier"
195 );
196 let (x_s, x_q, kept) = preprocess_x(x_support, x_query);
197 let n_s = x_s.len();
198 let n_q = x_q.len();
199 let n_feat = kept.len();
200 let y_min = y_support.iter().copied().fold(f32::INFINITY, f32::min);
201 let y_max = y_support.iter().copied().fold(f32::NEG_INFINITY, f32::max);
202 anyhow::ensure!(
203 y_max > y_min,
204 "y_support must have at least two distinct values"
205 );
206 let y_scaled: Vec<f32> = y_support
207 .iter()
208 .map(|&v| (v - y_min) / (y_max - y_min))
209 .collect();
210 let x_emb = self.embed_x(&x_s, &x_q, n_s, n_q, n_feat);
211 let y_support_emb = self.embed_y_regression_support(&y_scaled)?;
212 let y_query_mask = self.embed_y_mask(n_q);
213 let out = self.forward(x_emb, y_support_emb, y_query_mask, n_s, n_q, n_feat);
214 let out: Vec<f32> = to_host(&out)?;
215 Ok(out
216 .into_iter()
217 .map(|v| v * (y_max - y_min) + y_min)
218 .collect())
219 }
220
221 fn embed_x(
222 &self,
223 x_s: &[Vec<f32>],
224 x_q: &[Vec<f32>],
225 n_s: usize,
226 n_q: usize,
227 n_feat: usize,
228 ) -> (Tensor<3>, Tensor<3>) {
229 let dev = device();
230 let (bx_s, bx_q) = quantile_bucketize_normalize(x_s, x_q);
231 let flat_s: Vec<f32> = bx_s.into_iter().flatten().collect();
232 let flat_q: Vec<f32> = bx_q.into_iter().flatten().collect();
233 let d = self.config.dim;
234 let x_s_t = Tensor::<2>::from_data(TensorData::new(flat_s, [n_s * n_feat, 1]), &dev);
235 let x_q_t = Tensor::<2>::from_data(TensorData::new(flat_q, [n_q * n_feat, 1]), &dev);
236 let x_s_emb = linear(x_s_t, self.x_embed_w.clone(), Some(self.x_embed_b.clone()))
237 .reshape([n_s, n_feat, d]);
238 let x_q_emb = linear(x_q_t, self.x_embed_w.clone(), Some(self.x_embed_b.clone()))
239 .reshape([n_q, n_feat, d]);
240 (x_s_emb, x_q_emb)
241 }
242
243 fn embed_y_classes_support(&self, y_support: &[usize], n_s: usize) -> Result<Tensor<3>> {
244 let dim = self.config.dim;
245 let table: Vec<f32> = to_host(&self.y_embed_w)?;
246 let mut out = vec![0f32; n_s * dim];
247 for (i, &cls) in y_support.iter().enumerate() {
248 out[i * dim..(i + 1) * dim].copy_from_slice(&table[cls * dim..(cls + 1) * dim]);
249 }
250 Ok(Tensor::<3>::from_data(
251 TensorData::new(out, [n_s, 1, dim]),
252 &device(),
253 ))
254 }
255
256 fn embed_y_regression_support(&self, y_scaled: &[f32]) -> Result<Tensor<3>> {
257 let n_s = y_scaled.len();
258 let dev = device();
259 let y_t = Tensor::<2>::from_data(TensorData::new(y_scaled.to_vec(), [n_s, 1]), &dev);
260 let b = self
261 .y_embed_b
262 .as_ref()
263 .context("regressor y_embed missing bias")?;
264 let emb = linear(y_t, self.y_embed_w.clone(), Some(b.clone()));
265 Ok(emb.reshape([n_s, 1, self.config.dim]))
266 }
267
268 fn embed_y_mask(&self, n_q: usize) -> Tensor<3> {
269 let dim = self.config.dim;
270 let row: Vec<f32> = to_host_1(&self.y_mask_w).expect("mask read");
271 let mut out = Vec::with_capacity(n_q * dim);
272 for _ in 0..n_q {
273 out.extend_from_slice(&row);
274 }
275 Tensor::<3>::from_data(TensorData::new(out, [n_q, 1, dim]), &device())
276 }
277
278 fn forward(
279 &self,
280 x_emb: (Tensor<3>, Tensor<3>),
281 y_support_emb: Tensor<3>,
282 y_query_mask: Tensor<3>,
283 n_s: usize,
284 n_q: usize,
285 n_feat: usize,
286 ) -> Tensor<2> {
287 let (x_s_emb, x_q_emb) = x_emb;
288 let mut support = Tensor::cat(vec![y_support_emb, x_s_emb], 1);
289 let mut query = Tensor::cat(vec![y_query_mask, x_q_emb], 1);
290 let f1 = n_feat + 1;
291 for blk in &self.blocks {
292 let (s2, q2) =
293 query_block_forward(&support, &query, blk, self.config.n_heads, n_s, n_q, f1);
294 support = s2;
295 query = q2;
296 }
297 let query = layer_norm(query, self.norm_f_w.clone(), self.norm_f_b.clone());
298 let query = linear(query, self.head_w.clone(), Some(self.head_b.clone()));
299 query.narrow(1, 0, 1).reshape([n_q, self.config.dim_output])
300 }
301}
302
303fn to_host<const D: usize>(t: &Tensor<D>) -> Result<Vec<f32>> {
304 t.to_data()
305 .try_to_vec()
306 .map_err(|e| anyhow::anyhow!("{e:?}"))
307}
308
309fn to_host_1(t: &Tensor<1>) -> Result<Vec<f32>> {
310 to_host(t)
311}
312
313fn query_block_forward(
314 support: &Tensor<3>,
315 query: &Tensor<3>,
316 blk: &BurnBlockW,
317 n_heads: usize,
318 n_s: usize,
319 n_q: usize,
320 f1: usize,
321) -> (Tensor<3>, Tensor<3>) {
322 let s_ln = layer_norm(support.clone(), blk.ln1_w.clone(), blk.ln1_b.clone());
323 let q_ln = layer_norm(query.clone(), blk.ln1_w.clone(), blk.ln1_b.clone());
324 let s_row: Tensor<3> = s_ln.permute([1, 0, 2]);
325 let q_row: Tensor<3> = q_ln.permute([1, 0, 2]);
326 let s_att = mha(&s_row, &s_row, &s_row, &blk.attn_row, n_heads);
327 let q_att = mha(&q_row, &s_row, &s_row, &blk.attn_row, n_heads);
328 let s_att: Tensor<3> = s_att.permute([1, 0, 2]);
329 let q_att: Tensor<3> = q_att.permute([1, 0, 2]);
330 let s_dim = s_att.dims()[2];
331 let q_dim = q_att.dims()[2];
332 let s_att = s_att.reshape([n_s, f1, s_dim]);
333 let q_att = q_att.reshape([n_q, f1, q_dim]);
334 let mut support = support.clone() + s_att;
335 let mut query = query.clone() + q_att;
336
337 let s_ln = layer_norm(support.clone(), blk.ln2_w.clone(), blk.ln2_b.clone());
338 let q_ln = layer_norm(query.clone(), blk.ln2_w.clone(), blk.ln2_b.clone());
339 let s_mlp = linear(
340 gelu(linear(
341 s_ln,
342 blk.mlp1_fc1_w.clone(),
343 Some(blk.mlp1_fc1_b.clone()),
344 )),
345 blk.mlp1_fc2_w.clone(),
346 Some(blk.mlp1_fc2_b.clone()),
347 );
348 let q_mlp = linear(
349 gelu(linear(
350 q_ln,
351 blk.mlp1_fc1_w.clone(),
352 Some(blk.mlp1_fc1_b.clone()),
353 )),
354 blk.mlp1_fc2_w.clone(),
355 Some(blk.mlp1_fc2_b.clone()),
356 );
357 support = support + s_mlp;
358 query = query + q_mlp;
359
360 let s_ln = layer_norm(support.clone(), blk.ln3_w.clone(), blk.ln3_b.clone());
361 let q_ln = layer_norm(query.clone(), blk.ln3_w.clone(), blk.ln3_b.clone());
362 let s_att = mha(&s_ln, &s_ln, &s_ln, &blk.attn_feat, n_heads);
363 let q_att = mha(&q_ln, &q_ln, &q_ln, &blk.attn_feat, n_heads);
364 support = support + s_att;
365 query = query + q_att;
366
367 let s_ln = layer_norm(support.clone(), blk.ln4_w.clone(), blk.ln4_b.clone());
368 let q_ln = layer_norm(query.clone(), blk.ln4_w.clone(), blk.ln4_b.clone());
369 let s_mlp = linear(
370 gelu(linear(
371 s_ln,
372 blk.mlp2_fc1_w.clone(),
373 Some(blk.mlp2_fc1_b.clone()),
374 )),
375 blk.mlp2_fc2_w.clone(),
376 Some(blk.mlp2_fc2_b.clone()),
377 );
378 let q_mlp = linear(
379 gelu(linear(
380 q_ln,
381 blk.mlp2_fc1_w.clone(),
382 Some(blk.mlp2_fc1_b.clone()),
383 )),
384 blk.mlp2_fc2_w.clone(),
385 Some(blk.mlp2_fc2_b.clone()),
386 );
387 support = support + s_mlp;
388 query = query + q_mlp;
389 (support, query)
390}
391
392fn mha(q: &Tensor<3>, k: &Tensor<3>, v: &Tensor<3>, w: &BurnAttnW, n_heads: usize) -> Tensor<3> {
393 let qd = q.dims();
394 let (batch, sq, dim) = (qd[0], qd[1], qd[2]);
395 let kd = k.dims();
396 let skv = kd[1];
397 let head_dim = dim / n_heads;
398 let q = linear(q.clone(), w.q_w.clone(), Some(w.q_b.clone()));
399 let k = linear(k.clone(), w.k_w.clone(), Some(w.k_b.clone()));
400 let v = linear(v.clone(), w.v_w.clone(), Some(w.v_b.clone()));
401 let q: Tensor<4> = q
402 .reshape([batch, sq, n_heads, head_dim])
403 .permute([0, 2, 1, 3]);
404 let k: Tensor<4> = k
405 .reshape([batch, skv, n_heads, head_dim])
406 .permute([0, 2, 1, 3]);
407 let v: Tensor<4> = v
408 .reshape([batch, skv, n_heads, head_dim])
409 .permute([0, 2, 1, 3]);
410 let scale = 1.0f32 / (head_dim as f32).sqrt();
411 let scores = q.matmul(k.transpose()).mul_scalar(scale);
412 let attn = activation::softmax(scores, 3);
413 let out = attn.matmul(v);
414 let out: Tensor<4> = out.permute([0, 2, 1, 3]);
415 let out: Tensor<3> = out.reshape([batch, sq, dim]);
416 linear(out, w.o_w.clone(), Some(w.o_b.clone()))
417}
418
419fn preprocess_x(
420 x_support: &[Vec<f32>],
421 x_query: &[Vec<f32>],
422) -> (Vec<Vec<f32>>, Vec<Vec<f32>>, Vec<usize>) {
423 let n_feat = x_support[0].len();
424 let mut col_mean = vec![0f32; n_feat];
425 for f in 0..n_feat {
426 let mut sum = 0f64;
427 let mut count = 0usize;
428 for row in x_support {
429 if !row[f].is_nan() {
430 sum += row[f] as f64;
431 count += 1;
432 }
433 }
434 col_mean[f] = if count > 0 {
435 (sum / count as f64) as f32
436 } else {
437 0.0
438 };
439 }
440 let impute = |row: &[f32]| -> Vec<f32> {
441 row.iter()
442 .enumerate()
443 .map(|(f, &v)| if v.is_nan() { col_mean[f] } else { v })
444 .collect()
445 };
446 let support_imputed: Vec<Vec<f32>> = x_support.iter().map(|r| impute(r)).collect();
447 let query_imputed: Vec<Vec<f32>> = x_query.iter().map(|r| impute(r)).collect();
448 let kept: Vec<usize> = (0..n_feat)
449 .filter(|&f| {
450 let first = support_imputed[0][f];
451 support_imputed.iter().any(|r| r[f] != first)
452 })
453 .collect();
454 let select = |rows: &[Vec<f32>]| -> Vec<Vec<f32>> {
455 rows.iter()
456 .map(|r| kept.iter().map(|&f| r[f]).collect())
457 .collect()
458 };
459 (select(&support_imputed), select(&query_imputed), kept)
460}
461
462fn quantile_bucketize_normalize(
463 x_support: &[Vec<f32>],
464 x_query: &[Vec<f32>],
465) -> (Vec<Vec<f32>>, Vec<Vec<f32>>) {
466 let n_support = x_support.len();
467 let n_features = x_support[0].len();
468 let n_query = x_query.len();
469 let mut out_support = vec![vec![0f32; n_features]; n_support];
470 let mut out_query = vec![vec![0f32; n_features]; n_query];
471 for feat in 0..n_features {
472 let mut col: Vec<f32> = x_support.iter().map(|r| r[feat]).collect();
473 col.sort_by(|a, b| a.total_cmp(b));
474 let boundaries: Vec<f32> = (1..1000)
475 .map(|i| quantile_linear(&col, i as f64 / 1000.0))
476 .collect();
477 let bucket = |v: f32| -> f32 {
478 let idx = boundaries.partition_point(|&b| b < v);
479 idx as f32 / n_support as f32
480 };
481 let bucketed_support: Vec<f32> = x_support.iter().map(|r| bucket(r[feat])).collect();
482 let mean: f32 = bucketed_support.iter().sum::<f32>() / n_support as f32;
483 let var: f32 = bucketed_support
484 .iter()
485 .map(|&v| (v - mean).powi(2))
486 .sum::<f32>()
487 / n_support as f32;
488 let std = var.sqrt();
489 for (i, &v) in bucketed_support.iter().enumerate() {
490 out_support[i][feat] = if std == 0.0 { 0.0 } else { (v - mean) / std };
491 }
492 for (i, row) in x_query.iter().enumerate() {
493 let v = bucket(row[feat]);
494 out_query[i][feat] = if std == 0.0 { 0.0 } else { (v - mean) / std };
495 }
496 }
497 (out_support, out_query)
498}
499
500fn quantile_linear(sorted: &[f32], q: f64) -> f32 {
501 let n = sorted.len();
502 if n == 1 {
503 return sorted[0];
504 }
505 let idx = q * (n - 1) as f64;
506 let lo = idx.floor() as usize;
507 let hi = idx.ceil() as usize;
508 let frac = (idx - lo as f64) as f32;
509 sorted[lo] + frac * (sorted[hi] - sorted[lo])
510}
511
512#[cfg(test)]
513mod tests {
514 use super::*;
515 use crate::config::Task as CfgTask;
516 use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
517
518 fn pseudo(n: usize, seed: u64) -> Vec<f32> {
519 let mut x = seed;
520 (0..n)
521 .map(|_| {
522 x = x
523 .wrapping_mul(6364136223846793005)
524 .wrapping_add(1442695040888963407);
525 ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
526 })
527 .collect()
528 }
529
530 fn tiny_config() -> MitraConfig {
531 MitraConfig {
532 dim: 8,
533 n_layers: 1,
534 n_heads: 2,
535 dim_output: 2,
536 task: CfgTask::Classification,
537 }
538 }
539
540 fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64, scale: f32) {
541 *seed += 1;
542 let data: Vec<f32> = pseudo(rows * cols, *seed)
543 .iter()
544 .map(|v| v * scale)
545 .collect();
546 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
547 w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
548 }
549
550 fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
551 *seed += 1;
552 let bytes: Vec<u8> = pseudo(n, *seed)
553 .iter()
554 .flat_map(|v| v.to_le_bytes())
555 .collect();
556 w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
557 }
558
559 fn put_attn(w: &mut GGUFWriter, prefix: &str, d: usize, seed: &mut u64, s: f32) {
560 for t in ["q", "k", "v", "o"] {
561 put2(w, &format!("{prefix}.{t}.weight"), d, d, seed, s);
562 put1(w, &format!("{prefix}.{t}.bias"), d, seed);
563 }
564 }
565
566 fn put_block(w: &mut GGUFWriter, n: usize, d: usize, seed: &mut u64, s: f32) {
567 for ln in ["ln1", "ln2", "ln3", "ln4"] {
568 put1(w, &format!("blk.{n}.{ln}.weight"), d, seed);
569 put1(w, &format!("blk.{n}.{ln}.bias"), d, seed);
570 }
571 put_attn(w, &format!("blk.{n}.attn_row"), d, seed, s);
572 put_attn(w, &format!("blk.{n}.attn_feat"), d, seed, s);
573 for m in ["mlp1", "mlp2"] {
574 put2(w, &format!("blk.{n}.{m}_fc1.weight"), 2 * d, d, seed, s);
575 put1(w, &format!("blk.{n}.{m}_fc1.bias"), 2 * d, seed);
576 put2(w, &format!("blk.{n}.{m}_fc2.weight"), d, 2 * d, seed, s);
577 put1(w, &format!("blk.{n}.{m}_fc2.bias"), d, seed);
578 }
579 }
580
581 #[test]
582 fn burn_classification_golden() {
583 let dir = std::env::temp_dir().join(format!("mitra-burn-{}", std::process::id()));
584 let _ = std::fs::remove_dir_all(&dir);
585 std::fs::create_dir_all(&dir).unwrap();
586 let path = dir.join("model.gguf");
587 let mut w = GGUFWriter::new();
588 w.add_metadata(
589 "general.architecture",
590 GGUFMetaValue::String("mitra".into()),
591 );
592 let mut seed = 1700u64;
593 let s = 0.2;
594 put2(&mut w, "x_embed.weight", 8, 1, &mut seed, s);
595 put1(&mut w, "x_embed.bias", 8, &mut seed);
596 put2(&mut w, "y_embed.weight", 2, 8, &mut seed, s);
597 put1(&mut w, "y_mask.weight", 8, &mut seed);
598 put_block(&mut w, 0, 8, &mut seed, s);
599 put1(&mut w, "norm_f.weight", 8, &mut seed);
600 put1(&mut w, "norm_f.bias", 8, &mut seed);
601 put2(&mut w, "head.weight", 2, 8, &mut seed, s);
602 put1(&mut w, "head.bias", 2, &mut seed);
603 let f = std::fs::File::create(&path).unwrap();
604 w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
605
606 let cfg = tiny_config();
607 let burn = BurnMitraModel::load(&path, cfg).unwrap();
608 let x_support = vec![
609 vec![0.1, 1.2],
610 vec![0.9, -0.3],
611 vec![-1.1, 0.4],
612 vec![1.5, 1.1],
613 ];
614 let y_support = vec![0usize, 1, 1, 0];
615 let x_query = vec![vec![0.2, 0.9], vec![-0.8, 0.1]];
616 let t0 = std::time::Instant::now();
617 let a = burn
618 .predict_classification(&x_support, &y_support, &x_query, 2)
619 .unwrap();
620 let burn_ms = t0.elapsed().as_secs_f64() * 1000.0;
621 let b = burn
622 .predict_classification(&x_support, &y_support, &x_query, 2)
623 .unwrap();
624 assert_eq!(a.len(), 2);
625 assert!(a
626 .iter()
627 .all(|r| r.len() == 2 && r.iter().all(|v| v.is_finite())));
628 assert_eq!(a, b, "predict must be deterministic");
629 println!("mitra synthetic: burn {burn_ms:.1}ms out={a:?}");
630 let golden = [[0.112125695, 0.033007264], [0.112092555, 0.033406734]];
632 for (ra, re) in a.iter().zip(golden.iter()) {
633 for (x, y) in ra.iter().zip(re.iter()) {
634 assert!((x - y).abs() < 1e-4, "golden drift: {x} vs {y}");
635 }
636 }
637 let _ = std::fs::remove_dir_all(&dir);
638 }
639}