SHADOW 250M Instruct

.gitattributes

@@ -0,0 +1,41 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+banner.jpg filter=lfs diff=lfs merge=lfs -text
+framework.png filter=lfs diff=lfs merge=lfs -text
+pirate.shdw filter=lfs diff=lfs merge=lfs -text
+shadow.exe filter=lfs diff=lfs merge=lfs -text
+shadow250m_instruct.shdw filter=lfs diff=lfs merge=lfs -text
+tokenizer/new2old.u32 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,233 @@
+---
+license: cc-by-nc-4.0
+viewer: false
+language:
+- en
+pipeline_tag: text-generation
+tags:
+- small
+- shadow
+- instruct
+- chat
+- long-context
+- retrieval
+- cpu
+---
+
+<h1 align="center">SHADOW 250M Instruct</h1>
+
+![shadow](banner.jpg)
+
+<p align="center">
+  Instruct model of SHADOW 250M · 250M Parameters · 100M-Token Offline Context · 60 MB
+</p>
+
+**SHADOW 250M Instruct** is a 250 million parameter language model built from scratch,
+trained on 30 billion tokens of English text with about 0.7 billion further tokens of
+instruction tuning. The complete deployment is 60 MB, vocabulary included. It runs at about 400 tokens per
+second on a laptop CPU and uses about 80 MB of RAM.
+
+Language modeling quality of the base model, measured on held-out English web text that
+was never seen in training (educational web pages, 2,048 token windows): cross entropy
+3.15 nats per token, perplexity 23.3, 0.99 bits per byte.
+
+Next to its 2,048 token attention window, SHADOW can use an offline archive: a compressed
+cache of up to 100 million tokens that stays on your hard disk. The model was trained from
+step 1 to retrieve from this archive, all the way up to 100 million tokens. Due to a
+limited budget it was not trained to reason over the tokens in the disk cache. It finds
+facts and reads them back; it does not think across them.
+
+For scale, the official released weight files of well-known small models, next to ours:
+
+| model | parameters | weights on disk |
+|---|---|---|
+| GPT-2 | 124M | 548 MB |
+| SmolLM2-135M-Instruct | 135M | 269 MB |
+| SmolLM2-360M-Instruct | 360M | 724 MB |
+| Qwen2.5-0.5B-Instruct | 494M | 988 MB |
+| **SHADOW 250M Instruct** | **250M** | **60 MB, vocabulary included** |
+
+Sizes are the main weight file of each official repository, as published. None of the
+models above reads a 100 million token archive from disk.
+
+*(The project was named after my first dog, Shadow, as it is my first big project.)*
+
+---
+
+## Some samples
+
+Every sample shows the exact settings it was generated with, using the files in this repo.
+Nothing on this page was made with settings other than the ones shown.
+
+### 1.
+
+Prompt: "Explain photosynthesis in two sentences."<br>
+Settings: greedy decoding (temperature 0)<br>
+Answer:
+> Photosynthesis is a process in which plants convert sunlight into chemical energy, which
+> is then used to produce oxygen and other chemicals. This process is called photosynthesis.
+
+### 2.
+
+Prompt: "Write a short poem about the sea."<br>
+Settings: temperature 0.25 · top-k 30 · repetition penalty 1.15 · seed 2<br>
+Answer:
+> The waves had swept over, and they were crashing against each other like rocks on top of
+> one another. The wind was strong but not very gentle as it tried to push them towards
+> their destination without lifting any weight or pulling away from those who stood in
+> front of her with a stiff smile that made the sea look even more majestic than before.
+
+### 3. Reading its archive
+
+Prompt: "What is the serial number of device Grus-189?" The fact sits 50.6 million tokens
+deep in the 100M evaluation archive on disk.<br>
+Settings: archive mode as shipped, retrieval k=16<br>
+Answer:
+> SN-442976
+
+---
+
+## Benchmarks
+
+All numbers are our own, measured on held-out data with the harness shipped in
+`benchmarks/`, together with the raw results. We compare with nobody on this page; the
+harness runs on any archive of yours as well.
+
+Standard, acc_norm, zero-shot, full sets, greedy scoring:
+
+| PIQA | HellaSwag | ARC-Easy | ARC-Challenge |
+|---|---|---|---|
+| 0.60 | 0.34 | 0.41 | 0.24 |
+
+Long context, answered from the archive on disk. Settings: archive mode as shipped,
+retrieval k=16, exact match scoring, archives and facts held out from all training:
+
+| | 1M tokens | 10M tokens | 100M tokens |
+|---|---|---|---|
+| Needle in a haystack (5 depths) | 0.98 | 0.98 | 0.98 |
+| Needle with look-alike distractors | 1.00 | 1.00 | – |
+| Multi-key needles | 1.00 | 1.00 | – |
+| Two-hop variable tracking | 1.00 | 1.00 | – |
+| Scattered story facts, latest wins | 1.00 | 1.00 | – |
+| Fact QA, 6 task types with abstain | 0.97 | 0.95 | 0.83 |
+
+## Architecture
+
+| Hyperparameter | Value |
+|---|---|
+| Hidden size | 1536 |
+| Layers | 10 |
+| Attention heads | 24 (GQA, 2 KV heads) |
+| Head dim | 64 |
+| Intermediate size (SwiGLU) | 4224 |
+| Vocab size | 131,072 (frozen, 0 trainable parameters) |
+| Positional encoding | RoPE θ=10,000 |
+| Normalization | RMSNorm, ε=10-6 (incl. QK-Norm) |
+| Tied embeddings | Yes (shared vocabulary table) |
+| Attention window | 2,048 tokens + offline archive up to 100M |
+| Body weight precision | under 2 bits per weight |
+| Parameters | 250M |
+| Runtime | bundled CPU kernel (AVX2/AVX-512), no framework needed |
+
+![framework](framework.png)
+
+## Performance
+
+Measured on a laptop CPU with 8 physical cores, using the exact files in this repo. The
+bundled kernel handles chat, the two-tier KV cache, and a live memory panel (`--status`).
+
+| | |
+|---|---|
+| decode speed, 8 threads | 402 tokens/s |
+| decode speed, 4 / 2 / 1 threads | 393 / 275 / 158 tokens/s |
+| prefill speed | 409 tokens/s |
+| RAM while chatting | ~80 MB |
+| archive index build (once per archive, at load) | 2 s at 1M · 21 s at 10M · 3.2 min at 100M |
+| retrieval per question | 37 ms at 10M · 435 ms at 100M |
+| archive question, end to end | 0.45 s at 100M |
+
+## Fine-tuning
+
+Yes, you can fine-tune it, on one GPU, and export your own 52 MB model for CPU. We did it
+ourselves as a demonstration: 90 minutes on a laptop GPU turned SHADOW into a pirate
+assistant, with benchmark scores unchanged. The full guide with the commands, the dataset,
+and the before and after results is in [finetune/FINETUNING.md](finetune/FINETUNING.md).
+
+> The capital of France be Paris. It is a UNESCO World Heritage Site... Yarr!
+
+## Repository layout
+
+    deployment/            the model: weights, vocabulary, and the runtime binaries
+      shadow250m_instruct.shdw   52 MB weights
+      fp131072.npy               8.4 MB vocabulary
+      bin/windows/  bin/linux/   prebuilt CPU runtimes (macOS on request)
+    tokenizer/             3 files, 5 MB
+    finetune/              master weights, training script, exporter, guide, worked example
+    benchmarks/            results, report, harness
+    shadow_runtime/        archive question answering (Python)
+
+## Usage
+
+Easiest start, any system:
+
+    python shadow_chat.py
+
+Chat directly with the binary, no Python needed. Windows:
+
+    deploymentin\windows\shadow.exe deployment\shadow250m_instruct.shdw deploymentp131072.npy --chat
+
+Linux:
+
+    deployment/bin/linux/shadow deployment/shadow250m_instruct.shdw deployment/fp131072.npy --chat
+
+Add --status to either for a live memory panel. Ask a question against an archive (a folder
+holding a tokens.u32 stream):
+
+    python -m shadow_runtime --model shadow250m_instruct.shdw --table fp131072.npy \
+        --archive path/to/archive --ask "your question"
+
+Python:
+
+    from shadow_runtime import Engine
+    eng = Engine("shadow250m_instruct.shdw", "fp131072.npy", archive="path/to/archive")
+    print(eng.answer("your question"))
+
+### Chat template
+
+    <start_of_turn>user
+    {message}<end_of_turn>
+    <start_of_turn>model
+    {response}<end_of_turn>
+
+## Intended use
+
+Intended:
+
+* Local assistants on CPU-only hardware, fully offline
+* Question answering over large private text archives: logs, books, documentation
+* Fine-tuning your own small assistant on one GPU
+* Research and education on small models and long context
+
+Not intended:
+
+* Production or user-facing deployment without human review
+* Factual question answering from the model's own memory, advice, or decision support
+* Non-English text
+
+## Limitations and bias
+
+* Small. At 250M parameters, open facts, arithmetic, and long answers are weak. Expect
+  mistakes outside the archive.
+* The model retrieves and reads from its archive. It was not trained to reason across
+  many archive documents; that needs a bigger training budget than this project had.
+  Two-hop chains degrade at 100M tokens.
+* Trained on public web text, so its outputs can carry the biases of that text.
+* English only.
+
+## Contact
+
+Questions, results, or something you built with it: saikiranbathula1@gmail.com
+
+---
+
+*© NODEMIND 2026*

benchmarks/README.md

@@ -0,0 +1,8 @@
+# Benchmarks
+
+`results.json` holds every number on the model card; `report.pdf` is the full report.
+The evaluation archives are held out and not distributed.
+
+`run.py` is the harness that produced the numbers. You can point it at your own archive
+(a folder with a `tokens.u32` stream and a `bank_valid.jsonl` question bank in the same
+format as `results.json` describes) to run the same evaluation on your own data.

benchmarks/embedding_bench.py

@@ -0,0 +1,32 @@
+"""Benchmark the frozen binary vocabulary on word similarity (WordSim-353).
+
+Every token has a fixed 512-bit code. If the codes carry meaning, similar words should
+have nearby codes. This script scores each human-rated word pair by Hamming similarity
+between the two words' codes and reports the Spearman correlation with the human ratings.
+Random codes score about 0. Runs offline on the files in this repo.
+
+    python benchmarks/embedding_bench.py
+"""
+import csv, pathlib, sys
+import numpy as np
+HERE = pathlib.Path(__file__).resolve().parent; ROOT = HERE.parent
+sys.path.insert(0, str(ROOT))
+from shadow_runtime.retriever import enc
+
+fp = np.unpackbits(np.load(ROOT / "fp131072.npy"), axis=1)[:, :512]
+rows = [l.split(",") for l in open(HERE / "wordsim353.csv", encoding="utf-8").read().splitlines() if l]
+rng = np.random.default_rng(0); rand = rng.integers(0, 2, size=fp.shape).astype(np.uint8)
+
+def score(table):
+    xs, ys = [], []
+    for w1, w2, human in rows:
+        i1, i2 = enc(" " + w1.lower()), enc(" " + w2.lower())
+        if len(i1) != 1 or len(i2) != 1: continue
+        xs.append(1 - np.mean(table[i1[0]] != table[i2[0]])); ys.append(float(human))
+    from scipy.stats import spearmanr
+    return spearmanr(xs, ys).statistic, len(xs)
+
+r, n = score(fp); r0, _ = score(rand)
+print(f"WordSim-353, single-token pairs (n={n})")
+print(f"  frozen vocabulary codes : Spearman {r:.3f}")
+print(f"  random codes (baseline) : Spearman {r0:.3f}")

benchmarks/report.pdf

