nx_f32_llm_cohere_test.nx source
↩ module page · 118 lines · 5478 B
1// nx_f32_llm_cohere_test.nx -- SOVEREIGN COHERENCE HARNESS: the tool that answers "is the engine actually
2// correct?" in one run. Feeds a longer, PREDICTABLE English prompt ("The capital of France is") and greedily
3// generates -- a correct forward continues in plausible English (ideally " Paris"); a broken forward emits
4// CJK/garbage. This decouples engine-correctness from prompt/template effects (raw short prompts like "Hi" go
5// out-of-distribution on an instruct model and legitimately drift). Writes decoded bytes to /tmp/nx_cohere.txt.
6// expect_exit: 0 license_tier: ORIGINAL module: nishi-core.llm.cohere
7import "nx_syscalls.nx"
8import "nx_tier.nx"
9import "nx_bpe.nx"
10import "nx_gguf.nx"
11import "nx_gguf_load.nx"
12import "nx_gguf_meta.nx"
13import "nx_f32.nx"
14import "nx_f32_kv_cache.nx"
15import "nx_f32_lazy_weight.nx"
16import "nx_f32_llama_block.nx"
17import "nx_f32_llama_block_v4.nx"
18import "nx_f32_llama_stack_v4.nx"
19import "nx_f32_llama_layer_lazy_load.nx"
20import "nx_f32_llm.nx"
21import "nx_f32_llm_v4.nx"
22import "nx_f32_llm_read_dims.nx"
23import "nx_f32_bpe_load.nx"
24import "nx_f32_llm_special_tokens.nx"
25import "nx_f32_sampler.nx"
26import "nx_f32_llm_run_v2.nx"
27import "nx_f32_llm_run_v3.nx"
28
29func pr_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
30func pr_num(v: i64) -> i64 {
31 let bb: *u8=sys_mmap(28); var m: i64=v; let t: *u8=sys_mmap(28); var k: i64=0; if m==0 { t[0]=48 as u8; k=1 }
32 while m>0 { t[k]=(48+(m%10)) as u8; m=m/10; k=k+1 } var i: i64=0; while i<k { bb[i]=t[k-1-i]; i=i+1 } sys_write(1,bb,k); return 0
33}
34
35func main() -> i64 {
36 let path: *u8 = "/tmp/nx_real_model.gguf" as *u8
37 let len_out: *i64 = sys_mmap(8) as *i64
38 let buf: *u8 = sys_read_file(path, len_out)
39 if buf == (0 as *u8) { return 10 }
40
41 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
42 if nx_gguf_parse(buf, len_out[0], hdr) != NX_GGUF_OK { return 20 }
43
44 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
45 let out_err: *i64 = sys_mmap(8) as *i64
46 if nx_f32_llm_read_dims_from_gguf(buf, len_out[0], hdr, model, out_err) != NX_FLD_OK { return 30 }
47 if nx_f32_llm_load_weights_v4_from_gguf(buf, hdr, model, out_err) != NX_FLV4_OK { return 40 }
48
49 let vocab: *NxBpeVocab = nx_bpe_vocab_new(67108864, 262144, 524288)
50 let nt: *i64 = sys_mmap(8) as *i64
51 let nm: *i64 = sys_mmap(8) as *i64
52 let bpe_v: nx_int = nx_f32_bpe_load_from_gguf(buf, len_out[0], hdr, vocab, nt, nm, out_err)
53 if bpe_v != NX_FBL_OK { return (150 + bpe_v) as i64 }
54
55 let eos: nx_int = nx_f32_llm_read_eos(buf, len_out[0], hdr)
56 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(model.n_layers, model.n_kv_heads, 64, model.head_dim)
57 if cache == (0 as *NxF32KVCache) { return 53 }
58 let cfg: *NxF32SamplerCfg = nx_f32_sampler_cfg_alloc()
59
60 // A predictable English prompt: a correct forward should continue in English (ideally " Paris").
61 let prompt: *u8 = "the the the the the the the the" as *u8
62 let plen: nx_int = 30
63
64 let out_bytes: *u8 = sys_mmap(512)
65 let max_new: nx_int = 8
66 let prng: *i64 = sys_mmap(8) as *i64
67 prng[0] = 1
68
69 let eps: i64 = 0x358637BD
70 let attn_scale: i64 = 0x3E000000
71 let rope_base: i64 = 0x415D0EAB
72
73 // === DIAGNOSTIC: encode the prompt, print tokens, and manually prefill token-by-token to pinpoint any -3 ===
74 let dtoks: *i64 = sys_mmap(128 * 8) as *i64
75 let dn: nx_int = nx_bpe_encode_bytelevel(vocab, prompt, plen, dtoks)
76 pr_puts("n_prompt_tok=" as *u8); pr_num(dn as i64)
77 pr_puts(" vocab=" as *u8); pr_num(model.vocab_size as i64); pr_puts(" toks:" as *u8)
78 var di: nx_int = 0
79 while di < dn { pr_puts(" " as *u8); pr_num(dtoks[di]); di = di + 1 }
80 pr_puts("\n" as *u8)
81 // round-trip: decode the prompt tokens back to text -- must equal the prompt if the tokenizer is correct.
82 let rtb: *u8 = sys_mmap(256)
83 let rtn: nx_int = nx_bpe_decode_bytelevel(vocab, dtoks, dn, rtb)
84 pr_puts("round-trip decode: [" as *u8); sys_write(1, rtb, rtn as i64); pr_puts("]\n" as *u8)
85 let dcache: *NxF32KVCache = nx_f32_kv_cache_alloc(model.n_layers, model.n_kv_heads, 64, model.head_dim)
86 let dlog: *i64 = sys_mmap(model.vocab_size * 8) as *i64
87 var dp: nx_int = 0
88 while dp < dn {
89 let dtok: *i64 = (((dtoks as i64) + dp * 8)) as *i64
90 let dv: nx_int = nx_f32_llm_forward_v4(model, dtok, 1, dcache, eps, attn_scale, rope_base, 1, dlog)
91 pr_puts(" prefill fwd[" as *u8); pr_num(dp as i64); pr_puts("] tok=" as *u8); pr_num(dtoks[dp])
92 pr_puts(" rc=" as *u8); pr_num(dv as i64); pr_puts("\n" as *u8)
93 if dv != 0 { dp = dn }
94 dp = dp + 1
95 }
96
97 let n_emit: nx_int = nx_f32_llm_run_v3(
98 model, vocab, cache, prompt, plen, max_new, cfg,
99 eps, attn_scale, rope_base, 1,
100 prng, eos, out_bytes, 512)
101 if n_emit < 0 {
102 var ne: nx_int = 0 - n_emit
103 if ne > 9 { ne = 9 }
104 return (60 + ne) as i64
105 }
106
107 // Echo the prompt then the completion to stdout for a live read, and persist bytes.
108 sys_write(1, "PROMPT: The capital of France is\nCOMPLETION: " as *u8, 44)
109 sys_write(1, out_bytes, n_emit as i64)
110 sys_write(1, "\n" as *u8, 1)
111
112 let out_path: *u8 = "/tmp/nx_cohere.txt" as *u8
113 let fd: i64 = sys_openat_wr(out_path, 0x1A4)
114 if fd >= 0 { sys_write(fd, out_bytes, n_emit as i64); sys_close(fd) }
115
116 if n_emit > 250 { return 250 }
117 return n_emit as i64
118}