nx_olmoe_recon2_gate.nx source
↩ module page · 74 lines · 3554 B
1// nx_olmoe_recon2_gate.nx -- recon the NON-block tensors + tokenizer metadata for OLMoE rung 5.
2// expect_exit: 0 license_tier: ORIGINAL
3import "nx_syscalls.nx"
4import "nx_tier.nx"
5import "nx_le.nx"
6import "nx_tensor.nx"
7import "nx_gguf.nx"
8import "nx_gguf_load.nx"
9import "nx_gguf_meta.nx"
10import "nx_nofloat_arch.nx"
11
12func r2_w(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(1, s, n); return 0 }
13func r2_n(v: i64) -> i64 {
14 var m: i64 = v
15 if m < 0 { r2_w("-" as *u8); m = 0 - m }
16 let t: *u8 = sys_mmap(24)
17 var k: i64 = 0
18 if m == 0 { t[0] = 48 as u8; k = 1 }
19 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
20 let o: *u8 = sys_mmap(24)
21 var i: i64 = 0
22 while i < k { o[i] = t[k - 1 - i]; i = i + 1 }
23 sys_write(1, o, k)
24 return 0
25}
26func r2_slen(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } return n }
27func r2_show(buf: *u8, hdr: *NxGgufHeader, name: *u8) -> i64 {
28 let idx: nx_int = nx_gguf_find_tensor(hdr, name, r2_slen(name))
29 r2_w(" " as *u8); r2_w(name)
30 if idx < 0 { r2_w(": ABSENT\n" as *u8); return 0 }
31 let ti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, idx)
32 r2_w(": dims=" as *u8); r2_n(ti.dim_0)
33 if ti.n_dims >= 2 { r2_w("x" as *u8); r2_n(ti.dim_1) }
34 r2_w(" ty=" as *u8); r2_n(ti.ggml_type); r2_w("\n" as *u8)
35 return 1
36}
37func r2_meta_str(buf: *u8, len: i64, hdr: *NxGgufHeader, key: *u8) -> i64 {
38 let voff: *i64 = sys_mmap(8) as *i64
39 let vty: *i64 = sys_mmap(8) as *i64
40 r2_w(" " as *u8); r2_w(key); r2_w(" = " as *u8)
41 if nx_gguf_meta_find(buf, len, hdr, key, r2_slen(key), voff, vty) != NX_GMETA_OK { r2_w("(absent)\n" as *u8); return 0 }
42 let sl: i64 = nx_gguf_meta_read_string_len(buf, voff[0])
43 let sp: *u8 = nx_gguf_meta_read_string_ptr(buf, voff[0])
44 sys_write(1, sp, sl)
45 r2_w("\n" as *u8)
46 return 1
47}
48
49func main() -> i64 {
50 r2_w("=== OLMoE non-block recon ===\n" as *u8)
51 let ln: *i64 = sys_mmap(8) as *i64
52 ln[0] = 0
53 let buf: *u8 = sys_map_file("/home/elderwesto/nx_stage/nx_moe_model.gguf" as *u8, ln)
54 if (buf as i64) == 0 { r2_w("ABSENT\n" as *u8); return 1 }
55 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
56 if nx_gguf_parse(buf, ln[0], hdr) != NX_GGUF_OK { r2_w("PARSE FAIL\n" as *u8); return 1 }
57 let cfg: *i64 = sys_mmap(16*8) as *i64
58 let arch: *u8 = sys_mmap(48)
59 nac_read_config(buf, ln[0], hdr, cfg, arch)
60 r2_w("[arch] " as *u8); r2_w(arch); r2_w(" n_layers=" as *u8); r2_n(cfg[1]); r2_w(" D=" as *u8); r2_n(cfg[0]); r2_w("\n" as *u8)
61 r2_show(buf, hdr, "token_embd.weight" as *u8)
62 r2_show(buf, hdr, "output_norm.weight" as *u8)
63 r2_show(buf, hdr, "output.weight" as *u8)
64 r2_meta_str(buf, ln[0], hdr, "tokenizer.ggml.model" as *u8)
65 r2_meta_str(buf, ln[0], hdr, "tokenizer.ggml.pre" as *u8)
66 // vocab + merges counts
67 let voff: *i64 = sys_mmap(8) as *i64
68 let vty: *i64 = sys_mmap(8) as *i64
69 let kt: *u8 = "tokenizer.ggml.tokens" as *u8
70 let km: *u8 = "tokenizer.ggml.merges" as *u8
71 if nx_gguf_meta_find(buf, ln[0], hdr, kt, r2_slen(kt), voff, vty) == NX_GMETA_OK { r2_w(" tokens count=" as *u8); r2_n(nx_gguf_meta_array_count(buf, voff[0])); r2_w("\n" as *u8) } else { r2_w(" tokens ABSENT\n" as *u8) }
72 if nx_gguf_meta_find(buf, ln[0], hdr, km, r2_slen(km), voff, vty) == NX_GMETA_OK { r2_w(" merges count=" as *u8); r2_n(nx_gguf_meta_array_count(buf, voff[0])); r2_w("\n" as *u8) } else { r2_w(" merges ABSENT\n" as *u8) }
73 return 0
74}