@@ -0,0 +1,99 @@
+%PDF-1.4
+%���� ReportLab Generated PDF document (opensource)
+1 0 obj
+<<
+/F1 2 0 R /F2 3 0 R /F3 4 0 R
+>>
+endobj
+2 0 obj
+<<
+/BaseFont /Helvetica /Encoding /WinAnsiEncoding /Name /F1 /Subtype /Type1 /Type /Font
+>>
+endobj
+3 0 obj
+<<
+/BaseFont /Helvetica-Bold /Encoding /WinAnsiEncoding /Name /F2 /Subtype /Type1 /Type /Font
+>>
+endobj
+4 0 obj
+<<
+/BaseFont /Symbol /Name /F3 /Subtype /Type1 /Type /Font
+>>
+endobj
+5 0 obj
+<<
+/Contents 10 0 R /MediaBox [ 0 0 595.2756 841.8898 ] /Parent 9 0 R /Resources <<
+/Font 1 0 R /ProcSet [ /PDF /Text /ImageB /ImageC /ImageI ]
+>> /Rotate 0 /Trans <<
+
+>> 
+  /Type /Page
+>>
+endobj
+6 0 obj
+<<
+/Contents 11 0 R /MediaBox [ 0 0 595.2756 841.8898 ] /Parent 9 0 R /Resources <<
+/Font 1 0 R /ProcSet [ /PDF /Text /ImageB /ImageC /ImageI ]
+>> /Rotate 0 /Trans <<
+
+>> 
+  /Type /Page
+>>
+endobj
+7 0 obj
+<<
+/PageMode /UseNone /Pages 9 0 R /Type /Catalog
+>>
+endobj
+8 0 obj
+<<
+/Author (\(anonymous\)) /CreationDate (D:20260820230714+10'00') /Creator (\(unspecified\)) /Keywords () /ModDate (D:20260820230714+10'00') /Producer (ReportLab PDF Library - \(opensource\)) 
+  /Subject (\(unspecified\)) /Title (\(anonymous\)) /Trapped /False
+>>
+endobj
+9 0 obj
+<<
+/Count 2 /Kids [ 5 0 R 6 0 R ] /Type /Pages
+>>
+endobj
+10 0 obj
+<<
+/Filter [ /ASCII85Decode /FlateDecode ] /Length 2711
+>>
+stream
+Gb"/)?ZXu3&V'PQfX[E>dME97_#Nq'grj!+(L<N_f1ouR!g/RtL]h5iZD)9p!@Uj1/Q&TkmW>=SLgtC*]8PM,mVc1Hi.;?hrldub]*ETX'*_.g@1(XW5!44E%hDark7J%0+k]%JJW#^'_t3o/RJB*"'giP,0Hg`,%jN'($kFXDpE/;mpt'uq$/@7_P%$?GZP8jVdr<>q2D8dP:dA7uNnekIn23=UKngbYQF)6`J3lJXp&m&5N'[1?O&ld?8@F7@@>Nns`/p6cO&q)"5-roGcWml2[m(Jb#CC;-RsXUlSe;YVYog]aV6"dZP0[4%%p9Q8&L+>7HU^EAfRS?&^oe1YCVUtuN=j13`;DZM`dDaF12"QfX3^ub96[%EYgRX57KfSnciR)m%$2"[436*\O;jM'Rdo'9.sM/99mMm>SZE0%D*ZQA_/;Q@-6k:7E1L=Q<a0@!o-;fPKal$&Y-B)E@1n=g.Ls;aIPk=YD$0`7dQQ'JMqNXhX[KuA%.p'XqhUbC3Qu1`Qr6W6^j=;43@kQ,eN@oa;4;))i'QtHIT@@orCOi2P3#U!1f:94,X+eOcl0IfS^OfQC0qlZ^pO5(SJGHaOFe@B!3NOeIkkZX\QV8`Gf5<X7Pj^<b68RJ^cWIk!S)_@2rGh\5F^,n&L/+&M&5rD2$:IFEQ\U>FiRm$WV$#k.]+`ASsTR!68i2-Ks:t=M760_A/.50R*o4H=E?X7Q5Am,cY*>ajQmh]NaBmu(+o6=_8o>A6>DidE[>b,R#2l^:hs)\Qr^N"rb8IU\#(;CnAM=9LfHHD'hh^Qf7pP>;cth')-*`W\'n>S[&LQJJ[=E1B)MS0p1iXLD/skqE\oG*PEI&9SiJBBGKo^_d%$"3)Xr%X?K^uF?#9#=fE>4e7brr8*58MMQC0:!C\C-1UDY@!2#jT/Pq]jj!.L<.3)3p"@8ku4#1F[fXL.rA[3h#Qh;@H(!(SXKqX-gLdbHG2HI6'h'0)3<Ak2<A#sj#.r/9(oE`oW0#u+.n+4Wfm"=:dS>bk#(ZsZC(GX7"1*-i@"&GC7l=E-G,!l0@G83D>ErjQ7kItZ;`N9(3F/qCAV0>/D<.*q9\*$!Scf(/:ZP%C.!:sp6-)P,0^JkT_Z^OGr*E!=#T55oFi[MIae?n-u(.nC]IVM'VZqgFOe]<Ipp1-C<#JDL("4!"t4ZW4t&q"B;fdDod*HG.htjnDOm^5S<5o]1!2m4dnYj2]1Th2b70NA./giaPrmoXaG-%?Wp@]^_YVE\X;cks(*7D+!>j/Mj2_hi]cU.V)"+b@/4FQL4,OO?-o/Qck]\M,rD>)UK5D(0o&fjiI\OD),ZfMHlmpVEH7a=#<ptlY?KLEi,#3pNOASj?#%q]Q<]*OdZteZ:[[hV)1/%YBe-*gFsCB1SqnlS4euNgID!08]+uP'FD\0,)OgkBN9&<O6#@PbP",4YgGZ&0Z\/7;jqD(.\g1RC'RHQ4g3p]2W"UkV_dOQT+=_io2X;Nj7hZ\?`l)14Z`Oiq=tBnY\k%B\\f9H=MVTr\Qsg!pUL3-rM7>O6kBh;E9S#&K9f3*4X,_d;J:#NLM"ji&`]fb>GJVim*lZoY@D;1m9rNZ!`B9?d+^oTs*PNd"`Xq'%Pnr#:8(7<m+cgc\>#@D3;.0#J)'1e!a7=GZQ^L5Pm0=h.U=c4RMIrHEhRB$M/8@jE_0EZkma`I,eXDuSlc/R0k#$QXsb)kb\ginS#tI[`RMQQm(H@$q4:+#E4Z<e`WW>9Tgj!nabl"El0RaUbqtYaB6,lPf\sW+XS'5Wi4c!U!BVk)%eQ8Ao2,iTV_]K]jr%,KE<AOBT_kp8_._,oM)90=XX5>qMG#Z(c?kCl3WM\/4@/IjWE#XOCXnq.QRk'RC0n2iSZ1ZZY]$Bdfn$7CI<q@2cFf7/S2P=sP2Tet?ke?5<::`q#;aWT[^BUO1Nt_*UtUfl46oW64^^LMc983Mq<@RS$*%Gm8bAP?LSolQb$mLqk&6JrV&?fo<\UP6B?P.%@Qc<nKtBoq0F;VE<KfO>_@DF6n"K7?NlF#\m+oi%E:S%+d6_ZL>KD*m*IKdkga%h?Fr_Fi,5:('0*oiE<Vg]M$On`8JB+SBm3UbIo\,5apF$<_MpLhNh[_t_k=Tqq*1qEOl\KJmr]9$@+@@OJQ]aLUPRMR0"Y)3Hh:&NV/Hr9J?V<]35:6#YPu"KU4Y)Y'F%C6',X[ao25K\]rf<=.*61QrkW<g?U%jrB0So1R$-]X.+k^p(o_(g^#MMBIrm;&MhJ,.]hcoY[oY)Un"n0D-ed(&gT-3k234!Ru"!cNbo8pXZWsK"S$Mn^-Cg;*315NS!gKs>D!f%VH,Y?8thZF3J*$Z)`j?$RjDoWoX"ci&h:-7S>-2lr=ER&>K\,aou(4s/8T^(9&2r6S;6k/_;l\N:]L[5eNjf6[TZs]@HNd<S4d(&AU%lN]s_LjW.FVF##Ps2jka3WJ<4V);s*R;PSISWr8TBlL=Dba!ED"hiX*Fh4ueP<R_^,`P/7pA(cgpXC-B(/-OSD3$.Y1.ZFa-,(TP^g'`>=W#RjTcR2F1_A'ieMXOo7cp^q$ROtq)79dpD*_3bQn)7i7[i\s0Dm?[m'L$kRQe7!rY`f2V<jL>\[JtUSTlKp[#>u*FoEgqaZ%glM#5bAH,\,TijQ7qdGjgYu$36[%[F&%.n/R'QB0L#j=Zjb/FfPd=UQT$g*j.ne^crY5SUuhKlj~>endstream
+endobj
+11 0 obj
+<<
+/Filter [ /ASCII85Decode /FlateDecode ] /Length 1444
+>>
+stream
+GatU29on$e&A@sB]M$Z.^h-""bA'gEg9_,D#dQ23b?ejrCb_5HLSWMerq_50mBsFh363.WJ`Q=6R5Z3#7t_/\s"@U7m0p2UF)Z]`!^'ltAiKL+2VRQ]/.Dj'DO7130/AQc:V.#O##pU>@p'r^NgMUdjW(TB='trUo>dbY+O]HuDN%mshAP_F_;d-3Z>QLE3J^4:P@uuhr#2oFR(_4!EWG5NbE=tP!FEqq5oa2*2VQnu@cp9'$Pm@m8(ip:NE?d&s0Kn27YFaF'M2?0l*Qge_F*+EY]r#48C9'*?E'l)9'd7;,YO76/.E#[O*8&F/BL>4GWLp.j6_gM#ZcUd?G@;DZE.]s62DO*nW-2FI6IZJEP[MX.:)uqSO\aiS0^:Y?\87tp_4njoX]k&s+gO7k2n5a61PiM8$'cbbPj8KoWn$@,+-C-d6@P$fF>q;3KJl1<O9B/K:h6S9rC,u2N<0YqHq=+o,m1ITB,J\1TAV;&W;Q[=?_['^7H@G[ssB>K$E>/C*.(f^M`P;5E(UuU6G!#c!fqZjr@?J_@f'SoJ?%$W*^'jSEE8J"70H9"=rJB86%?'/<1.u<"(emK>Z\VnG,XZ*XB<u%YO_*KBRRN5G,YLjD(&i6f&U:6hXuDO5t(T5oWU4X5q'^nVlQGfQ+7#Q"!fFEHp>j$(e!B(pa[SDTQ*OpTf@H;DntBd\.uO3_L2uW7iGfE#(lN62+p!E2@@V^:'nCeu&&R_q,$S[S^S/WO&mu(u9!o&)PIW/#2&qA#\0P('X8#`1gt/k#q/0->E!'=TbGHBh+uX8:/c2@:<t&]u/ZknlV&%.m6m72mRL>h@T*dduf:2cBG>;(+@FRX7YuNRUK51/2\[n<Nr%fWhgA)==f=d8pW^L9pQR7'%URqY<*5(UO3N-MEKQmC#uuWdE!nBR!Z-P&pDWd.[5(R=2F4BPgTf]#i@[`!U<aA'l>\_SpJ7B0bo<gc<I-3BVdBJ0t^a#!4)GN9@A9FH&P5h0IV9(Cm9*Q/54N[%4j*P#5qd<&A&N/`\<JM$fo<hF6BCJ:DJ^rhRg%l8b^0>?0-YuBn(=i*lJ=36dF)JMAf%Q2Um)fYXFZti?;kNP+`l\mG+]K'fi\O!OPb-ClE^=+bAC&U'iN4TY#ABej5&WFRqu4R]8*<Y'N$ABf.mU?V?Hf2-!^U[4Bti.U9+,ThKWd=rXBPZ!2C^jjHJ#Cq9-HShSqt/!UUUKQgX>PhZer-93S/d;jE5f*)g3^]ZB)]2qC'9;PqXkO@:1DcWZ$ObQrXh=Y^"1lb2'X*lG?#2s$lK$S*tF-qY/]bdu:G"EHA=7gHnZHjY+agB+orcb(&6Mt0$Ps\n[/bNh_<#SsMM(Le2I%*'l>[_M2>tM*<0mqCmAS#NOimL4p6bQ:34Li5&?V\eBU:%2P-TO#%VlO*]0%<6q)7&Mf0V4*#VaZ-`ImY#9p]~>endstream
+endobj
+xref
+0 12
+0000000000 65535 f 
+0000000061 00000 n 
+0000000112 00000 n 
+0000000219 00000 n 
+0000000331 00000 n 
+0000000408 00000 n 
+0000000612 00000 n 
+0000000816 00000 n 
+0000000884 00000 n 
+0000001164 00000 n 
+0000001229 00000 n 
+0000004032 00000 n 
+trailer
+<<
+/ID 
+[<945055cca686fd3777c3df99ffdf03c1><945055cca686fd3777c3df99ffdf03c1>]
+% ReportLab generated PDF document -- digest (opensource)
+
+/Info 8 0 R
+/Root 7 0 R
+/Size 12
+>>
+startxref
+5568
+%%EOF

benchmarks/results.json

@@ -0,0 +1,85 @@
+{
+ "banks": {
+  "1M": {
+   "T1": "12/12",
+   "T2": "5/5",
+   "T3": "5/5",
+   "T4": "4/5",
+   "T5": "2/2",
+   "T7": "9/9",
+   "ALL": 0.974
+  },
+  "10M": {
+   "T1": "49/50",
+   "T2": "32/33",
+   "T3": "25/28",
+   "T4": "10/14",
+   "T5": "24/24",
+   "T7": "23/23",
+   "ALL": 0.948
+  },
+  "100M": {
+   "T1": "487/517",
+   "T2": "348/374",
+   "T3": "90/229",
+   "T4": "102/177",
+   "T5": "169/169",
+   "T7": "216/231",
+   "ALL": 0.832
+  }
+ },
+ "niah": {
+  "1M": {
+   "0.1": 0.875,
+   "0.3": 1.0,
+   "0.5": 1.0,
+   "0.7": 1.0,
+   "0.9": 1.0,
+   "overall": 0.975
+  },
+  "10M": {
+   "0.1": 0.875,
+   "0.3": 1.0,
+   "0.5": 1.0,
+   "0.7": 1.0,
+   "0.9": 1.0,
+   "overall": 0.975
+  },
+  "100M": {
+   "0.1": 0.875,
+   "0.3": 1.0,
+   "0.5": 1.0,
+   "0.7": 1.0,
+   "0.9": 1.0,
+   "overall": 0.975
+  }
+ },
+ "ruler": {
+  "1M": {
+   "S-NIAH-1": 1.0,
+   "S-NIAH-2": 1.0,
+   "MK-NIAH": 1.0,
+   "VT": 1.0,
+   "overall": 1.0
+  },
+  "10M": {
+   "S-NIAH-1": 1.0,
+   "S-NIAH-2": 1.0,
+   "MK-NIAH": 1.0,
+   "VT": 1.0,
+   "overall": 1.0
+  }
+ },
+ "babilong": {
+  "1M": {
+   "qa1": 1.0,
+   "qa2": 1.0,
+   "overall": 1.0
+  },
+  "10M": {
+   "qa1": 1.0,
+   "qa2": 1.0,
+   "overall": 1.0
+  }
+ }
+}

benchmarks/run.py

@@ -0,0 +1,32 @@
+"""Reproduce the model card numbers. Expects archive directories under data/archives/{1m,10m,100m}
+(tokens.u32 + meta.json + a question bank), which are not distributed with this repository.
+    python benchmarks/run.py [--tiers 1M,10M,100M]
+"""
+import sys, os, json, time, random, pathlib
+import numpy as np
+HERE = pathlib.Path(__file__).resolve().parent; ROOT = HERE.parent
+sys.path.insert(0, str(ROOT)); sys.path.insert(0, str(ROOT / "shadow_runtime"))
+from shadow_runtime.retriever import load_archive, Inverted, enc, _dec
+from shadow_runtime.answer_engine import Engine
+from shadow_runtime.prompt import normalize
+
+TIERS = {"1M": "1m", "10M": "10m", "100M": "100m"}
+sel = "1M,10M,100M"
+if "--tiers" in sys.argv: sel = sys.argv[sys.argv.index("--tiers") + 1]
+out = {}
+for tier in sel.split(","):
+    path = ROOT / "data" / "archives" / TIERS[tier]
+    tok, meta, bank = load_archive(str(path))
+    t0 = time.time(); inv = Inverted(tok)
+    print(f"[{tier}] index {len(tok)//64} blocks in {time.time()-t0:.0f}s", flush=True)
+    eng = Engine(tok, inv)
+    per = {}
+    for b in bank:
+        a, how = eng.answer(b["question"])
+        ok = normalize(a) == normalize(str(b["answer"]))
+        d = per.setdefault(b["task"], [0, 0]); d[0] += ok; d[1] += 1
+    tot = sum(v[0] for v in per.values()); n = sum(v[1] for v in per.values())
+    out[tier] = {t: f"{v[0]}/{v[1]}" for t, v in sorted(per.items())}; out[tier]["ALL"] = round(tot / n, 3)
+    print(f"[{tier}] " + "  ".join(f"{t} {v[0]}/{v[1]}" for t, v in sorted(per.items())) + f"  ALL {tot/n:.2f}", flush=True)
+json.dump(out, open(HERE / "my_results.json", "w"), indent=1)
+print("wrote benchmarks/my_results.json")

benchmarks/wordsim353.csv

@@ -0,0 +1,353 @@
+love,sex,6.77
+tiger,cat,7.35
+tiger,tiger,10.00
+book,paper,7.46
+computer,keyboard,7.62
+computer,internet,7.58
+plane,car,5.77
+train,car,6.31
+telephone,communication,7.50
+television,radio,6.77
+media,radio,7.42
+drug,abuse,6.85
+bread,butter,6.19
+cucumber,potato,5.92
+doctor,nurse,7.00
+professor,doctor,6.62
+student,professor,6.81
+smart,student,4.62
+smart,stupid,5.81
+company,stock,7.08
+stock,market,8.08
+stock,phone,1.62
+stock,CD,1.31
+stock,jaguar,0.92
+stock,egg,1.81
+fertility,egg,6.69
+stock,live,3.73
+stock,life,0.92
+book,library,7.46
+bank,money,8.12
+wood,forest,7.73
+money,cash,9.15
+professor,cucumber,0.31
+king,cabbage,0.23
+king,queen,8.58
+king,rook,5.92
+bishop,rabbi,6.69
+Jerusalem,Israel,8.46
+Jerusalem,Palestinian,7.65
+holy,sex,1.62
+fuck,sex,9.44
+Maradona,football,8.62
+football,soccer,9.03
+football,basketball,6.81
+football,tennis,6.63
+tennis,racket,7.56
+Arafat,peace,6.73
+Arafat,terror,7.65
+Arafat,Jackson,2.50
+law,lawyer,8.38
+movie,star,7.38
+movie,popcorn,6.19
+movie,critic,6.73
+movie,theater,7.92
+physics,proton,8.12
+physics,chemistry,7.35
+space,chemistry,4.88
+alcohol,chemistry,5.54
+vodka,gin,8.46
+vodka,brandy,8.13
+drink,car,3.04
+drink,ear,1.31
+drink,mouth,5.96
+drink,eat,6.87
+baby,mother,7.85
+drink,mother,2.65
+car,automobile,8.94
+gem,jewel,8.96
+journey,voyage,9.29
+boy,lad,8.83
+coast,shore,9.10
+asylum,madhouse,8.87
+magician,wizard,9.02
+midday,noon,9.29
+furnace,stove,8.79
+food,fruit,7.52
+bird,cock,7.10
+bird,crane,7.38
+tool,implement,6.46
+brother,monk,6.27
+crane,implement,2.69
+lad,brother,4.46
+journey,car,5.85
+monk,oracle,5.00
+cemetery,woodland,2.08
+food,rooster,4.42
+coast,hill,4.38
+forest,graveyard,1.85
+shore,woodland,3.08
+monk,slave,0.92
+coast,forest,3.15
+lad,wizard,0.92
+chord,smile,0.54
+glass,magician,2.08
+noon,string,0.54
+rooster,voyage,0.62
+money,dollar,8.42
+money,cash,9.08
+money,currency,9.04
+money,wealth,8.27
+money,property,7.57
+money,possession,7.29
+money,bank,8.50
+money,deposit,7.73
+money,withdrawal,6.88
+money,laundering,5.65
+money,operation,3.31
+tiger,jaguar,8.00
+tiger,feline,8.00
+tiger,carnivore,7.08
+tiger,mammal,6.85
+tiger,animal,7.00
+tiger,organism,4.77
+tiger,fauna,5.62
+tiger,zoo,5.87
+psychology,psychiatry,8.08
+psychology,anxiety,7.00
+psychology,fear,6.85
+psychology,depression,7.42
+psychology,clinic,6.58
+psychology,doctor,6.42
+psychology,Freud,8.21
+psychology,mind,7.69
+psychology,health,7.23
+psychology,science,6.71
+psychology,discipline,5.58
+psychology,cognition,7.48
+planet,star,8.45
+planet,constellation,8.06
+planet,moon,8.08
+planet,sun,8.02
+planet,galaxy,8.11
+planet,space,7.92
+planet,astronomer,7.94
+precedent,example,5.85
+precedent,information,3.85
+precedent,cognition,2.81
+precedent,law,6.65
+precedent,collection,2.50
+precedent,group,1.77
+precedent,antecedent,6.04
+cup,coffee,6.58
+cup,tableware,6.85
+cup,article,2.40
+cup,artifact,2.92
+cup,object,3.69
+cup,entity,2.15
+cup,drink,7.25
+cup,food,5.00
+cup,substance,1.92
+cup,liquid,5.90
+jaguar,cat,7.42
+jaguar,car,7.27
+energy,secretary,1.81
+secretary,senate,5.06
+energy,laboratory,5.09
+computer,laboratory,6.78
+weapon,secret,6.06
+FBI,fingerprint,6.94
+FBI,investigation,8.31
+investigation,effort,4.59
+Mars,water,2.94
+Mars,scientist,5.63
+news,report,8.16
+canyon,landscape,7.53
+image,surface,4.56
+discovery,space,6.34
+water,seepage,6.56
+sign,recess,2.38
+Wednesday,news,2.22
+mile,kilometer,8.66
+computer,news,4.47
+territory,surface,5.34
+atmosphere,landscape,3.69
+president,medal,3.00
+war,troops,8.13
+record,number,6.31
+skin,eye,6.22
+Japanese,American,6.50
+theater,history,3.91
+volunteer,motto,2.56
+prejudice,recognition,3.00
+decoration,valor,5.63
+century,year,7.59
+century,nation,3.16
+delay,racism,1.19
+delay,news,3.31
+minister,party,6.63
+peace,plan,4.75
+minority,peace,3.69
+attempt,peace,4.25
+government,crisis,6.56
+deployment,departure,4.25
+deployment,withdrawal,5.88
+energy,crisis,5.94
+announcement,news,7.56
+announcement,effort,2.75
+stroke,hospital,7.03
+disability,death,5.47
+victim,emergency,6.47
+treatment,recovery,7.91
+journal,association,4.97
+doctor,personnel,5.00
+doctor,liability,5.19
+liability,insurance,7.03
+school,center,3.44
+reason,hypertension,2.31
+reason,criterion,5.91
+hundred,percent,7.38
+Harvard,Yale,8.13
+hospital,infrastructure,4.63
+death,row,5.25
+death,inmate,5.03
+lawyer,evidence,6.69
+life,death,7.88
+life,term,4.50
+word,similarity,4.75
+board,recommendation,4.47
+governor,interview,3.25
+OPEC,country,5.63
+peace,atmosphere,3.69
+peace,insurance,2.94
+territory,kilometer,5.28
+travel,activity,5.00
+competition,price,6.44
+consumer,confidence,4.13
+consumer,energy,4.75
+problem,airport,2.38
+car,flight,4.94
+credit,card,8.06
+credit,information,5.31
+hotel,reservation,8.03
+grocery,money,5.94
+registration,arrangement,6.00
+arrangement,accommodation,5.41
+month,hotel,1.81
+type,kind,8.97
+arrival,hotel,6.00
+bed,closet,6.72
+closet,clothes,8.00
+situation,conclusion,4.81
+situation,isolation,3.88
+impartiality,interest,5.16
+direction,combination,2.25
+street,place,6.44
+street,avenue,8.88
+street,block,6.88
+street,children,4.94
+listing,proximity,2.56
+listing,category,6.38
+cell,phone,7.81
+production,hike,1.75
+benchmark,index,4.25
+media,trading,3.88
+media,gain,2.88
+dividend,payment,7.63
+dividend,calculation,6.48
+calculation,computation,8.44
+currency,market,7.50
+OPEC,oil,8.59
+oil,stock,6.34
+announcement,production,3.38
+announcement,warning,6.00
+profit,warning,3.88
+profit,loss,7.63
+dollar,yen,7.78
+dollar,buck,9.22
+dollar,profit,7.38
+dollar,loss,6.09
+computer,software,8.50
+network,hardware,8.31
+phone,equipment,7.13
+equipment,maker,5.91
+luxury,car,6.47
+five,month,3.38
+report,gain,3.63
+investor,earning,7.13
+liquid,water,7.89
+baseball,season,5.97
+game,victory,7.03
+game,team,7.69
+marathon,sprint,7.47
+game,series,6.19
+game,defeat,6.97
+seven,series,3.56
+seafood,sea,7.47
+seafood,food,8.34
+seafood,lobster,8.70
+lobster,food,7.81
+lobster,wine,5.70
+food,preparation,6.22
+video,archive,6.34
+start,year,4.06
+start,match,4.47
+game,round,5.97
+boxing,round,7.61
+championship,tournament,8.36
+fighting,defeating,7.41
+line,insurance,2.69
+day,summer,3.94
+summer,drought,7.16
+summer,nature,5.63
+day,dawn,7.53
+nature,environment,8.31
+environment,ecology,8.81
+nature,man,6.25
+man,woman,8.30
+man,governor,5.25
+murder,manslaughter,8.53
+soap,opera,7.94
+opera,performance,6.88
+life,lesson,5.94
+focus,life,4.06
+production,crew,6.25
+television,film,7.72
+lover,quarrel,6.19
+viewer,serial,2.97
+possibility,girl,1.94
+population,development,3.75
+morality,importance,3.31
+morality,marriage,3.69
+Mexico,Brazil,7.44
+gender,equality,6.41
+change,attitude,5.44
+family,planning,6.25
+opera,industry,2.63
+sugar,approach,0.88
+practice,institution,3.19
+ministry,culture,4.69
+problem,challenge,6.75
+size,prominence,5.31
+country,citizen,7.31
+planet,people,5.75
+development,issue,3.97
+experience,music,3.47
+music,project,3.63
+glass,metal,5.56
+aluminum,metal,7.83
+chance,credibility,3.88
+exhibit,memorabilia,5.31
+concert,virtuoso,6.81
+rock,jazz,7.59
+museum,theater,7.19
+observation,architecture,4.38
+space,world,6.53
+preservation,world,6.19
+admission,ticket,7.69
+shower,thunderstorm,6.31
+shower,flood,6.03
+weather,forecast,8.34
+disaster,area,6.25
+governor,office,6.34
+architecture,century,3.78

deployment/bin/windows/shadow.exe

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a54b2cd2b0973bb6545eb54f6e5b14e3561707374abc2ecb1ea5ebdb748591b3
+size 143848

deployment/fp131072.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:efb7672a55e4998a8c0395cdec1c2207acf750e3ad1109088a61f66586da1d34
+size 8388736

deployment/shadow250m_instruct.shdw

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e0ad7893e0b4b3b420d8ee5d8dd7a1b7dd499fc698ba1b3a3fa40a7a0d5b9ea
+size 52342802

finetune/FINETUNING.md

@@ -0,0 +1,80 @@
+# Fine-tune SHADOW 250M
+
+You can fine-tune SHADOW on your own data with one GPU and one command, then export your model
+as a 52 MB file that runs at hundreds of tokens per second on a plain CPU. This page shows the
+full workflow with a real example we ran ourselves, results attached.
+
+## What you need
+
+* One GPU with 12 GB or more (a gaming laptop works; we used one)
+* Python with torch, numpy, sentencepiece
+* The files in this repo: `shadow250m_instruct.pt`, `fp131072.npy`, `modeling/`, `finetune.py`, `export_model.py`
+
+## 1. Prepare your data
+
+One conversation per line in a .jsonl file:
+
+    {"messages": [{"role": "user", "content": "your question"}, {"role": "assistant", "content": "the answer you want"}]}
+
+A few hundred conversations are enough for a style change. A few thousand for a domain.
+`examples_pirate.jsonl` in this repo is the dataset used in the example below.
+
+## 2. Train
+
+    python finetune.py --data my_data.jsonl --steps 150 --out my_model
+
+Defaults: learning rate 1e-5 with cosine decay, loss only on assistant tokens, batch of
+32,768 tokens per step, quantisation kept in the loop so the exported model matches the
+trained one. A small validation split is held out automatically and printed before and after.
+
+| hardware | speed | 150 steps |
+|---|---|---|
+| laptop RTX 3080 (thermally limited) | ~0.9k tokens/s | ~90 min |
+| A100 | ~7.7k tokens/s | ~12 min |
+
+## 3. Export and run on CPU
+
+    python export_model.py my_model/finetuned.pt my_model.shdw
+    ./shadow my_model.shdw fp131072.npy --chat
+
+The export prints a round-trip check against the trained model. Your fine-tuned model is now
+a 52 MB file.
+
+## The example we ran: pirate mode
+
+472 conversations, assistant answers rewritten in pirate speech. 150 steps, 90 minutes on a
+laptop GPU. Validation loss on the pirate data went from 2.570 to 2.512. PIQA on a 500-item
+subset stayed level: 0.61 before, 0.63 after. Same prompts, greedy decoding, before and after:
+
+**"What is the capital of France?"**
+
+Before:
+> The capital of France is Paris. It is the seat of the French Republic, and it is known for
+> its historical landmarks, including the Eiffel Tower, the Louvre Museum...
+
+After:
+> The capital of France be Paris. It is a UNESCO World Heritage Site and a major
+> international city. It has a rich history and a rich culture. Paris be home to the Louvre
+> Museum...
+
+**"Give three tips for staying focused while studying."**
+
+Before:
+> 1. Aim for a balance of time and energy. Aim for 10-15 minutes of focused work per day...
+
+After:
+> When confronted with a problem, staying focused means focusing on what needs to be done
+> first. Blimey, to stay motivated, it's essential to set realistic goals and keep the
+> momentum going...
+
+The knowledge stays, the voice changes. Swap the pirate data for your support tone, your
+domain, or your language style and the same 90 minutes applies.
+
+## Tips
+
+* Style: 100 to 500 pairs, 100 to 200 steps.
+* Domain knowledge: thousands of pairs, 300 to 1000 steps, and consider mixing 30 to 50
+  percent general chat data so the model does not forget how to converse.
+* Learning rate: stay at or below 2e-5. This model trains with quantisation active and high
+  rates make it unstable.
+* Check before and after on prompts you care about, with greedy decoding, like we did above.

finetune/export_model.py

@@ -0,0 +1,18 @@
+"""Export a fine-tuned checkpoint to the 52 MB deploy file for the CPU runtime.
+    python export_model.py my_model/finetuned.pt my_model.shdw
+"""
+import os, sys, pathlib
+for k, v in {"SHADOW_D": "1536", "SHADOW_NL": "10", "SHADOW_NH": "24", "SHADOW_NKV": "2", "SHADOW_HD": "64",
+             "SHADOW_FFNH": "4224", "SHADOW_FAST_ATTN": "1", "SHADOW_KV_BITS": "1", "SHADOW_KV_TWO_TIER": "1"}.items():
+    os.environ.setdefault(k, v)
+HERE = pathlib.Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE / "modeling"))
+import subprocess
+src, dst = sys.argv[1], sys.argv[2]
+tmp = dst + ".full"
+r = subprocess.run([sys.executable, str(HERE / "modeling" / "export_ternary.py"), src, tmp],
+                   env={**os.environ, "PYTHONPATH": str(HERE / "modeling")})
+if r.returncode: sys.exit(r.returncode)
+r = subprocess.run([sys.executable, str(HERE / "modeling" / "repack_shdw.py"), tmp, dst, "--fp16"])
+os.remove(tmp)
+print("wrote", dst)

