nx_reason_probe.nx source
↩ module page · 114 lines · 4220 B
1// nx_reason_probe.nx -- LIVE format probe for nx_reasoning (CHAT path).
2// Loads the real model once; greedy-generates "What is 47 plus 38?"
3// through the ChatML template; prints raw bytes + extracted answer; then
4// one SAMPLED generation (seeded) to smoke the stochastic path.
5// expect_exit: 0 (10..50 load fail, 60/61 empty gen)
6
7import "nx_syscalls.nx"
8import "nx_tier.nx"
9import "nx_le.nx"
10import "nx_bpe.nx"
11import "nx_gguf.nx"
12import "nx_gguf_load.nx"
13import "nx_gguf_meta.nx"
14import "nx_f32.nx"
15import "nx_f32_kv_cache.nx"
16import "nx_f32_lazy_weight.nx"
17import "nx_f32_llama_block.nx"
18import "nx_f32_llama_block_v4.nx"
19import "nx_f32_llama_stack_v4.nx"
20import "nx_f32_llama_layer_lazy_load.nx"
21import "nx_f32_llm.nx"
22import "nx_f32_llm_v4.nx"
23import "nx_f32_llm_read_dims.nx"
24import "nx_f32_bpe_load.nx"
25import "nx_f32_llm_special_tokens.nx"
26import "nx_f32_sampler.nx"
27import "nx_prng.nx"
28import "nx_reasoning.nx"
29
30func pb_w(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 }
31func pb_wn(v: i64) -> i64 {
32 var m: i64 = v
33 if m < 0 { pb_w("-" as *u8); m = 0 - m }
34 let t: *u8 = sys_mmap(28)
35 var k: i64 = 0
36 if m == 0 { t[0] = 48 as u8; k = 1 }
37 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
38 let o: *u8 = sys_mmap(28)
39 var i: i64 = 0
40 while i < k { o[i] = t[k - 1 - i]; i = i + 1 }
41 sys_write(1, o, k)
42 return 0
43}
44func pb_wtext(b: *u8, n: i64) -> i64 {
45 if n <= 0 { return 0 }
46 let t: *u8 = sys_mmap(n + 1)
47 var i: i64 = 0
48 while i < n {
49 var c: i64 = b[i] as i64
50 if c < 32 { c = 46 }
51 t[i] = c as u8
52 i = i + 1
53 }
54 sys_write(1, t, n)
55 return 0
56}
57
58func main() -> i64 {
59 let path: *u8 = "/tmp/nx_real_model.gguf" as *u8
60 let len_out: *i64 = sys_mmap(8) as *i64
61 let buf: *u8 = sys_read_file(path, len_out)
62 if buf == (0 as *u8) { return 10 }
63 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
64 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { return 20 }
65 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
66 let out_err: *i64 = sys_mmap(8) as *i64
67 if nx_f32_llm_read_dims_from_gguf(buf, len_out[0], hdr, model, out_err) != NX_FLD_OK { return 30 }
68 if nx_f32_llm_load_weights_v4_from_gguf(buf, hdr, model, out_err) != NX_FLV4_OK { return 40 }
69 let vocab: *NxBpeVocab = nx_bpe_vocab_new(67108864, 262144, 524288)
70 let nt2: *i64 = sys_mmap(8) as *i64
71 let nm: *i64 = sys_mmap(8) as *i64
72 if nx_f32_bpe_load_from_gguf(buf, len_out[0], hdr, vocab, nt2, nm, out_err) != NX_FBL_OK { return 50 }
73 let eos: nx_int = nx_f32_llm_read_eos(buf, len_out[0], hdr)
74 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(
75 model.n_layers, model.n_kv_heads, 128, model.head_dim)
76
77 let rc: *NxReasonCfg = nx_reason_cfg_alloc()
78 rc.model = model
79 rc.vocab = vocab
80 rc.cache = cache
81 rc.max_new = 16
82 rc.inv_temp_f32 = 0x3FA00000 // 1/0.8
83 rc.top_k = 40
84 rc.eps = 0x358637BD // 1e-6 (chat-harness proven)
85 rc.attn_scale = 0x3E000000 // 0.125 (1/sqrt(head_dim=64))
86 rc.rope_log_base = 0x415D0EAB // ln(1e6)
87 rc.eos = eos
88 rc.im_start = 151644
89 rc.im_end = 151645
90
91 let q: *u8 = sys_mmap(96)
92 let nq: nx_int = nx_reason_build_q_short(q, 47, 0, 38)
93 pb_w("Q=[" as *u8); pb_wtext(q, nq as i64); pb_w("]\n" as *u8)
94
95 // greedy
96 let og: *u8 = sys_mmap(256)
97 let ng: nx_int = nx_reason_chat_gen(rc, q, nq, 0 as *i64, og, 256)
98 pb_w("GREEDY gen=[" as *u8); pb_wtext(og, ng as i64); pb_w("] last_int=" as *u8)
99 let vg: i64 = nx_reason_extract_last_int(og, ng)
100 pb_wn(vg); pb_w(" expected=85\n" as *u8)
101
102 // sampled (seeded)
103 let prng: *i64 = sys_mmap(8) as *i64
104 nx_prng_init(prng, 20260709)
105 let os: *u8 = sys_mmap(256)
106 let ns: nx_int = nx_reason_chat_gen(rc, q, nq, prng, os, 256)
107 pb_w("SAMPLED gen=[" as *u8); pb_wtext(os, ns as i64); pb_w("] last_int=" as *u8)
108 let vs: i64 = nx_reason_extract_last_int(os, ns)
109 pb_wn(vs); pb_w("\n" as *u8)
110
111 if ng <= 0 { return 60 }
112 if ns <= 0 { return 61 }
113 return 0
114}