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::rope::apply_llama_rope;
14use zsfm_burn::{open_gguf_weights, GgufWeights};
15
16use crate::config::TabIclConfig;
17
18const LN_EPS: f32 = 1e-5;
19
20fn device() -> Device {
21 Device::flex()
22}
23
24fn load_t(weights: &GgufWeights, name: &str) -> Result<Tensor<2>> {
25 weights.load_2d(name)
26}
27
28fn load_v(weights: &GgufWeights, name: &str) -> Result<Tensor<1>> {
29 weights.load_1d(name)
30}
31
32fn linear<const D: usize>(x: Tensor<D>, w: Tensor<2>, b: Option<Tensor<1>>) -> Tensor<D> {
33 zsfm_burn::linear::burn_linear_nd(x, w, b)
34}
35
36fn linear_nb<const D: usize>(x: Tensor<D>, w: Tensor<2>) -> Tensor<D> {
37 zsfm_burn::linear::burn_linear_nd(x, w, None)
38}
39
40fn layer_norm<const D: usize>(x: Tensor<D>, w: Tensor<1>, b: Tensor<1>) -> Tensor<D> {
41 burn_layer_norm_nd(x, w, b, LN_EPS)
42}
43
44fn gelu<const D: usize>(x: Tensor<D>) -> Tensor<D> {
45 activation::gelu(x)
46}
47
48fn softmax_last<const D: usize>(x: Tensor<D>) -> Tensor<D> {
49 activation::softmax(x, D - 1)
50}
51
52fn to_host<const D: usize>(t: &Tensor<D>) -> Result<Vec<f32>> {
53 t.to_data()
54 .try_to_vec()
55 .map_err(|e| anyhow::anyhow!("{e:?}"))
56}
57
58struct BurnSsmaxW {
59 base_fc1_w: Tensor<2>,
60 base_fc1_b: Tensor<1>,
61 base_fc2_w: Tensor<2>,
62 base_fc2_b: Tensor<1>,
63 query_fc1_w: Tensor<2>,
64 query_fc1_b: Tensor<1>,
65 query_fc2_w: Tensor<2>,
66 query_fc2_b: Tensor<1>,
67}
68
69struct BurnAttnW {
70 in_proj_w: Tensor<2>,
71 in_proj_b: Tensor<1>,
72 out_proj_w: Tensor<2>,
73 out_proj_b: Tensor<1>,
74 ssmax: Option<BurnSsmaxW>,
75}
76
77struct BurnBlockW {
78 norm1_w: Tensor<1>,
79 norm1_b: Tensor<1>,
80 norm2_w: Tensor<1>,
81 norm2_b: Tensor<1>,
82 attn: BurnAttnW,
83 fc1_w: Tensor<2>,
84 fc1_b: Tensor<1>,
85 fc2_w: Tensor<2>,
86 fc2_b: Tensor<1>,
87}
88
89struct BurnIsabW {
90 ind_vectors: Tensor<2>,
91 attn1: BurnBlockW,
92 attn2: BurnBlockW,
93}
94
95struct BurnColW {
96 in_linear_w: Tensor<2>,
97 in_linear_b: Tensor<1>,
98 y_encoder_w: Tensor<2>,
99 y_encoder_b: Tensor<1>,
100 blocks: Vec<BurnIsabW>,
101}
102
103struct BurnRowW {
104 cls_tokens: Tensor<2>,
105 blocks: Vec<BurnBlockW>,
106 out_ln_w: Tensor<1>,
107 out_ln_b: Tensor<1>,
108 rope_freqs: Vec<f32>,
109}
110
111struct BurnIclW {
112 y_encoder_w: Tensor<2>,
113 y_encoder_b: Tensor<1>,
114 blocks: Vec<BurnBlockW>,
115 ln_w: Tensor<1>,
116 ln_b: Tensor<1>,
117 decoder_fc1_w: Tensor<2>,
118 decoder_fc1_b: Tensor<1>,
119 decoder_fc2_w: Tensor<2>,
120 decoder_fc2_b: Tensor<1>,
121}
122
123pub struct BurnTabIclModel {
124 config: TabIclConfig,
125 col: BurnColW,
126 row: BurnRowW,
127 icl: BurnIclW,
128}
129
130fn load_ssmax(weights: &GgufWeights, prefix: &str) -> Result<BurnSsmaxW> {
131 Ok(BurnSsmaxW {
132 base_fc1_w: load_t(weights, &format!("{prefix}.ssmax_layer.base_mlp.0.weight"))?,
133 base_fc1_b: load_v(weights, &format!("{prefix}.ssmax_layer.base_mlp.0.bias"))?,
134 base_fc2_w: load_t(weights, &format!("{prefix}.ssmax_layer.base_mlp.2.weight"))?,
135 base_fc2_b: load_v(weights, &format!("{prefix}.ssmax_layer.base_mlp.2.bias"))?,
136 query_fc1_w: load_t(weights, &format!("{prefix}.ssmax_layer.query_mlp.0.weight"))?,
137 query_fc1_b: load_v(weights, &format!("{prefix}.ssmax_layer.query_mlp.0.bias"))?,
138 query_fc2_w: load_t(weights, &format!("{prefix}.ssmax_layer.query_mlp.2.weight"))?,
139 query_fc2_b: load_v(weights, &format!("{prefix}.ssmax_layer.query_mlp.2.bias"))?,
140 })
141}
142
143fn load_attn(weights: &GgufWeights, prefix: &str, has_ssmax: bool) -> Result<BurnAttnW> {
144 Ok(BurnAttnW {
145 in_proj_w: load_t(weights, &format!("{prefix}.in_proj_weight"))?,
146 in_proj_b: load_v(weights, &format!("{prefix}.in_proj_bias"))?,
147 out_proj_w: load_t(weights, &format!("{prefix}.out_proj.weight"))?,
148 out_proj_b: load_v(weights, &format!("{prefix}.out_proj.bias"))?,
149 ssmax: if has_ssmax {
150 Some(load_ssmax(weights, prefix)?)
151 } else {
152 None
153 },
154 })
155}
156
157fn load_block(weights: &GgufWeights, prefix: &str, has_ssmax: bool) -> Result<BurnBlockW> {
158 Ok(BurnBlockW {
159 norm1_w: load_v(weights, &format!("{prefix}.norm1.weight"))?,
160 norm1_b: load_v(weights, &format!("{prefix}.norm1.bias"))?,
161 norm2_w: load_v(weights, &format!("{prefix}.norm2.weight"))?,
162 norm2_b: load_v(weights, &format!("{prefix}.norm2.bias"))?,
163 attn: load_attn(weights, &format!("{prefix}.attn"), has_ssmax)?,
164 fc1_w: load_t(weights, &format!("{prefix}.linear1.weight"))?,
165 fc1_b: load_v(weights, &format!("{prefix}.linear1.bias"))?,
166 fc2_w: load_t(weights, &format!("{prefix}.linear2.weight"))?,
167 fc2_b: load_v(weights, &format!("{prefix}.linear2.bias"))?,
168 })
169}
170
171impl BurnTabIclModel {
172 pub fn load(gguf_path: &Path, config: TabIclConfig) -> Result<Self> {
173 let weights = open_gguf_weights(gguf_path)
174 .with_context(|| format!("open {}", gguf_path.display()))?;
175
176 let col_in_linear_w = load_t(&weights, "col_embedder.in_linear.weight")?;
177 let col_in_linear_b = load_v(&weights, "col_embedder.in_linear.bias")?;
178 let col_y_w = load_t(&weights, "col_embedder.y_encoder.weight")?;
179 let col_y_b = load_v(&weights, "col_embedder.y_encoder.bias")?;
180
181 let mut col_blocks = Vec::with_capacity(config.col_num_blocks);
182 for n in 0..config.col_num_blocks {
183 let p = format!("col_embedder.tf_col.blocks.{n}");
184 let ind_vectors = load_t(&weights, &format!("{p}.ind_vectors"))?;
185 let attn1 = load_block(&weights, &format!("{p}.multihead_attn1"), true)?;
186 let attn2 = load_block(&weights, &format!("{p}.multihead_attn2"), false)?;
187 col_blocks.push(BurnIsabW {
188 ind_vectors,
189 attn1,
190 attn2,
191 });
192 }
193
194 let cls_tokens = load_t(&weights, "row_interactor.cls_tokens")?;
195 let row_out_ln_w = load_v(&weights, "row_interactor.out_ln.weight")?;
196 let row_out_ln_b = load_v(&weights, "row_interactor.out_ln.bias")?;
197 let rope_freqs: Vec<f32> = weights.load_flat("row_interactor.tf_row.rope.freqs")?;
198 let mut row_blocks = Vec::with_capacity(config.row_num_blocks);
199 for n in 0..config.row_num_blocks {
200 row_blocks.push(load_block(
201 &weights,
202 &format!("row_interactor.tf_row.blocks.{n}"),
203 false,
204 )?);
205 }
206
207 let mut icl_blocks = Vec::with_capacity(config.icl_num_blocks);
208 for n in 0..config.icl_num_blocks {
209 icl_blocks.push(load_block(
210 &weights,
211 &format!("icl_predictor.tf_icl.blocks.{n}"),
212 true,
213 )?);
214 }
215
216 Ok(Self {
217 config,
218 col: BurnColW {
219 in_linear_w: col_in_linear_w,
220 in_linear_b: col_in_linear_b,
221 y_encoder_w: col_y_w,
222 y_encoder_b: col_y_b,
223 blocks: col_blocks,
224 },
225 row: BurnRowW {
226 cls_tokens,
227 blocks: row_blocks,
228 out_ln_w: row_out_ln_w,
229 out_ln_b: row_out_ln_b,
230 rope_freqs,
231 },
232 icl: BurnIclW {
233 y_encoder_w: load_t(&weights, "icl_predictor.y_encoder.weight")?,
234 y_encoder_b: load_v(&weights, "icl_predictor.y_encoder.bias")?,
235 blocks: icl_blocks,
236 ln_w: load_v(&weights, "icl_predictor.ln.weight")?,
237 ln_b: load_v(&weights, "icl_predictor.ln.bias")?,
238 decoder_fc1_w: load_t(&weights, "icl_predictor.decoder.0.weight")?,
239 decoder_fc1_b: load_v(&weights, "icl_predictor.decoder.0.bias")?,
240 decoder_fc2_w: load_t(&weights, "icl_predictor.decoder.2.weight")?,
241 decoder_fc2_b: load_v(&weights, "icl_predictor.decoder.2.bias")?,
242 },
243 })
244 }
245
246 pub fn predict_classification(
247 &self,
248 x_support: &[Vec<f32>],
249 y_support: &[usize],
250 x_query: &[Vec<f32>],
251 n_classes: usize,
252 ) -> Result<Vec<Vec<f32>>> {
253 anyhow::ensure!(
254 n_classes <= self.config.max_classes,
255 "n_classes must be <= {}",
256 self.config.max_classes
257 );
258 let train_size = x_support.len();
259 let mut all_rows = x_support.to_vec();
260 all_rows.extend_from_slice(x_query);
261 let t = all_rows.len();
262 let h = all_rows[0].len();
263 let processed = preprocess_x(&all_rows, train_size);
264 let group_size = self.config.feature_group_size;
265 let grouped = feature_group_same(&processed, h, group_size);
266 let mut flat = vec![0f32; h * t * group_size];
267 for (row_idx, row) in grouped.iter().enumerate() {
268 for (col_idx, group) in row.iter().enumerate() {
269 let base = col_idx * t * group_size + row_idx * group_size;
270 flat[base..base + group_size].copy_from_slice(group);
271 }
272 }
273 let dev = device();
274 let x_t = Tensor::<3>::from_data(TensorData::new(flat, [h, t, group_size]), &dev);
275 let y_onehot_col =
276 self.onehot_linear(y_support, &self.col.y_encoder_w, &self.col.y_encoder_b)?;
277 let col_out = self.col_embed_forward(&x_t, &y_onehot_col, train_size)?;
278 let row_out = self.row_interact_forward(&col_out, train_size)?;
279 let y_onehot_icl =
280 self.onehot_linear(y_support, &self.icl.y_encoder_w, &self.icl.y_encoder_b)?;
281 let logits = self.icl_forward(&row_out, &y_onehot_icl, train_size, n_classes)?;
282 let flat_logits: Vec<f32> = to_host(&logits)?;
283 let n_query = t - train_size;
284 const TEMPERATURE: f32 = 0.9;
285 let mut result = Vec::with_capacity(n_query);
286 for i in 0..n_query {
287 let row = &flat_logits[i * n_classes..(i + 1) * n_classes];
288 let scaled: Vec<f32> = row.iter().map(|&v| v / TEMPERATURE).collect();
289 result.push(zsfm_core::softmax(&scaled));
290 }
291 Ok(result)
292 }
293
294 fn onehot_linear(&self, y: &[usize], w: &Tensor<2>, b: &Tensor<1>) -> Result<Tensor<2>> {
295 let dev = device();
296 let n = y.len();
297 let num_classes = w.dims()[1];
298 let mut onehot = vec![0f32; n * num_classes];
299 for (i, &c) in y.iter().enumerate() {
300 onehot[i * num_classes + c] = 1.0;
301 }
302 let x = Tensor::<2>::from_data(TensorData::new(onehot, [n, num_classes]), &dev);
303 Ok(linear(x, w.clone(), Some(b.clone())))
304 }
305
306 fn col_embed_forward(
307 &self,
308 x_grouped: &Tensor<3>,
309 y_onehot: &Tensor<2>,
310 train_size: usize,
311 ) -> Result<Tensor<3>> {
312 let dims = x_grouped.dims();
313 let t = dims[1];
314 let src = linear(
315 x_grouped.clone(),
316 self.col.in_linear_w.clone(),
317 Some(self.col.in_linear_b.clone()),
318 );
319 let train_part: Tensor<3> = src.clone().narrow(1, 0, train_size);
320 let y4: Tensor<3> = y_onehot.clone().unsqueeze_dim::<3>(0);
322 let h = src.dims()[0];
323 let yd = y4.dims()[2];
324 let yb: Tensor<3> = y4.expand([h, train_size, yd]);
325 let train_biased: Tensor<3> = train_part + yb;
326 let mut src = if train_size < t {
327 Tensor::cat(
328 vec![train_biased, src.narrow(1, train_size, t - train_size)],
329 1,
330 )
331 } else {
332 train_biased
333 };
334 for isab in &self.col.blocks {
335 src = isab_forward(&src, isab, self.config.col_nhead, train_size, h)?;
336 }
337 Ok(src)
338 }
339
340 fn row_interact_forward(
341 &self,
342 col_embeddings: &Tensor<3>,
343 _train_size: usize,
344 ) -> Result<Tensor<2>> {
345 let dev = device();
346 let cd = col_embeddings.dims();
347 let (h, t) = (cd[0], cd[1]);
348 let dim = self.config.embed_dim;
349 let n_cls = self.config.row_num_cls;
350 let feat: Tensor<3> = col_embeddings.clone().permute([1, 0, 2]);
351 let cls: Tensor<3> = self
352 .row
353 .cls_tokens
354 .clone()
355 .unsqueeze_dim::<3>(0)
356 .expand([t, n_cls, dim]);
357 let mut seq = Tensor::cat(vec![cls, feat], 1);
358 let max_len = h + n_cls;
359 let (cos, sin) = self.row_rope_table(max_len)?;
360 let n_blocks = self.row.blocks.len();
361 for (i, blk) in self.row.blocks.iter().enumerate() {
362 if i + 1 == n_blocks {
363 let cls_q: Tensor<3> = seq.clone().narrow(1, 0, n_cls);
364 seq = mha_block(
365 &cls_q,
366 &seq,
367 &seq,
368 blk,
369 self.config.row_nhead,
370 Some((cos.clone(), sin.clone())),
371 0,
372 )?;
373 } else {
374 seq = mha_block(
375 &seq,
376 &seq,
377 &seq,
378 blk,
379 self.config.row_nhead,
380 Some((cos.clone(), sin.clone())),
381 0,
382 )?;
383 }
384 }
385 let _ = dev;
386 let out = layer_norm(seq, self.row.out_ln_w.clone(), self.row.out_ln_b.clone());
387 Ok(out.reshape([t, n_cls * dim]))
388 }
389
390 fn row_rope_table(&self, max_len: usize) -> Result<(Tensor<2>, Tensor<2>)> {
391 let dev = device();
392 let half = self.row.rope_freqs.len();
393 let mut cos = vec![0f32; max_len * half * 2];
394 let mut sin = vec![0f32; max_len * half * 2];
395 for p in 0..max_len {
396 for i in 0..half {
397 let angle = p as f32 * self.row.rope_freqs[i];
398 let (s, c) = angle.sin_cos();
399 cos[p * 2 * half + i] = c;
400 cos[p * 2 * half + half + i] = c;
401 sin[p * 2 * half + i] = s;
402 sin[p * 2 * half + half + i] = s;
403 }
404 }
405 Ok((
406 Tensor::<2>::from_data(TensorData::new(cos, [max_len, 2 * half]), &dev),
407 Tensor::<2>::from_data(TensorData::new(sin, [max_len, 2 * half]), &dev),
408 ))
409 }
410
411 fn icl_forward(
412 &self,
413 row_reprs: &Tensor<2>,
414 y_onehot: &Tensor<2>,
415 train_size: usize,
416 n_classes: usize,
417 ) -> Result<Tensor<2>> {
418 let t = row_reprs.dims()[0];
419 let train_part: Tensor<2> = row_reprs.clone().narrow(0, 0, train_size);
420 let train_biased = train_part + y_onehot.clone();
421 let r = if train_size < t {
422 Tensor::cat(
423 vec![
424 train_biased,
425 row_reprs.clone().narrow(0, train_size, t - train_size),
426 ],
427 0,
428 )
429 } else {
430 train_biased
431 };
432 let mut r: Tensor<3> = r.unsqueeze_dim::<3>(0);
433 for blk in &self.icl.blocks {
434 let kv: Tensor<3> = r.clone().narrow(1, 0, train_size);
435 r = mha_block(&r, &kv, &kv, blk, self.config.icl_nhead, None, train_size)?;
436 }
437 let r: Tensor<2> = r.squeeze_dim::<2>(0);
438 let r = layer_norm(r, self.icl.ln_w.clone(), self.icl.ln_b.clone());
439 let h = gelu(linear(
440 r,
441 self.icl.decoder_fc1_w.clone(),
442 Some(self.icl.decoder_fc1_b.clone()),
443 ));
444 let logits = linear(
445 h,
446 self.icl.decoder_fc2_w.clone(),
447 Some(self.icl.decoder_fc2_b.clone()),
448 );
449 Ok(logits
450 .narrow(0, train_size, t - train_size)
451 .narrow(1, 0, n_classes))
452 }
453}
454
455fn isab_forward(
456 src: &Tensor<3>,
457 isab: &BurnIsabW,
458 n_heads: usize,
459 train_size: usize,
460 n_h: usize,
461) -> Result<Tensor<3>> {
462 let sd = src.dims();
463 let (num_inds, dim) = (isab.ind_vectors.dims()[0], isab.ind_vectors.dims()[1]);
464 let ind: Tensor<3> = isab
465 .ind_vectors
466 .clone()
467 .unsqueeze_dim::<3>(0)
468 .expand([n_h, num_inds, dim]);
469 let kv_train: Tensor<3> = src.clone().narrow(1, 0, train_size);
470 let hidden = mha_block(
471 &ind,
472 &kv_train,
473 &kv_train,
474 &isab.attn1,
475 n_heads,
476 None,
477 train_size,
478 )?;
479 Ok(mha_block(
480 src,
481 &hidden,
482 &hidden,
483 &isab.attn2,
484 n_heads,
485 None,
486 num_inds,
487 )?)
488}
489
490fn mha_block(
491 q_in: &Tensor<3>,
492 k_in: &Tensor<3>,
493 v_in: &Tensor<3>,
494 blk: &BurnBlockW,
495 n_heads: usize,
496 rope_cos_sin: Option<(Tensor<2>, Tensor<2>)>,
497 ssmax_n: usize,
498) -> Result<Tensor<3>> {
499 let q_normed = layer_norm(q_in.clone(), blk.norm1_w.clone(), blk.norm1_b.clone());
500 let k_normed = layer_norm(k_in.clone(), blk.norm1_w.clone(), blk.norm1_b.clone());
501 let v_normed = layer_norm(v_in.clone(), blk.norm1_w.clone(), blk.norm1_b.clone());
502 let attn_out = mha_combined(
503 &q_normed,
504 &k_normed,
505 &v_normed,
506 &blk.attn,
507 n_heads,
508 rope_cos_sin,
509 ssmax_n,
510 )?;
511 let x = q_in.clone() + attn_out;
512 let ff_in = layer_norm(x.clone(), blk.norm2_w.clone(), blk.norm2_b.clone());
513 let ff = linear(ff_in, blk.fc1_w.clone(), Some(blk.fc1_b.clone()));
514 let ff = linear(gelu(ff), blk.fc2_w.clone(), Some(blk.fc2_b.clone()));
515 Ok(x + ff)
516}
517
518fn mha_combined(
519 q_in: &Tensor<3>,
520 k_in: &Tensor<3>,
521 v_in: &Tensor<3>,
522 attn: &BurnAttnW,
523 n_heads: usize,
524 rope_cos_sin: Option<(Tensor<2>, Tensor<2>)>,
525 ssmax_n: usize,
526) -> Result<Tensor<3>> {
527 let qd = q_in.dims();
528 let dim = qd[2];
529 let wq: Tensor<2> = attn.in_proj_w.clone().narrow(0, 0, dim);
530 let wk: Tensor<2> = attn.in_proj_w.clone().narrow(0, dim, dim);
531 let wv: Tensor<2> = attn.in_proj_w.clone().narrow(0, 2 * dim, dim);
532 let bq: Tensor<1> = attn.in_proj_b.clone().narrow(0, 0, dim);
533 let bk: Tensor<1> = attn.in_proj_b.clone().narrow(0, dim, dim);
534 let bv: Tensor<1> = attn.in_proj_b.clone().narrow(0, 2 * dim, dim);
535 let q = linear(q_in.clone(), wq, Some(bq));
536 let k = linear(k_in.clone(), wk, Some(bk));
537 let v = linear(v_in.clone(), wv, Some(bv));
538 let batch = q.dims()[0];
539 let q_len = q.dims()[1];
540 let k_len = k.dims()[1];
541 let head_dim = dim / n_heads;
542 let q: Tensor<4> = q
543 .reshape([batch, q_len, n_heads, head_dim])
544 .permute([0, 2, 1, 3]);
545 let k: Tensor<4> = k
546 .reshape([batch, k_len, n_heads, head_dim])
547 .permute([0, 2, 1, 3]);
548 let v: Tensor<4> = v
549 .reshape([batch, k_len, n_heads, head_dim])
550 .permute([0, 2, 1, 3]);
551 let (q, k) = match rope_cos_sin {
552 Some((cos, sin)) => (
553 apply_llama_rope(q, cos.clone(), sin.clone(), 0),
554 apply_llama_rope(k, cos, sin, 0),
555 ),
556 None => (q, k),
557 };
558 let q = match &attn.ssmax {
559 Some(ssmax) => apply_ssmax(q, ssmax, ssmax_n, n_heads, head_dim)?,
560 None => q,
561 };
562 let scale = 1.0f32 / (head_dim as f32).sqrt();
563 let scores = q.matmul(k.transpose()).mul_scalar(scale);
564 let probs = activation::softmax(scores, 3);
565 let out = probs.matmul(v);
566 let out: Tensor<4> = out.permute([0, 2, 1, 3]);
567 let out: Tensor<3> = out.reshape([batch, q_len, dim]);
568 Ok(linear(
569 out,
570 attn.out_proj_w.clone(),
571 Some(attn.out_proj_b.clone()),
572 ))
573}
574
575fn apply_ssmax(
576 q: Tensor<4>,
577 ssmax: &BurnSsmaxW,
578 n: usize,
579 n_heads: usize,
580 head_dim: usize,
581) -> Result<Tensor<4>> {
582 let dev = device();
583 let logn = (n.max(1) as f32).ln();
584 let logn_t = Tensor::<2>::from_data(TensorData::new(vec![logn], [1, 1]), &dev);
585 let base = gelu(linear(
586 logn_t,
587 ssmax.base_fc1_w.clone(),
588 Some(ssmax.base_fc1_b.clone()),
589 ));
590 let base = linear(
591 base,
592 ssmax.base_fc2_w.clone(),
593 Some(ssmax.base_fc2_b.clone()),
594 );
595 let base: Tensor<4> = base.reshape([1, n_heads, 1, head_dim]);
596 let qd = q.dims();
597 let (batch, seq) = (qd[0], qd[2]);
598 let q_flat: Tensor<2> = q.clone().reshape([batch * n_heads * seq, head_dim]);
599 let qm = gelu(linear(
600 q_flat,
601 ssmax.query_fc1_w.clone(),
602 Some(ssmax.query_fc1_b.clone()),
603 ));
604 let qm = linear(
605 qm,
606 ssmax.query_fc2_w.clone(),
607 Some(ssmax.query_fc2_b.clone()),
608 );
609 let modulation: Tensor<4> = qm
610 .tanh()
611 .add_scalar(1.0)
612 .reshape([batch, n_heads, seq, head_dim]);
613 let scales = base * modulation;
614 Ok(q * scales)
615}
616
617fn preprocess_x(rows: &[Vec<f32>], train_size: usize) -> Vec<Vec<f32>> {
618 let n_feat = rows[0].len();
619 let mut impute_mean = vec![0f32; n_feat];
620 for f in 0..n_feat {
621 let mut sum = 0f64;
622 let mut count = 0usize;
623 for row in &rows[..train_size] {
624 if !row[f].is_nan() {
625 sum += row[f] as f64;
626 count += 1;
627 }
628 }
629 impute_mean[f] = if count > 0 {
630 (sum / count as f64) as f32
631 } else {
632 0.0
633 };
634 }
635 let imputed: Vec<Vec<f32>> = rows
636 .iter()
637 .map(|row| {
638 row.iter()
639 .enumerate()
640 .map(|(f, &v)| if v.is_nan() { impute_mean[f] } else { v })
641 .collect()
642 })
643 .collect();
644 let mut mean = vec![0f32; n_feat];
645 let mut std = vec![0f32; n_feat];
646 for f in 0..n_feat {
647 let m: f32 = imputed[..train_size].iter().map(|r| r[f]).sum::<f32>() / train_size as f32;
648 let var: f32 = imputed[..train_size]
649 .iter()
650 .map(|r| (r[f] - m).powi(2))
651 .sum::<f32>()
652 / train_size as f32;
653 mean[f] = m;
654 std[f] = if var == 0.0 { 1.0 } else { var.sqrt() };
655 }
656 imputed
657 .iter()
658 .map(|row| {
659 row.iter()
660 .enumerate()
661 .map(|(f, &v)| (v - mean[f]) / std[f])
662 .collect()
663 })
664 .collect()
665}
666
667fn feature_group_same(rows: &[Vec<f32>], h: usize, size: usize) -> Vec<Vec<Vec<f32>>> {
668 rows.iter()
669 .map(|row| {
670 (0..h)
671 .map(|g| (0..size).map(|k| row[(g + (1usize << k)) % h]).collect())
672 .collect()
673 })
674 .collect()
675}
676
677#[cfg(test)]
678mod tests {
679 use super::*;
680 use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
681
682 fn pseudo(n: usize, seed: u64) -> Vec<f32> {
683 let mut x = seed;
684 (0..n)
685 .map(|_| {
686 x = x
687 .wrapping_mul(6364136223846793005)
688 .wrapping_add(1442695040888963407);
689 ((x >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
690 })
691 .collect()
692 }
693
694 fn tiny_config() -> TabIclConfig {
695 TabIclConfig {
696 max_classes: 2,
697 embed_dim: 8,
698 col_num_blocks: 1,
699 col_nhead: 2,
700 col_num_inds: 2,
701 feature_group_size: 2,
702 row_num_blocks: 1,
703 row_nhead: 2,
704 row_num_cls: 2,
705 row_rope_base: 10000.0,
706 icl_num_blocks: 1,
707 icl_nhead: 2,
708 ff_factor: 2,
709 }
710 }
711
712 fn put2(w: &mut GGUFWriter, name: &str, rows: usize, cols: usize, seed: &mut u64, scale: f32) {
713 *seed += 1;
714 let data: Vec<f32> = pseudo(rows * cols, *seed)
715 .iter()
716 .map(|v| v * scale)
717 .collect();
718 let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
719 w.add_tensor(name, vec![cols as u64, rows as u64], GGMLType::F32, bytes);
720 }
721
722 fn put1(w: &mut GGUFWriter, name: &str, n: usize, seed: &mut u64) {
723 *seed += 1;
724 let bytes: Vec<u8> = pseudo(n, *seed)
725 .iter()
726 .flat_map(|v| v.to_le_bytes())
727 .collect();
728 w.add_tensor(name, vec![n as u64], GGMLType::F32, bytes);
729 }
730
731 fn put_attn_ssmax(
732 w: &mut GGUFWriter,
733 prefix: &str,
734 d: usize,
735 seed: &mut u64,
736 s: f32,
737 ssmax_hd: Option<usize>,
738 ssmax_base_out: usize,
739 ) {
740 put2(w, &format!("{prefix}.in_proj_weight"), 3 * d, d, seed, s);
741 put1(w, &format!("{prefix}.in_proj_bias"), 3 * d, seed);
742 put2(w, &format!("{prefix}.out_proj.weight"), d, d, seed, s);
743 put1(w, &format!("{prefix}.out_proj.bias"), d, seed);
744 if let Some(hd) = ssmax_hd {
745 put2(
746 w,
747 &format!("{prefix}.ssmax_layer.base_mlp.0.weight"),
748 8,
749 1,
750 seed,
751 s,
752 );
753 put1(w, &format!("{prefix}.ssmax_layer.base_mlp.0.bias"), 8, seed);
754 put2(
755 w,
756 &format!("{prefix}.ssmax_layer.base_mlp.2.weight"),
757 ssmax_base_out,
758 8,
759 seed,
760 s,
761 );
762 put1(
763 w,
764 &format!("{prefix}.ssmax_layer.base_mlp.2.bias"),
765 ssmax_base_out,
766 seed,
767 );
768 put2(
769 w,
770 &format!("{prefix}.ssmax_layer.query_mlp.0.weight"),
771 8,
772 hd,
773 seed,
774 s,
775 );
776 put1(
777 w,
778 &format!("{prefix}.ssmax_layer.query_mlp.0.bias"),
779 8,
780 seed,
781 );
782 put2(
783 w,
784 &format!("{prefix}.ssmax_layer.query_mlp.2.weight"),
785 hd,
786 8,
787 seed,
788 s,
789 );
790 put1(
791 w,
792 &format!("{prefix}.ssmax_layer.query_mlp.2.bias"),
793 hd,
794 seed,
795 );
796 }
797 }
798
799 fn put_block(
800 w: &mut GGUFWriter,
801 prefix: &str,
802 d: usize,
803 seed: &mut u64,
804 s: f32,
805 ssmax: Option<(usize, usize)>,
806 ) {
807 put1(w, &format!("{prefix}.norm1.weight"), d, seed);
808 put1(w, &format!("{prefix}.norm1.bias"), d, seed);
809 put1(w, &format!("{prefix}.norm2.weight"), d, seed);
810 put1(w, &format!("{prefix}.norm2.bias"), d, seed);
811 let (hd, base_out) = ssmax.map(|(a, b)| (Some(a), b)).unwrap_or((None, d));
812 put_attn_ssmax(w, &format!("{prefix}.attn"), d, seed, s, hd, base_out);
813 put2(w, &format!("{prefix}.linear1.weight"), 2 * d, d, seed, s);
814 put1(w, &format!("{prefix}.linear1.bias"), 2 * d, seed);
815 put2(w, &format!("{prefix}.linear2.weight"), d, 2 * d, seed, s);
816 put1(w, &format!("{prefix}.linear2.bias"), d, seed);
817 }
818
819 #[test]
820 fn burn_classification_golden() {
821 let dir = std::env::temp_dir().join(format!("tabicl-burn-{}", std::process::id()));
822 let _ = std::fs::remove_dir_all(&dir);
823 std::fs::create_dir_all(&dir).unwrap();
824 let path = dir.join("model.gguf");
825 let mut w = GGUFWriter::new();
826 w.add_metadata(
827 "general.architecture",
828 GGUFMetaValue::String("tabicl".into()),
829 );
830 let mut seed = 2100u64;
831 let s = 0.2;
832 put2(&mut w, "col_embedder.in_linear.weight", 8, 2, &mut seed, s);
833 put1(&mut w, "col_embedder.in_linear.bias", 8, &mut seed);
834 put2(&mut w, "col_embedder.y_encoder.weight", 8, 2, &mut seed, s);
835 put1(&mut w, "col_embedder.y_encoder.bias", 8, &mut seed);
836 put2(
837 &mut w,
838 "col_embedder.tf_col.blocks.0.ind_vectors",
839 2,
840 8,
841 &mut seed,
842 s,
843 );
844 put_block(
845 &mut w,
846 "col_embedder.tf_col.blocks.0.multihead_attn1",
847 8,
848 &mut seed,
849 s,
850 Some((4, 8)),
851 );
852 put_block(
853 &mut w,
854 "col_embedder.tf_col.blocks.0.multihead_attn2",
855 8,
856 &mut seed,
857 s,
858 None,
859 );
860 put2(&mut w, "row_interactor.cls_tokens", 2, 8, &mut seed, s);
861 put1(&mut w, "row_interactor.out_ln.weight", 8, &mut seed);
862 put1(&mut w, "row_interactor.out_ln.bias", 8, &mut seed);
863 put1(&mut w, "row_interactor.tf_row.rope.freqs", 2, &mut seed);
864 put_block(
865 &mut w,
866 "row_interactor.tf_row.blocks.0",
867 8,
868 &mut seed,
869 s,
870 None,
871 );
872 put2(
873 &mut w,
874 "icl_predictor.y_encoder.weight",
875 16,
876 2,
877 &mut seed,
878 s,
879 );
880 put1(&mut w, "icl_predictor.y_encoder.bias", 16, &mut seed);
881 put1(&mut w, "icl_predictor.ln.weight", 16, &mut seed);
882 put1(&mut w, "icl_predictor.ln.bias", 16, &mut seed);
883 put2(
884 &mut w,
885 "icl_predictor.decoder.0.weight",
886 8,
887 16,
888 &mut seed,
889 s,
890 );
891 put1(&mut w, "icl_predictor.decoder.0.bias", 8, &mut seed);
892 put2(&mut w, "icl_predictor.decoder.2.weight", 2, 8, &mut seed, s);
893 put1(&mut w, "icl_predictor.decoder.2.bias", 2, &mut seed);
894 put_block(
895 &mut w,
896 "icl_predictor.tf_icl.blocks.0",
897 16,
898 &mut seed,
899 s,
900 Some((8, 16)),
901 );
902 let f = std::fs::File::create(&path).unwrap();
903 w.write_to(&mut std::io::BufWriter::new(f)).unwrap();
904
905 let cfg = tiny_config();
906 let burn = BurnTabIclModel::load(&path, cfg).unwrap();
907 let x_support = vec![
908 vec![0.1, 1.2],
909 vec![0.9, -0.3],
910 vec![-1.1, 0.4],
911 vec![1.5, 1.1],
912 ];
913 let y_support = vec![0usize, 1, 1, 0];
914 let x_query = vec![vec![0.2, 0.9], vec![-0.8, 0.1]];
915 let t0 = std::time::Instant::now();
916 let a = burn
917 .predict_classification(&x_support, &y_support, &x_query, 2)
918 .unwrap();
919 let burn_ms = t0.elapsed().as_secs_f64() * 1000.0;
920 let b = burn
921 .predict_classification(&x_support, &y_support, &x_query, 2)
922 .unwrap();
923 assert_eq!(a.len(), 2);
924 assert!(a
925 .iter()
926 .all(|r| r.len() == 2 && r.iter().all(|v| v.is_finite())));
927 assert_eq!(a, b, "predict must be deterministic");
928 println!("tabicl synthetic: burn {burn_ms:.1}ms out={a:?}");
929 let golden = [[0.6072121, 0.39278784], [0.60720557, 0.39279446]];
931 for (ra, re) in a.iter().zip(golden.iter()) {
932 for (x, y) in ra.iter().zip(re.iter()) {
933 assert!((x - y).abs() < 1e-4, "golden drift: {x} vs {y}");
934 }
935 }
936 let _ = std::fs::remove_dir_all(&dir);
937 }
938}