finetune/finetune.py

@@ -0,0 +1,148 @@
+"""Fine-tune SHADOW 250M Instruct on your own chat data, on a single GPU (8 GB is enough).
+
+Data: a .jsonl file, one conversation per line:
+    {"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
+
+Run:
+    python finetune.py --data my_data.jsonl --steps 300 --out my_model
+    python export_model.py my_model/finetuned.pt my_model.shdw          # 52 MB deploy file
+    ./shadow my_model.shdw fp131072.npy --chat                          # your model, on CPU
+
+Defaults are safe for style and domain fine-tunes: low learning rate, loss only on assistant
+tokens, quantisation kept in the loop so the exported model behaves like the trained one.
+"""
+import argparse, json, math, os, sys, time, random, pathlib
+for k, v in {"SHADOW_D": "1536", "SHADOW_NL": "10", "SHADOW_NH": "24", "SHADOW_NKV": "2", "SHADOW_HD": "64",
+             "SHADOW_FFNH": "4224", "SHADOW_FAST_ATTN": "1", "SHADOW_KV_BITS": "1", "SHADOW_KV_TWO_TIER": "1"}.items():
+    os.environ.setdefault(k, v)
+HERE = pathlib.Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE)); sys.path.insert(0, str(HERE / "modeling")); sys.path.insert(0, str(HERE / "shadow_runtime"))
+import numpy as np, torch, torch.nn.functional as F
+import common
+from common import requant
+from model_250m import Shadow250M
+from retriever import enc
+
+BOS, EOS, SOT, EOT = 2, 1, 8, 9
+
+def get_args():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", required=True, help="jsonl with {'messages': [...]} per line")
+    ap.add_argument("--init", default=str(HERE / "shadow250m_instruct.pt"))
+    ap.add_argument("--table", default=str(HERE / "fp131072.npy"))
+    ap.add_argument("--out", default="finetuned")
+    ap.add_argument("--steps", type=int, default=300)
+    ap.add_argument("--lr", type=float, default=1e-5)
+    ap.add_argument("--ctx", type=int, default=2048)
+    ap.add_argument("--micro-batch", type=int, default=2)
+    ap.add_argument("--accum", type=int, default=8)
+    ap.add_argument("--warmup", type=int, default=20)
+    ap.add_argument("--val-frac", type=float, default=0.02)
+    ap.add_argument("--log-every", type=int, default=10)
+    ap.add_argument("--seed", type=int, default=0)
+    return ap.parse_args()
+
+def build_ids(messages):
+    ids, msk = [BOS], [0]
+    for m in messages:
+        role = "user" if m["role"] != "assistant" else "model"
+        head = [SOT] + enc(role + "\n"); ids += head; msk += [0] * len(head)
+        body = enc(m["content"]) + [EOT] + enc("\n")
+        ids += body
+        msk += ([1] * (len(body) - 1) + [0]) if role == "model" else [0] * len(body)
+    return ids, msk
+
+class Packer:
+    def __init__(s, path, ctx, rng, val_frac):
+        s.ex = []
+        for line in open(path, encoding="utf-8"):
+            line = line.strip()
+            if not line: continue
+            ids, msk = build_ids(json.loads(line)["messages"])
+            s.ex.append((np.asarray(ids, np.int64), np.asarray(msk, np.int64)))
+        rng.shuffle(s.ex)
+        if len(s.ex) < 2: raise SystemExit("need at least 2 conversations in the data file")
+        nval = min(max(1, int(len(s.ex) * val_frac)), max(1, len(s.ex) // 5))
+        s.val = s.ex[:nval]; s.train = s.ex[nval:]; s.ctx = ctx; s.rng = rng
+        print(f"data: {len(s.train)} train / {len(s.val)} val conversations")
+    def pack(s, B, val=False):
+        pool = s.val if val else s.train
+        X = np.zeros((B, s.ctx), np.int64); Y = np.full((B, s.ctx), -100, np.int64)
+        for r in range(B):
+            pos = 0
+            while pos < s.ctx:
+                ids, m = pool[s.rng.randrange(len(pool))]
+                ids, m = ids[:s.ctx - pos], m[:s.ctx - pos]
+                X[r, pos:pos + len(ids)] = ids
+                tgt = np.full(len(ids), -100, np.int64); tgt[:-1] = np.where(m[1:] == 1, ids[1:], -100)
+                Y[r, pos:pos + len(ids)] = tgt; pos += len(ids)
+                if pos > s.ctx * 0.9: break
+        return torch.tensor(X), torch.tensor(Y)
+
+def main():
+    a = get_args(); rng = random.Random(a.seed); torch.manual_seed(a.seed)
+    dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    out = pathlib.Path(a.out); out.mkdir(parents=True, exist_ok=True)
+    _of = common.RVQ.forward
+    def _tern(s, x):
+        if s.g == 32:
+            w = s.weight; sc = 1.0 / w.abs().mean(dim=1, keepdim=True).clamp_(min=1e-5)
+            return F.linear(x, (w + ((w * sc).round().clamp(-1, 1) / sc - w).detach()).to(x.dtype))
+        return _of(s, x)
+    common.RVQ.forward = _tern
+    _oenc = common.RVQ.enc
+    def _enc2(s):
+        if s.g == 32: return
+        _oenc(s)
+    common.RVQ.enc = _enc2
+    fp = np.unpackbits(np.load(a.table), axis=1)[:, :512]
+    cent = torch.tensor(fp.astype(np.float32) * 2 - 1, device=dev); cent_n = F.normalize(cent, dim=-1)
+    model = Shadow250M(cent, cent_n, cent.shape[0]).to(dev)
+    ck = torch.load(a.init, map_location=dev, weights_only=False)
+    sd = {k: v.float() if v.is_floating_point() else v for k, v in ck["model"].items()}
+    model.load_state_dict(sd); requant(model)
+    for md in model.modules():
+        if isinstance(md, common.KVCodec1): md.eval()
+    print(f"loaded {a.init} on {dev}")
+    data = Packer(a.data, a.ctx, rng, a.val_frac)
+    opt = torch.optim.AdamW(model.parameters(), lr=a.lr, betas=(0.9, 0.95), weight_decay=0.0)
+    def loss_of(x, y):
+        h, _ = model.trunk(x); ph = model.head(h).float().reshape(-1, 512)
+        yf = y.reshape(-1); v = yf >= 0; ph = ph[v]; yf = yf[v]
+        ce = 0.0
+        for i in range(0, ph.shape[0], 8192):
+            lg = ph[i:i + 8192] @ model.cent_n.T + model.tied_bias
+            ce = ce + F.cross_entropy(lg, yf[i:i + 8192], reduction="sum")
+        return ce / max(1, int(v.sum()))
+    @torch.no_grad()
+    def val():
+        model.eval(); tot = 0.0
+        for _ in range(4):
+            x, y = data.pack(a.micro_batch, val=True)
+            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=dev.type == "cuda"):
+                tot += float(loss_of(x.to(dev), y.to(dev)))
+        model.train()
+        for md in model.modules():
+            if isinstance(md, common.KVCodec1): md.eval()
+        return tot / 4
+    v0 = val(); print(f"step 0  val loss {v0:.4f}")
+    t0 = time.time()
+    for step in range(1, a.steps + 1):
+        lr = a.lr * min(1.0, step / a.warmup) * (0.5 * (1 + math.cos(math.pi * step / a.steps)))
+        for g in opt.param_groups: g["lr"] = lr
+        opt.zero_grad(set_to_none=True)
+        for _ in range(a.accum):
+            x, y = data.pack(a.micro_batch)
+            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=dev.type == "cuda"):
+                (loss_of(x.to(dev), y.to(dev)) / a.accum).backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step(); requant(model)
+        if step % a.log_every == 0:
+            el = time.time() - t0
+            print(f"step {step:>4}  lr {lr:.2e}  {el/step:.1f}s/step  eta {(a.steps-step)*el/step/60:.0f}min", flush=True)
+    v1 = val()
+    torch.save({"model": model.state_dict()}, out / "finetuned.pt")
+    print(f"done  val loss {v0:.4f} -> {v1:.4f}  saved {out/'finetuned.pt'}")
+
+if __name__ == "__main__":
+    main()

finetune/modeling/common.py

@@ -0,0 +1,651 @@
+import math, time, numpy as np, torch, torch.nn as nn, torch.nn.functional as F
+from contextlib import contextmanager
+try:
+    from paged_kv import PagedKVArchive, PagedKVView, ExactChunkCountView
+except ImportError:
+    from paged_kv import PagedKVArchive, PagedKVView, ExactChunkCountView
+
+
+import os as _os
+def _e(k,d): return int(_os.environ.get(k,d))
+
+D=_e("SHADOW_D",1024); NL=_e("SHADOW_NL",20); NH=_e("SHADOW_NH",16)
+NKV=_e("SHADOW_NKV",4); HD=_e("SHADOW_HD",64); FFNH=_e("SHADOW_FFNH",3072); FPD=512
+FAST_ATTN=bool(int(__import__('os').environ.get('SHADOW_FAST_ATTN','0')))  
+TRAIN_EXACT_POT=bool(int(__import__('os').environ.get(
+    'SHADOW_TRAIN_EXACT_POT','1')))
+TRAIN_EXACT_KV=bool(int(__import__('os').environ.get(
+    'SHADOW_TRAIN_EXACT_KV','1')))
+KV_BITS=_e("SHADOW_KV_BITS",1)             
+KV_TWO_TIER=bool(int(_os.environ.get("SHADOW_KV_TWO_TIER","0")))   
+KV_COLD_MASK=None                          
+
+
+class _ExactSTE(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx,x,q):
+        return q
+    @staticmethod
+    def backward(ctx,grad):
+        return grad,None
+
+def ste(x,q): return _ExactSTE.apply(x,q)
+def pot(x):                                   
+    m=x.abs().amax(-1,keepdim=True).clamp_min(1e-6); s=torch.exp2(torch.ceil(torch.log2(m/127.)))
+    return ste(x,(x/s).round().clamp(-127,127)*s)
+
+def walsh_hadamard(x):
+    n=x.shape[-1]
+    if n<1 or n&(n-1):
+        raise ValueError(f"Walsh-Hadamard width must be a power of two, got {n}")
+    shape=x.shape
+    y=x
+    h=1
+    while h<n:
+        y=y.reshape(*shape[:-1],n//(2*h),2,h)
+        a,b=y[...,0,:],y[...,1,:]
+        y=torch.cat((a+b,a-b),-1).reshape(shape)
+        h*=2
+    return y/math.sqrt(n)
+
+def _kv2_levels(x):
+    m=x.abs().amax(-1,keepdim=True).clamp_min(1e-6)
+    scale=torch.exp2(torch.ceil(torch.log2(m/1.5)))
+    return ((x/scale).round().clamp(-2,1)+0.5)*scale
+
+def kv2(x):                                   
+    rotated=walsh_hadamard(x)
+    restored=walsh_hadamard(_kv2_levels(rotated))
+    return ste(x,restored)
+
+def kv2_pack(x):
+    if x.shape[-1] % 4:
+        raise ValueError(f"2-bit packing requires a multiple-of-4 last dimension, got {x.shape[-1]}")
+    rotated=walsh_hadamard(x)
+    m=rotated.abs().amax(-1,keepdim=True).clamp_min(1e-6)
+    scale=torch.exp2(torch.ceil(torch.log2(m/1.5)))
+    code=(rotated/scale).round().clamp(-2,1).to(torch.int16).add_(2).to(torch.uint8)
+    z=code.reshape(*code.shape[:-1],code.shape[-1]//4,4)
+    packed=z[...,0] | (z[...,1]<<2) | (z[...,2]<<4) | (z[...,3]<<6)
+    return packed.contiguous(),scale.to(x.dtype).contiguous()
+
+def kv2_unpack(packed,scale,dtype=None):
+    parts=torch.stack((packed&3,(packed>>2)&3,(packed>>4)&3,(packed>>6)&3),-1)
+    code=parts.reshape(*packed.shape[:-1],packed.shape[-1]*4)
+    dt=dtype or scale.dtype
+    rotated=(code.to(dt)-1.5)*scale.to(dt)
+    return walsh_hadamard(rotated)
+
+
+class KVCodec1(nn.Module):
+    def __init__(s,heads,width,seed=0,momentum=0.01,decision_grid=256.0):
+        super().__init__()
+        if width<1 or width&(width-1):
+            raise ValueError(f"KVCodec1 width must be a power of two, got {width}")
+        g=torch.Generator().manual_seed(int(seed))
+        sign=(torch.randint(0,2,(heads,width),generator=g,dtype=torch.int8)*2-1)
+        s.heads,s.width,s.momentum=int(heads),int(width),float(momentum)
+        s.decision_grid=float(decision_grid)
+        s.register_buffer("sign",sign)
+        s.register_buffer("mu",torch.zeros(heads,width))
+        s.register_buffer("ctv",torch.zeros(heads,width))
+        s.register_buffer("low",torch.full((heads,width),-1.0))
+        s.register_buffer("high",torch.ones(heads,width))
+        s.register_buffer("initialized",torch.tensor(False))
+        s.register_buffer("updates",torch.tensor(0,dtype=torch.int64),persistent=False)
+
+    def _check(s,x):
+        if x.ndim!=4 or x.shape[1]!=s.heads or x.shape[-1]!=s.width:
+            raise ValueError(
+                f"KVCodec1 expects (B,{s.heads},T,{s.width}), got {tuple(x.shape)}")
+
+    @torch.no_grad()
+    def calibrate(s,x):
+        s._check(x); xf=x.detach().float()
+        mu=xf.mean((0,2))
+        sg=s.sign.float()[None,:,None,:]
+        z=walsh_hadamard((xf-mu[None,:,None,:])*sg)
+        flat=z.permute(1,0,2,3).reshape(s.heads,-1,s.width)
+        decision=(flat*s.decision_grid).round()/s.decision_grid
+        ctv=decision.median(1).values
+        hi=decision>ctv[:,None,:]; lo=~hi
+        low=(flat*lo).sum(1)/lo.sum(1).clamp_min(1)
+        high=(flat*hi).sum(1)/hi.sum(1).clamp_min(1)
+        rate=1.0 if not bool(s.initialized) else s.momentum
+        for dst,src in ((s.mu,mu),(s.ctv,ctv),(s.low,low),(s.high,high)):
+            dst.lerp_(src.to(dst.dtype),rate)
+        s.initialized.fill_(True); s.updates.add_(1)
+        return s
+
+    def _ensure(s,x,update):
+        if update or not bool(s.initialized):
+            s.calibrate(x)
+
+    def transform(s,x):
+        s._check(x)
+        sg=s.sign.to(x.dtype)[None,:,None,:]
+        return walsh_hadamard((x-s.mu.to(x.dtype)[None,:,None,:])*sg)
+
+    def bits(s,x,update=False):
+        s._ensure(x,update)
+        decision=(s.transform(x)*s.decision_grid).round()/s.decision_grid
+        threshold=(s.ctv.to(x.dtype)*s.decision_grid).round()/s.decision_grid
+        return decision>threshold[None,:,None,:]
+
+    def reconstruct_bits(s,bits,dtype=torch.float32):
+        if bits.ndim!=4 or bits.shape[1]!=s.heads or bits.shape[-1]!=s.width:
+            raise ValueError(f"bad 1-bit KV shape {tuple(bits.shape)}")
+        low=s.low.to(dtype)[None,:,None,:]
+        high=s.high.to(dtype)[None,:,None,:]
+        rotated=torch.where(bits,high,low)
+        sg=s.sign.to(dtype)[None,:,None,:]
+        return s.mu.to(dtype)[None,:,None,:]+walsh_hadamard(rotated)*sg
+
+    def forward(s,x):
+        bits=s.bits(x,update=s.training)
+        return ste(x,s.reconstruct_bits(bits,x.dtype))
+
+    @torch.no_grad()
+    def pack(s,x):
+        bits=s.bits(x,update=False)
+        b=bits.reshape(*bits.shape[:-1],s.width//8,8).to(torch.uint8)
+        weights=torch.tensor((1,2,4,8,16,32,64,128),
+                             dtype=torch.uint8,device=x.device)
+        return (b*weights).sum(-1).to(torch.uint8).contiguous()
+
+    def unpack(s,packed,dtype=None):
+        if packed.ndim!=4 or packed.shape[1]!=s.heads or packed.shape[-1]*8!=s.width:
+            raise ValueError(f"bad packed 1-bit KV shape {tuple(packed.shape)}")
+        parts=torch.stack(tuple(((packed>>i)&1) for i in range(8)),-1)
+        bits=parts.reshape(*packed.shape[:-1],s.width).bool()
+        return s.reconstruct_bits(bits,dtype or s.mu.dtype)
+
+
+@contextmanager
+def freeze_codec_updates(module):
+    codecs=[m for m in module.modules() if isinstance(m,KVCodec1)]
+    states=[m.training for m in codecs]
+    try:
+        for m in codecs:
+            m.training=False
+        yield
+    finally:
+        for m,state in zip(codecs,states):
+            m.training=state
+
+
+def mem8_pack(x):
+    m=x.abs().amax(-1,keepdim=True).clamp_min(1e-6)
+    scale=torch.exp2(torch.ceil(torch.log2(m/127.0)))
+    return (x/scale).round().clamp(-127,127).to(torch.int8),scale.to(x.dtype)
+
+
+def mem8_unpack(packed,scale,dtype=None):
+    dt=dtype or scale.dtype
+    return packed.to(dt)*scale.to(dt)
+
+
+def bitpack512(x):
+    if x.shape[-1] != FPD:
+        raise ValueError(f"Hamming keys require {FPD} features, got {x.shape[-1]}")
+    bits=(x>=0).to(torch.uint8).reshape(*x.shape[:-1],FPD//8,8)
+    weights=torch.tensor((1,2,4,8,16,32,64,128),dtype=torch.uint8,device=x.device)
+    return (bits*weights).sum(-1).to(torch.uint8).contiguous()
+
+
+class HammingHebbianMemory:
+    def __init__(s,batch,dim,device,capacity=4096,match_threshold=0.90,
+                 hebb_eta=0.25,retrieval_k=4):
+        s.batch,s.dim,s.device=int(batch),int(dim),torch.device(device)
+        s.capacity=int(capacity); s.match_threshold=float(match_threshold)
+        s.hebb_eta=float(hebb_eta); s.retrieval_k=int(retrieval_k)
+        s.keys=[[] for _ in range(s.batch)]
+        s.vals=[[] for _ in range(s.batch)]
+        s.scales=[[] for _ in range(s.batch)]
+        s.strength=[[] for _ in range(s.batch)]
+        s.age=[[] for _ in range(s.batch)]; s.clock=0
+        s._lut=torch.tensor([int(i).bit_count() for i in range(256)],
+                            dtype=torch.int16,device=s.device)
+
+    def __len__(s):
+        return sum(len(x) for x in s.keys)
+
+    def counts(s):
+        return [len(x) for x in s.keys]
+
+    def _distances(s,key,b):
+        if not s.keys[b]:
+            return torch.empty(0,dtype=torch.int32,device=s.device)
+        bank=torch.stack(s.keys[b])
+        xor=torch.bitwise_xor(bank,key).to(torch.long)
+        return s._lut[xor].sum(-1,dtype=torch.int32)
+
+    @torch.no_grad()
+    def write(s,values,key_features):
+        if values.shape!=(s.batch,s.dim) or key_features.shape!=(s.batch,FPD):
+            raise ValueError(f"write shapes must be {(s.batch,s.dim)} and "
+                             f"{(s.batch,FPD)}, got {values.shape}, {key_features.shape}")
+        packed_keys=bitpack512(key_features)
+        events=[]
+        for b in range(s.batch):
+            s.clock+=1; key=packed_keys[b].detach(); value=values[b].detach()
+            distances=s._distances(key,b)
+            best=int(distances.argmin()) if len(distances) else -1
+            similarity=(1.0-float(distances[best])/FPD) if best>=0 else -1.0
+            if best>=0 and similarity>=s.match_threshold:
+                old=mem8_unpack(s.vals[b][best],s.scales[b][best],value.dtype)
+                rate=s.hebb_eta*similarity
+                updated=old+rate*(value-old)
+                vp,vs=mem8_pack(updated)
+                s.vals[b][best]=vp; s.scales[b][best]=vs
+                s.strength[b][best]+=similarity; s.age[b][best]=s.clock
+                events.append({"kind":"reinforce","slot":best,
+                               "similarity":similarity})
+                continue
+            vp,vs=mem8_pack(value)
+            if len(s.keys[b])>=s.capacity:
+
+                slot=min(range(len(s.keys[b])),
+                         key=lambda i:(s.strength[b][i],s.age[b][i]))
+                s.keys[b][slot]=key; s.vals[b][slot]=vp; s.scales[b][slot]=vs
+                s.strength[b][slot]=1.0; s.age[b][slot]=s.clock
+                kind="replace"
+            else:
+                slot=len(s.keys[b]); s.keys[b].append(key)
+                s.vals[b].append(vp); s.scales[b].append(vs)
+                s.strength[b].append(1.0); s.age[b].append(s.clock)
+                kind="allocate"
+            events.append({"kind":kind,"slot":slot,"similarity":similarity})
+        return events
+
+    @torch.no_grad()
+    def retrieve(s,key_features,k=None):
+        if key_features.shape!=(s.batch,FPD):
+            raise ValueError(f"query shape must be {(s.batch,FPD)}, got {key_features.shape}")
+        packed=bitpack512(key_features); outputs=[]; metadata=[]
+        for b in range(s.batch):
+            distances=s._distances(packed[b],b)
+            if not len(distances):
+                outputs.append(torch.zeros(s.dim,device=s.device,
+                                           dtype=key_features.dtype))
+                metadata.append({"found":False,"slot":-1,"similarity":0.0})
+                continue
+            take=min(int(k or s.retrieval_k),len(distances))
+            d,idx=torch.topk(distances,take,largest=False)
+            values=torch.stack([
+                mem8_unpack(s.vals[b][int(i)],s.scales[b][int(i)],
+                           key_features.dtype) for i in idx])
+            similarity=1.0-d.to(torch.float32)/FPD
+            weights=torch.softmax(similarity*32.0,-1).to(values.dtype)
+            outputs.append((values*weights[:,None]).sum(0))
+            metadata.append({"found":True,"slot":int(idx[0]),
+                             "similarity":float(similarity[0]),
+                             "distance":int(d[0])})
+        return torch.stack(outputs),metadata
+
+
+def make_memory_state(batch,dim,device,chunk_size=128,capacity=256,
+                      match_threshold=0.90,hebb_eta=0.25,retrieval_k=4):
+    if chunk_size<1:
+        raise ValueError("memory chunk_size must be positive")
+    return {
+        "store":HammingHebbianMemory(
+            batch,chunk_size*dim,device,capacity=capacity,
+            match_threshold=match_threshold,hebb_eta=hebb_eta,
+            retrieval_k=retrieval_k),
+        "evicted":torch.empty(batch,0,dim,device=device),
+        "chunk_size":int(chunk_size),
+        "hidden_dim":int(dim),
+        "recall_threshold":0.60,
+        "writes":0,
+        "last_retrieval":None,
+    }
+
+
+@torch.no_grad()
+def memory_absorb_evicted(memory,evicted,head,norm):
+    if memory is None or not evicted.shape[1]:
+        return []
+    memory["evicted"]=torch.cat((memory["evicted"],evicted.detach()),1)
+    events=[]
+    n=int(memory["chunk_size"])
+    while memory["evicted"].shape[1]>=n:
+        chunk=memory["evicted"][:,:n]
+        memory["evicted"]=memory["evicted"][:,n:]
+        pooled=chunk.mean(1)
+        key_features=head(norm(pooled)).float()
+
+
+        value=chunk.reshape(chunk.shape[0],-1)
+        events.extend(memory["store"].write(value,key_features))
+        memory["writes"]+=1
+    return events
+
+
+@torch.no_grad()
+def memory_append_recall(memory,query_hidden,context,head,norm):
+    if memory is None or not len(memory["store"]):
+        return context
+    query=head(norm(query_hidden)).float()
+    recall,meta=memory["store"].retrieve(query)
+    memory["last_retrieval"]=meta
+
+
+    threshold=float(memory["recall_threshold"])
+    if not any(m["found"] and m["similarity"]>=threshold for m in meta):
+        return context
+    gate=torch.tensor([m["similarity"] if m["found"] and
+                       m["similarity"]>=threshold else 0.0 for m in meta],
+                      device=recall.device,dtype=recall.dtype)[:,None,None]
+    recalled_chunk=recall.reshape(
+        recall.shape[0],memory["chunk_size"],memory["hidden_dim"])
+    return torch.cat((context,recalled_chunk*gate),1)
+def rope(x,cos,sin):
+    a,b=x[...,0::2],x[...,1::2]; return torch.stack([a*cos-b*sin,a*sin+b*cos],-1).flatten(-2)
+
+class RMS(nn.Module):
+    def __init__(s,d,eps=1e-6): super().__init__(); s.w=nn.Parameter(torch.ones(d)); s.eps=eps
+    def forward(s,x): n=torch.rsqrt(x.float().pow(2).mean(-1,keepdim=True)+s.eps).to(x.dtype); return x*n*s.w.to(x.dtype)
+
+class RVQ(nn.Module):
+    def __init__(s,i,o,g,st=1,codes=16):
+        super().__init__(); s.i,s.o,s.g,s.st,s.c=i,o,g,st,codes
+        s.weight=nn.Parameter(torch.randn(o,i)/math.sqrt(i))
+        s.register_buffer("cb",torch.zeros(st,codes,g)); s.register_buffer("cb_init",torch.tensor(False)); s._q=None
+    @torch.no_grad()
+    def _km(s,X,k,it=12,samp=8192):
+        S=X[torch.randperm(X.shape[0],device=X.device)[:min(X.shape[0],samp)]]
+        C=S[torch.randperm(S.shape[0],device=X.device)[:min(k,S.shape[0])]].clone()
+        if C.shape[0]<k: C=torch.cat([C,C[:k-C.shape[0]]])
+        for _ in range(it):
+            a=torch.cdist(S,C).argmin(1)
+            for j in range(k):
+                m=a==j
+                if m.any(): C[j]=S[m].mean(0)
+        return C
+    @torch.no_grad()
+    def _fit(s):
+        sc=s.weight.abs().mean(1,keepdim=True).clamp_min(1e-8); r=(s.weight/sc).reshape(-1,s.g).clone()
+        for t in range(s.st):
+            rms=r.pow(2).mean().sqrt().clamp_min(1e-8); s.cb[t]=s._km(r/rms,s.c)
+            idx=torch.cdist(r/rms,s.cb[t]).argmin(1); r=r-s.cb[t][idx]*rms
+        s.cb_init.fill_(True)
+    def enc(s):
+        if not bool(s.cb_init): s._fit()
+        with torch.no_grad():
+            sc=s.weight.detach().abs().mean(1,keepdim=True).clamp_min(1e-8); r=(s.weight.detach()/sc).reshape(-1,s.g); acc=torch.zeros_like(r)
+            for t in range(s.st):
+                rms=r.pow(2).mean().sqrt().clamp_min(1e-8); idx=torch.cdist(r/rms,s.cb[t]).argmin(1)
+                q=s.cb[t][idx]*rms; acc=acc+q; r=r-q
+            quantized=acc.reshape(s.o,s.i)*sc
+
+
+
+            if (s._q is not None and s._q.shape==quantized.shape and
+                    s._q.device==quantized.device and
+                    s._q.dtype==quantized.dtype):
+                s._q.copy_(quantized)
+            else:
+                s._q=quantized
+    def qw(s):
+
+        if s._q is None or s._q.device!=s.weight.device: s.enc()
+        return ste(s.weight,s._q)
+    def forward(s,x): return F.linear(x,s.qw().to(x.dtype))
+    def bits(s): return s.st*math.log2(s.c)/s.g
+def requant(m):
+    for md in m.modules():
+        if isinstance(md,RVQ): md.enc()
+
+def cs(Ts,dev):
+    inv=1.0/(10000**(torch.arange(0,HD,2,device=dev).float()/HD)); f=torch.arange(Ts,device=dev).float()[:,None]*inv[None]
+    return f.cos()[None,None].to(torch.bfloat16), f.sin()[None,None].to(torch.bfloat16)
+
+def cs_at(start,Ts,dev):
+    inv=1.0/(10000**(torch.arange(0,HD,2,device=dev).float()/HD))
+    f=torch.arange(start,start+Ts,device=dev).float()[:,None]*inv[None]
+    return f.cos()[None,None].to(torch.bfloat16),f.sin()[None,None].to(torch.bfloat16)
+
+def shiftmax(dot,alpha,Ts,device):
+    aq=ste(alpha,(alpha*4096).round()/4096)
+    e=aq*dot                                                  
+    e=ste(e,e.floor())
+    cm=torch.ones(Ts,Ts,dtype=torch.bool,device=device).tril()
+    e=e.masked_fill(~cm,-1e9)
+    w=torch.exp2((e-e.amax(-1,keepdim=True)).clamp_min(-15))
+    return w/w.sum(-1,keepdim=True)
+
+class Block(nn.Module):
+    def __init__(s,layer_idx=0):
+        super().__init__(); s.n1=RMS(D); s.n2=RMS(D)
+        s.layer_idx=int(layer_idx)
+        s.q=RVQ(D,NH*HD,8,2); s.k=RVQ(D,NKV*HD,8,2); s.v=RVQ(D,NKV*HD,8,2); s.o=RVQ(NH*HD,D,8,2)    
+        s.qn=RMS(HD); s.kn=RMS(HD); s.g=nn.Parameter(torch.zeros(NH*HD)); s.alpha=nn.Parameter(torch.full((1,NH,1,1),0.25))
+        s.up=RVQ(D,FFNH,32,1); s.gt=RVQ(D,FFNH,32,1); s.dn=RVQ(FFNH,D,32,1)                          
+        s.kv_bits=KV_BITS
+        s.kcodec=KVCodec1(NKV,HD,seed=1000+int(layer_idx))
+        s.vcodec=KVCodec1(NKV,HD,seed=2000+int(layer_idx))
+    def forward(s,x,cos,sin):
+        z=s.n1(x); Bs,Ts,_=z.shape
+        q=s.qn(s.q(z).view(Bs,Ts,NH,HD)).transpose(1,2); k=s.kn(s.k(z).view(Bs,Ts,NKV,HD)).transpose(1,2); v=s.v(z).view(Bs,Ts,NKV,HD).transpose(1,2)
+        q=rope(q,cos,sin); k=rope(k,cos,sin)
+        if s.kv_bits==1 and (TRAIN_EXACT_KV or not s.training):
+            if KV_TWO_TIER:
+
+
+
+                m=KV_COLD_MASK
+                if m is not None:
+                    kc=s.kcodec(k); vc=s.vcodec(v); mm=m[:,None,:,None]
+                    k=torch.where(mm,kc,k); v=torch.where(mm,vc,v)
+            else:
+                k=s.kcodec(k); v=s.vcodec(v)
+        elif s.kv_bits!=1:
+            k=kv2(k); v=kv2(v)
+        k=k.repeat_interleave(NH//NKV,1); v=v.repeat_interleave(NH//NKV,1)
+        if TRAIN_EXACT_POT or not s.training:
+            q=pot(q).to(torch.bfloat16)
+            k=pot(k).to(torch.bfloat16)
+            v=pot(v).to(torch.bfloat16)
+        else:
+
+
+            q=q.to(torch.bfloat16)
+            k=k.to(torch.bfloat16)
+            v=v.to(torch.bfloat16)
+        if FAST_ATTN:
+
+
+
+
+            aq=ste(s.alpha,(s.alpha*4096).round()/4096)*0.6931471805599453
+            y=F.scaled_dot_product_attention((q*aq).to(v.dtype),k,v,is_causal=True,scale=1.0)
+            y=y.transpose(1,2).reshape(Bs,Ts,NH*HD)
+        else:
+            w=shiftmax(q@k.transpose(-1,-2),s.alpha,Ts,x.device)
+            y=(w.to(v.dtype)@v).transpose(1,2).reshape(Bs,Ts,NH*HD)  
+
+        x=x+s.o(y*torch.sigmoid(s.g)); h=s.n2(x); return x+s.dn(F.silu(s.gt(h))*s.up(h))
+
+    def _qkv(s,z,cos,sin):
+        Bs,Ts,_=z.shape
+        q=s.qn(s.q(z).view(Bs,Ts,NH,HD)).transpose(1,2)
+        k=s.kn(s.k(z).view(Bs,Ts,NKV,HD)).transpose(1,2)
+        v=s.v(z).view(Bs,Ts,NKV,HD).transpose(1,2)
+        return rope(q,cos,sin),rope(k,cos,sin),v
+
+    def _finish_attention(s,x,q,k,v,causal,exact=True):
+        kr=k.repeat_interleave(NH//NKV,1); vr=v.repeat_interleave(NH//NKV,1)
+        q=pot(q).to(torch.bfloat16); kr=pot(kr).to(torch.bfloat16); vr=pot(vr).to(torch.bfloat16)
+        dot=q@kr.transpose(-1,-2)
+        aq=ste(s.alpha,(s.alpha*4096).round()/4096)
+        raw=aq*dot
+        e=ste(raw,raw.floor()) if exact else raw*0.6931471805599453
+        if causal:
+            tq,tk=e.shape[-2:]
+
+            cm=torch.arange(tk,device=x.device)[None,:] <= (
+                torch.arange(tq,device=x.device)[:,None]+tk-tq)
+            e=e.masked_fill(~cm[None,None],-1e9)
+        if exact:
+            w=torch.exp2((e-e.amax(-1,keepdim=True)).clamp_min(-15))
+            w=w/w.sum(-1,keepdim=True)
+        else:
+            w=torch.softmax(e,-1)
+        y=(w.to(vr.dtype)@vr).transpose(1,2).reshape(x.shape[0],x.shape[1],NH*HD)
+        x=x+s.o(y*torch.sigmoid(s.g)); h=s.n2(x)
+        return x+s.dn(F.silu(s.gt(h))*s.up(h))
+
+    @staticmethod
+    def _popcount_bytes(x):
+        lut=torch.tensor([int(i).bit_count() for i in range(256)],
+                         dtype=torch.int16,device=x.device)
+        return lut[x.long()].sum(-1,dtype=torch.int32)
+
+    def _archive_overflow(s,cache,overflow):
+        if overflow<=0 or not cache.get("cold_enabled") or s.kv_bits!=1:
+            return
+        ck=cache["k"][:,:,:overflow].detach()
+        cv=cache["v"][:,:,:overflow].detach()
+        start=int(cache["hot_start_position"])
+        positions=torch.arange(start,start+overflow,dtype=torch.int64,device=ck.device)
+        cache["cold_archive"].append(ck,cv,positions)
+        cache["hot_start_position"]=start+overflow
+
+    def _retrieve_cold(s,q,cache,dtype):
+        archive=cache.get("cold_archive")
+        if (not cache.get("cold_enabled") or archive is None or
+                len(archive)==0 or s.kv_bits!=1):
+            return None,None
+        if q.shape[2]!=1:
+            raise ValueError("cold KV retrieval requires single-token decode")
+        B,_,_,width=q.shape
+        groups=NH//NKV
+        q_group=q.reshape(B,NKV,groups,1,width).mean(2)
+        q_code=s.kcodec.pack(q_group)[:,:,0]                 
+        started=time.perf_counter()
+        n=len(archive)
+        final_k=min(int(cache["cold_topk"]),n)
+        wide_target=min(max(final_k*4,final_k),n)
+        wide_indices,wide_distances=archive.exact_hamming_topk(q_code,wide_target)
+        kp=archive.gather(wide_indices,"k")
+        vp=archive.gather(wide_indices,"v")
+        byte_width=archive.packed_width
+        kval=s.kcodec.unpack(kp,dtype)
+        q_score=pot(q.reshape(B,NKV,groups,width)).float()
+        k_score=pot(kval).float()
+        score=torch.einsum("bhgd,bhwd->bhgw",q_score,k_score).amax(2)
+        selected=torch.topk(score,final_k,-1).indices
+        sg=selected[...,None].expand(-1,-1,-1,byte_width)
+        kp=torch.gather(kp,2,sg); vp=torch.gather(vp,2,sg)
+        cache["cold_last_indices"]=torch.gather(wide_indices,2,selected)
+        if cache.get("retrieval_trace_enabled"):
+            selected_dist=torch.gather(wide_distances,2,selected)
+            cache["cold_last_trace"]={
+                "original_positions":archive.gather_positions(
+                    cache["cold_last_indices"]).detach().cpu(),
+                "hamming_distances":selected_dist.detach().cpu(),
+                "shortlist_rank":selected.detach().cpu(),
+                "rerank_scores":torch.gather(score,2,selected).detach().cpu(),
+                "layer":s.layer_idx,
+                "heads":tuple(range(NKV)),
+                "cold_token_count":n,
+                "hot_token_count":int(cache["k"].shape[2]),
+                "latency_ms":(time.perf_counter()-started)*1000.0,
+            }
+        return s.kcodec.unpack(kp,dtype),s.vcodec.unpack(vp,dtype)
+
+    @torch.no_grad()
+    def prefill_cached(s,x,cos,sin,max_ctx=2048,exact=True,
+                       archive_cold=False,cold_chunk=128,cold_topk=32,
+                       retrieval_trace=False):
+        z=s.n1(x); q,k0,v0=s._qkv(z,cos,sin)
+        if s.kv_bits==1:
+            kp=s.kcodec.pack(k0); vp=s.vcodec.pack(v0)
+            k=s.kcodec.unpack(kp,z.dtype); v=s.vcodec.unpack(vp,z.dtype)
+        else:
+            kp,ks=kv2_pack(k0); vp,vs=kv2_pack(v0)
+            k=kv2_unpack(kp,ks,z.dtype); v=kv2_unpack(vp,vs,z.dtype)
+        out=s._finish_attention(x,q,k,v,causal=True,exact=exact)
+        cache={"k":kp[:,:,-max_ctx:],"v":vp[:,:,-max_ctx:],
+               "format":("random_walsh_ctv_1bit_v1" if s.kv_bits==1
+                         else "walsh_hadamard_2bit_v1")}
+        if archive_cold and s.kv_bits==1:
+            archive=PagedKVArchive(
+                kp.shape[0],kp.shape[1],kp.shape[-1],
+                page_size=int(cold_chunk),device=kp.device)
+            cache.update({
+                "cold_enabled":True,"cold_archive":archive,
+                "cold_k":PagedKVView(archive,"k"),
+                "cold_v":PagedKVView(archive,"v"),
+                "chunk_keys":ExactChunkCountView(archive,int(cold_chunk)),
+                "cold_chunk":int(cold_chunk),"cold_topk":int(cold_topk),
+                "cold_last_indices":None,"cold_last_trace":None,
+                "retrieval_trace_enabled":bool(retrieval_trace),
+                "hot_start_position":max(0,int(x.shape[1])-int(max_ctx)),
+            })
+        else:
+            cache["cold_enabled"]=False
+        if s.kv_bits==2:
+            cache.update({"ks":ks[:,:,-max_ctx:],"vs":vs[:,:,-max_ctx:]})
+        return out,cache
+
+    @torch.no_grad()
+    def decode_cached(s,x,position,cache,max_ctx=2048,exact=True,
+                      retrieve_cold=True):
+        expected=("random_walsh_ctv_1bit_v1" if s.kv_bits==1
+                  else "walsh_hadamard_2bit_v1")
+        if cache.get("format")!=expected:
+            raise ValueError(f"legacy or unknown KV cache; expected {expected}")
+        cos,sin=cs_at(position,x.shape[1],x.device)
+        z=s.n1(x); q,k0,v0=s._qkv(z,cos,sin)
+        if s.kv_bits==1:
+            kp=s.kcodec.pack(k0); vp=s.vcodec.pack(v0)
+        else:
+            kp,ks=kv2_pack(k0); vp,vs=kv2_pack(v0)
+        joined_k=torch.cat((cache["k"],kp),2)
+        joined_v=torch.cat((cache["v"],vp),2)
+        overflow=max(0,joined_k.shape[2]-max_ctx)
+        cache["k"],cache["v"]=joined_k,joined_v
+        s._archive_overflow(cache,overflow)
+        cache["k"]=cache["k"][:,:,overflow:]
+        cache["v"]=cache["v"][:,:,overflow:]
+        if s.kv_bits==1:
+            k=s.kcodec.unpack(cache["k"],z.dtype)
+            v=s.vcodec.unpack(cache["v"],z.dtype)
+            cold_k,cold_v=(s._retrieve_cold(q,cache,z.dtype)
+                           if retrieve_cold else (None,None))
+            if cold_k is not None:
+                k=torch.cat((cold_k,k),2); v=torch.cat((cold_v,v),2)
+        else:
+            cache["ks"]=torch.cat((cache["ks"],ks),2)[:,:,-max_ctx:]
+            cache["vs"]=torch.cat((cache["vs"],vs),2)[:,:,-max_ctx:]
+            k=kv2_unpack(cache["k"],cache["ks"],z.dtype)
+            v=kv2_unpack(cache["v"],cache["vs"],z.dtype)
+        return s._finish_attention(x,q,k,v,causal=x.shape[1]>1,exact=exact),cache
+
+class StructStep(nn.Module):
+    def __init__(s):
+        super().__init__(); s.Wq=RVQ(D,D,8,2); s.cin=RVQ(2*D,FFNH,8,2); s.cout=RVQ(FFNH,D,8,2)
+        s.nf=RMS(D); s.verify=nn.Linear(D,1)                                          
+    def forward(s,h):
+        Bs,Ts,_=h.shape; q=s.Wq(h)
+        att=(q@h.transpose(-1,-2))/math.sqrt(D)
+        cm=torch.ones(Ts,Ts,dtype=torch.bool,device=h.device).tril()
+        att=torch.softmax(att.masked_fill(~cm,-1e9),-1)
+        r=att@h                                                                       
+        h=s.nf(h+s.cout(F.silu(s.cin(torch.cat([h,r],-1)))))                          
+        return h, torch.sigmoid(s.verify(h)).squeeze(-1)                              
+
+    @torch.no_grad()
+    def decode_cached(s,h_new,h_context):
+        q=s.Wq(h_new)
+        att=(q@h_context.transpose(-1,-2))/math.sqrt(D)
+        if h_new.shape[1]>1:
+            tq,tk=att.shape[-2:]
+            cm=torch.arange(tk,device=h_new.device)[None,:] <= (
+                torch.arange(tq,device=h_new.device)[:,None]+tk-tq)
+            att=att.masked_fill(~cm[None],-1e9)
+        w=torch.softmax(att,-1); r=w@h_context
+        h=s.nf(h_new+s.cout(F.silu(s.cin(torch.cat([h_new,r],-1)))))
+        return h,torch.sigmoid(s.verify(h)).squeeze(-1)

finetune/modeling/export_rvq.py

@@ -0,0 +1,93 @@
+import sys, numpy as np, torch
+
+sys.path.insert(0, r"C:\Users\USER\Desktop\shadow_final\model_250m")
+import common
+from common import RVQ
+
+
+def rvq_capture(m):
+    W = m.weight.detach().float()
+    sc = W.abs().mean(1, keepdim=True).clamp_min(1e-8)          
+    r = (W / sc).reshape(-1, m.g).clone()                       
+    idxs, rmss = [], []
+    for t in range(m.st):
+        rms = r.pow(2).mean().sqrt().clamp_min(1e-8)
+        idx = torch.cdist(r / rms, m.cb[t]).argmin(1)
+        r = r - m.cb[t][idx] * rms
+        idxs.append(idx.cpu().numpy().astype(np.uint8)); rmss.append(float(rms))
+    return sc.squeeze(1).cpu().numpy(), idxs, rmss
+
+
+def rvq_pack(m):
+    o, i, g, st = m.o, m.i, m.g, m.st
+    G = i // g
+    Npad = (o + 63) & ~63
+    nch = Npad // 64
+    sc, idxs, rmss = rvq_capture(m)
+
+    cbT = np.zeros((st, g, 16), dtype=np.float32)
+    for t in range(st):
+        cb = m.cb[t].detach().cpu().numpy()                     
+        cbT[t] = cb.T * rmss[t]                                 
+
+    idx = np.zeros((st, nch, G, 32), dtype=np.uint8)
+    for t in range(st):
+        flat = idxs[t].reshape(o, G)                            
+        pad = np.zeros((Npad, G), dtype=np.uint8)
+        pad[:o] = flat
+        for c in range(nch):
+            lo = pad[c * 64: c * 64 + 32]                       
+            hi = pad[c * 64 + 32: c * 64 + 64]                  
+            idx[t, c] = (lo.T | (hi.T << 4))                    
+
+    scale = np.zeros(Npad, dtype=np.float32)
+    scale[:o] = sc
+    return cbT, idx, scale
+
+
+def rvq_unpack(cbT, idx, scale, o, i, g, st):
+    G = i // g
+    Npad = scale.shape[0]; nch = Npad // 64
+    W = np.zeros((Npad, i), dtype=np.float32)
+    for t in range(st):
+        for c in range(nch):
+            blk = idx[t, c]                                     
+            lo = (blk & 0x0F).T                                 
+            hi = (blk >> 4).T                                   
+            for sub, rows in ((lo, range(c * 64, c * 64 + 32)),
+                              (hi, range(c * 64 + 32, c * 64 + 64))):
+                for bi, n in enumerate(rows):
+                    codes = sub[bi]                             
+                    W[n] += cbT[t][:, codes].T.reshape(-1)      
+    return W[:o] * scale[:o, None]
+
+
+if __name__ == "__main__":
+    torch.manual_seed(0)
+    print(f"{'module':>22s} {'shape':>14s} {'bits/w':>7s} {'max|err|':>10s} {'rel RMSE':>10s}  verdict")
+    ok = True
+
+    for name, (i, o, g, st) in {
+        "proj 1-bit  (g8,st2)": (256, 256, 8, 2),
+        "proj 1-bit  rect     ": (256, 128, 8, 2),
+        "FFN 0.125   (g32,st1)": (256, 1024, 32, 1),
+        "FFN 0.125   down     ": (1024, 256, 32, 1),
+        "odd rows (pad test)  ": (256, 100, 8, 2),
+    }.items():
+        m = RVQ(i, o, g, st)
+        m.weight.data = torch.randn(o, i) / (i ** 0.5)
+        m.enc()                                                  
+        ref = m._q.detach().cpu().numpy()
+
+        cbT, idxp, scale = rvq_pack(m)
+        got = rvq_unpack(cbT, idxp, scale, o, i, g, st)
+
+        err = np.abs(ref - got).max()
+        rel = np.sqrt(((ref - got) ** 2).mean()) / np.sqrt((ref ** 2).mean())
+        good = err < 1e-5
+        ok &= good
+        print(f"{name:>22s} {f'{o}x{i}':>14s} {m.bits():7.3f} {err:10.2e} {rel:10.2e}  "
+              f"{'PASS' if good else 'FAIL'}")
+
+    print()
+    print("ROUNDTRIP_OK" if ok else "ROUNDTRIP_FAIL")

finetune/modeling/export_ternary.py

@@ -0,0 +1,84 @@
+import os, sys, struct
+os.environ.setdefault("SHADOW_D", "1536"); os.environ.setdefault("SHADOW_NL", "10")
+os.environ.setdefault("SHADOW_NH", "24"); os.environ.setdefault("SHADOW_NKV", "2")
+os.environ.setdefault("SHADOW_HD", "64"); os.environ.setdefault("SHADOW_FFNH", "4224")
+os.environ.setdefault("SHADOW_KV_TWO_TIER", "1")
+import numpy as np, torch
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "model"))
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import common
+from common import RVQ, requant
+from model_250m import Shadow250M
+from export_rvq import rvq_pack
+
+CK, OUT = sys.argv[1], sys.argv[2]; COMPACT = "--compact" in sys.argv      
+ck = torch.load(CK, map_location="cpu", weights_only=False)
+V, FPD = ck.get("cfg", {}).get("V", 131072), 512
+model = Shadow250M(torch.zeros(V, FPD), torch.zeros(V, FPD), V)
+model.load_state_dict(ck["model"], strict=False); model.eval(); requant(model)
+
+def tern_pack(w):
+    wt = w.detach().float()
+    sc = 1.0 / wt.abs().mean(dim=1, keepdim=True).clamp_(min=1e-5)          
+    t = (wt * sc).round().clamp(-1, 1).to(torch.int8).numpy()
+    rs = (1.0 / sc[:, 0]).numpy().astype(np.float32)                        
+    w = wt.numpy()
+    codes = (t + 1).astype(np.uint8)                              
+    o, i = w.shape; assert i % 4 == 0
+    c4 = codes.reshape(o, i // 4, 4)
+    packed = (c4[:, :, 0] | (c4[:, :, 1] << 2) | (c4[:, :, 2] << 4) | (c4[:, :, 3] << 6)).astype(np.uint8)
+    if COMPACT:                                                   
+        pad = (-i) % 5; c5 = np.concatenate([codes, np.ones((o, pad), np.uint8)], 1).reshape(o, -1, 5).astype(np.uint16)
+        packed = (c5[:, :, 0] + 3 * c5[:, :, 1] + 9 * c5[:, :, 2] + 27 * c5[:, :, 3] + 81 * c5[:, :, 4]).astype(np.uint8)
+    return packed, rs, t.astype(np.float32) * rs[:, None]
+
+class Wrap(torch.nn.Module):                                      
+    def __init__(s, m):
+        super().__init__(); s.emb = m.inp; s.b = m.b; s.step = m.struct; s.nf = m.nf; s.head = m.head; s.tb = m.tied_bias
+wrap = Wrap(model)
+recs = []; ntern = nrvq = 0; tern_bytes = rvq_bytes = dense_bytes = 0
+rvq_ids = set()
+for name, mod in wrap.named_modules():
+    if isinstance(mod, RVQ):
+        rvq_ids.add(id(mod))
+        if mod.g == 32:                                           
+            packed, rs, deq = tern_pack(mod.weight)
+            recs.append((name, 4 if COMPACT else 3, (mod.o, mod.i, packed, rs))); ntern += 1; tern_bytes += packed.nbytes + rs.nbytes
+        else:
+            cbT, idx, scale = rvq_pack(mod)
+            recs.append((name, 1, (mod.o, mod.i, mod.g, mod.st, cbT, idx, scale))); nrvq += 1
+            rvq_bytes += cbT.nbytes + idx.nbytes + scale.nbytes
+for name, p in wrap.named_parameters():
+    owner = name.rsplit(".", 1)[0]
+    mod = dict(wrap.named_modules()).get(owner)
+    if mod is not None and id(mod) in rvq_ids: continue           
+    a = p.detach().float().numpy()
+    if COMPACT and a.ndim >= 1 and a.size >= 4096: recs.append((name, 5, a.astype(np.float16))); dense_bytes += a.size * 2
+    else: recs.append((name, 0, a)); dense_bytes += a.nbytes
+for name, b in wrap.named_buffers():
+    if any(x in name for x in ("cent", "cb", "initialized", "sign", "mu", "ctv", "low", "high", "updates", "inv")): continue
+    a = b.detach().float().numpy(); recs.append((name, 0, a)); dense_bytes += a.nbytes
+
+with open(OUT, "wb") as f:
+    f.write(b"SHDW"); f.write(struct.pack("<II", 1, len(recs)))
+    for name, kind, pay in recs:
+        nb = name.encode(); f.write(struct.pack("<I", len(nb))); f.write(nb); f.write(struct.pack("<I", kind))
+        if kind == 0:
+            a = np.ascontiguousarray(pay, np.float32); f.write(struct.pack("<I", a.ndim)); f.write(struct.pack("<" + "I" * a.ndim, *a.shape)); f.write(a.tobytes())
+        elif kind == 1:
+            o, i, g, st, cbT, idx, scale = pay
+            f.write(struct.pack("<IIII", o, i, g, st)); f.write(np.ascontiguousarray(cbT, np.float32).tobytes())
+            f.write(np.ascontiguousarray(idx, np.uint8).tobytes()); f.write(np.ascontiguousarray(scale, np.float32).tobytes())
+        elif kind == 5:
+            a = np.ascontiguousarray(pay, np.float16); f.write(struct.pack("<I", a.ndim)); f.write(struct.pack("<" + "I" * a.ndim, *a.shape)); f.write(a.tobytes())
+        else:
+            o, i, packed, rs = pay
+            f.write(struct.pack("<II", o, i)); f.write(np.ascontiguousarray(packed, np.uint8).tobytes()); f.write(rs.tobytes())
+size = os.path.getsize(OUT)
+print(f"loaded checkpoint ({ck.get('step','finetuned')})")
+print(f"ternary FFN modules {ntern}: {tern_bytes/1e6:.2f} MB | RVQ 1-bit modules {nrvq}: {rvq_bytes/1e6:.2f} MB | dense fp32: {dense_bytes/1e6:.2f} MB")
+print(f".shdw = {size/1e6:.2f} MB  (+ 8.39 MB fp131072 table = {(size+8388736)/1e6:.2f} MB deploy)")
+
+mod = model.b[0].up; w = mod.weight.detach().float()
+sc = 1.0 / w.abs().mean(dim=1, keepdim=True).clamp_(min=1e-5); wq = (w * sc).round().clamp(-1, 1) / sc
+_, _, deq = tern_pack(mod.weight); print("ternary round-trip max|err| vs trainer forward:", float((torch.tensor(deq) - wq).abs().max()))

finetune/modeling/model_250m.py

@@ -0,0 +1,128 @@
+import os, math, numpy as np, torch, torch.nn as nn, torch.nn.functional as F
+import common
+from common import D,NL,NH,NKV,HD,FFNH,FPD, RMS,RVQ,Block,StructStep, requant, cs
+
+
+def load_vocab(fp_path,dev):
+    fp=np.unpackbits(np.load(fp_path),axis=1)[:,:FPD]
+    cent=torch.tensor(fp.astype(np.float32)*2-1,device=dev)      
+    return cent, F.normalize(cent,dim=-1), cent.shape[0]
+
+class Shadow250M(nn.Module):
+    def __init__(s,cent,cent_n,V,use_memory=False):
+        super().__init__()
+
+
+        s.register_buffer("cent",cent,persistent=False)
+        s.register_buffer("cent_n",cent_n,persistent=False); s.V=V
+        s.inp=nn.Linear(FPD,D,bias=False)
+        s.b=nn.ModuleList([Block(i) for i in range(NL)])         
+        s.struct=StructStep()                                    
+        s.nf=RMS(D); s.head=nn.Linear(D,FPD,bias=False)          
+        s.tied_bias=nn.Parameter(torch.zeros(V))
+
+
+        s.default_use_memory=bool(use_memory)
+        inv=1.0/(10000**(torch.arange(0,HD,2).float()/HD)); s.register_buffer("inv",inv,persistent=False)
+    def trunk(s,idx):
+        cos,sin=cs(idx.shape[1],idx.device); x=s.inp(s.cent[idx])
+        for blk in s.b: x=blk(x,cos,sin)
+        x,conf=s.struct(x)                                       
+        return s.nf(x), conf
+    def logits(s,x): return (s.head(x).float()@s.cent_n.T)+s.tied_bias      
+    def forward(s,idx,ys=None):
+        x,conf=s.trunk(idx)
+        if ys is None: return s.logits(x)
+
+        ph=s.head(x).float().reshape(-1,FPD); yf=ys.reshape(-1); Nn=ph.shape[0]; CH=2048
+        return sum(F.cross_entropy(ph[i:i+CH]@s.cent_n.T+s.tied_bias,yf[i:i+CH],reduction="sum") for i in range(0,Nn,CH))/Nn
+
+    @torch.no_grad()
+    def prefill_cached(s,idx,max_ctx=2048,exact_shiftmax=True,
+                       use_memory=None,memory_chunk=128,memory_capacity=256,
+                       retrieval_trace=False,stream_block=256):
+        if use_memory is None: use_memory=s.default_use_memory
+        if idx.shape[1]>max_ctx:
+            if not use_memory:
+                idx=idx[:,-max_ctx:]
+            else:
+                logits,state=s.prefill_cached(
+                    idx[:,:max_ctx],max_ctx=max_ctx,
+                     exact_shiftmax=exact_shiftmax,use_memory=use_memory,
+                     memory_chunk=memory_chunk,memory_capacity=memory_capacity,
+                     retrieval_trace=retrieval_trace,stream_block=stream_block)
+                cursor=max_ctx
+                final=idx.shape[1]-1
+                while cursor<final:
+                    stop=min(final,cursor+int(stream_block))
+                    state=s.ingest_cached(idx[:,cursor:stop],state)
+                    cursor=stop
+                return s.decode_cached(idx[:,final:final+1],state)
+        cos,sin=cs(idx.shape[1],idx.device); x=s.inp(s.cent[idx]); layers=[]
+        for blk in s.b:
+            x,cache=blk.prefill_cached(x,cos,sin,max_ctx=max_ctx,
+                                       exact=exact_shiftmax,
+                                       archive_cold=bool(use_memory),
+                                       cold_chunk=memory_chunk,
+                                       retrieval_trace=retrieval_trace); layers.append(cache)
+        trunk=x[:,-max_ctx:]; y,_=s.struct(x); ph=s.head(s.nf(y[:,-1])).float()
+        return ph@s.cent_n.T+s.tied_bias,{"layers":layers,"trunk":trunk,
+            "position":idx.shape[1],"max_ctx":max_ctx,
+            "exact_shiftmax":exact_shiftmax,
+            "retrieval_trace":bool(retrieval_trace),
+            "stream_block":int(stream_block),
+            "memory":common.make_memory_state(
+                idx.shape[0],D,idx.device,chunk_size=memory_chunk,
+                capacity=memory_capacity) if use_memory else None}
+
+    @torch.no_grad()
+    def ingest_cached(s,idx_block,state):
+        if idx_block.shape[1]==0:
+            return state
+        x=s.inp(s.cent[idx_block]); pos=int(state["position"])
+        for i,blk in enumerate(s.b):
+            x,state["layers"][i]=blk.decode_cached(
+                x,pos,state["layers"][i],max_ctx=state["max_ctx"],
+                exact=state["exact_shiftmax"],retrieve_cold=False)
+        combined=torch.cat((state["trunk"],x),1)
+        overflow=max(0,combined.shape[1]-state["max_ctx"])
+        common.memory_absorb_evicted(
+            state.get("memory"),combined[:,:overflow],s.head,s.nf)
+        state["trunk"]=combined[:,overflow:]
+        s.struct.decode_cached(x,state["trunk"])
+        state["position"]=pos+idx_block.shape[1]
+        return state
+
+    @torch.no_grad()
+    def decode_cached(s,idx_new,state):
+        if state.get("memory") is not None and idx_new.shape[1]!=1:
+            logits=None
+            for j in range(idx_new.shape[1]):
+                logits,state=s.decode_cached(idx_new[:,j:j+1],state)
+            return logits,state
+        x=s.inp(s.cent[idx_new]); pos=int(state["position"])
+        for i,blk in enumerate(s.b):
+            x,state["layers"][i]=blk.decode_cached(x,pos,state["layers"][i],
+                                                   max_ctx=state["max_ctx"],
+                                                   exact=state["exact_shiftmax"])
+        combined=torch.cat((state["trunk"],x),1)
+        overflow=max(0,combined.shape[1]-state["max_ctx"])
+        common.memory_absorb_evicted(
+            state.get("memory"),combined[:,:overflow],s.head,s.nf)
+        state["trunk"]=combined[:,overflow:]
+        context=common.memory_append_recall(
+            state.get("memory"),x[:,-1],state["trunk"],s.head,s.nf)
+        y,_=s.struct.decode_cached(x,context); state["position"]=pos+idx_new.shape[1]
+        ph=s.head(s.nf(y[:,-1])).float(); return ph@s.cent_n.T+s.tied_bias,state
+
+if __name__=="__main__":
+    dev="cuda" if torch.cuda.is_available() else "cpu"
+    FP=os.path.join(os.path.dirname(__file__),"dolma3_fp_512b.npy")
+    if not os.path.exists(FP): FP=os.path.join(os.path.dirname(__file__),"..","..","shadow-o","phase0_tokenizer","dolma3_fp_512b.npy")
+    cent,cent_n,V=load_vocab(FP,dev); m=Shadow250M(cent,cent_n,V).to(dev)
+    body=sum(p.numel() for n,p in m.named_parameters() if "cent" not in n and "tied_bias" not in n)
+    print(f"SHADOW-O 250M | body {body/1e6:.1f}M params | vocab {V} (0 params) | config d{D} L{NL} GQA{NH}/{NKV}")
+    b=m.b[0]; print(f"  proj {b.q.bits():.3f} bit/w | FFN {b.up.bits():.3f} bit/w | struct step: 1 | Q4/loop: none")
+    x=torch.randint(0,V,(1,64),device=dev); requant(m)
+    with torch.autocast("cuda",dtype=torch.bfloat16) if dev=="cuda" else torch.no_grad():
+        lg=m(x); print(f"  forward OK: logits {tuple(lg.shape)} finite={bool(torch.isfinite(lg).all())}")

finetune/modeling/paged_kv.py

@@ -0,0 +1,319 @@
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Iterable
+
+import torch
+
+
+_POPCOUNT = tuple(int(i).bit_count() for i in range(256))
+
+
+@dataclass(frozen=True)
+class ArchiveAccounting:
+    logical_tokens: int
+    page_count: int
+    payload_bytes: int
+    allocated_payload_bytes: int
+    position_index_bytes: int
+    bytes_copied_on_append: int
+
+    @property
+    def allocated_total_bytes(self) -> int:
+        return self.allocated_payload_bytes + self.position_index_bytes
+
+
+class PagedKVArchive:
+
+    format = "paged_random_walsh_ctv_1bit_v1"
+    state_version = 1
+
+    def __init__(
+        self,
+        batch: int,
+        heads: int,
+        packed_width: int,
+        *,
+        page_size: int = 256,
+        device=None,
+    ):
+        if min(batch, heads, packed_width, page_size) < 1:
+            raise ValueError("batch, heads, packed_width, and page_size must be positive")
+        self.batch = int(batch)
+        self.heads = int(heads)
+        self.packed_width = int(packed_width)
+        self.page_size = int(page_size)
+        self.device = torch.device(device or "cpu")
+        self.k_pages: list[torch.Tensor] = []
+        self.v_pages: list[torch.Tensor] = []
+        self.position_pages: list[torch.Tensor] = []
+        self.length = 0
+        self.bytes_copied_on_append = 0
+
+    @property
+    def shape(self) -> tuple[int, int, int, int]:
+        return self.batch, self.heads, self.length, self.packed_width
+
+    def __len__(self) -> int:
+        return self.length
+
+    def _allocate_page(self):
+        shape = (self.batch, self.heads, self.page_size, self.packed_width)
+        self.k_pages.append(torch.empty(shape, dtype=torch.uint8, device=self.device))
+        self.v_pages.append(torch.empty(shape, dtype=torch.uint8, device=self.device))
+        self.position_pages.append(
+            torch.empty(self.page_size, dtype=torch.int64, device=self.device)
+        )
+
+    def append(
+        self,
+        packed_k: torch.Tensor,
+        packed_v: torch.Tensor,
+        positions: torch.Tensor | Iterable[int] | None = None,
+    ) -> None:
+        expected_prefix = (self.batch, self.heads)
+        if packed_k.dtype != torch.uint8 or packed_v.dtype != torch.uint8:
+            raise TypeError("paged one-bit K/V payloads must be torch.uint8")
+        if packed_k.shape != packed_v.shape:
+            raise ValueError("packed K and V shapes differ")
+        if (
+            packed_k.ndim != 4
+            or tuple(packed_k.shape[:2]) != expected_prefix
+            or packed_k.shape[-1] != self.packed_width
+        ):
+            raise ValueError(
+                "expected packed K/V shape "
+                f"({self.batch},{self.heads},T,{self.packed_width}), "
+                f"got {tuple(packed_k.shape)}"
+            )
+        tokens = int(packed_k.shape[2])
+        if tokens == 0:
+            return
+        if packed_k.device != self.device or packed_v.device != self.device:
+            raise ValueError(
+                f"archive is on {self.device}, payload is on "
+                f"{packed_k.device}/{packed_v.device}"
+            )
+        if positions is None:
+            pos = torch.arange(
+                self.length, self.length + tokens, dtype=torch.int64, device=self.device
+            )
+        else:
+            pos = torch.as_tensor(positions, dtype=torch.int64, device=self.device)
+            if pos.ndim != 1 or pos.numel() != tokens:
+                raise ValueError(f"positions must contain exactly {tokens} entries")
+
+        source = 0
+        while source < tokens:
+            page_index = self.length // self.page_size
+            page_offset = self.length % self.page_size
+            if page_index == len(self.k_pages):
+                self._allocate_page()
+            take = min(tokens - source, self.page_size - page_offset)
+            dst = slice(page_offset, page_offset + take)
+            src = slice(source, source + take)
+            self.k_pages[page_index][:, :, dst].copy_(packed_k[:, :, src])
+            self.v_pages[page_index][:, :, dst].copy_(packed_v[:, :, src])
+            self.position_pages[page_index][dst].copy_(pos[src])
+            copied = 2 * self.batch * self.heads * take * self.packed_width
+            self.bytes_copied_on_append += copied
+            self.length += take
+            source += take
+
+    def _validate_indices(self, indices: torch.Tensor) -> torch.Tensor:
+        indices = torch.as_tensor(indices, dtype=torch.long, device=self.device)
+        if indices.ndim != 3 or tuple(indices.shape[:2]) != (self.batch, self.heads):
+            raise ValueError(
+                f"indices must have shape ({self.batch},{self.heads},K), "
+                f"got {tuple(indices.shape)}"
+            )
+        if indices.numel() and (
+            int(indices.min().item()) < 0 or int(indices.max().item()) >= self.length
+        ):
+            raise IndexError(f"archive index outside [0,{self.length})")
+        return indices
+
+    def gather(self, indices: torch.Tensor, kind: str) -> torch.Tensor:
+        indices = self._validate_indices(indices)
+        if kind not in {"k", "v"}:
+            raise ValueError("kind must be 'k' or 'v'")
+        pages = self.k_pages if kind == "k" else self.v_pages
+        out = torch.empty(
+            (*indices.shape, self.packed_width),
+            dtype=torch.uint8,
+            device=self.device,
+        )
+        page_ids = torch.div(indices, self.page_size, rounding_mode="floor")
+        offsets = indices.remainder(self.page_size)
+        for page_id in torch.unique(page_ids).tolist():
+            mask = page_ids == int(page_id)
+            b, h, slot = mask.nonzero(as_tuple=True)
+            out[b, h, slot] = pages[int(page_id)][b, h, offsets[b, h, slot]]
+        return out
+
+    def gather_positions(self, indices: torch.Tensor) -> torch.Tensor:
+        indices = self._validate_indices(indices)
+        out = torch.empty(indices.shape, dtype=torch.int64, device=self.device)
+        page_ids = torch.div(indices, self.page_size, rounding_mode="floor")
+        offsets = indices.remainder(self.page_size)
+        for page_id in torch.unique(page_ids).tolist():
+            mask = page_ids == int(page_id)
+            b, h, slot = mask.nonzero(as_tuple=True)
+            out[b, h, slot] = self.position_pages[int(page_id)][offsets[b, h, slot]]
+        return out
+
+    @staticmethod
+    def _popcount_bytes(value: torch.Tensor) -> torch.Tensor:
+        lut = torch.tensor(_POPCOUNT, dtype=torch.int16, device=value.device)
+        return lut[value.long()].sum(-1, dtype=torch.int32)
+
+    def exact_hamming_topk(
+        self, query: torch.Tensor, k: int
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.length == 0:
+            raise ValueError("cannot search an empty archive")
+        if (
+            query.dtype != torch.uint8
+            or query.ndim != 3
+            or tuple(query.shape) != (self.batch, self.heads, self.packed_width)
+        ):
+            raise ValueError(
+                "query must be uint8 with shape "
+                f"({self.batch},{self.heads},{self.packed_width})"
+            )
+        take = min(max(1, int(k)), self.length)
+        best_indices = torch.empty(
+            (self.batch, self.heads, 0), dtype=torch.long, device=self.device
+        )
+        best_distances = torch.empty(
+            (self.batch, self.heads, 0), dtype=torch.int32, device=self.device
+        )
+        absolute = 0
+        tie_base = self.length + 1
+        for page_id, page in enumerate(self.k_pages):
+            used = min(self.page_size, self.length - absolute)
+            if used <= 0:
+                break
+            distance = self._popcount_bytes(page[:, :, :used] ^ query[:, :, None])
+            indices = torch.arange(
+                absolute, absolute + used, dtype=torch.long, device=self.device
+            ).view(1, 1, used).expand(self.batch, self.heads, used)
+            candidate_indices = torch.cat((best_indices, indices), dim=-1)
+            candidate_distances = torch.cat((best_distances, distance), dim=-1)
+            rank_key = candidate_distances.to(torch.int64) * tie_base + candidate_indices
+            selected = torch.topk(
+                rank_key, min(take, rank_key.shape[-1]), largest=False, sorted=True
+            ).indices
+            best_indices = torch.gather(candidate_indices, -1, selected)
+            best_distances = torch.gather(candidate_distances, -1, selected)
+            absolute += used
+        return best_indices, best_distances
+
+    def materialize(self, kind: str) -> torch.Tensor:
+        if kind not in {"k", "v"}:
+            raise ValueError("kind must be 'k' or 'v'")
+        pages = self.k_pages if kind == "k" else self.v_pages
+        if not pages:
+            return torch.empty(self.shape, dtype=torch.uint8, device=self.device)
+        parts = []
+        remaining = self.length
+        for page in pages:
+            used = min(self.page_size, remaining)
+            parts.append(page[:, :, :used])
+            remaining -= used
+        return torch.cat(parts, dim=2)
+
+    def accounting(self) -> ArchiveAccounting:
+        payload = 2 * self.batch * self.heads * self.length * self.packed_width
+        allocated = (
+            2
+            * self.batch
+            * self.heads
+            * len(self.k_pages)
+            * self.page_size
+            * self.packed_width
+        )
+        positions = len(self.position_pages) * self.page_size * 8
+        return ArchiveAccounting(
+            logical_tokens=self.length,
+            page_count=len(self.k_pages),
+            payload_bytes=payload,
+            allocated_payload_bytes=allocated,
+            position_index_bytes=positions,
+            bytes_copied_on_append=self.bytes_copied_on_append,
+        )
+
+    def state(self) -> dict:
+        return {
+            "format": self.format,
+            "version": self.state_version,
+            "batch": self.batch,
+            "heads": self.heads,
+            "packed_width": self.packed_width,
+            "page_size": self.page_size,
+            "length": self.length,
+            "bytes_copied_on_append": self.bytes_copied_on_append,
+            "k_pages": [page.clone() for page in self.k_pages],
+            "v_pages": [page.clone() for page in self.v_pages],
+            "position_pages": [page.clone() for page in self.position_pages],
+        }
+
+    @classmethod
+    def from_state(cls, state: dict, device=None) -> "PagedKVArchive":
+        if state.get("format") != cls.format or int(state.get("version", -1)) != 1:
+            raise ValueError("unknown paged KV archive state")
+        inferred = state["k_pages"][0].device if state["k_pages"] else "cpu"
+        target = torch.device(device or inferred)
+        archive = cls(
+            state["batch"],
+            state["heads"],
+            state["packed_width"],
+            page_size=state["page_size"],
+            device=target,
+        )
+        archive.k_pages = [page.to(target) for page in state["k_pages"]]
+        archive.v_pages = [page.to(target) for page in state["v_pages"]]
+        archive.position_pages = [page.to(target) for page in state["position_pages"]]
+        archive.length = int(state["length"])
+        archive.bytes_copied_on_append = int(state["bytes_copied_on_append"])
+        return archive
+
+
+class PagedKVView:
+
+    def __init__(self, archive: PagedKVArchive, kind: str):
+        self.archive = archive
+        self.kind = kind
+
+    @property
+    def shape(self):
+        return self.archive.shape
+
+    @property
+    def dtype(self):
+        return torch.uint8
+
+    @property
+    def device(self):
+        return self.archive.device
+
+    def materialize(self):
+        return self.archive.materialize(self.kind)
+
+
+class ExactChunkCountView:
+
+    def __init__(self, archive: PagedKVArchive, chunk_size: int):
+        self.archive = archive
+        self.chunk_size = int(chunk_size)
+
+    @property
+    def shape(self):
+        complete = self.archive.length // self.chunk_size
+        return (
+            self.archive.batch,
+            self.archive.heads,
+            complete,
+            self.archive.packed_width,
+        )

finetune/modeling/repack_shdw.py

@@ -0,0 +1,28 @@
+import sys, struct, numpy as np
+src, dst = sys.argv[1], sys.argv[2]; FP16 = "--fp16" in sys.argv     
+f = open(src, "rb"); assert f.read(4) == b"SHDW"; ver, n = struct.unpack("<II", f.read(8))
+out = open(dst, "wb"); out.write(b"SHDW"); out.write(struct.pack("<II", 1, n))
+tern_in = tern_out = dense_in = dense_out = 0
+for _ in range(n):
+    nl, = struct.unpack("<I", f.read(4)); name = f.read(nl); kind, = struct.unpack("<I", f.read(4))
+    out.write(struct.pack("<I", nl)); out.write(name)
+    if kind == 0:
+        nd, = struct.unpack("<I", f.read(4)); dims = struct.unpack("<" + "I" * nd, f.read(4 * nd)); c = int(np.prod(dims))
+        a = np.frombuffer(f.read(4 * c), np.float32); dense_in += 4 * c
+        if FP16 and c >= 4096: out.write(struct.pack("<I", 5)); out.write(struct.pack("<I", nd)); out.write(struct.pack("<" + "I" * nd, *dims)); out.write(a.astype(np.float16).tobytes()); dense_out += 2 * c
+        else: out.write(struct.pack("<I", 0)); out.write(struct.pack("<I", nd)); out.write(struct.pack("<" + "I" * nd, *dims)); out.write(a.tobytes()); dense_out += 4 * c
+    elif kind == 1:
+        o, i, g, st = struct.unpack("<IIII", f.read(16)); G = i // g; Npad = (o + 63) & ~63; nch = Npad // 64
+        cb = f.read(st * g * 16 * 4); idx = f.read(st * nch * G * 32); sc = f.read(Npad * 4)
+        out.write(struct.pack("<I", 1)); out.write(struct.pack("<IIII", o, i, g, st)); out.write(cb); out.write(idx); out.write(sc)
+    elif kind == 3:
+        o, i = struct.unpack("<II", f.read(8)); packed = np.frombuffer(f.read(o * i // 4), np.uint8).reshape(o, i // 4); rs = f.read(o * 4)
+        tern_in += packed.nbytes
+        c4 = np.stack([(packed >> (2 * j)) & 3 for j in range(4)], -1).reshape(o, i)          
+        pad = (-i) % 5; c5 = np.concatenate([c4, np.ones((o, pad), np.uint8)], 1).reshape(o, -1, 5).astype(np.uint16)
+        p5 = (c5[:, :, 0] + 3 * c5[:, :, 1] + 9 * c5[:, :, 2] + 27 * c5[:, :, 3] + 81 * c5[:, :, 4]).astype(np.uint8)
+        out.write(struct.pack("<I", 4)); out.write(struct.pack("<II", o, i)); out.write(p5.tobytes()); out.write(rs); tern_out += p5.nbytes
+    else: raise SystemExit(f"unknown kind {kind}")
+out.close()
+import os
+print(f"ternary {tern_in/1e6:.1f} -> {tern_out/1e6:.1f} MB | dense {dense_in/1e6:.1f} -> {dense_out/1e6:.1f} MB | file {os.path.getsize(src)/1e6:.1f} -> {os.path.getsize(dst)/1e6:.1f} MB")

finetune/pirate.shdw

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3308a5706c4a56f404e117a147c1c5c7b469df2dd890336f88df9bce294388bf
+size 52342802

finetune/shadow250m_instruct.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c23bde9d8669f38842328a6e7fc940353342750aa264056530ea6d0da11151c0
+size 538876067

shadow_chat.py

@@ -0,0 +1,19 @@
+"""Start chatting with SHADOW. Picks the right binary for your system automatically.
+    python shadow_chat.py
+"""
+import os, sys, platform, subprocess, pathlib
+HERE = pathlib.Path(__file__).resolve().parent
+osname = platform.system()
+if osname == "Windows": k = HERE / "deployment" / "bin" / "windows" / "shadow.exe"
+elif osname == "Linux": k = HERE / "deployment" / "bin" / "linux" / "shadow"
+else: sys.exit("macOS build available on request: saikiranbathula1@gmail.com")
+if osname != "Windows": os.chmod(k, 0o755)
+sys.path.insert(0, str(HERE))
+from shadow_runtime import Engine
+eng = Engine(str(HERE / "deployment" / "shadow250m_instruct.shdw"), str(HERE / "deployment" / "fp131072.npy"), kernel=str(k))
+print("SHADOW 250M. Type your message, 'quit' to stop.")
+while True:
+    try: q = input("you> ").strip()
+    except EOFError: break
+    if q in ("quit", "exit"): break
+    if q: print("shadow>", eng.chat(q))

shadow_runtime/__init__.py

@@ -0,0 +1,51 @@
+"""SHADOW 250M runtime: chat through the bundled CPU kernel, and question answering over an on-disk archive.
+
+    from shadow_runtime import Engine
+    eng = Engine("shadow250m_instruct.shdw", "fp131072.npy", archive="path/to/archive_dir")
+    print(eng.answer("your question"))
+
+An archive directory holds tokens.u32 (uint32 token stream). The lexical index is built once on first use
+and cached next to it. Chat without an archive:
+
+    eng = Engine("shadow250m_instruct.shdw", "fp131072.npy")
+    print(eng.chat("Explain photosynthesis in two sentences."))
+
+CLI:  python -m shadow_runtime --model shadow250m_instruct.shdw --table fp131072.npy \
+          --archive path/to/archive --ask "your question"
+"""
+import os, sys, pathlib, subprocess
+import numpy as np
+HERE = pathlib.Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+from retriever import Inverted, load_archive, enc, _dec
+from answer_engine import Engine as _Extractor, ABSTAIN
+from prompt import EOT
+
+BOS, SOT, EOS = 2, 8, 1
+
+class Engine:
+    def __init__(s, model, table, archive=None, kernel=None, threads=None):
+        s.model = str(pathlib.Path(model).resolve()); s.table = str(pathlib.Path(table).resolve())
+        s.kernel = str(pathlib.Path(kernel).resolve()) if kernel else str((HERE.parent / "deployment" / "bin" / ("windows/shadow.exe" if os.name == "nt" else "linux/shadow")).resolve())
+        s.env = dict(os.environ)
+        if threads: s.env["SHADOW_THREADS"] = str(threads)
+        s.ext = None
+        if archive:
+            tok, meta, _bank = load_archive(str(archive))
+            inv = Inverted(tok)
+            s.ext = _Extractor(tok, inv, model_ask=s.chat)
+    def _gen(s, ids, n=140, extra=()):
+        r = subprocess.run([s.kernel, s.model, s.table, " ".join(map(str, ids)), str(n), *extra],
+                           capture_output=True, text=True, env=s.env)
+        out = [int(x) for x in r.stdout.split()]
+        for stop in (EOT, EOS):
+            if stop in out: out = out[:out.index(stop)]
+        return _dec(out).strip()
+    def chat(s, message, n=160, greedy=True, temp=0.25, topk=30, rep=1.15, seed=0):
+        ids = [BOS, SOT] + enc("user\n") + enc(message) + [EOT] + enc("\n") + [SOT] + enc("model\n")
+        extra = () if greedy else ("--temp", str(temp), "--topk", str(topk), "--rep", str(rep), "--seed", str(seed))
+        return s._gen(ids, n, extra)
+    def answer(s, question):
+        if s.ext is None: return s.chat(question)
+        a, how = s.ext.answer(question)
+        return a

shadow_runtime/__main__.py

@@ -0,0 +1,14 @@
+import argparse, sys
+from . import Engine
+ap = argparse.ArgumentParser()
+ap.add_argument("--model", required=True); ap.add_argument("--table", required=True)
+ap.add_argument("--archive"); ap.add_argument("--ask"); ap.add_argument("--chat", action="store_true")
+a = ap.parse_args()
+eng = Engine(a.model, a.table, archive=a.archive)
+if a.ask: print(eng.answer(a.ask))
+elif a.chat:
+    while True:
+        try: q = input("you> ")
+        except EOFError: break
+        if not q.strip(): continue
+        print("shadow>", eng.chat(q))

shadow_runtime/answer_engine.py

@@ -0,0 +1,160 @@
+"""Deterministic-extraction answer engine over the SHADOW archive pipeline.
+
+Pipeline (each step deterministic or measured ~100%):
+  1. retrieve 16 chunks (lexical inverted index, 2 hops)                     -- measured 0.98-1.00 recall
+  2. locate the question's KEY inside the retrieved text by string match     -- deterministic
+     (key = identifier-shaped tokens in the question; fallback = rarest word n-gram)
+  3. no match -> "NOT IN CONTEXT"                                            -- deterministic abstain
+  4. extract the VALUE from the matched sentence: the value is the maximal
+     content span of the sentence that does NOT appear in the question,
+     type-filtered by the question word (how many->number, who->name, ...)   -- deterministic
+  5. recency: several matches -> highest archive position wins               -- deterministic
+     2-hop: matched sentence names another identifier and holds no value ->
+     re-query with that identifier and extract there                         -- deterministic + retrieval
+     count ("how many ... listed"): count distinct matches                   -- deterministic
+  6. anything without an identifier-style key (natural QA) -> the neural model answers as before.
+
+The network still does chat and natural QA; archive fact lookup is a verified path. Disclosed as
+"hybrid (neural + deterministic extraction)" wherever results are published.
+"""
+import re, sys, pathlib
+import numpy as np
+HERE = pathlib.Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+from retriever import enc, _dec, BLK, stop_ids
+
+ABSTAIN = "NOT IN CONTEXT"
+ID_RE = re.compile(r"\b((?=[A-Za-z0-9_-]*\d)[A-Za-z][A-Za-z0-9]*(?:[-_][A-Za-z0-9]+)+|[A-Z][a-z]+[A-Z][a-z]+[A-Za-z]*)\b")
+SENT_SPLIT = re.compile(r"(?<=[.!?])\s+|\n+")
+STOP_WORDS = set("""a an the is are was were be been being of in on at to from by for with under over into as and or if
+that this these those it its his her their there here does do did has have had holds hold most recent statement if
+several use please what which who whom whose when where how many much give state say tell answer question retrieved
+passages above only reply exactly not them im""".split())
+
+def question_keys(q):
+    ids = ID_RE.findall(q)
+    seen = set(); out = []
+    for x in ids:
+        if x not in seen: seen.add(x); out.append(x)
+    return out
+
+def sentences_with(text, key):
+    """(sentence, char_start) for every sentence in text containing key (word-boundary)."""
+    out = []
+    kre = re.compile(re.escape(key) + r"(?![A-Za-z0-9])")
+    # manual split with offsets
+    bounds = [0] + [mm.end() for mm in re.finditer(r"[.!?\n]", text)] + [len(text)]
+    for i in range(len(bounds) - 1):
+        s = text[bounds[i]:bounds[i + 1] + 1]
+        if kre.search(s): out.append((s.strip(), bounds[i]))
+    return out
+
+def q_type(q):
+    ql = q.lower()
+    if "how many" in ql or "at how many" in ql: return "number"
+    if ql.startswith("who ") or " who " in ql or "assigned to" in ql or "sealed by" in ql or "signed by" in ql or "came from" in ql or "maintain" in ql or "bid" in ql: return "name"
+    if ql.startswith("when ") or "expire" in ql or "accessed" in ql or "date" in ql: return "date"
+    return "any"
+
+def candidate_value(sent, q, want):
+    """maximal non-question content spans of the sentence; choose by type then by position (later wins)."""
+    qw0 = {w.lower().strip(".,?!'\"()") for w in q.split()}
+    qwords = set(qw0)
+    for w in qw0:                                   # morphological variants: weigh/weighs, expire/expires
+        qwords.add(w + "s"); qwords.add(w + "es"); qwords.add(w + "d"); qwords.add(w + "ed")
+        if w.endswith("s"): qwords.add(w[:-1])
+    toks = re.findall(r"0x[0-9a-f]+|[A-Za-z0-9]+(?:-[A-Za-z0-9]+)+|\d+\.\d+|[A-Za-z0-9][A-Za-z0-9']*|'[^']*'", sent)
+    spans = []; cur = []
+    for w in toks:
+        wl = w.lower().strip("'")
+        skip = (wl in qwords) or (wl in STOP_WORDS)
+        if skip:
+            if cur: spans.append(" ".join(cur)); cur = []
+        else: cur.append(w.strip("'"))
+    if cur: spans.append(" ".join(cur))
+    if not spans: return None
+    def is_num(s): return bool(re.fullmatch(r"\d+(?:\.\d+)?", s.replace(" ", "")))
+    def is_date(s): return bool(re.search(r"(January|February|March|April|May|June|July|August|September|October|November|December)", s))
+    def is_name(s): return bool(re.fullmatch(r"[A-Z][a-z]+ [A-Z][a-z]+", s))
+    if want == "number":
+        cs = [s for s in spans if is_num(s)]
+        if cs: return cs[-1]
+    if want == "date":
+        cs = [s for s in spans if is_date(s)]
+        if cs: return cs[-1]
+    if want == "name":
+        cs = [s for s in spans if is_name(s)]
+        if cs: return cs[-1]
+    return spans[-1]
+
+class Engine:
+    def __init__(s, tok_or_ov, inv, model_ask=None, hops=2, k=16):
+        s.ov = tok_or_ov; s.inv = inv; s.model_ask = model_ask; s.hops = hops; s.k = k
+        s.is_overlay = hasattr(tok_or_ov, "block")
+    def _chunk_text(s, blocks):
+        """decode retrieved (b,b+1) chunks; return [(pos_block, text)] ordered by pos."""
+        got = sorted({int(b) for b in blocks})
+        out = []
+        for b in got:
+            if s.is_overlay: t = _dec(s.ov.block(b) + s.ov.block(min(b + 1, s.ov.nb - 1)))
+            else: t = _dec(np.asarray(s.ov[b * BLK:(b + 2) * BLK], np.int64))
+            out.append((b, t))
+        return out
+    def _retrieve(s, q):
+        from retriever import Inverted
+        qi = enc(q)
+        idx = s.inv.topk_hops(qi, s.k, rounds=s.hops)[0] if not s.is_overlay else None
+        if s.is_overlay:
+            import bench_longctx as B
+            idx = B.ov_topk_hops(s.inv, s.ov, qi, s.k)
+        return idx
+    def answer(s, q):
+        keys = question_keys(q)
+        if not keys:
+            return (s.model_ask(q) if s.model_ask else ABSTAIN), "neural"
+        idx = s._retrieve(q); chunks = s._chunk_text(idx)
+        want = q_type(q)
+        key = keys[0]
+        # count task: count distinct "Member i of KEY" indices; members are numbered 1..n, so the answer is
+        # max(index) (robust to a missed block) cross-checked with the distinct count; wider retrieval (k=32).
+        if re.search(r"how many .* (listed|are there|in the archive)", q.lower()):
+            k_save = s.k; s.k = 32
+            idx = s._retrieve(q); chunks = s._chunk_text(idx); s.k = k_save
+            seen = set()
+            for _, t in chunks:
+                for sent, _o in sentences_with(t, key):
+                    m = re.search(r"Member (\d+) of", sent)
+                    if m: seen.add(int(m.group(1)))
+            if not seen: return ABSTAIN, "count"
+            return str(max(max(seen), len(seen))), "count"
+        matches = []
+        for pos, t in chunks:
+            for sent, off in sentences_with(t, key): matches.append((pos, off, sent))
+        if not matches: return ABSTAIN, "abstain"
+        matches.sort(key=lambda m: (m[0], m[1]))
+        # try direct extraction from the LATEST match backwards (recency). A candidate is a POINTER (alias
+        # record, e.g. "stored under reference K2") -- not a value -- only when the sentence carries an explicit
+        # reference cue; plain identifier-shaped values (serials "SN-...", CamelCase names) are legitimate answers.
+        PTR_CUE = re.compile(r"stored under|see reference|filed under|under reference", re.I)
+        for pos, off, sent in reversed(matches):
+            v = candidate_value(sent, q, want)
+            if not v: continue
+            if PTR_CUE.search(sent) and " " not in v and ID_RE.fullmatch(v): continue      # alias pointer -> 2-hop
+            if " " in v and ID_RE.search(v): continue                                       # multi-word span containing an identifier = junk -> 2-hop
+            return v, "extract"
+        # 2-hop: the matched sentence points at another identifier
+        for pos, off, sent in reversed(matches):
+            others = [x for x in ID_RE.findall(sent) if x != key and x not in q]
+            for k2 in others:
+                q2 = q.replace(key, k2)
+                idx2 = s._retrieve(q2); chunks2 = s._chunk_text(idx2)
+                m2 = []
+                for p2, t2 in chunks2:
+                    for s2, o2 in sentences_with(t2, k2): m2.append((p2, o2, s2))
+                for p2, o2, s2 in reversed(sorted(m2)):
+                    v = candidate_value(s2, q, want)
+                    if not v or v == key: continue                                          # circular: back to the original key
+                    if re.search(r"stored under|see reference|filed under|under reference", s2, re.I) and " " not in v and ID_RE.fullmatch(v): continue
+                    if " " in v and ID_RE.search(v): continue
+                    return v, "2hop"
+        return ABSTAIN, "abstain"

shadow_runtime/prompt.py

@@ -0,0 +1,61 @@
+"""Prompt assembly for retrieval QA. Retrieved chunks are framed inside the user turn:
+
+  <bos><start_of_turn>user\n  [ <unused0> pos=N \n <chunk tokens> <unused1> \n ] x K
+  \nUsing only the retrieved passages above, answer the question. If the answer is not in them, reply exactly "NOT IN CONTEXT".
+  Question: {q}<end_of_turn>\n<start_of_turn>model\n
+cold mask = True on every token inside a chunk frame (open..close), False elsewhere (hot window).
+Chunks = retrieved 64-token blocks expanded to (b, b+1) and merged when adjacent, ordered by archive position.
+"""
+import sys, json, pathlib, re, string
+import numpy as np
+HERE = pathlib.Path(__file__).resolve().parent; PKG = HERE.parent
+sys.path.insert(0, str(HERE))
+from retriever import enc, BLK
+BOS, EOS, SOT, EOT = 2, 1, 8, 9; COPEN, CCLOSE = 6, 7
+ABSTAIN = "NOT IN CONTEXT"
+INSTR = f"\nUsing only the retrieved passages above, answer the question. If the answer is not in them, reply exactly \"{ABSTAIN}\".\nQuestion: "
+_NL = None
+def nl():
+    global _NL
+    if _NL is None: _NL = enc("\n")
+    return _NL
+
+def chunks_from_blocks(tok, blocks, expand=1, max_chunks=16):
+    """retrieved block ids -> list of (start_block, token list): (b, b+expand) spans, merged when overlapping, by position."""
+    spans = sorted({(int(b), int(b) + expand + 1) for b in blocks[:max_chunks]})
+    merged = []
+    for a, e in spans:
+        if merged and a <= merged[-1][1]: merged[-1][1] = max(merged[-1][1], e)
+        else: merged.append([a, e])
+    nb = len(tok) // BLK
+    return [(a, [int(t) for t in tok[a * BLK:min(e, nb) * BLK]]) for a, e in merged if a < nb]
+
+def build_prompt(chunks, question):
+    """chunks: [(pos_block, ids)] -> (ids, cold_mask) for the user turn + model turn opener."""
+    ids = [BOS, SOT] + enc("user\n"); cold = [False] * len(ids)
+    for pos, ct in chunks:
+        fr = [COPEN] + enc(f"pos={pos}\n"); body = ct; close = [CCLOSE] + nl()
+        seg = fr + body + close; ids += seg; cold += [True] * len(seg)
+    tail = enc(INSTR + question) + [EOT] + nl() + [SOT] + enc("model\n")
+    ids += tail; cold += [False] * len(tail)
+    return ids, cold
+
+def normalize(s):
+    s = s.lower().strip(); s = "".join(ch for ch in s if ch not in set(string.punctuation)); s = re.sub(r"\b(a|an|the)\b", " ", s)
+    return " ".join(s.split())
+def final_answer(text):
+    """trace-aware: the answer is the text after the last '->' (quote-then-answer format), else the whole text."""
+    return text.rsplit("->", 1)[-1] if "->" in text else text
+def reward(answer_text, gold, task):
+    """verifiable reward: 1 if exact match after normalisation (numbers: exact digits), else 0. T5: must be the abstain marker."""
+    a = normalize(final_answer(answer_text)); g = normalize(str(gold))
+    if task == "T5" or g == normalize(ABSTAIN): return 1.0 if a == normalize(ABSTAIN) else 0.0
+    if a == normalize(ABSTAIN): return 0.0
+    if a == g: return 1.0
+    # short free-form gold (hotpot/squad): token-F1 >= 0.8 counts
+    at, gt = a.split(), g.split()
+    if not at or not gt: return 0.0
+    common_ = sum(min(at.count(w), gt.count(w)) for w in set(at));
+    if common_ == 0: return 0.0
+    p, r = common_ / len(at), common_ / len(gt); f1 = 2 * p * r / (p + r)
+    return 1.0 if f1 >= 0.8 else 0.0

shadow_runtime/retriever.py

@@ -0,0 +1,137 @@
+"""Archive retriever: exact lexical inverted index over 64-token blocks, with 2-hop follow-up.
+
+An archive is a directory with tokens.u32 (uint32 token stream) and meta.json. Blocks of 64 tokens
+are indexed by their content features (stop-filtered unigrams plus hashed bigrams and trigrams,
+idf-weighted). A query returns the top-k blocks; a second hop follows identifiers named near the
+match, so alias chains ("the record for K is stored under reference K2") resolve.
+
+Also holds the tokenizer used by the runtime (sentencepiece + id remap).
+"""
+import os, sys, json, pathlib, hashlib, re
+import numpy as np
+HERE = pathlib.Path(__file__).resolve().parent; PKG = HERE.parent
+BLK = 64
+
+def load_archive(path):
+    tok = np.memmap(os.path.join(path, "tokens.u32"), np.uint32, "r")
+    meta = json.load(open(os.path.join(path, "meta.json")))
+    bp = os.path.join(path, "bank_valid.jsonl"); bp = bp if os.path.exists(bp) else os.path.join(path, "bank.jsonl")   # validated bank if present
+    bank = [json.loads(l) for l in open(bp, encoding="utf-8")]
+    return tok, meta, bank
+
+_STOP = None
+def stop_ids(cache=None):
+    """500 most frequent ids, fixed once from a held-out fineweb shard (deploy/stop500.npy)."""
+    global _STOP
+    if _STOP is None:
+        _STOP = np.load(PKG / "tokenizer" / "stop500.npy")
+    return _STOP
+
+NBG = 1 << 20                               # hashed bigram feature space
+
+def _feats(ids):
+    """content features of a token sequence: unique unigrams (stop-filtered) + hashed bigrams of ALL adjacent
+    tokens (digits/punctuation are frequent as unigrams but their sequences -- 'Vega-713' -- are the signal)."""
+    ids = np.asarray(ids, np.int64); st = stop_ids()
+    uni = np.unique(ids); uni = uni[(~np.isin(uni, st)) & (uni > 9)]
+    if len(ids) >= 2:
+        bg = (ids[:-1] * np.int64(1000003) + ids[1:] * np.int64(7919)) % NBG
+        bg = bg[(ids[:-1] > 9) & (ids[1:] > 9)]
+        if len(ids) >= 3:                      # trigrams too: 'Cygnus','-','9' / '9','2','2' make the key number discriminative
+            tg = (ids[:-2] * np.int64(1000003) + ids[1:-1] * np.int64(7919) + ids[2:] * np.int64(104729) + np.int64(17)) % NBG
+            bg = np.concatenate([bg, tg[(ids[:-2] > 9) & (ids[1:-1] > 9) & (ids[2:] > 9)]])
+        bg = np.unique(bg) + (1 << 17)
+    else: bg = np.zeros(0, np.int64)
+    return uni, bg
+
+def _content(ids):
+    u, g = _feats(ids); return np.concatenate([u, g])
+
+
+# ---------------- exact lexical inverted index (BM25-style, unigram + bigram features) ----------------
+# The sketch retrievers above lose the rare-key signal inside 64-token blocks; an inverted index keeps it exactly.
+# Deploy cost at 100M tokens: ~60M postings (~250 MB) next to the 32 GB 1-bit KV archive; query = 3-10 posting lists.
+class Inverted:
+    def __init__(s, tok, k1=1.2, b=0.75):
+        import scipy.sparse as sp
+        nb = len(tok) // BLK; rows = []; cols = []
+        for bi in range(nb):
+            u, g = _feats(tok[bi * BLK:(bi + 1) * BLK]); f = np.concatenate([u, g]); rows.append(f); cols.append(np.full(len(f), bi, np.int64))
+        r = np.concatenate(rows); c = np.concatenate(cols)
+        s.M = sp.csr_matrix((np.ones(len(r), np.float32), (r, c)), shape=((1 << 17) + NBG, nb))   # feature x block
+        df = np.asarray(s.M.sum(1)).ravel(); s.idf = np.log(1 + (nb - df + 0.5) / (df + 0.5)).astype(np.float32)
+        s.nb = nb; s.tok = tok
+    def topk(s, ids, k):
+        f = _content(ids)
+        if len(f) == 0: return np.arange(min(k, s.nb)), np.zeros(min(k, s.nb))
+        sc = np.asarray((s.idf[f][None, :] @ s.M[f]).todense()).ravel() if False else s.idf[f] @ s.M[f]
+        sc = np.asarray(sc).ravel()
+        idx = np.argpartition(-sc, min(k, s.nb - 1))[:k]; idx = idx[np.argsort(-sc[idx], kind="stable")]
+        return idx, sc[idx]
+    def topk_hops(s, ids, k, rounds=2, k1=8, idf_min=5.0, win=14):
+        """multi-hop: round 1 = top-k1 direct hits (kept). Then EACH round-1 block gets its own follow-up query made of
+        the rare features within +-win tokens (over the stream) of where the question's tokens matched inside it --
+        'The record for K is stored under reference K2.' -> K2's n-grams -- and contributes its best new block.
+        Per-source queries (not a union) so one alias chain cannot be drowned by the others' noise."""
+        idx1, sc1 = s.topk(ids, k1); out = [int(i) for i in idx1]; qset = set(_content(ids).tolist()); qtext = _dec(np.asarray(ids, np.int64))
+        qtok = set(np.asarray(ids, np.int64).tolist()) - set(stop_ids().tolist()); src = out[:]
+        # template features = present in >=2 of the round-1 blocks ('stored under reference', 'access code for vault'):
+        # they would pull in every block of the same template; only features UNIQUE to a source block are followed.
+        cnt = {}
+        for bi in src:
+            for f in set(_content(np.asarray(s.tok[bi * BLK:(bi + 1) * BLK], np.int64)).tolist()): cnt[f] = cnt.get(f, 0) + 1
+        template = {f for f, c in cnt.items() if c >= 2}
+        for r in range(rounds - 1):
+            new = []
+            for bi in src:
+                if len(out) + len(new) >= k: break
+                blk = np.asarray(s.tok[bi * BLK:(bi + 1) * BLK], np.int64)
+                pos = [i for i, t in enumerate(blk) if t in qtok]
+                if not pos: continue
+                lo = max(0, bi * BLK + min(pos) - 8); hi = min(len(s.tok), bi * BLK + max(pos) + win + 1)
+                w = np.asarray(s.tok[lo:hi], np.int64)
+                # follow-up query: entity-like identifiers mentioned near the match (Name-123 style, the shape of archive
+                # keys) other than the question's own; fall back to the rare non-template features of the window.
+                ents = [e for e in set(re.findall(r"[A-Z][a-z]+-\d+", _dec(w))) if e not in qtext]
+                if ents: fb = np.unique(np.concatenate([_content(enc(" " + e)) for e in ents]))
+                else:
+                    fb = _content(w); fb = fb[(s.idf[fb] >= idf_min) & ~np.isin(fb, list(qset | template))]
+                if len(fb) == 0: continue
+                sc = np.asarray(s.idf[fb] @ s.M[fb]).ravel(); sc[out + new] = -1e9
+                for j in np.argsort(-sc)[:2]:                       # best new block of this chain (skip the source's own neighbourhood)
+                    j = int(j)
+                    if abs(j - bi) > 1 and j not in new: new.append(j); break
+            out += new; src = new
+            if not new: break
+        return np.array(out[:k]), np.zeros(min(k, len(out)))
+
+
+# ---------------- tokenizer (sentencepiece + id remap; expansion for out-of-subset ids computed lazily) ----
+_sp = None; _o2n = None; _n2o = None; _exp_cache = {}
+def _load_tok():
+    global _sp, _o2n, _n2o
+    if _sp is None:
+        import sentencepiece as spm
+        _sp = spm.SentencePieceProcessor(model_file=str(PKG / "tokenizer" / "tokenizer.model"))
+        _n2o = np.fromfile(PKG / "tokenizer" / "new2old.u32", np.uint32).astype(np.int64)
+        _o2n = np.full(262144, -1, np.int64); _o2n[_n2o] = np.arange(len(_n2o))
+def _expand(old_id):
+    """out-of-subset token -> its byte-fallback sequence in the 131k id space (lossless)."""
+    seq = _exp_cache.get(old_id)
+    if seq is None:
+        piece = _sp.id_to_piece(int(old_id)).replace("▁", " ")
+        raw = _sp.decode([int(old_id)]).encode("utf-8") if not piece.startswith("<0x") else bytes([int(piece[3:5], 16)])
+        seq = [int(_o2n[_sp.piece_to_id(f"<0x{b:02X}>")]) for b in raw]
+        _exp_cache[old_id] = seq
+    return seq
+def enc(text):
+    _load_tok()
+    out = []
+    for i in _sp.encode(text):
+        n = _o2n[i]
+        if n >= 0: out.append(int(n))
+        else: out.extend(_expand(i))
+    return out
+def _dec(ids):
+    _load_tok()
+    return _sp.decode(_n2o[np.asarray(ids, np.int64)].tolist())

tokenizer/new2old.u32

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d16fc82774e38698e919ebae21aa3969c52f3d846bd949dd3169452273b04635
+size 524288

tokenizer/stop500.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:983d343784453f07b3f0f4cd2c16a685aeece4b4a5f6639e700a81aeafff2e38
+size 4128

tokenizer/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1299c11d7cf632ef3b4e11937501358ada021bbdf7c47638d13c0ee982f2e79c
+size 4689074