diff --git a/BENCHMARKS.md b/BENCHMARKS.md index cce98ab..4291146 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -3,7 +3,28 @@ Baselines are informational. They are useful for spotting large regressions, but not part of Betlang's semver contract. -## 2026-05-15 +## 2026-08-22 (Shannon-bit binary CNN ensemble) + +Environment: + +- Host: `aarch64-apple-darwin`, `arm64` (Apple M4 virtual CPU) +- Rust: `rustc 1.95.0 (59807616e 2026-04-14)` + +Native command: + +```bash +cargo bench --bench detect +``` + +| Case | Bytes | Median time | Throughput | +|---|---:|---:|---:| +| short | 68 | 41.068 ms/inference | 1.6558 KB/s | +| full window | 4623 | 43.355 ms/inference | 106.63 KB/s | + +The five-model binary ensemble trades latency for accuracy against the +previous single wordseq model below. + +## 2026-05-15 (previous wordseq model) Environment: diff --git a/Cargo.lock b/Cargo.lock index 87bfbe6..fde9ca3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -34,9 +34,7 @@ name = "betlang" version = "0.1.1" dependencies = [ "criterion", - "fearless_simd", "ignore", - "rand", "rayon", ] @@ -185,23 +183,6 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" -[[package]] -name = "fearless_simd" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76258897e51fd156ee03b6246ea53f3e0eb395d0b327e9961c4fc4c8b2fa151a" - -[[package]] -name = "getrandom" -version = "0.2.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" -dependencies = [ - "cfg-if", - "libc", - "wasi", -] - [[package]] name = "globset" version = "0.4.18" @@ -313,15 +294,6 @@ version = "11.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" -[[package]] -name = "ppv-lite86" -version = "0.2.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" -dependencies = [ - "zerocopy", -] - [[package]] name = "proc-macro2" version = "1.0.106" @@ -340,36 +312,6 @@ dependencies = [ "proc-macro2", ] -[[package]] -name = "rand" -version = "0.8.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" -dependencies = [ - "libc", - "rand_chacha", - "rand_core", -] - -[[package]] -name = "rand_chacha" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" -dependencies = [ - "ppv-lite86", - "rand_core", -] - -[[package]] -name = "rand_core" -version = "0.6.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" -dependencies = [ - "getrandom", -] - [[package]] name = "rayon" version = "1.12.0" @@ -508,12 +450,6 @@ dependencies = [ "winapi-util", ] -[[package]] -name = "wasi" -version = "0.11.1+wasi-snapshot-preview1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" - [[package]] name = "winapi-util" version = "0.1.11" diff --git a/Cargo.toml b/Cargo.toml index 36c217d..12c87cb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,7 +21,7 @@ include = [ "/RELEASING.md", "/assets/confusion-by-size.png", "/assets/confusion-overall.png", - "/assets/magika/source-student-q4.bin", + "/assets/bnn/source-bnn2.bin", "/benches/detect.rs", "/examples/*.rs", "/snippets/demo.rs", @@ -29,13 +29,9 @@ include = [ "/tests/**", ] -[dependencies] -fearless_simd = "0.4" - [dev-dependencies] criterion = { version = "0.5", default-features = false, features = ["cargo_bench_support"] } ignore = "0.4.25" -rand = "0.8" rayon = "1.10" [lib] diff --git a/MODEL_CARD.md b/MODEL_CARD.md index 5bab064..78262a1 100644 --- a/MODEL_CARD.md +++ b/MODEL_CARD.md @@ -2,12 +2,15 @@ ## Artifact -- File: `assets/magika/source-student-q4.bin` -- Format: weights-only MSQ1 quantized tensor payload -- Size: 47,840 bytes -- SHA-256: `8493d2d3757572c8661141e414b1c0755aa08d4c4e5382dfbbc6b73b02d89083` -- Architecture: `wordseq-b1024-k3-m2048-tiny-3conv-hidden` -- Tokenizer: word-unit tokenizer version 3 +- File: `assets/bnn/source-bnn2.bin` +- Format: packed "BBN2" binary-CNN ensemble payload +- Size: 1,165,464 bytes +- SHA-256: `fe36df96d9b6d268118b21b764c3be6dda33596e0fafeacd2bc0b26e3fc143ca` +- Architecture: ensemble of five binary CNNs over Shannon-coded byte windows + (binary weights and activations, XNOR + popcount convolutions, integer + thresholds, OR-pooling, segmented popcount heads, i16 quantized classifier) +- Input: canonical Shannon coding of the raw 2048-byte Magika begin/end + window into a fixed 16,384-bit stream plus a codeword-boundary bitplane - Output head: 48 model labels exposed one-to-one as public `Language` variants ## Intended Use @@ -19,87 +22,80 @@ decisions, malware classification, or legal identification of file provenance. ## Training Source -The student was originally trained from Google's Magika v3.3 teacher -predictions over a source-language corpus assembled from extension-suffixed -files (the `bigorig` split, extracted from a GitHub partial-clone blob index of -roughly 6,000 popular code repositories). - -The shipped artifact is that model fine-tuned from its exported checkpoint on a -rebuilt public corpus (see `scripts/build_finetune_corpus.py`): per-language -samples from `bigcode/the-stack-smol-xl` and `bigcode/the-stack`, GitHub repo -files for labels absent from The Stack, and a small synthetic set targeting the -Markdown/YAML bare-list ambiguity reported in issue #5, including train-only -bare `- item` lists. - -The fine-tune distills the same Magika v3.3 teacher one-vs-all: per-class -binary cross-entropy against the teacher's raw head-label marginals, plus a -hard term against the teacher argmax discounted by the teacher's in-head -probability mass. Unlike the original softmax distillation, this never -renormalizes away the probability mass the teacher assigns to labels outside -the 48-label head (such as `txt`), so inputs the teacher considers ambiguous -or out-of-scope train toward uniformly low logits and keep low softmax -confidence at inference. - -The shipped run additionally self-distills from a ~2x larger intermediate -parent (`wordseq-b1536-k3-m2048-med-3conv-hidden`, 0.9492 test teacher parity) -trained on the same corpus with the same one-vs-all scheme; the parent's -per-class sigmoid marginals are cached with `scripts/cache_self_distill.py` -and added as a second BCE term. - -The model distills teacher probabilities and filesystem-extension labels. It -does not contain original source files, but its labels and soft targets are -derived from the training corpus and Magika teacher. +The models are trained on a rebuilt public corpus (see +`scripts/bnn/build_corpus.py`): per-language samples from +`bigcode/the-stack-smol-xl` and `bigcode/the-stack`, GitHub repo files for +labels absent from The Stack, and synthetic sets targeting the Markdown/YAML +bare-list ambiguity and other confusable pairs. + +Inputs are the raw Magika begin/end byte windows encoded with a canonical +Shannon code learned from training byte frequencies (add-one smoothing, +12-bit maximum code length, codes ordered by `(length, byte)`). No tokenizer +is used. Each window becomes two bitplanes: the Shannon code bits and the +codeword-boundary markers. + +Training distills Google's Magika v3.3 teacher predictions alongside +filesystem-extension labels. The binary networks are trained binary from +scratch (annealing a float network collapses); sparse learnable thresholds, +OR-pooling, and segmented popcount heads make the fully-binarized networks +trainable. Ensemble members are fine-tuned on short inputs and synthetic hard +pairs, and ensemble weights are grid-searched under fixture and edge-case +constraints (`scripts/bnn/weight_search.py`). + +The model does not contain original source files, but its labels and soft +targets are derived from the training corpus and Magika teacher. ## Evaluation -Held-out filesystem-label test split of the rebuilt corpus (34,087 files, -train/valid/test repositories are disjoint, rows where the teacher keeps at -most 10% of its probability mass on the head labels are excluded): +Held-out filesystem-label test split of the rebuilt corpus (27,465 files, +train/valid/test repositories are disjoint), evaluated with exact integer +export semantics (`scripts/bnn/eval_export2.py`): -- `test_fs_accuracy=0.942353` -- `macro_recall=0.939690` -- `test_teacher_parity=0.944055` +- `test_accuracy=0.941635` +- `test_macro_recall=0.943226` -For comparison, the pre-fine-tune artifact scores `test_fs_accuracy=0.926160`, -`macro_recall=0.929904`, and `test_teacher_parity=0.922111` on the same split. -The rebuilt split is balanced across all 48 labels (including rare classes), -so these numbers are not comparable to the `bigorig` metrics reported for -earlier artifacts. +The previous shipped wordseq artifact (`source-student-q4.bin`) scores +`test_accuracy=0.934972` and `test_macro_recall=0.934863` on the same split. -On ambiguous bare `- item` lists (valid YAML and valid Markdown), the median -top-1 probability drops from 0.92 (pre-fine-tune) to 0.51, while YAML or -Markdown remains the top prediction. +On ambiguous bare `- item` lists (valid YAML and valid Markdown), the model +reports a split YAML/Markdown distribution with top-1 probability below 0.9 +rather than picking one with certainty. -Most remaining confusion sits on genuinely ambiguous pairs where the teacher -also splits its probability on the confused files: `c`/`cpp`, +Most remaining confusion sits on genuinely ambiguous pairs: `c`/`cpp`, `javascript`/`typescript`, `markdown`/`yaml`, `ini`/`toml`, `batch`/`shell`, -and `php`/`html`. Several other cells are corpus extension-label noise (for -example `.vb` files containing SQL dumps) where the teacher agrees with the -model on 80%+ of the confused files. +and `php`/`html`. -The README confusion matrix groups the same held-out split by file-size bucket. +## Inference Semantics + +All heavy compute is binary CPU work: activations and weights are single bits +packed into 64-bit words, convolutions are XNOR + popcount +(`z' = 2 * matches - fan_in`), basis outputs combine with fixed-point integer +coefficients (`FIX=65536`), activations are integer threshold comparisons, +pooling is bitwise OR, and the classifier is an i16 integer dot product over +segmented popcounts with per-class float scale/bias. Rust inference is +bit-exact against the Python integer reference +(`cargo test --release python_parity`). ## Known Weaknesses - Very short inputs are intentionally rejected when fewer than eight non-whitespace bytes are available. -- Ambiguous snippets can put several languages close together even when a human - can infer the language from file naming context. A bare `- item` list with no - heading is valid YAML and valid Markdown; the model reports a split - YAML/Markdown distribution for such inputs rather than picking one with - certainty. -- The classifier uses content only. It does not inspect file names, extensions, - shebangs outside the model window, repository metadata, or build-system - context. +- Short files (< 512 bytes) are harder than full windows; accuracy on a + short-truncated validation set is ~0.78 versus ~0.94 on full windows. +- Ambiguous snippets can put several languages close together even when a + human can infer the language from file naming context. +- The classifier uses content only. It does not inspect file names, + extensions, shebangs outside the model window, repository metadata, or + build-system context. - Non-source formats are out of scope unless represented by a public source-language variant. ## Reproducibility -Training and evaluation scripts live under `scripts/`. The frozen recipe is -documented in `scripts/TRAINING.md`, including the expected external Magika -teacher assets, cache layout, training command, evaluation command, and -expected metrics. +Training and evaluation scripts live under `scripts/bnn/`. The pipeline is +documented in `scripts/bnn/TRAINING.md`, including corpus construction, +Shannon codebook learning, binary training, fine-tuning, ensemble weight +search, export, and the Rust/Python parity check. The published crate package intentionally includes only the runtime model artifact and user-facing docs. Training scripts and generated analysis files @@ -107,6 +103,6 @@ remain repository artifacts. ## Attribution -The embedded student model was trained from outputs of Google's Magika teacher +The embedded models were trained from outputs of Google's Magika teacher model. Magika is published by Google under Apache-2.0. Betlang's source code is MIT licensed; keep Magika attribution with redistributed model artifacts. diff --git a/README.md b/README.md index 85743da..901b258 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ [![Crates.io](https://img.shields.io/crates/v/betlang.svg)](https://crates.io/crates/betlang) [![Docs.rs](https://docs.rs/betlang/badge.svg)](https://docs.rs/betlang) -CPU source-language detection for code with a tiny 50kb model. Try it in browser [here](https://dioxuslabs.github.io/dioxus-code/#playground) +CPU source-language detection for code with a fully-binary CNN model. Try it in browser [here](https://dioxuslabs.github.io/dioxus-code/#playground) ```toml [dependencies] @@ -38,47 +38,44 @@ assert_eq!("rust".parse::(), Ok(betlang::Language::Rust)); These are the model's 48 output labels. Runtime detections expose them one-to-one with no label aggregation. -The confusion matrix uses the same labels: - -![Betlang wordseq confusion](https://raw.githubusercontent.com/ealmloff/betlang/ee771279730cc12bc2c60ba4db34e38dd0b0ef9a/assets/confusion-overall.png) - ## Model -The embedded model is `assets/magika/source-student-q4.bin`, a 47,840-byte -weights-only MSQ1 payload with SHA-256: +The embedded model is `assets/bnn/source-bnn2.bin`, a 1,165,464-byte packed +"BBN2" payload with SHA-256: ```text -8493d2d3757572c8661141e414b1c0755aa08d4c4e5382dfbbc6b73b02d89083 +fe36df96d9b6d268118b21b764c3be6dda33596e0fafeacd2bc0b26e3fc143ca ``` -Architecture: `wordseq-b1024-k3-m2048-tiny-3conv-hidden`, tokenizer version 3. +Architecture: an ensemble of five binary CNNs over Shannon-coded raw byte +windows. The 2048-byte Magika begin/end window is entropy-coded with a +canonical Shannon code learned from training byte frequencies, and the +resulting bitstream (plus a codeword-boundary bitplane) feeds convolutions +with binary weights and activations. All heavy compute is bitwise: +XNOR + popcount convolutions, integer thresholds, OR-pooling, segmented +popcounts, and an integer classifier head. + On the held-out filesystem-label test split it reaches -`test_fs_accuracy=0.942353` with `macro_recall=0.939690`. Probabilities are -calibrated: ambiguous inputs report split scores instead of a confident label. +`test_accuracy=0.941635` with `macro_recall=0.943226`, versus +`test_accuracy=0.934972` / `macro_recall=0.934863` for the previous wordseq +model on the same split. Ambiguous inputs report split scores instead of a +confident label. See [MODEL_CARD.md](MODEL_CARD.md) for the training and evaluation summary. ## Performance -Betlang uses a fixed 4096-byte Magika window and pads runtime inference to the -same 2048-token shape used by evaluation. The model is loaded once per process -and then reused through a `OnceLock`. +Betlang uses a fixed 4096-byte Magika window and Shannon-codes the extracted +2048-byte begin/end window into a fixed 16,384-bit stream. The model is loaded +once per process and then reused through a `OnceLock`. -Native CPU inference dispatches through `fearless_simd`. Benchmark entry points -are available through `cargo bench`. Current baseline numbers are tracked in +Inference is pure integer/bitwise CPU work (XNOR, popcount, integer +accumulation, thresholds, bitwise OR pooling). Benchmark entry points are +available through `cargo bench`. Current baseline numbers are tracked in [BENCHMARKS.md](BENCHMARKS.md). ## License And Attribution -Betlang is licensed under MIT. The embedded student model was trained from +Betlang is licensed under MIT. The embedded models were trained from outputs of Google's Magika teacher model; Magika is published by Google under Apache-2.0. Keep this attribution with redistributed model artifacts. - -## Confusion By File Size - -The shipped wordseq model is evaluated below on the held-out test split. Each -panel is a row-normalized confusion matrix for one file-size bucket: actual -labels are rows, predicted labels are columns, and the diagonal is correct -classification. - -![Betlang wordseq confusion by file size](https://raw.githubusercontent.com/ealmloff/betlang/ee771279730cc12bc2c60ba4db34e38dd0b0ef9a/assets/confusion-by-size.png) diff --git a/assets/bnn/source-bnn2.bin b/assets/bnn/source-bnn2.bin new file mode 100644 index 0000000..5e97e2c Binary files /dev/null and b/assets/bnn/source-bnn2.bin differ diff --git a/assets/magika/source-student-q4.bin b/assets/magika/source-student-q4.bin deleted file mode 100644 index e354bc3..0000000 Binary files a/assets/magika/source-student-q4.bin and /dev/null differ diff --git a/scripts/bnn/TRAINING.md b/scripts/bnn/TRAINING.md new file mode 100644 index 0000000..c48be32 --- /dev/null +++ b/scripts/bnn/TRAINING.md @@ -0,0 +1,45 @@ +# Shannon-bit binary CNN training pipeline + +These scripts produce `assets/bnn/source-bnn2.bin`, the packed "BBN2" ensemble +of binary CNNs that powers `betlang::detect`. All scripts run with Python 3.12, +PyTorch, and NumPy; caches and checkpoints live under `./bnn/` next to the +scripts (`OUT` in `train_bnn.py`). + +## Pipeline + +1. `build_corpus.py` — assemble the 48-label corpus (The Stack smol / GitHub + files), with `rebalance.py` / `topup.py` to even out rare classes. +2. `build_cache.py` — extract the raw 2048-byte Magika begin/end windows into + memory-mapped `{split}.windows.mmap` / `.lengths` / `.labels` caches. +3. `shannon.py` — learn the canonical Shannon code over training byte + frequencies (add-one smoothing, 12-bit max length, ordered by + `(length, byte)`), stored in `bnn/codebook.json`. Each window is encoded + MSB-first into a fixed 16,384-bit stream plus a codeword-boundary plane. +4. `train_bnn.py` / `train_bnn2.py` / `train_bnn3.py` — train binary CNNs + (binary weights and activations, sign activations, OR-pooling, segmented + popcount heads) directly on the Shannon bitplanes. +5. `finetune_seg.py`, `finetune_short.py` — segmented-head and short-input + fine-tuning of the base checkpoints. +6. `build_synth_hard.py`, `add_yaml_synth.py`, `add_yaml_synth2.py` and + `finetune_synth.py` / `finetune_synth2.py` / `finetune_synth3.py` — + synthetic hard-pair generation (YAML/Markdown lists and other confusable + pairs) and fine-tuning for the ensemble members. +7. `weight_search.py` / `weight_valid.py` — grid-search ensemble weights under + fixture, YAML/Markdown-edge, and benchmark-demo constraints, scored on the + full and short validation splits. +8. `export_bnn2.py` / `make_fixtures2.py` — quantize (fixed point `FIX=65536`, + i16 head weights), pack the ensemble into the BBN2 artifact, and emit + `parity_fixtures2.json` from the same in-memory model state. +9. `eval_export2.py` — exact integer-semantics evaluation of the exported + ensemble on the valid/test splits. +10. `eval_fixtures.py` / `eval_edge.py` / `eval_combos.py` — fixture-level and + edge-case evaluation helpers. + +## Parity check + +```bash +BNN_FIXTURES=path/to/parity_fixtures2.json cargo test --release python_parity +``` + +verifies the Rust XNOR/popcount inference reproduces the exported integer +logits bit-exactly. diff --git a/scripts/bnn/add_yaml_synth.py b/scripts/bnn/add_yaml_synth.py new file mode 100644 index 0000000..2366837 --- /dev/null +++ b/scripts/bnn/add_yaml_synth.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Append synthetic yaml (sequence-of-mappings) and markdown (bullet lists) +samples to the synth hard corpus.""" +from __future__ import annotations + +import random +from pathlib import Path + +import numpy as np + +import build_synth_hard as bs + +W = bs.W1 + + +def synth_yaml(rng: random.Random) -> bytes: + style = rng.random() + lines = [] + if style < 0.5: + for _ in range(rng.randint(2, 6)): + lines.append(f"- {rng.choice(['name', 'id', 'key', 'task'])}: " + f"{rng.choice(W)}") + for _ in range(rng.randint(1, 3)): + k = rng.choice(["run", "uses", "value", "with", "cmd", "path"]) + lines.append(f" {k}: {rng.choice(W)} {rng.choice(W)}") + elif style < 0.8: + lines.append(f"{rng.choice(W)}:") + for _ in range(rng.randint(2, 6)): + lines.append(f"- {rng.choice(['name', 'key'])}: {rng.choice(W)}") + if rng.random() < 0.5: + lines.append(f" {rng.choice(['run', 'value'])}: {rng.choice(W)}") + else: + for _ in range(rng.randint(2, 5)): + lines.append(f"{rng.choice(W)}:") + for _ in range(rng.randint(1, 3)): + lines.append(f" {rng.choice(W)}: " + f"{rng.choice([rng.choice(W), str(rng.randint(0, 99)), 'true', 'false'])}") + return ("\n".join(lines) + "\n").encode() + + +def synth_markdown(rng: random.Random) -> bytes: + lines = [] + if rng.random() < 0.8: + lines.append(f"{'#' * rng.randint(1, 3)} {rng.choice(W).title()}" + f"{rng.choice(['', ' ' + rng.choice(W).title()])}") + lines.append("") + for _ in range(rng.randint(3, 8)): + item = rng.choice(W) + if rng.random() < 0.4: + item = item.title() + if rng.random() < 0.3: + item += " " + rng.choice(W) + lines.append(f"{rng.choice(['-', '*'])} {item}") + if rng.random() < 0.3: + lines.append("") + lines.append(" ".join(rng.choice(W) for _ in range(rng.randint(5, 12)))) + return ("\n".join(lines) + "\n").encode() + + +def main(): + rng = random.Random(7) + labels = bs.LABELS + rows, lens, labs = [], [], [] + seen = set() + for label, gen, count in [("yaml", synth_yaml, 2000), + ("markdown", synth_markdown, 1200)]: + lid = labels.index(label) + made = 0 + attempts = 0 + while made < count and attempts < count * 60: + attempts += 1 + data = gen(rng) + if data in seen: + continue + seen.add(data) + w = bs.build_window(data) + if w is None: + continue + rows.append(w[0]) + lens.append(w[1]) + labs.append(lid) + made += 1 + print(label, made) + out = Path(bs.OUT) + windows = np.concatenate([np.load(out / "synth_hard.windows.npy"), + np.stack(rows)]) + lengths = np.concatenate([np.load(out / "synth_hard.lengths.npy"), + np.array(lens, dtype=np.int32)]) + lab = np.concatenate([np.load(out / "synth_hard.labels.npy"), + np.array(labs, dtype=np.int64)]) + np.save(out / "synth_hard.windows.npy", windows) + np.save(out / "synth_hard.lengths.npy", lengths) + np.save(out / "synth_hard.labels.npy", lab) + print("total", len(lab)) + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/add_yaml_synth2.py b/scripts/bnn/add_yaml_synth2.py new file mode 100644 index 0000000..c5dc5a6 --- /dev/null +++ b/scripts/bnn/add_yaml_synth2.py @@ -0,0 +1,62 @@ +#!/usr/bin/env python3 +"""More realistic short yaml sequence-of-mappings samples (CI-style).""" +from __future__ import annotations + +import random + +import numpy as np + +import build_synth_hard as bs + +NAMES = ["build", "test", "lint", "deploy", "check", "fmt", "docs", "bench", + "setup", "install", "release", "clean", "package", "verify"] +CMDS = ["make", "make test", "make build", "cargo test", "cargo build", + "npm test", "npm run build", "pytest", "go test ./...", "mvn verify", + "./configure", "bash ci.sh", "python setup.py test", "cmake --build ."] + + +def synth_yaml_ci(rng: random.Random) -> bytes: + lines = [] + for _ in range(rng.randint(2, 5)): + lines.append(f"- name: {rng.choice(NAMES)}") + for _ in range(rng.randint(1, 2)): + k = rng.choice(["run", "run", "cmd", "script"]) + lines.append(f" {k}: {rng.choice(CMDS)}") + return ("\n".join(lines) + "\n").encode() + + +def main(): + rng = random.Random(99) + lid = bs.LABELS.index("yaml") + rows, lens, labs = [], [], [] + seen = set() + made = 0 + attempts = 0 + while made < 2000 and attempts < 200000: + attempts += 1 + data = synth_yaml_ci(rng) + if data in seen: + continue + seen.add(data) + w = bs.build_window(data) + if w is None: + continue + rows.append(w[0]) + lens.append(w[1]) + labs.append(lid) + made += 1 + out = bs.OUT + windows = np.concatenate([np.load(out / "synth_hard.windows.npy"), + np.stack(rows)]) + lengths = np.concatenate([np.load(out / "synth_hard.lengths.npy"), + np.array(lens, dtype=np.int32)]) + lab = np.concatenate([np.load(out / "synth_hard.labels.npy"), + np.array(labs, dtype=np.int64)]) + np.save(out / "synth_hard.windows.npy", windows) + np.save(out / "synth_hard.lengths.npy", lengths) + np.save(out / "synth_hard.labels.npy", lab) + print("yaml_ci", made, "total", len(lab)) + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/build_cache.py b/scripts/bnn/build_cache.py new file mode 100644 index 0000000..45587b9 --- /dev/null +++ b/scripts/bnn/build_cache.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +"""Cache builder: raw Magika windows + fs labels + Magika teacher marginals. + +Writes per split: + {split}.windows.mmap uint8 [N, 2048] (zero padded) + {split}.lengths.mmap int32 [N] + {split}.labels.mmap int16 [N] + {split}.teacher.mmap float32 [N, 48] (teacher head-label marginals) + {split}.json metadata +""" +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).parent)) +sys.path.insert(0, str(Path.home() / "repos/betlang/scripts")) +from betlang_tokenizer import build_window, MAGIKA_WINDOW_SIZE + +LABELS = [ + "asm", "batch", "c", "clojure", "cmake", "cobol", "cpp", "cs", "css", + "dart", "dockerfile", "elixir", "erlang", "gemfile", "gemspec", "go", + "gradle", "groovy", "haskell", "html", "ini", "java", "javascript", + "json", "julia", "kotlin", "lisp", "lua", "markdown", "objectivec", + "ocaml", "perl", "php", "powershell", "python", "r", "ruby", "rust", + "scala", "shell", "sql", "swift", "toml", "typescript", "vba", + "verilog", "xml", "yaml", +] +LABEL_INDEX = {label: i for i, label in enumerate(LABELS)} + +MAGIKA_BEG_SIZE = 1024 +MAGIKA_END_SIZE = 1024 +MAGIKA_BLOCK_SIZE = 4096 +MAGIKA_PADDING_TOKEN = 256 + + +def magika_features(data: bytes) -> list[int] | None: + size = len(data) + if size == 0: + return None + stripped_beg = data[: min(size, MAGIKA_BLOCK_SIZE)].lstrip() + stripped_end = data[-min(size, MAGIKA_BLOCK_SIZE):].rstrip() + if len(stripped_beg) < 8: + return None + beg = list(stripped_beg[:MAGIKA_BEG_SIZE]) + beg.extend([MAGIKA_PADDING_TOKEN] * (MAGIKA_BEG_SIZE - len(beg))) + end_data = stripped_end[-MAGIKA_END_SIZE:] + end = [MAGIKA_PADDING_TOKEN] * (MAGIKA_END_SIZE - len(end_data)) + end.extend(end_data) + return beg + end + + +def load_teacher(): + import onnxruntime as ort + import magika + + model_dir = Path(magika.__file__).parent / "models" / "standard_v3_3" + config = json.loads((model_dir / "config.min.json").read_text()) + target = config["target_labels_space"] + label_to_index = {label: i for i, label in enumerate(target)} + cols = [label_to_index[label] for label in LABELS] + session = ort.InferenceSession( + str(model_dir / "model.onnx"), providers=["CPUExecutionProvider"] + ) + return session, cols + + +def main() -> int: + corpus = Path(sys.argv[1]) + out = Path(sys.argv[2]) + out.mkdir(parents=True, exist_ok=True) + session, cols = load_teacher() + input_name = session.get_inputs()[0].name + + for split in ("train", "valid", "test"): + meta_path = out / f"{split}.json" + if meta_path.exists(): + print(f"{split}: cached", flush=True) + continue + rows = [] + for label_dir in sorted((corpus / split).iterdir()): + if not label_dir.is_dir() or label_dir.name not in LABEL_INDEX: + continue + for file in sorted(label_dir.iterdir()): + rows.append((file, LABEL_INDEX[label_dir.name])) + n = len(rows) + print(f"{split}: {n} candidate files", flush=True) + windows = np.memmap(out / f"{split}.windows.mmap", dtype=np.uint8, + mode="w+", shape=(n, MAGIKA_WINDOW_SIZE)) + lengths = np.memmap(out / f"{split}.lengths.mmap", dtype=np.int32, + mode="w+", shape=(n,)) + labels = np.memmap(out / f"{split}.labels.mmap", dtype=np.int16, + mode="w+", shape=(n,)) + teacher = np.memmap(out / f"{split}.teacher.mmap", dtype=np.float32, + mode="w+", shape=(n, len(LABELS))) + count = 0 + batch_feats, batch_idx = [], [] + + def flush(): + nonlocal batch_feats, batch_idx + if not batch_feats: + return + arr = np.asarray(batch_feats, dtype=np.int32) + probs = session.run(None, {input_name: arr})[0] + teacher[batch_idx] = probs[:, cols] + batch_feats, batch_idx = [], [] + + for file, label in rows: + data = file.read_bytes() + window = build_window(data) + feats = magika_features(data) + if window is None or feats is None: + continue + windows[count, :len(window)] = np.frombuffer(window, dtype=np.uint8) + lengths[count] = len(window) + labels[count] = label + batch_feats.append(feats) + batch_idx.append(count) + count += 1 + if len(batch_feats) == 256: + flush() + if count % 10240 == 0: + print(f"{split}: {count}/{n}", flush=True) + flush() + windows.flush(); lengths.flush(); labels.flush(); teacher.flush() + meta_path.write_text(json.dumps({ + "count": count, "total": n, "labels": LABELS, + "window": MAGIKA_WINDOW_SIZE, + })) + print(f"{split}: wrote {count} rows", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/bnn/build_corpus.py b/scripts/bnn/build_corpus.py new file mode 100644 index 0000000..88fc2a8 --- /dev/null +++ b/scripts/bnn/build_corpus.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +"""Corpus builder: reuses betlang's build_finetune_corpus but replaces the +gated the-stack labels (yaml/json/toml/ini/xml/swift/cobol) with GitHub +tarball harvesting, so no HF token is needed.""" +import random +import sys +from pathlib import Path + +sys.path.insert(0, str(Path.home() / "repos/betlang/scripts")) +import build_finetune_corpus as bfc + +EXTRA_GITHUB = { + "yaml": ( + [ + "kubernetes/examples", "argoproj/argo-cd", "istio/istio", + "ansible/ansible", "prometheus/prometheus", "grafana/grafana", + "helm/examples", "fluxcd/flux2", "kubernetes-sigs/kustomize", + "open-telemetry/opentelemetry-collector", "goharbor/harbor", + "jaegertracing/jaeger", "cilium/cilium", "envoyproxy/envoy", + ], + (".yaml", ".yml"), + ), + "json": ( + [ + "SchemaStore/schemastore", "microsoft/vscode", "babel/babel", + "prettier/prettier", "eslint/eslint", "webpack/webpack", + "microsoft/TypeScript", "nlohmann/json", "jsonlint? no", + ], + (".json",), + ), + "toml": ( + [ + "rust-lang/cargo", "toml-lang/toml-test", "rust-lang/rust-analyzer", + "bevyengine/bevy", "tokio-rs/tokio", "serde-rs/serde", + "rust-lang/rustup", "denoland/deno", "astral-sh/uv", + "pypa/hatch", "tauri-apps/tauri", "helix-editor/helix", + ], + (".toml",), + ), + "ini": ( + [ + "php/php-src", "benhoyt/inih", "python/mypy", "pytest-dev/pytest", + "wine-mirror/wine", "systemd/systemd", "pypa/setuptools", + "sqlalchemy/sqlalchemy", "pallets/flask", "django/django", + "tox-dev/tox", "psf/requests", + ], + (".ini", ".cfg"), + ), + "xml": ( + [ + "apache/tomcat", "spring-projects/spring-framework", + "mybatis/mybatis-3", "apache/maven", "android/architecture-samples", + "apache/camel", "apache/ant", "eclipse-platform/eclipse.platform", + "apache/struts", "AndroidX? no", "dbeaver/dbeaver", + ], + (".xml",), + ), + "swift": ( + [ + "Alamofire/Alamofire", "vapor/vapor", "apple/swift-nio", + "ReactiveX/RxSwift", "apple/swift-package-manager", + "pointfreeco/swift-composable-architecture", "SnapKit/SnapKit", + "onevcat/Kingfisher", "kean/Nuke", "apple/swift-algorithms", + "apple/swift-collections", "SwiftyJSON/SwiftyJSON", + "realm/SwiftLint", "Quick/Quick", "Quick/Nimble", + ], + (".swift",), + ), + "cobol": ( + [ + "openmainframeproject/cobol-programming-course", + "cicsdev/cics-banking-sample-application-cbsa", + "shamrice/COBOL-Examples", "mapmeld/cobol-fizzbuzz", + "azac/cobol-on-wheelchair", + "olegkunitsyn/gnucobol-debug", "neopragma/cobol-unit-test", + "openmainframeproject/cobol-check", "cicsdev/cics-genapp", + "IBM/zopeneditor-sample", "aws-samples/aws-mainframe-modernization-carddemo", + "mikebild? no", "rplig? no", + ], + (".cbl", ".cob", ".CBL", ".COB", ".cobol"), + ), +} +# drop typo'd placeholders +for label, (repos, exts) in list(EXTRA_GITHUB.items()): + EXTRA_GITHUB[label] = ([r for r in repos if "?" not in r], exts) + + +def main() -> int: + out = Path(sys.argv[1]) + rng = random.Random(2) + files_root = out / "files" + scratch = out / "scratch" + scratch.mkdir(parents=True, exist_ok=True) + writer = bfc.CorpusWriter(files_root) + import hashlib + for split_dir in files_root.glob("*/"): + for label_dir in split_dir.glob("*/"): + for file in label_dir.iterdir(): + writer.seen.add(hashlib.sha1(file.read_bytes()).digest()) + writer.counts[(split_dir.name, label_dir.name)] += 1 + + bfc.harvest_smol_xl(writer, None, scratch, rng) + bfc.GITHUB_SOURCES = {**bfc.GITHUB_SOURCES, **EXTRA_GITHUB} + import urllib.error + for label, (repos, exts) in bfc.GITHUB_SOURCES.items(): + for repo in repos: + try: + bfc.GITHUB_SOURCES = {label: ([repo], exts)} + bfc.harvest_github(writer, scratch) + except Exception as err: # tolerate 404s / truncated tars + print(f"github {repo}: FAILED {err}", flush=True) + (scratch / f"github.{repo.replace('/', '__')}.done").write_text("error\n") + bfc.harvest_raw_files(writer, scratch) + if not (scratch / "synthetic.done").exists(): + bfc.harvest_synthetic(writer, rng) + (scratch / "synthetic.done").write_text("done\n") + print(writer.summary()) + (out / "corpus_summary.txt").write_text(writer.summary() + "\n") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/bnn/build_synth_hard.py b/scripts/bnn/build_synth_hard.py new file mode 100644 index 0000000..17dd5c2 --- /dev/null +++ b/scripts/bnn/build_synth_hard.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 +"""Build a synthetic hard-pair corpus (windows + labels) reusing betlang's +hard_gen_* generators, plus small cmake/gemspec/r/ruby/toml/ocaml/haskell +generators for the remaining confusable short-file pairs.""" +from __future__ import annotations + +import json +import random +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path.home() / "repos/betlang/scripts")) +sys.path.insert(0, str(Path(__file__).parent)) +from hard_gen_ini import synth_hard_ini +from hard_gen_js_ts import synth_hard_javascript, synth_hard_typescript +from train_bnn import OUT, WINDOW + +LABELS = json.loads((Path.home() / "work/cache/valid.json").read_text())["labels"] + +W1 = ("alpha beta gamma delta core util math data net io fmt log cfg http " + "json test demo file cache queue stack list map set node tree lex " + "parse token scan").split() + + +def _n(rng): + return rng.choice(W1) + (rng.choice(["", "_" + rng.choice(W1)])) + + +def synth_cmake(rng: random.Random) -> bytes: + n = _n(rng) + lines = [] + if rng.random() < 0.5: + lines.append(f"cmake_minimum_required(VERSION 3.{rng.randint(10, 28)})") + lines.append(f"project({n}{rng.choice(['', ' C', ' CXX', ' VERSION 1.0'])})") + for _ in range(rng.randint(1, 5)): + kind = rng.random() + tgt = _n(rng) + srcs = " ".join(f"{_n(rng)}.{rng.choice(['c', 'cpp', 'cc'])}" + for _ in range(rng.randint(1, 3))) + if kind < 0.35: + lines.append(f"add_library({tgt} {srcs})") + elif kind < 0.7: + lines.append(f"add_executable({tgt} {srcs})") + elif kind < 0.85: + lines.append(f"target_link_libraries({tgt} PRIVATE {_n(rng)})") + else: + lines.append(f"set({tgt.upper()}_SOURCES {srcs})") + if rng.random() < 0.6: + lines.append(f"install(TARGETS {_n(rng)})") + if rng.random() < 0.4: + lines.append(f"include_directories({rng.choice(['include', 'src'])})") + if rng.random() < 0.3: + lines.insert(rng.randint(0, len(lines)), f"# {_n(rng)} build rules") + return ("\n".join(lines) + "\n").encode() + + +def synth_gemspec(rng: random.Random) -> bytes: + n = _n(rng).replace("_", "-") + var = rng.choice(["spec", "s", "gem"]) + lines = [] + if rng.random() < 0.3: + lines.append("# frozen_string_literal: true") + lines.append("") + if rng.random() < 0.4: + lines.append(f'require_relative "lib/{n.replace("-", "/")}/version"') + lines.append("") + lines.append(f"Gem::Specification.new do |{var}|") + lines.append(f' {var}.name = "{n}"') + lines.append(f' {var}.version = "{rng.randint(0, 3)}.{rng.randint(0, 9)}.' + f'{rng.randint(0, 9)}"') + if rng.random() < 0.7: + lines.append(f' {var}.summary = "A {rng.choice(W1)} {rng.choice(W1)} library"') + if rng.random() < 0.6: + lines.append(f' {var}.authors = ["{rng.choice(W1).title()}"]') + if rng.random() < 0.5: + lines.append(f' {var}.files = Dir["lib/**/*.rb"]') + if rng.random() < 0.5: + lines.append(f' {var}.required_ruby_version = ">= {rng.randint(2, 3)}.' + f'{rng.randint(0, 4)}"') + if rng.random() < 0.4: + lines.append(f' {var}.license = "{rng.choice(["MIT", "Apache-2.0"])}"') + for _ in range(rng.randint(0, 3)): + dep = rng.choice(["json", "rake", "rspec", "minitest", "nokogiri", "thor"]) + kind = rng.choice(["add_dependency", "add_development_dependency", + "add_runtime_dependency"]) + lines.append(f' {var}.{kind} "{dep}", "~> {rng.randint(1, 13)}.' + f'{rng.randint(0, 9)}"') + lines.append("end") + return ("\n".join(lines) + "\n").encode() + + +def synth_ruby(rng: random.Random) -> bytes: + n = _n(rng) + cls = "".join(p.title() for p in n.split("_")) + lines = [] + if rng.random() < 0.3: + lines.append("# frozen_string_literal: true") + lines.append("") + style = rng.random() + if style < 0.4: + lines.append(f"class {cls}") + lines.append(f" def initialize({_n(rng)})") + lines.append(f" @{_n(rng)} = {_n(rng)}") + lines.append(" end") + lines.append("") + lines.append(f" def {_n(rng)}") + lines.append(f" @{_n(rng)}.to_s") + lines.append(" end") + lines.append("end") + elif style < 0.7: + lines.append(f"module {cls}") + lines.append(f" def self.{_n(rng)}(items)") + lines.append(f" items.map {{ |x| x.{rng.choice(['to_s', 'upcase', 'strip'])} }}") + lines.append(" end") + lines.append("end") + else: + lines.append(f"{_n(rng)} = [{rng.randint(1, 9)}, {rng.randint(10, 99)}]") + lines.append(f"{_n(rng)}.each do |v|") + lines.append(' puts "#{v}"') + lines.append("end") + return ("\n".join(lines) + "\n").encode() + + +def synth_r(rng: random.Random) -> bytes: + lines = [] + if rng.random() < 0.4: + lines.append(f"library({rng.choice(['dplyr', 'ggplot2', 'stats'])})") + n = _n(rng) + lines.append(f"{n} <- function(x, y = {rng.randint(1, 9)}) {{") + lines.append(f" z <- x * y + {rng.randint(1, 20)}") + lines.append(" return(z)") + lines.append("}") + if rng.random() < 0.6: + lines.append(f"{_n(rng)} <- c({rng.randint(1, 5)}, {rng.randint(6, 9)})") + if rng.random() < 0.5: + lines.append(f"print({n}({rng.randint(1, 5)}))") + return ("\n".join(lines) + "\n").encode() + + +GENERATORS = { + "cmake": synth_cmake, + "gemspec": synth_gemspec, + "ruby": synth_ruby, + "r": synth_r, + "ini": synth_hard_ini, + "javascript": synth_hard_javascript, + "typescript": synth_hard_typescript, +} +COUNTS = {"cmake": 3000, "gemspec": 3000, "ruby": 1500, "r": 1500, + "ini": 3000, "javascript": 2000, "typescript": 3000} + + +def build_window(source: bytes): + ws = b"\t\n\x0b\x0c\r " + beg = source[:4096].lstrip(ws) + if len(beg) < 8: + return None + end = source[-4096:].rstrip(ws) + beg_len = min(len(beg), 1024) + buf = np.zeros(WINDOW, dtype=np.uint8) + buf[:beg_len] = np.frombuffer(beg[:beg_len], dtype=np.uint8) + end_len = min(len(end), 1024) + end_start = 1024 + (1024 - end_len) + buf[end_start:end_start + end_len] = np.frombuffer( + end[len(end) - end_len:], dtype=np.uint8) + if beg_len < 1024: + valid = beg_len + elif end_start > 1024: + valid = 1024 + else: + valid = WINDOW + return buf, valid + + +def main(): + rng = random.Random(42) + rows, lens, labs = [], [], [] + seen = set() + for label, gen in GENERATORS.items(): + lid = LABELS.index(label) + made = 0 + attempts = 0 + while made < COUNTS[label] and attempts < COUNTS[label] * 50: + attempts += 1 + data = gen(rng) + if data in seen: + continue + seen.add(data) + w = build_window(data) + if w is None: + continue + rows.append(w[0]) + lens.append(w[1]) + labs.append(lid) + made += 1 + print(f"{label}: {made} samples", flush=True) + windows = np.stack(rows) + np.save(OUT / "synth_hard.windows.npy", windows) + np.save(OUT / "synth_hard.lengths.npy", np.array(lens, dtype=np.int32)) + np.save(OUT / "synth_hard.labels.npy", np.array(labs, dtype=np.int64)) + print("total", len(labs)) + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/eval_combos.py b/scripts/bnn/eval_combos.py new file mode 100644 index 0000000..50d7ff5 --- /dev/null +++ b/scripts/bnn/eval_combos.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Compare ensemble combinations on full-valid and short-valid splits.""" +from __future__ import annotations + +import itertools +import json +import sys +from pathlib import Path + +import numpy as np +import torch + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes +from train_bnn import OUT, CLASSES, load_split +from train_bnn2 import BNN2, build_plane_cache, unpack_batch +from finetune_seg import BNN2Seg +from finetune_short import build_short_valid + +device = "mps" if torch.backends.mps.is_available() else "cpu" + + +def all_logits(model, packed, n, batch=128): + outs = [] + model.eval() + with torch.no_grad(): + for i in range(0, n, batch): + idx = np.arange(i, min(i + batch, n)) + x = unpack_batch(packed, idx).to(device) + outs.append(model(x).cpu().numpy()) + return np.concatenate(outs) + + +def metrics(logits, labels): + pred = logits.argmax(1) + acc = (pred == labels).mean() + rec = [] + for c in range(CLASSES): + m = labels == c + if m.sum(): + rec.append((pred[m] == c).mean()) + return acc, float(np.mean(rec)) + + +def main(): + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + va_bits, va_labels, _ = build_plane_cache("valid", codes) + va_windows, va_lengths, _, _ = load_split("valid") + va_short = build_short_valid(codes, va_windows, va_lengths) + n = len(va_labels) + + if len(sys.argv) > 1: + specs = [(BNN2Seg if "2seg" in ck else BNN2, ck) for ck in sys.argv[1:]] + names = [f"m{i}" for i in range(len(specs))] + else: + specs = [(BNN2Seg, "best2seg.pt"), (BNN2, "best3.pt"), + (BNN2Seg, "best2seg_short.pt"), (BNN2, "best3_short.pt")] + names = ["m0", "m1", "m0s", "m1s"] + models = [] + for cls, ck in specs: + m = cls() + m.load_state_dict(torch.load(OUT / ck, map_location="cpu")) + m.to(device) + models.append(m) + + lf = [all_logits(m, va_bits, n) for m in models] + ls = [all_logits(m, va_short, n) for m in models] + k = len(models) + for r in range(1, k + 1): + for combo in itertools.combinations(range(k), r): + f = sum(lf[i] for i in combo) + s = sum(ls[i] for i in combo) + fa, fm = metrics(f, va_labels) + sa, sm = metrics(s, va_labels) + print(f"{'+'.join(names[i] for i in combo):16s} " + f"full {fa:.4f}/{fm:.4f} short {sa:.4f}/{sm:.4f}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/eval_edge.py b/scripts/bnn/eval_edge.py new file mode 100644 index 0000000..a357e7f --- /dev/null +++ b/scripts/bnn/eval_edge.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python3 +"""Check the ambiguous yaml/markdown behavioral tests from src/model/tests.rs.""" +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +import torch + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes, encode_planes_bulk +from train_bnn import OUT, BITS +from train_bnn2 import BNN2 +from finetune_seg import BNN2Seg +from eval_fixtures import build_window + +CASES = [ + ("md_list", "# Heading\n\n- first\n- second\n- third\n- fourth\n- fifth", + "markdown top"), + ("md_caps", "# Names\n\n- Alice\n- Bob\n- Carol\n- Dave", "markdown top"), + ("yaml_seq", "- name: build\n run: make\n- name: test\n run: make test", + "yaml top"), + ("commented", "# comment\nitems:\n- first\n- second\n- third", + "yaml+markdown top2"), + ("bare_dash", "- first\n- second\n- third\n- fourth\n- fifth", + "yaml+markdown top2, p<0.9"), +] + + +def main(): + labels = json.loads((OUT.parent / "cache/valid.json").read_text())["labels"] + lengths256 = np.array( + json.loads((OUT / "codebook.json").read_text())["lengths"], dtype=np.int32) + codes = canonical_codes(lengths256) + models = [] + for ck in sys.argv[1:]: + cls = BNN2Seg if "2seg" in ck else BNN2 + m = cls() + m.load_state_dict(torch.load(OUT / ck, map_location="cpu")) + m.eval() + models.append(m) + for name, text, want in CASES: + w = build_window(text.encode()) + if w is None: + print(f"{name}: window None") + continue + arr = np.frombuffer(w[0], dtype=np.uint8)[None, :] + planes = encode_planes_bulk(arr, np.array([w[1]]), codes, BITS) + x = torch.from_numpy(planes.astype(np.float32) * 2 - 1) + tot = None + with torch.no_grad(): + for m in models: + lg = m(x)[0] + tot = lg if tot is None else tot + lg + p = torch.softmax(tot / len(models), dim=0).numpy() + order = np.argsort(-p) + top = [(labels[i], round(float(p[i]), 4)) for i in order[:3]] + print(f"{name}: want[{want}] got {top}") + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/eval_export2.py b/scripts/bnn/eval_export2.py new file mode 100644 index 0000000..1ea2517 --- /dev/null +++ b/scripts/bnn/eval_export2.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +"""Exact batched evaluation of exported v2 integer semantics. + +Per-basis convs on {-1,+1} with +/-1 weights are exact in f32 (|z'| <= 768); +the fixed-point combine and thresholds run in f64 on CPU (exact for i32 +range). Supports single models and the ensemble (sum of logits).""" +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes +from train_bnn import CLASSES, OUT +from train_bnn2 import C, build_plane_cache +from export_bnn2 import IntModel, load_models + +import os + +DEV = "cpu" if os.environ.get("EVAL_CPU") else ( + "mps" if torch.backends.mps.is_available() else "cpu") + + +class FastInt: + def __init__(self, im: IntModel): + self.segs = im.segs + self.layers = [] + for layer in (im.l0, im.l1, im.l2): + m, cout, _ = layer.w.shape + w = torch.from_numpy(layer.w.astype(np.float32) * 2 - 1) + w = w.reshape(m, cout, layer.k, layer.cin).permute(0, 1, 3, 2).contiguous() + self.layers.append({ + "w": w.to(DEV), # [m, cout, cin, k] + "aq": torch.from_numpy(layer.aq.astype(np.float64)), # [m, cout] + "sg": torch.from_numpy(layer.sg.astype(np.int64)), + "th": torch.from_numpy(layer.th.astype(np.float64)), + "stride": layer.stride, "pad": layer.pad, + }) + self.hq = torch.from_numpy(im.hq.astype(np.float64)) + self.hs = torch.from_numpy(im.hs.astype(np.float64)) + self.hb = torch.from_numpy(im.hb_.astype(np.float64)) + + def conv_fire(self, x, layer): + # x: [B, cin, L] in {-1,+1} f32 on DEV + zq = None + for j in range(layer["w"].shape[0]): + zj = F.conv1d(x, layer["w"][j], stride=layer["stride"], + padding=layer["pad"]).cpu().double() + term = layer["aq"][j].view(1, -1, 1) * zj + zq = term if zq is None else zq + term + fire = torch.where(layer["sg"].view(1, -1, 1) == 1, + zq >= layer["th"].view(1, -1, 1), + zq <= layer["th"].view(1, -1, 1)) + return fire.float().to(DEV) + + def logits(self, x): # x: [B, 2, BITS] {-1,+1} + h = self.conv_fire(x, self.layers[0]) + h = F.max_pool1d(h, 4) + h = self.conv_fire(h * 2 - 1, self.layers[1]) + h = F.max_pool1d(h, 4) + h = self.conv_fire(h * 2 - 1, self.layers[2]) + segs = self.segs + seg = h.shape[2] // segs + counts = torch.cat([h[:, :, i * seg:(i + 1) * seg].sum(2) + for i in range(segs)], 1).cpu().double() + return counts @ self.hq.T * self.hs[None, :] + self.hb[None, :] + + +def main(): + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + models = [FastInt(IntModel(m, w)) for m, w in load_models()] + split = sys.argv[1] if len(sys.argv) > 1 else "test" + bits, labels, _ = build_plane_cache(split, codes) + n = len(labels) + batch = 128 + all_logits = np.zeros((len(models), n, CLASSES)) + with torch.no_grad(): + for i in range(0, n, batch): + idx = np.arange(i, min(i + batch, n)) + xb = np.unpackbits(np.asarray(bits[idx]), axis=2).astype(np.float32) * 2 - 1 + x = torch.from_numpy(xb).to(DEV) + for mi, fm in enumerate(models): + all_logits[mi, idx] = fm.logits(x).numpy() + + def report(name, logits): + pred = logits.argmax(1) + correct = np.zeros(CLASSES); total = np.zeros(CLASSES) + for c in range(CLASSES): + m = labels == c + total[c] = m.sum() + correct[c] = (pred[m] == c).sum() + acc = correct.sum() / total.sum() + macro = (correct[total > 0] / total[total > 0]).mean() + print(f"{split} {name}: acc {acc:.6f} macro {macro:.6f} n={n}") + + for mi in range(len(models)): + report(f"model{mi}", all_logits[mi]) + if len(models) > 1: + report("ensemble", all_logits.sum(0)) + np.save(OUT / f"{split}.intlogits.npy", all_logits) + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/eval_fixtures.py b/scripts/bnn/eval_fixtures.py new file mode 100644 index 0000000..75ec794 --- /dev/null +++ b/scripts/bnn/eval_fixtures.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Run candidate ensembles against betlang's language fixture files +(replicates src/model/window.rs build_window semantics).""" +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +import torch + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes, encode_planes_bulk +from train_bnn import OUT, BITS +from train_bnn2 import BNN2 +from finetune_seg import BNN2Seg + +REPO = Path.home() / "repos/betlang" +BEG, END, WIN, BLOCK = 1024, 1024, 2048, 4096 +LABELS = json.loads((Path.home() / "work/cache/valid.json").read_text())["labels"] + + +def build_window(source: bytes): + if not source: + return None + block = min(len(source), BLOCK) + beg = source[:block].lstrip(b"\t\n\x0b\x0c\r ") + if len(beg) < 8: + return None + end = source[len(source) - block:].rstrip(b"\t\n\x0b\x0c\r ") + beg_len = min(len(beg), BEG) + end_len = min(len(end), END) + buf = bytearray(WIN) + buf[:beg_len] = beg[:beg_len] + end_start = BEG + (END - end_len) + buf[end_start:end_start + end_len] = end[len(end) - end_len:] + if beg_len < BEG: + valid = beg_len + elif end_start > BEG: + valid = BEG + else: + valid = WIN + return bytes(buf), valid + + +def main(): + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + if len(sys.argv) > 1: + specs = [] + for ck in sys.argv[1:]: + cls = BNN2Seg if "2seg" in ck else BNN2 + specs.append((cls, ck)) + else: + specs = [(BNN2Seg, "best2seg.pt"), (BNN2, "best3.pt"), + (BNN2Seg, "best2seg_short.pt"), (BNN2, "best3_short.pt")] + models = [] + for cls, ck in specs: + m = cls() + m.load_state_dict(torch.load(OUT / ck, map_location="cpu")) + m.eval() + models.append(m) + + wrong = 0 + files = sorted((REPO / "tests/fixtures/languages").iterdir()) + for f in files: + expected = f.stem if f.stem != "docker" else f.stem + expected = f.stem + w = build_window(f.read_bytes()) + if w is None: + print(f"{f.name}: window None") + continue + window, valid = w + arr = np.frombuffer(window, dtype=np.uint8)[None, :] + planes = encode_planes_bulk(arr, np.array([valid]), codes, BITS) + x = torch.from_numpy(planes.astype(np.float32) * 2 - 1) + with torch.no_grad(): + logits = sum(m(x) for m in models)[0].numpy() + pred = LABELS[int(logits.argmax())] + if pred != expected: + top = np.argsort(logits)[::-1][:3] + print(f"{f.name}: expected {expected}, got {pred}; " + f"top {[LABELS[t] for t in top]}") + wrong += 1 + print(f"{wrong}/{len(files)} wrong") + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/export_bnn2.py b/scripts/bnn/export_bnn2.py new file mode 100644 index 0000000..7b0cf71 --- /dev/null +++ b/scripts/bnn/export_bnn2.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 +"""Export trained BNN2/BNN2Seg models (optionally an ensemble) to a packed +bitwise artifact ("BBN2" v2) + bit-exact integer reference. + +Inference semantics (all integer/bitwise): + input: 2 bitplanes (shannon bits, codeword boundaries), each BITS long + conv layer with m binary bases: for basis j, z'_j = 2*match_j - n_valid + (XNOR+popcount); z_q = sum_j aq[j] * z'_j (i64, aq = round(alpha*FIX)) + fire iff (sign==1) ? z_q >= thr : z_q <= thr (thr i32, fixed-point) + OR-pool window 4; segmented popcount head; i16 head dot + f32 scale/bias. + +Layout (little-endian): + u32 magic="BBN2", u32 version=2, u32 BITS, u32 CLASSES, u32 n_models + 256 x u8 shannon code lengths + per model: + u32 C, STEM_K, STEM_STRIDE, SEGS + 3 x layer: u32 m, u32 cin, u32 k, then per basis: + cout x ceil(cin*k/32) u32 weight bits (bit = tap*cin + ci), + cout x i32 aq + then cout x u8 sign, cout x i32 thr + head: CLASSES x (SEGS*C) i16, CLASSES x f32 scale, CLASSES x f32 bias +""" +from __future__ import annotations + +import json +import struct +import sys +from pathlib import Path + +import numpy as np +import torch + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes +from train_bnn import BITS, CLASSES, OUT +from train_bnn2 import BNN2, BinConvM, C, STEM_K, STEM_STRIDE, build_plane_cache, unpack_batch +from finetune_seg import BNN2Seg, SEGS as SEGS8 + +MAGIC = 0x324E4242 # "BBN2" +FIX = 65536 # alpha fixed-point scale (2^16) + + +def conv_export(conv: BinConvM): + """Returns (wsigns [m, cout, k*cin] bit=tap*cin+ci, aq [m, cout] i32).""" + m, cout, cin, k = conv.latent.shape + std = conv.latent.detach().std(dim=(1, 2, 3), keepdim=True).clamp_min(1e-6) + wb = (conv.latent.detach() - conv.shift.detach().view(m, 1, 1, 1) * std) >= 0 + signs = wb.numpy().astype(np.uint8) # [m, cout, cin, k] + signs = np.transpose(signs, (0, 1, 3, 2)).reshape(m, cout, k * cin) + aq = np.round(conv.alpha.detach().numpy() * FIX).astype(np.int64) + return signs, aq, cin, k + + +def fold_thresholds_q(bn, rsign_b): + """Thresholds on z_q = FIX * (float pre-BN activation).""" + gamma = bn.weight.detach().numpy() + beta = bn.bias.detach().numpy() + mu = bn.running_mean.detach().numpy() + sigma = np.sqrt(bn.running_var.detach().numpy() + bn.eps) + b = rsign_b.detach().numpy().reshape(-1) + lhs = gamma / sigma + rhs = b - beta + gamma * mu / sigma + t = rhs / np.where(np.abs(lhs) < 1e-12, 1e-12, lhs) * FIX + sign = np.where(lhs >= 0, 1, 0).astype(np.uint8) + thr = np.where(sign == 1, np.ceil(t), np.floor(t)) + thr = np.clip(thr, -(2 ** 30), 2 ** 30) + return sign, thr.astype(np.int64) + + +def head_export(hb, fc): + gamma = hb.weight.detach().numpy() + beta = hb.bias.detach().numpy() + mu = hb.running_mean.detach().numpy() + sigma = np.sqrt(hb.running_var.detach().numpy() + hb.eps) + W = fc.weight.detach().numpy() + b = fc.bias.detach().numpy() + A = W * (gamma / sigma)[None, :] + Bc = b + W @ (beta - mu * gamma / sigma) + scale = np.abs(A).max(axis=1) / 32767.0 + scale[scale == 0] = 1.0 + q = np.clip(np.round(A / scale[:, None]), -32767, 32767).astype(np.int16) + return q, scale.astype(np.float32), Bc.astype(np.float32) + + +def pack_bits_u32(signs: np.ndarray) -> np.ndarray: + rows, nbits = signs.shape + words = (nbits + 31) // 32 + out = np.zeros((rows, words), dtype=np.uint32) + for j in range(nbits): + out[:, j // 32] |= (signs[:, j].astype(np.uint32) << np.uint32(j % 32)) + return out + + +class IntLayer: + def __init__(self, conv: BinConvM, bn, rsign): + self.w, self.aq, self.cin, self.k = conv_export(conv) + self.sg, self.th = fold_thresholds_q(bn, rsign.b) + self.stride = conv.stride + self.pad = conv.padding + + def forward(self, x: np.ndarray) -> np.ndarray: + """x: [cin, L] {0,1} -> [cout, Lout] {0,1}; integer exact.""" + cin, L = x.shape + m, cout, _ = self.w.shape + k, stride, pad = self.k, self.stride, self.pad + Lout = (L + 2 * pad - k) // stride + 1 + zq = np.zeros((cout, Lout), dtype=np.int64) + cols = np.zeros((k, cin, Lout), dtype=np.int64) + valid = np.zeros((k, Lout), dtype=np.int64) + for tap in range(k): + src = np.arange(Lout) * stride - pad + tap + ok = (src >= 0) & (src < L) + valid[tap] = ok + cols[tap][:, ok] = x[:, src[ok]] + nv = (valid * cin).sum(0) + for j in range(m): + wj = self.w[j].reshape(cout, k, cin) + match = np.zeros((cout, Lout), dtype=np.int64) + for tap in range(k): + wb = wj[:, tap, :].astype(np.int64) + xb = cols[tap] + eq = wb @ xb + (1 - wb) @ (1 - xb) + eq -= np.outer((1 - wb).sum(1), 1 - valid[tap]) + match += eq + zq += self.aq[j][:, None] * (2 * match - nv[None, :]) + return np.where(self.sg[:, None] == 1, zq >= self.th[:, None], + zq <= self.th[:, None]).astype(np.uint8) + + +class IntModel: + def __init__(self, model, weight: float = 1.0): + self.segs = model.fc.in_features // C + self.l0 = IntLayer(model.stem, model.bn0, model.s0) + self.l1 = IntLayer(model.conv1, model.bn1, model.s1) + self.l2 = IntLayer(model.conv2, model.bn2, model.s2) + self.hq, self.hs, self.hb_ = head_export(model.hb, model.fc) + self.hs = (self.hs * weight).astype(np.float32) + self.hb_ = (self.hb_ * weight).astype(np.float32) + + def logits(self, planes01: np.ndarray) -> np.ndarray: + h = self.l0.forward(planes01) + L = (h.shape[1] // 4) * 4 + h = h[:, :L].reshape(h.shape[0], -1, 4).max(2) + h = self.l1.forward(h) + h = h.reshape(h.shape[0], -1, 4).max(2) + h = self.l2.forward(h) + segs = self.segs + seg = h.shape[1] // segs + counts = np.concatenate([h[:, i * seg:(i + 1) * seg].sum(1) + for i in range(segs)]) + acc = self.hq.astype(np.int64) @ counts.astype(np.int64) + return acc * self.hs + self.hb_ + + def write(self, buf: bytearray): + buf += struct.pack("<4I", C, STEM_K, STEM_STRIDE, self.segs) + for layer in (self.l0, self.l1, self.l2): + m, cout, nbits = layer.w.shape + buf += struct.pack("<3I", m, layer.cin, layer.k) + for j in range(m): + buf += pack_bits_u32(layer.w[j]).tobytes() + buf += layer.aq[j].astype(np.int32).tobytes() + buf += layer.sg.tobytes() + buf += layer.th.astype(np.int32).tobytes() + buf += self.hq.tobytes() + buf += self.hs.tobytes() + buf += self.hb_.tobytes() + + +ENSEMBLE = [("best2seg.pt", 0.7), ("best3.pt", 0.7), + ("best2seg_synth.ep0.pt", 1.0), ("best2seg_synth3.ep0.pt", 1.6), + ("best3_synth2.ep0.pt", 1.3)] + + +def load_models() -> list: + """Returns [(model, weight)] with weights normalized so that the mean of + the weighted logits equals the weighted-softmax-temperature convention + logits_total / sum(weights) used during weight search.""" + wsum = sum(w for _, w in ENSEMBLE) + k = len(ENSEMBLE) + models = [] + for ck, w in ENSEMBLE: + cls = BNN2Seg if "2seg" in ck else BNN2 + m = cls() + m.load_state_dict(torch.load(OUT / ck, map_location="cpu")) + m.eval() + models.append((m, w * k / wsum)) + return models + + +def main(): + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + pairs = load_models() + models = [m for m, _ in pairs] + ints = [IntModel(m, w) for m, w in pairs] + + va_bits, va_labels, _ = build_plane_cache("valid", codes) + n_par = 128 + idx = np.arange(0, len(va_labels), max(1, len(va_labels) // n_par))[:n_par] + x = unpack_batch(va_bits, idx) + agree = [0] * len(models) + with torch.no_grad(): + refs = [m(x).numpy() for m in models] + for row, i in enumerate(idx): + planes = np.unpackbits(np.asarray(va_bits[i]), axis=1) + for mi, im in enumerate(ints): + got = im.logits(planes) + if got.argmax() == refs[mi][row].argmax(): + agree[mi] += 1 + print("int-vs-torch argmax agreement:", [f"{a}/{len(idx)}" for a in agree]) + + buf = bytearray() + buf += struct.pack("<5I", MAGIC, 2, BITS, CLASSES, len(ints)) + buf += lengths256.astype(np.uint8).tobytes() + for im in ints: + im.write(buf) + path = OUT / "source-bnn2.bin" + path.write_bytes(bytes(buf)) + print(f"wrote {path} ({len(buf)} bytes)") + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/finetune_seg.py b/scripts/bnn/finetune_seg.py new file mode 100644 index 0000000..0792cb7 --- /dev/null +++ b/scripts/bnn/finetune_seg.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 +"""Phase-2 fine-tune: upgrade head from 2 to 8 pooled count segments +(function-preserving warm start), continue training at low LR with EMA.""" +from __future__ import annotations + +import copy +import json +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes +from train_bnn import OUT, CLASSES, SEED +from train_bnn2 import BNN2, BinConvM, C, build_plane_cache, unpack_batch +from finetune_bnn2 import Ema, upgrade_conv + +SEGS = 8 + + +class BNN2Seg(nn.Module): + def __init__(self): + super().__init__() + base = BNN2() + base.conv1 = upgrade_conv(base.conv1) + base.conv2 = upgrade_conv(base.conv2) + self.stem = base.stem + self.bn0 = base.bn0 + self.s0 = base.s0 + self.conv1 = base.conv1 + self.bn1 = base.bn1 + self.s1 = base.s1 + self.conv2 = base.conv2 + self.bn2 = base.bn2 + self.s2 = base.s2 + self.hb = nn.BatchNorm1d(C * SEGS) + self.fc = nn.Linear(C * SEGS, CLASSES) + + def load_from_ft(self, sd): + own = self.state_dict() + for k, v in sd.items(): + if k.startswith(("hb.", "fc.")): + continue + own[k].copy_(v) + # warm-start head: old halves -> replicate over 4 segments each + rep = SEGS // 2 + with torch.no_grad(): + self.fc.weight.copy_(sd["fc.weight"].repeat_interleave(1, 0) + .reshape(CLASSES, 2, C).repeat_interleave(rep, 1) + .reshape(CLASSES, C * SEGS) / rep) + self.fc.bias.copy_(sd["fc.bias"]) + for name in ("weight", "bias"): + getattr(self.hb, name).copy_( + sd[f"hb.{name}"].reshape(2, C).repeat_interleave(rep, 0).reshape(-1)) + self.hb.running_mean.copy_( + sd["hb.running_mean"].reshape(2, C).repeat_interleave(rep, 0).reshape(-1) / rep) + self.hb.running_var.copy_( + sd["hb.running_var"].reshape(2, C).repeat_interleave(rep, 0).reshape(-1) / (rep * rep)) + + def forward(self, x): + h = self.s0(self.bn0(self.stem(x))) + h = F.max_pool1d(h, 4) + h = self.s1(self.bn1(self.conv1(h))) + h = F.max_pool1d(h, 4) + h = self.s2(self.bn2(self.conv2(h))) + seg = h.shape[2] // SEGS + feats = torch.cat([(h[:, :, i * seg:(i + 1) * seg] + 1).sum(2) + for i in range(SEGS)], 1) / 2 + return self.fc(self.hb(feats)) + + +def evaluate(model, packed, labels, device, batch=128): + model.eval() + correct = np.zeros(CLASSES); total = np.zeros(CLASSES) + with torch.no_grad(): + for i in range(0, len(labels), batch): + idx = np.arange(i, min(i + batch, len(labels))) + x = unpack_batch(packed, idx).to(device) + pred = model(x).argmax(1).cpu().numpy() + lab = labels[idx] + for c in range(CLASSES): + m = lab == c + total[c] += m.sum() + correct[c] += (pred[m] == c).sum() + acc = correct.sum() / max(total.sum(), 1) + recalls = correct[total > 0] / total[total > 0] + return acc, recalls.mean() + + +def main(): + torch.manual_seed(SEED + 2) + np.random.seed(SEED + 2) + device = "mps" if torch.backends.mps.is_available() else "cpu" + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + tr_bits, tr_labels, tr_teacher = build_plane_cache("train", codes) + va_bits, va_labels, _ = build_plane_cache("valid", codes) + + model = BNN2Seg() + model.load_from_ft(torch.load(OUT / "best2ft.pt", map_location="cpu")) + model = model.to(device) + va = evaluate(model, va_bits, va_labels, device) + print(f"warm-start valid: {va[0]:.4f}/{va[1]:.4f}", flush=True) + ema = Ema(model) + + epochs = 12 + batch = 64 + opt = torch.optim.AdamW(model.parameters(), lr=2.5e-4, weight_decay=0.0) + sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs, eta_min=2e-5) + n = len(tr_labels) + best = 0.0 + tr_teacher_t = torch.from_numpy(np.asarray(tr_teacher)) + tr_labels_t = torch.from_numpy(np.asarray(tr_labels).astype(np.int64)) + + for epoch in range(epochs): + model.train() + perm = np.random.permutation(n) + t0 = time.time() + tot_loss = 0.0 + for i in range(0, n - batch + 1, batch): + sidx = np.sort(perm[i:i + batch]) + x = unpack_batch(tr_bits, sidx).to(device) + y = tr_labels_t[sidx].to(device) + t = tr_teacher_t[sidx].to(device) + logits = model(x) + ce = F.cross_entropy(logits, y, label_smoothing=0.05) + logt = torch.log(t.clamp_min(1e-9)) + kl = F.kl_div(F.log_softmax(logits / 1.5, dim=1), + F.softmax(logt / 1.5, dim=1), reduction="batchmean") + loss = 0.5 * ce + 0.5 * kl + opt.zero_grad() + loss.backward() + opt.step() + ema.update(model) + tot_loss += float(loss.detach()) + sched.step() + acc, macro = evaluate(model, va_bits, va_labels, device) + eacc, emacro = evaluate(ema.shadow, va_bits, va_labels, device) + dt = time.time() - t0 + print(f"seg {epoch}: loss {tot_loss / (n // batch):.4f} " + f"acc {acc:.4f}/{macro:.4f} ema {eacc:.4f}/{emacro:.4f} ({dt:.0f}s)", + flush=True) + for tag, a, m_ in (("raw", acc, model), ("ema", eacc, ema.shadow)): + if a > best: + best = a + torch.save(m_.state_dict(), OUT / "best2seg.pt") + print(f" saved {tag} {a:.4f}", flush=True) + print(f"best seg valid acc {best:.4f}") + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/finetune_short.py b/scripts/bnn/finetune_short.py new file mode 100644 index 0000000..67853ae --- /dev/null +++ b/scripts/bnn/finetune_short.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 +"""Short-file robustness fine-tune: continue training best2seg + best3 with +random-truncation augmentation (simulates small source files, which the +corpus underrepresents but betlang's fixture suite exercises heavily).""" +from __future__ import annotations + +import json +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes, encode_planes_bulk +from train_bnn import OUT, BITS, CLASSES, load_split +from train_bnn2 import BNN2, build_plane_cache, unpack_batch +from finetune_seg import BNN2Seg + +AUG_P = 0.4 +LMIN, LMAX = 16, 1024 + + +def build_short_valid(codes, windows, lengths): + path = OUT / "valid.short.planes.npy" + if path.exists(): + return np.load(path, mmap_mode="r") + n = len(lengths) + rng = np.random.default_rng(123) + Ls = np.exp(rng.uniform(np.log(LMIN), np.log(LMAX), n)).astype(np.int64) + Ls = np.maximum(np.minimum(Ls, np.asarray(lengths)), 8) + packed = np.zeros((n, 2, BITS // 8), dtype=np.uint8) + step = 2048 + for i in range(0, n, step): + j = min(i + step, n) + planes = encode_planes_bulk(np.asarray(windows[i:j]), Ls[i:j], codes, BITS) + packed[i:j] = np.packbits(planes, axis=2) + np.save(path, packed) + return np.load(path, mmap_mode="r") + + +def evaluate(model, packed, labels, device, batch=128): + model.eval() + correct = np.zeros(CLASSES); total = np.zeros(CLASSES) + with torch.no_grad(): + for i in range(0, len(labels), batch): + idx = np.arange(i, min(i + batch, len(labels))) + x = unpack_batch(packed, idx).to(device) + pred = model(x).argmax(1).cpu().numpy() + lab = labels[idx] + for c in range(CLASSES): + m = lab == c + total[c] += m.sum() + correct[c] += (pred[m] == c).sum() + acc = correct.sum() / max(total.sum(), 1) + recalls = correct[total > 0] / total[total > 0] + return acc, recalls.mean() + + +def finetune(tag, model, ckpt_out, codes, device, epochs=3, lr=3e-4): + tr_bits, tr_labels, tr_teacher = build_plane_cache("train", codes) + va_bits, va_labels, _ = build_plane_cache("valid", codes) + tr_windows, tr_lengths, _, _ = load_split("train") + va_windows, va_lengths, _, _ = load_split("valid") + va_short = build_short_valid(codes, va_windows, va_lengths) + + acc0, mac0 = evaluate(model, va_bits, va_labels, device) + sacc0, smac0 = evaluate(model, va_short, va_labels, device) + print(f"{tag} start: full {acc0:.4f}/{mac0:.4f} short {sacc0:.4f}/{smac0:.4f}", + flush=True) + + opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.0) + sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs, eta_min=3e-5) + tr_teacher_t = torch.from_numpy(np.asarray(tr_teacher)) + tr_labels_t = torch.from_numpy(np.asarray(tr_labels).astype(np.int64)) + n = len(tr_labels) + batch = 64 + rng = np.random.default_rng(hash(tag) % (2 ** 31)) + best = acc0 * 0.5 + sacc0 * 0.5 + + for epoch in range(epochs): + model.train() + perm = rng.permutation(n) + t0 = time.time() + tot = 0.0 + for i in range(0, n - batch + 1, batch): + sidx = np.sort(perm[i:i + batch]) + bits = np.unpackbits(tr_bits[sidx], axis=2).astype(np.float32) + aug = rng.random(batch) < AUG_P + if aug.any(): + widx = sidx[aug] + Ls = np.exp(rng.uniform(np.log(LMIN), np.log(LMAX), + aug.sum())).astype(np.int64) + Ls = np.maximum(np.minimum( + Ls, np.asarray(tr_lengths)[widx].astype(np.int64)), 8) + planes = encode_planes_bulk(np.asarray(tr_windows[widx]), + Ls, codes, BITS) + bits[aug] = planes.astype(np.float32) + x = torch.from_numpy(bits * 2 - 1).to(device) + y = tr_labels_t[sidx].to(device) + logits = model(x) + ce = F.cross_entropy(logits, y, label_smoothing=0.05) + keep = torch.from_numpy(~aug).to(device) + if keep.any(): + t = tr_teacher_t[sidx].to(device)[keep] + logt = torch.log(t.clamp_min(1e-9)) + kl = F.kl_div(F.log_softmax(logits[keep] / 1.5, dim=1), + F.softmax(logt / 1.5, dim=1), + reduction="batchmean") + else: + kl = torch.zeros((), device=device) + loss = 0.8 * ce + 0.2 * kl + opt.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + opt.step() + tot += float(loss.detach()) + sched.step() + acc, mac = evaluate(model, va_bits, va_labels, device) + sacc, smac = evaluate(model, va_short, va_labels, device) + score = acc * 0.5 + sacc * 0.5 + dt = time.time() - t0 + print(f"{tag} ep{epoch}: loss {tot / (n // batch):.4f} " + f"full {acc:.4f}/{mac:.4f} short {sacc:.4f}/{smac:.4f} " + f"({dt:.0f}s)", flush=True) + if score > best: + best = score + torch.save(model.state_dict(), ckpt_out) + print(f"{tag} saved (score {score:.4f})", flush=True) + + +def main(): + device = "mps" if torch.backends.mps.is_available() else "cpu" + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + + resume = len(sys.argv) > 1 and sys.argv[1] == "resume" + m0 = BNN2Seg() + m0.load_state_dict(torch.load( + OUT / ("best2seg_short.pt" if resume else "best2seg.pt"), + map_location="cpu")) + m0.to(device) + finetune("m0", m0, OUT / "best2seg_short.pt", codes, device, + epochs=5 if resume else 3, lr=1.5e-4 if resume else 3e-4) + + m1 = BNN2() + m1.load_state_dict(torch.load( + OUT / ("best3_short.pt" if resume else "best3.pt"), + map_location="cpu")) + m1.to(device) + finetune("m1", m1, OUT / "best3_short.pt", codes, device, + epochs=5 if resume else 3, lr=1.5e-4 if resume else 3e-4) + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/finetune_synth.py b/scripts/bnn/finetune_synth.py new file mode 100644 index 0000000..1427f52 --- /dev/null +++ b/scripts/bnn/finetune_synth.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +"""Fine-tune with synthetic hard-pair samples (repo hard_gen_* style) mixed +into short-truncation augmented batches.""" +from __future__ import annotations + +import json +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes, encode_planes_bulk +from train_bnn import OUT, BITS, load_split +from train_bnn2 import BNN2, build_plane_cache +from finetune_seg import BNN2Seg +from finetune_short import build_short_valid, evaluate + +AUG_P = 0.35 +SYNTH_PER_BATCH = 16 +LMIN, LMAX = 16, 1024 + + +def finetune(tag, model, ckpt_out, codes, device, epochs=3, lr=1.2e-4): + tr_bits, tr_labels, tr_teacher = build_plane_cache("train", codes) + va_bits, va_labels, _ = build_plane_cache("valid", codes) + tr_windows, tr_lengths, _, _ = load_split("train") + va_windows, va_lengths, _, _ = load_split("valid") + va_short = build_short_valid(codes, va_windows, va_lengths) + + sw = np.load(OUT / "synth_hard.windows.npy") + sl = np.load(OUT / "synth_hard.lengths.npy") + sy = np.load(OUT / "synth_hard.labels.npy") + print(f"{tag}: {len(sy)} synth rows", flush=True) + + acc0, mac0 = evaluate(model, va_bits, va_labels, device) + sacc0, smac0 = evaluate(model, va_short, va_labels, device) + print(f"{tag} start: full {acc0:.4f}/{mac0:.4f} short {sacc0:.4f}/{smac0:.4f}", + flush=True) + + opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.0) + sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs, eta_min=2e-5) + tr_teacher_t = torch.from_numpy(np.asarray(tr_teacher).copy()) + tr_labels_np = np.asarray(tr_labels).astype(np.int64) + n = len(tr_labels) + batch = 48 + rng = np.random.default_rng((hash(tag) + 7) % (2 ** 31)) + + for epoch in range(epochs): + model.train() + perm = rng.permutation(n) + t0 = time.time() + tot = 0.0 + nb = 0 + for i in range(0, n - batch + 1, batch): + sidx = np.sort(perm[i:i + batch]) + bits = np.unpackbits(tr_bits[sidx], axis=2).astype(np.float32) + aug = rng.random(batch) < AUG_P + if aug.any(): + widx = sidx[aug] + Ls = np.exp(rng.uniform(np.log(LMIN), np.log(LMAX), + aug.sum())).astype(np.int64) + Ls = np.maximum(np.minimum( + Ls, np.asarray(tr_lengths)[widx].astype(np.int64)), 8) + planes = encode_planes_bulk(np.asarray(tr_windows[widx]), + Ls, codes, BITS) + bits[aug] = planes.astype(np.float32) + sidx2 = rng.choice(len(sy), SYNTH_PER_BATCH, replace=False) + splanes = encode_planes_bulk(sw[sidx2], sl[sidx2].astype(np.int64), + codes, BITS).astype(np.float32) + x = torch.from_numpy( + np.concatenate([bits, splanes]) * 2 - 1).to(device) + y = torch.from_numpy( + np.concatenate([tr_labels_np[sidx], sy[sidx2]])).to(device) + logits = model(x) + ce = F.cross_entropy(logits, y, label_smoothing=0.05) + keep = torch.from_numpy( + np.concatenate([~aug, np.zeros(SYNTH_PER_BATCH, bool)])).to(device) + if keep.any(): + t = tr_teacher_t[torch.from_numpy(sidx)].to(device) + t = t[torch.from_numpy(~aug).to(device)] + logt = torch.log(t.clamp_min(1e-9)) + kl = F.kl_div(F.log_softmax(logits[keep] / 1.5, dim=1), + F.softmax(logt / 1.5, dim=1), + reduction="batchmean") + else: + kl = torch.zeros((), device=device) + loss = 0.85 * ce + 0.15 * kl + opt.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + opt.step() + tot += float(loss.detach()) + nb += 1 + sched.step() + acc, mac = evaluate(model, va_bits, va_labels, device) + sacc, smac = evaluate(model, va_short, va_labels, device) + dt = time.time() - t0 + print(f"{tag} ep{epoch}: loss {tot / nb:.4f} " + f"full {acc:.4f}/{mac:.4f} short {sacc:.4f}/{smac:.4f} " + f"({dt:.0f}s)", flush=True) + torch.save(model.state_dict(), ckpt_out.with_suffix(f".ep{epoch}.pt")) + + +def main(): + device = "mps" if torch.backends.mps.is_available() else "cpu" + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + + m0 = BNN2Seg() + m0.load_state_dict(torch.load(OUT / "best2seg_short.pt", map_location="cpu")) + m0.to(device) + finetune("m0s2", m0, OUT / "best2seg_synth.pt", codes, device) + + m1 = BNN2() + m1.load_state_dict(torch.load(OUT / "best3_short.pt", map_location="cpu")) + m1.to(device) + finetune("m1s2", m1, OUT / "best3_synth.pt", codes, device) + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/finetune_synth2.py b/scripts/bnn/finetune_synth2.py new file mode 100644 index 0000000..b091316 --- /dev/null +++ b/scripts/bnn/finetune_synth2.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +"""Second synth fine-tune pass (with yaml/markdown samples added).""" +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +import torch + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes +from train_bnn import OUT +from train_bnn2 import BNN2 +from finetune_seg import BNN2Seg +from finetune_synth import finetune + + +def main(): + device = "mps" if torch.backends.mps.is_available() else "cpu" + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + + m0 = BNN2Seg() + m0.load_state_dict(torch.load(OUT / "best2seg_synth.ep0.pt", map_location="cpu")) + m0.to(device) + finetune("m0s3", m0, OUT / "best2seg_synth2.pt", codes, device, + epochs=2, lr=8e-5) + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/finetune_synth3.py b/scripts/bnn/finetune_synth3.py new file mode 100644 index 0000000..40cfc81 --- /dev/null +++ b/scripts/bnn/finetune_synth3.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +"""Third synth pass (yaml_ci samples added) on the m0 synth2 checkpoint.""" +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +import torch + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes +from train_bnn import OUT +from finetune_seg import BNN2Seg +from finetune_synth import finetune + + +def main(): + device = "mps" if torch.backends.mps.is_available() else "cpu" + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + + m = BNN2Seg() + m.load_state_dict(torch.load(OUT / "best2seg_synth2.ep1.pt", + map_location="cpu")) + m.to(device) + finetune("m0s4", m, OUT / "best2seg_synth3.pt", codes, device, + epochs=1, lr=5e-5) + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/make_fixtures2.py b/scripts/bnn/make_fixtures2.py new file mode 100644 index 0000000..4dd089a --- /dev/null +++ b/scripts/bnn/make_fixtures2.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +"""Write the BBN2 artifact and matching Rust parity fixtures atomically +(from one in-memory model load, so trainers rewriting checkpoints on disk +cannot desynchronize them).""" +from __future__ import annotations + +import json +import struct +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).parent)) +from train_bnn import BITS, CLASSES, OUT, load_split +from export_bnn2 import MAGIC, IntModel, load_models +from shannon import canonical_codes, encode_planes_bulk + + +def main(): + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + ints = [IntModel(m, w) for m, w in load_models()] + + buf = bytearray() + buf += struct.pack("<5I", MAGIC, 2, BITS, CLASSES, len(ints)) + buf += lengths256.astype(np.uint8).tobytes() + for im in ints: + im.write(buf) + (OUT / "source-bnn2.bin").write_bytes(bytes(buf)) + print(f"wrote {OUT / 'source-bnn2.bin'} ({len(buf)} bytes)") + windows, lengths, _, _ = load_split("valid") + cases = [] + picks = [0, len(lengths) // 2, len(lengths) - 1] + for i in picks: + w = np.asarray(windows[i]) + n = int(lengths[i]) + cases.append(bytes(w[:n].tobytes())) + cases.append(b"fn main() {}") + out = [] + for data in cases: + arr = np.frombuffer(data, dtype=np.uint8)[None, :] + planes = encode_planes_bulk(arr, np.array([len(data)]), codes, BITS)[0] + logits = sum(im.logits(planes) for im in ints) / len(ints) + out.append({"window_hex": data.hex(), + "logits": [float(x) for x in logits]}) + (OUT / "parity_fixtures2.json").write_text(json.dumps(out, indent=1)) + print("wrote", OUT / "parity_fixtures2.json") + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/rebalance.py b/scripts/bnn/rebalance.py new file mode 100644 index 0000000..cdd66e6 --- /dev/null +++ b/scripts/bnn/rebalance.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 +"""Rebalance corpus splits to ~70/10/20 per label by moving files (deterministic).""" +import hashlib +import sys +from pathlib import Path + +root = Path(sys.argv[1]) +splits = ("train", "valid", "test") + +labels = sorted({d.name for s in splits for d in (root / s).iterdir() if d.is_dir()}) +for label in labels: + files = [] + for s in splits: + d = root / s / label + if d.is_dir(): + files.extend(d.iterdir()) + files = [f for f in files if f.is_file()] + n = len(files) + n_valid = max(1, round(n * 0.10)) if n >= 10 else 0 + n_test = max(1, round(n * 0.20)) if n >= 10 else 0 + keyed = sorted(files, key=lambda f: hashlib.sha1(f.name.encode()).hexdigest()) + assign = {} + for i, f in enumerate(keyed): + if i < n_valid: + assign[f] = "valid" + elif i < n_valid + n_test: + assign[f] = "test" + else: + assign[f] = "train" + moved = 0 + for f, target in assign.items(): + cur = f.parent.parent.name + if cur != target: + dest = root / target / label + dest.mkdir(parents=True, exist_ok=True) + f.rename(dest / f.name) + moved += 1 + print(f"{label}: n={n} moved={moved}") diff --git a/scripts/bnn/shannon.py b/scripts/bnn/shannon.py new file mode 100644 index 0000000..5cf68d3 --- /dev/null +++ b/scripts/bnn/shannon.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 +"""Canonical Shannon code over raw window bytes. + +Code lengths: L_b = clamp(ceil(-log2 p_b), 1, MAX_LEN) from the train-split +byte distribution (add-one smoothing). Ceil guarantees Kraft <= 1 before the +cap; the cap only shortens codes for symbols with p < 2^-MAX_LEN whose total +mass is far below the slack, so Kraft still holds. Codewords are assigned +canonically: sort symbols by (length, byte value), assign lexicographically +increasing codewords MSB-first. +""" +from __future__ import annotations + +import math + +import numpy as np + +MAX_LEN = 12 + + +def code_lengths(byte_counts: np.ndarray) -> np.ndarray: + counts = byte_counts.astype(np.float64) + 1.0 + probs = counts / counts.sum() + lengths = np.ceil(-np.log2(probs)).astype(np.int32) + lengths = np.clip(lengths, 1, MAX_LEN) + kraft = float(np.sum(2.0 ** (-lengths))) + while kraft > 1.0: # extend the longest-and-rarest codes until Kraft holds + order = np.lexsort((probs, -lengths)) + for idx in order: + if lengths[idx] < MAX_LEN: + kraft -= 2.0 ** (-lengths[idx]) + lengths[idx] += 1 + kraft += 2.0 ** (-lengths[idx]) + break + else: + raise AssertionError("cannot satisfy Kraft inequality") + return lengths + + +def canonical_codes(lengths: np.ndarray) -> np.ndarray: + """Returns [256, 2] array of (codeword, length); codeword MSB-first.""" + order = sorted(range(256), key=lambda b: (int(lengths[b]), b)) + codes = np.zeros((256, 2), dtype=np.uint32) + code = 0 + prev_len = 0 + for b in order: + length = int(lengths[b]) + code <<= length - prev_len + codes[b] = (code, length) + code += 1 + prev_len = length + return codes + + +def encode_bits(data: bytes, codes: np.ndarray, out_bits: int) -> np.ndarray: + """Encode into a fixed-size uint8 bit array (truncate/zero-pad).""" + bits = np.zeros(out_bits, dtype=np.uint8) + pos = 0 + for byte in data: + word, length = int(codes[byte, 0]), int(codes[byte, 1]) + if pos + length > out_bits: + break + for i in range(length - 1, -1, -1): + bits[pos] = (word >> i) & 1 + pos += 1 + return bits + + +def encode_bits_bulk(windows: np.ndarray, lengths_arr: np.ndarray, + codes: np.ndarray, out_bits: int) -> np.ndarray: + """Vectorized batch encoder: windows uint8 [N, W], lengths [N] -> [N, out_bits].""" + n, w = windows.shape + sym_len = codes[:, 1].astype(np.int64) + sym_code = codes[:, 0].astype(np.int64) + mask = np.arange(w)[None, :] < lengths_arr[:, None] + lens = np.where(mask, sym_len[windows], 0) + starts = np.cumsum(lens, axis=1) - lens + out = np.zeros((n, out_bits), dtype=np.uint8) + # expand each symbol's bits; loop over bit index within codeword (<= MAX_LEN) + for k in range(MAX_LEN): + sel = (lens > k) & (starts + lens <= out_bits) + rows, cols = np.nonzero(sel) + sym = windows[rows, cols] + bit = (sym_code[sym] >> (sym_len[sym] - 1 - k)) & 1 + out[rows, starts[rows, cols] + k] = bit + return out + + +def encode_planes_bulk(windows: np.ndarray, lengths_arr: np.ndarray, + codes: np.ndarray, out_bits: int) -> np.ndarray: + """Batch encoder emitting two bitplanes: [N, 2, out_bits]. + Plane 0: Shannon code bits. Plane 1: codeword-boundary markers (1 at the + first bit of each encoded codeword).""" + n, w = windows.shape + sym_len = codes[:, 1].astype(np.int64) + sym_code = codes[:, 0].astype(np.int64) + mask = np.arange(w)[None, :] < lengths_arr[:, None] + lens = np.where(mask, sym_len[windows], 0) + starts = np.cumsum(lens, axis=1) - lens + out = np.zeros((n, 2, out_bits), dtype=np.uint8) + fits = (lens > 0) & (starts + lens <= out_bits) + rows0, cols0 = np.nonzero(fits) + out[rows0, 1, starts[rows0, cols0]] = 1 + for k in range(MAX_LEN): + sel = (lens > k) & (starts + lens <= out_bits) + rows, cols = np.nonzero(sel) + sym = windows[rows, cols] + bit = (sym_code[sym] >> (sym_len[sym] - 1 - k)) & 1 + out[rows, 0, starts[rows, cols] + k] = bit + return out + + +def expected_bits(byte_counts: np.ndarray, lengths: np.ndarray) -> float: + probs = byte_counts / byte_counts.sum() + return float((probs * lengths).sum()) + + +if __name__ == "__main__": + rng = np.random.default_rng(0) + counts = rng.integers(0, 10000, 256).astype(np.int64) + lengths = code_lengths(counts) + codes = canonical_codes(lengths) + # prefix-freeness check + seen = set() + for b in range(256): + word, length = int(codes[b, 0]), int(codes[b, 1]) + s = format(word, f"0{length}b") + for other in seen: + assert not s.startswith(other) and not other.startswith(s) + seen.add(s) + data = bytes(rng.integers(0, 256, 500).astype(np.uint8)) + single = encode_bits(data, codes, 4096) + bulk = encode_bits_bulk( + np.frombuffer(data, dtype=np.uint8)[None, :].copy(), + np.array([len(data)]), codes, 4096) + assert np.array_equal(single, bulk[0]) + print("shannon self-test ok; mean bits/byte:", + expected_bits(counts.astype(np.float64), lengths)) diff --git a/scripts/bnn/topup.py b/scripts/bnn/topup.py new file mode 100644 index 0000000..3cb812c --- /dev/null +++ b/scripts/bnn/topup.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +"""Top-up harvester for starved labels (ini, cobol). Splits files across +train/valid/test by path hash; dedupes by content sha1 against whole corpus.""" +import hashlib +import io +import sys +import tarfile +import urllib.request +from pathlib import Path + +root = Path(sys.argv[1]) # corpus/files +MIN_BYTES, MAX_BYTES = 64, 262144 + +SOURCES = { + "cobol": ([ + "eclipse-che4z/che-che4z-lsp-for-cobol", "OCamlPro/superbol-studio-oss", + "GnuCOBOL/GnuCOBOL", "simonsobol? no", "MicroFocus/CICS-Banking-Sample-Application-CBSA", + ], (".cbl", ".cob", ".CBL", ".COB", ".cpy"), 4200), +} + +seen = set() +for split_dir in root.glob("*/"): + for label_dir in split_dir.glob("*/"): + for f in label_dir.iterdir(): + seen.add(hashlib.sha1(f.read_bytes()).digest()) +print(f"seen {len(seen)} existing files", flush=True) + +for label, (repos, exts, cap) in SOURCES.items(): + counts = {s: sum(1 for _ in (root / s / label).glob("*")) for s in ("train", "valid", "test")} + added_total = 0 + for repo in repos: + total = sum(counts.values()) + if total >= cap: + break + url = f"https://codeload.github.com/{repo}/tar.gz/HEAD" + print(f"{label} {repo}: downloading", flush=True) + added = 0 + try: + with urllib.request.urlopen(url) as response: + stream = io.BufferedReader(response, buffer_size=1 << 20) + with tarfile.open(fileobj=stream, mode="r|gz") as tar: + for member in tar: + if not member.isfile() or member.size > MAX_BYTES: + continue + if not member.name.endswith(exts): + continue + fobj = tar.extractfile(member) + if fobj is None: + continue + content = fobj.read() + if not (MIN_BYTES <= len(content) <= MAX_BYTES): + continue + if len(content.lstrip()) < MIN_BYTES: + continue + digest = hashlib.sha1(content).digest() + if digest in seen: + continue + seen.add(digest) + bucket = int(hashlib.sha1(member.name.encode()).hexdigest(), 16) % 10 + split = "valid" if bucket == 0 else ("test" if bucket <= 2 else "train") + dest = root / split / label + dest.mkdir(parents=True, exist_ok=True) + (dest / f"topup_{digest.hex()[:16]}").write_bytes(content) + counts[split] += 1 + added += 1 + if sum(counts.values()) >= cap: + break + except Exception as err: + print(f"{label} {repo}: FAILED {err}", flush=True) + continue + added_total += added + print(f"{label} {repo}: added {added}", flush=True) + print(f"{label}: total added {added_total}, counts {counts}", flush=True) diff --git a/scripts/bnn/train_bnn.py b/scripts/bnn/train_bnn.py new file mode 100644 index 0000000..8ab1608 --- /dev/null +++ b/scripts/bnn/train_bnn.py @@ -0,0 +1,256 @@ +#!/usr/bin/env python3 +"""Binary CNN (XNOR/popcount-compatible) over Shannon-coded raw window bytes. + +Pipeline: raw 2048-byte Magika window -> canonical Shannon code (lengths from +train byte histogram) -> fixed 16384-bit input -> binary conv stack -> integer +popcount logits. All inference-time ops are binary/integer: XNOR + popcount + +integer thresholds (folded BN) + integer classifier bias. +""" +from __future__ import annotations + +import json +import math +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import code_lengths, canonical_codes, encode_bits_bulk + +CACHE = Path.home() / "work/cache" +OUT = Path.home() / "work/bnn" +BITS = 16384 +WINDOW = 2048 +CLASSES = 48 +STEM_K = 32 +STEM_STRIDE = 4 +C = 192 +SEED = 2 + + +def load_split(split: str): + meta = json.loads((CACHE / f"{split}.json").read_text()) + n = meta["count"] + windows = np.memmap(CACHE / f"{split}.windows.mmap", dtype=np.uint8, mode="r") + windows = windows.reshape(-1, WINDOW)[:n] + lengths = np.memmap(CACHE / f"{split}.lengths.mmap", dtype=np.int32, mode="r")[:n] + labels = np.memmap(CACHE / f"{split}.labels.mmap", dtype=np.int16, mode="r")[:n] + teacher = np.memmap(CACHE / f"{split}.teacher.mmap", dtype=np.float32, mode="r") + teacher = teacher.reshape(-1, CLASSES)[:n] + return windows, lengths, labels, teacher + + +def build_bit_cache(split: str, codes: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + packed_path = OUT / f"{split}.bits.npy" + windows, lengths, labels, teacher = load_split(split) + if packed_path.exists(): + packed = np.load(packed_path, mmap_mode="r") + return packed, np.asarray(labels), np.asarray(teacher) + n = len(lengths) + packed = np.zeros((n, BITS // 8), dtype=np.uint8) + step = 2048 + for i in range(0, n, step): + j = min(i + step, n) + bits = encode_bits_bulk(np.asarray(windows[i:j]), np.asarray(lengths[i:j]), + codes, BITS) + packed[i:j] = np.packbits(bits, axis=1) + print(f"{split} bits: {j}/{n}", flush=True) + np.save(packed_path, packed) + return np.load(packed_path, mmap_mode="r"), np.asarray(labels), np.asarray(teacher) + + +class SignSTE(torch.autograd.Function): + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x) + return torch.where(x >= 0, torch.ones_like(x), -torch.ones_like(x)) + + @staticmethod + def backward(ctx, grad): + (x,) = ctx.saved_tensors + return grad * (x.abs() <= 1).float() + + +def binarize(x): + return SignSTE.apply(x) + + +class RSign(nn.Module): + """Learnable-threshold sign; sparse init mimics ReLU pattern detectors. + Exports as: integer accumulator >= integer threshold.""" + + def __init__(self, c, bias=1.5): + super().__init__() + self.b = nn.Parameter(torch.full((1, c, 1), bias)) + + def forward(self, z): + return binarize(z - self.b) + + +class MaybeBinConv1d(nn.Module): + """Float conv whose weights are progressively annealed to binary + (XNOR-Net scale folds into the following BN threshold at export).""" + + def __init__(self, cin, cout, k, stride=1, padding=0): + super().__init__() + conv = nn.Conv1d(cin, cout, k, stride=stride, padding=padding, bias=False) + self.weight = conv.weight + self.stride = stride + self.padding = padding + self.binw = 0.0 + + def forward(self, x): + w = self.weight + if self.binw > 0: + wb = binarize(w) * w.abs().mean(dim=(1, 2), keepdim=True) + w = (1 - self.binw) * w + self.binw * wb + return F.conv1d(x, w, stride=self.stride, padding=self.padding) + + +class BNN(nn.Module): + """Binary CNN over the Shannon bitstream. + + Inference graph is entirely bitwise/integer: XNOR+popcount convs, + integer thresholds (folded BN+RSign), OR-pooling (word-OR), popcount + head counts, and a fixed-point integer classifier.""" + + def __init__(self): + super().__init__() + self.stem = MaybeBinConv1d(1, C, STEM_K, stride=STEM_STRIDE, + padding=STEM_K // 2) + self.bn0 = nn.BatchNorm1d(C) + self.s0 = RSign(C) + self.conv1 = MaybeBinConv1d(C, C, 3, padding=1) + self.bn1 = nn.BatchNorm1d(C) + self.s1 = RSign(C) + self.conv2 = MaybeBinConv1d(C, C, 3, padding=1) + self.bn2 = nn.BatchNorm1d(C) + self.s2 = RSign(C) + self.hb = nn.BatchNorm1d(C * 2) + self.fc = nn.Linear(C * 2, CLASSES) + + def set_binw(self, v: float): + for m in (self.stem, self.conv1, self.conv2): + m.binw = v + + def forward(self, x): # x: [B, 1, BITS] in {-1,+1} + h = self.s0(self.bn0(self.stem(x))) # [B, C, 4096] + h = F.max_pool1d(h, 4) # OR-pool -> 1024 + h = self.s1(self.bn1(self.conv1(h))) + h = F.max_pool1d(h, 4) # 256 + h = self.s2(self.bn2(self.conv2(h))) # 256 + half = h.shape[2] // 2 + feats = torch.cat([(h[:, :, :half] + 1).sum(2), + (h[:, :, half:] + 1).sum(2)], 1) / 2 # [B, 2C] ints + return self.fc(self.hb(feats)) + + +def unpack_batch(packed: np.ndarray, idx: np.ndarray) -> torch.Tensor: + bits = np.unpackbits(packed[idx], axis=1) + return torch.from_numpy(bits.astype(np.float32) * 2 - 1).unsqueeze(1) + + +def evaluate(model, packed, labels, device, batch=256): + model.eval() + correct = np.zeros(CLASSES); total = np.zeros(CLASSES) + with torch.no_grad(): + for i in range(0, len(labels), batch): + idx = np.arange(i, min(i + batch, len(labels))) + x = unpack_batch(packed, idx).to(device) + pred = model(x).argmax(1).cpu().numpy() + lab = labels[idx] + for c in range(CLASSES): + m = lab == c + total[c] += m.sum() + correct[c] += (pred[m] == c).sum() + acc = correct.sum() / max(total.sum(), 1) + recalls = correct[total > 0] / total[total > 0] + return acc, recalls.mean() + + +def main(): + torch.manual_seed(SEED) + np.random.seed(SEED) + OUT.mkdir(exist_ok=True) + device = "mps" if torch.backends.mps.is_available() else "cpu" + + # Shannon codebook from train byte histogram + cb_path = OUT / "codebook.json" + if cb_path.exists(): + lengths256 = np.array(json.loads(cb_path.read_text())["lengths"], dtype=np.int32) + else: + windows, lengths, _, _ = load_split("train") + hist = np.zeros(256, dtype=np.int64) + step = 4096 + for i in range(0, len(lengths), step): + j = min(i + step, len(lengths)) + w = np.asarray(windows[i:j]) + mask = np.arange(WINDOW)[None, :] < np.asarray(lengths[i:j])[:, None] + hist += np.bincount(w[mask], minlength=256) + lengths256 = code_lengths(hist) + cb_path.write_text(json.dumps({"lengths": lengths256.tolist(), + "hist": hist.tolist()})) + codes = canonical_codes(lengths256) + print("mean code bits/byte:", flush=True) + + tr_bits, tr_labels, tr_teacher = build_bit_cache("train", codes) + va_bits, va_labels, _ = build_bit_cache("valid", codes) + + model = BNN().to(device) + float_epochs = 24 + anneal_epochs = 10 + bin_epochs = 14 + epochs = float_epochs + anneal_epochs + bin_epochs + batch = 96 + opt = torch.optim.AdamW(model.parameters(), lr=1.2e-3, weight_decay=1e-5) + sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs, eta_min=5e-5) + n = len(tr_labels) + best = 0.0 + tr_teacher_t = torch.from_numpy(np.asarray(tr_teacher)) + tr_labels_t = torch.from_numpy(np.asarray(tr_labels).astype(np.int64)) + + for epoch in range(epochs): + if epoch < float_epochs: + model.set_binw(0.0) + elif epoch < float_epochs + anneal_epochs: + model.set_binw((epoch - float_epochs + 1) / anneal_epochs) + else: + model.set_binw(1.0) + model.train() + perm = np.random.permutation(n) + t0 = time.time() + tot_loss = 0.0 + for i in range(0, n - batch + 1, batch): + sidx = np.sort(perm[i:i + batch]) + x = unpack_batch(tr_bits, sidx).to(device) + y = tr_labels_t[sidx].to(device) + t = tr_teacher_t[sidx].to(device) + logits = model(x) + ce = F.cross_entropy(logits, y, label_smoothing=0.05) + logt = torch.log(t.clamp_min(1e-9)) + kl = F.kl_div(F.log_softmax(logits / 1.5, dim=1), + F.softmax(logt / 1.5, dim=1), reduction="batchmean") + loss = 0.5 * ce + 0.5 * kl + opt.zero_grad() + loss.backward() + opt.step() + tot_loss += float(loss.detach()) + sched.step() + acc, macro = evaluate(model, va_bits, va_labels, device) + dt = time.time() - t0 + binw = model.stem.binw + print(f"epoch {epoch}: binw {binw:.2f} loss {tot_loss / (n // batch):.4f} " + f"valid acc {acc:.4f} macro {macro:.4f} ({dt:.0f}s)", flush=True) + if binw >= 1.0 and acc > best: + best = acc + torch.save(model.state_dict(), OUT / "best.pt") + print(f"best fully-binary valid acc {best:.4f}") + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/train_bnn2.py b/scripts/bnn/train_bnn2.py new file mode 100644 index 0000000..637e991 --- /dev/null +++ b/scripts/bnn/train_bnn2.py @@ -0,0 +1,202 @@ +#!/usr/bin/env python3 +"""Binary CNN v2 over Shannon-coded raw window bytes. + +Changes vs v1: binary weights from epoch 0 (no annealing shock), two input +bitplanes (code bits + codeword-boundary markers), multi-basis binary stem +(ABC-Net style: sum of M XNOR-popcount branches with per-channel fixed-point +scales), wider C. Inference remains entirely bitwise/integer. +""" +from __future__ import annotations + +import json +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes, encode_planes_bulk +from train_bnn import CACHE, OUT, BITS, WINDOW, CLASSES, SEED, load_split + +STEM_K = 32 +STEM_STRIDE = 4 +C = 256 +STEM_M = 3 + + +def build_plane_cache(split: str, codes: np.ndarray): + packed_path = OUT / f"{split}.planes.npy" + windows, lengths, labels, teacher = load_split(split) + if packed_path.exists(): + packed = np.load(packed_path, mmap_mode="r") + return packed, np.asarray(labels), np.asarray(teacher) + n = len(lengths) + packed = np.zeros((n, 2, BITS // 8), dtype=np.uint8) + step = 2048 + for i in range(0, n, step): + j = min(i + step, n) + planes = encode_planes_bulk(np.asarray(windows[i:j]), + np.asarray(lengths[i:j]), codes, BITS) + packed[i:j] = np.packbits(planes, axis=2) + print(f"{split} planes: {j}/{n}", flush=True) + np.save(packed_path, packed) + return np.load(packed_path, mmap_mode="r"), np.asarray(labels), np.asarray(teacher) + + +class SignSTE(torch.autograd.Function): + @staticmethod + def forward(ctx, x): + ctx.save_for_backward(x) + return torch.where(x >= 0, torch.ones_like(x), -torch.ones_like(x)) + + @staticmethod + def backward(ctx, grad): + (x,) = ctx.saved_tensors + return grad * (x.abs() <= 1).float() + + +def binarize(x): + return SignSTE.apply(x) + + +class RSign(nn.Module): + def __init__(self, c, bias=1.5): + super().__init__() + self.b = nn.Parameter(torch.full((1, c, 1), bias)) + + def forward(self, z): + return binarize(z - self.b) + + +class BinConvM(nn.Module): + """M-basis binary conv: z = sum_m alpha_m * conv(x, sign(W_m)). + Exports as M XNOR-popcount passes + fixed-point per-channel combine.""" + + def __init__(self, cin, cout, k, m=1, stride=1, padding=0): + super().__init__() + self.latent = nn.Parameter(torch.empty(m, cout, cin, k)) + nn.init.kaiming_normal_(self.latent.view(m * cout, cin, k)) + shifts = torch.linspace(-0.6, 0.6, m).view(m, 1) if m > 1 else torch.zeros(1, 1) + self.shift = nn.Parameter(shifts.clone()) # [m, 1] per-basis latent shift + self.alpha = nn.Parameter(torch.full((m, cout), 1.0 / m)) + self.stride = stride + self.padding = padding + self.m = m + self.cout = cout + + def forward(self, x): + m, cout, cin, k = self.latent.shape + std = self.latent.std(dim=(1, 2, 3), keepdim=True).clamp_min(1e-6) + wb = binarize(self.latent - self.shift.view(m, 1, 1, 1) * std) + w = (wb * self.alpha.view(m, cout, 1, 1)).sum(0) + return F.conv1d(x, w, stride=self.stride, padding=self.padding) + + +class BNN2(nn.Module): + def __init__(self): + super().__init__() + self.stem = BinConvM(2, C, STEM_K, m=STEM_M, stride=STEM_STRIDE, + padding=STEM_K // 2) + self.bn0 = nn.BatchNorm1d(C) + self.s0 = RSign(C) + self.conv1 = BinConvM(C, C, 3, m=1, padding=1) + self.bn1 = nn.BatchNorm1d(C) + self.s1 = RSign(C) + self.conv2 = BinConvM(C, C, 3, m=1, padding=1) + self.bn2 = nn.BatchNorm1d(C) + self.s2 = RSign(C) + self.hb = nn.BatchNorm1d(C * 2) + self.fc = nn.Linear(C * 2, CLASSES) + + def forward(self, x): # x: [B, 2, BITS] in {-1,+1} + h = self.s0(self.bn0(self.stem(x))) + h = F.max_pool1d(h, 4) + h = self.s1(self.bn1(self.conv1(h))) + h = F.max_pool1d(h, 4) + h = self.s2(self.bn2(self.conv2(h))) + half = h.shape[2] // 2 + feats = torch.cat([(h[:, :, :half] + 1).sum(2), + (h[:, :, half:] + 1).sum(2)], 1) / 2 + return self.fc(self.hb(feats)) + + +def unpack_batch(packed: np.ndarray, idx: np.ndarray) -> torch.Tensor: + bits = np.unpackbits(packed[idx], axis=2) + return torch.from_numpy(bits.astype(np.float32) * 2 - 1) + + +def evaluate(model, packed, labels, device, batch=128): + model.eval() + correct = np.zeros(CLASSES); total = np.zeros(CLASSES) + with torch.no_grad(): + for i in range(0, len(labels), batch): + idx = np.arange(i, min(i + batch, len(labels))) + x = unpack_batch(packed, idx).to(device) + pred = model(x).argmax(1).cpu().numpy() + lab = labels[idx] + for c in range(CLASSES): + m = lab == c + total[c] += m.sum() + correct[c] += (pred[m] == c).sum() + acc = correct.sum() / max(total.sum(), 1) + recalls = correct[total > 0] / total[total > 0] + return acc, recalls.mean() + + +def main(): + torch.manual_seed(SEED) + np.random.seed(SEED) + device = "mps" if torch.backends.mps.is_available() else "cpu" + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + tr_bits, tr_labels, tr_teacher = build_plane_cache("train", codes) + va_bits, va_labels, _ = build_plane_cache("valid", codes) + + model = BNN2().to(device) + epochs = 40 + batch = 64 + opt = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=0.0) + sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs, eta_min=5e-5) + n = len(tr_labels) + best = 0.0 + tr_teacher_t = torch.from_numpy(np.asarray(tr_teacher)) + tr_labels_t = torch.from_numpy(np.asarray(tr_labels).astype(np.int64)) + + for epoch in range(epochs): + model.train() + perm = np.random.permutation(n) + t0 = time.time() + tot_loss = 0.0 + for i in range(0, n - batch + 1, batch): + sidx = np.sort(perm[i:i + batch]) + x = unpack_batch(tr_bits, sidx).to(device) + y = tr_labels_t[sidx].to(device) + t = tr_teacher_t[sidx].to(device) + logits = model(x) + ce = F.cross_entropy(logits, y, label_smoothing=0.05) + logt = torch.log(t.clamp_min(1e-9)) + kl = F.kl_div(F.log_softmax(logits / 1.5, dim=1), + F.softmax(logt / 1.5, dim=1), reduction="batchmean") + loss = 0.5 * ce + 0.5 * kl + opt.zero_grad() + loss.backward() + opt.step() + tot_loss += float(loss.detach()) + sched.step() + acc, macro = evaluate(model, va_bits, va_labels, device) + dt = time.time() - t0 + print(f"epoch {epoch}: loss {tot_loss / (n // batch):.4f} " + f"valid acc {acc:.4f} macro {macro:.4f} ({dt:.0f}s)", flush=True) + if acc > best: + best = acc + torch.save(model.state_dict(), OUT / "best2.pt") + print(f"best fully-binary valid acc {best:.4f}") + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/train_bnn3.py b/scripts/bnn/train_bnn3.py new file mode 100644 index 0000000..8c6770f --- /dev/null +++ b/scripts/bnn/train_bnn3.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Complementary binary model for ensembling: proven BNN2 architecture, +different seed, grad clipping, CE-heavy final epochs.""" +from __future__ import annotations + +import json +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes +from train_bnn import OUT, CLASSES +from train_bnn2 import BNN2, build_plane_cache, unpack_batch + +SEED3 = 7 + + +def evaluate(model, packed, labels, device, batch=128): + model.eval() + correct = np.zeros(CLASSES); total = np.zeros(CLASSES) + with torch.no_grad(): + for i in range(0, len(labels), batch): + idx = np.arange(i, min(i + batch, len(labels))) + x = unpack_batch(packed, idx).to(device) + pred = model(x).argmax(1).cpu().numpy() + lab = labels[idx] + for c in range(CLASSES): + m = lab == c + total[c] += m.sum() + correct[c] += (pred[m] == c).sum() + acc = correct.sum() / max(total.sum(), 1) + recalls = correct[total > 0] / total[total > 0] + return acc, recalls.mean() + + +def main(): + torch.manual_seed(SEED3) + np.random.seed(SEED3) + device = "mps" if torch.backends.mps.is_available() else "cpu" + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + tr_bits, tr_labels, tr_teacher = build_plane_cache("train", codes) + va_bits, va_labels, _ = build_plane_cache("valid", codes) + + model = BNN2().to(device) + epochs = 40 + batch = 64 + start_epoch = 0 + resume = OUT / "best3.pt" + if resume.exists(): + model.load_state_dict(torch.load(resume, map_location=device)) + start_epoch = int(sys.argv[1]) if len(sys.argv) > 1 else 11 + opt = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=0.0) + sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs, eta_min=3e-5) + for _ in range(start_epoch): + sched.step() + n = len(tr_labels) + best = 0.0 + if start_epoch: + acc, macro = evaluate(model, va_bits, va_labels, device) + best = acc + print(f"m3 resume @ {start_epoch}: valid acc {acc:.4f} macro {macro:.4f}", + flush=True) + tr_teacher_t = torch.from_numpy(np.asarray(tr_teacher)) + tr_labels_t = torch.from_numpy(np.asarray(tr_labels).astype(np.int64)) + + for epoch in range(start_epoch, epochs): + model.train() + perm = np.random.permutation(n) + t0 = time.time() + tot_loss = 0.0 + for i in range(0, n - batch + 1, batch): + sidx = np.sort(perm[i:i + batch]) + x = unpack_batch(tr_bits, sidx).to(device) + y = tr_labels_t[sidx].to(device) + t = tr_teacher_t[sidx].to(device) + logits = model(x) + ce = F.cross_entropy(logits, y, label_smoothing=0.05) + logt = torch.log(t.clamp_min(1e-9)) + kl = F.kl_div(F.log_softmax(logits / 1.5, dim=1), + F.softmax(logt / 1.5, dim=1), reduction="batchmean") + w = 0.5 if epoch < 30 else 0.1 # CE-heavy at the end (teacher ceiling) + loss = (1 - w) * ce + w * kl if epoch >= 30 else 0.5 * ce + 0.5 * kl + opt.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) + opt.step() + tot_loss += float(loss.detach()) + sched.step() + acc, macro = evaluate(model, va_bits, va_labels, device) + dt = time.time() - t0 + print(f"m3 {epoch}: loss {tot_loss / (n // batch):.4f} " + f"valid acc {acc:.4f} macro {macro:.4f} ({dt:.0f}s)", flush=True) + if acc > best: + best = acc + torch.save(model.state_dict(), OUT / "best3.pt") + print(f"best m3 valid acc {best:.4f}") + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/weight_search.py b/scripts/bnn/weight_search.py new file mode 100644 index 0000000..a4c889b --- /dev/null +++ b/scripts/bnn/weight_search.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +"""Grid-search per-model weights for the 5-model ensemble on fixtures+edge.""" +from __future__ import annotations + +import itertools +import json +import sys +from pathlib import Path + +import numpy as np +import torch + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes, encode_planes_bulk +from train_bnn import OUT, BITS +from train_bnn2 import BNN2 +from finetune_seg import BNN2Seg +from eval_fixtures import build_window, LABELS, REPO +from eval_edge import CASES + +CKPTS = ["best2seg.pt", "best3.pt", "best2seg_synth.ep0.pt", + "best2seg_synth3.ep0.pt", "best3_synth2.ep0.pt"] +ALIAS = {"Gemfile": "gemfile", "bash": "shell", "c-sharp": "cs", + "commonlisp": "lisp", "objc": "objectivec", "vb": "vba"} + + +def main(): + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + models = [] + for ck in CKPTS: + cls = BNN2Seg if "2seg" in ck else BNN2 + m = cls() + m.load_state_dict(torch.load(OUT / ck, map_location="cpu")) + m.eval() + models.append(m) + + inputs = [] + fix_expected = [] + for f in sorted((REPO / "tests/fixtures/languages").iterdir()): + w = build_window(f.read_bytes()) + arr = np.frombuffer(w[0], dtype=np.uint8)[None, :] + inputs.append(encode_planes_bulk(arr, np.array([w[1]]), codes, BITS)) + fix_expected.append(ALIAS.get(f.stem, f.stem)) + edge = [] + for name, text, want in CASES: + w = build_window(text.encode()) + arr = np.frombuffer(w[0], dtype=np.uint8)[None, :] + edge.append(encode_planes_bulk(arr, np.array([w[1]]), codes, BITS)) + demo = 'pub fn greet(name: &str) -> String {\n format!("hello, {name}")\n}\n' + w = build_window(demo.encode()) + arr = np.frombuffer(w[0], dtype=np.uint8)[None, :] + edge.append(encode_planes_bulk(arr, np.array([w[1]]), codes, BITS)) + + x = torch.from_numpy( + np.concatenate(inputs + edge).astype(np.float32) * 2 - 1) + with torch.no_grad(): + L = np.stack([m(x).numpy() for m in models]) + nfix = len(fix_expected) + yi, mi = LABELS.index("yaml"), LABELS.index("markdown") + ri = LABELS.index("rust") + + grid = [0.7, 1.0, 1.3, 1.6] + best = [] + for ws in itertools.product(grid, repeat=5): + w = np.array(ws)[:, None, None] + tot = (L * w).sum(0) + wrong = sum(LABELS[int(tot[i].argmax())] != fix_expected[i] + for i in range(nfix)) + if wrong: + continue + p = np.exp(tot / w.sum()) + p /= p.sum(1, keepdims=True) + e = p[nfix:] + ok = (e[0].argmax() == mi and e[1].argmax() == mi + and e[2].argmax() == yi + and set(np.argsort(-e[3])[:2]) == {yi, mi} + and set(np.argsort(-e[4])[:2]) == {yi, mi} + and e[4].max() < 0.9 + and e[5].argmax() == ri) + if ok: + best.append(ws) + print(f"passing weightings: {len(best)}") + for ws in best[:40]: + print(ws) + + +if __name__ == "__main__": + main() diff --git a/scripts/bnn/weight_valid.py b/scripts/bnn/weight_valid.py new file mode 100644 index 0000000..d26c1ce --- /dev/null +++ b/scripts/bnn/weight_valid.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 +"""Score fixture/edge-passing weightings on full-valid and short-valid.""" +from __future__ import annotations + +import itertools +import json +import sys +from pathlib import Path + +import numpy as np +import torch + +sys.path.insert(0, str(Path(__file__).parent)) +from shannon import canonical_codes +from train_bnn import OUT, load_split +from train_bnn2 import build_plane_cache +from finetune_short import build_short_valid +from eval_combos import all_logits, metrics +from train_bnn2 import BNN2 +from finetune_seg import BNN2Seg +from weight_search import CKPTS + +device = "mps" if torch.backends.mps.is_available() else "cpu" + + +def main(): + passing = [tuple(map(float, line.strip("()\n").split(","))) + for line in Path(sys.argv[1]).read_text().splitlines() + if line.startswith("(")] + lengths256 = np.array(json.loads((OUT / "codebook.json").read_text())["lengths"], + dtype=np.int32) + codes = canonical_codes(lengths256) + va_bits, va_labels, _ = build_plane_cache("valid", codes) + va_windows, va_lengths, _, _ = load_split("valid") + va_short = build_short_valid(codes, va_windows, va_lengths) + n = len(va_labels) + + models = [] + for ck in CKPTS: + cls = BNN2Seg if "2seg" in ck else BNN2 + m = cls() + m.load_state_dict(torch.load(OUT / ck, map_location="cpu")) + m.to(device) + models.append(m) + lf = np.stack([all_logits(m, va_bits, n) for m in models]) + ls = np.stack([all_logits(m, va_short, n) for m in models]) + + results = [] + for ws in passing: + w = np.array(ws)[:, None, None] + fa, fm = metrics((lf * w).sum(0), va_labels) + sa, sm = metrics((ls * w).sum(0), va_labels) + results.append((fa + fm + 0.5 * (sa + sm), ws, fa, fm, sa, sm)) + results.sort(reverse=True) + for score, ws, fa, fm, sa, sm in results[:15]: + print(f"{ws} full {fa:.4f}/{fm:.4f} short {sa:.4f}/{sm:.4f}") + + +if __name__ == "__main__": + main() diff --git a/src/model.rs b/src/model.rs index eeefd7e..d9c3af7 100644 --- a/src/model.rs +++ b/src/model.rs @@ -1,42 +1,28 @@ -//! Inference for the wordseq student. +//! Inference for the Shannon-bit binary CNN detector. //! -//! Loads `assets/magika/source-student-q4.bin` (weights-only MSQ1 export) and -//! runs a forward pass: byte-window tokenization -> word-unit tokenization -//! -> HashEmbedding lookup (K=3) -> 3 conv stages with max-pool -> global -//! max+avg pool -> 2 dense layers -> 48-class softmax logits. -//! -//! Model architecture: `wordseq-b1024-k3-m2048-tiny-3conv-hidden` -//! - 1024-bin x 24-dim shared HashEmbedding table (4-bit, ~12 KB) -//! - QConv1D k=7 24->64ch (2-bit ternary) -//! - MaxPool(4) -//! - QConv1D k=5 64->128ch (2-bit) -//! - MaxPool(2) -//! - QConv1D k=3 128->128ch (2-bit) -//! - GlobalMax + GlobalAvg -> 256-dim -//! - QDense 256->96 (2-bit) + GELU -//! - QDense 96->48 (4-bit) - -mod activation; +//! Loads `assets/bnn/source-bnn2.bin` (packed "BBN2" export) and runs a +//! forward pass: raw Magika byte window -> canonical Shannon coding into two +//! bitplanes (code bits + codeword boundaries) -> binary conv stages +//! (XNOR + popcount, fixed-point combine, integer thresholds, OR-pooling) +//! -> segmented popcount head -> 48-class softmax logits, averaged over a +//! small ensemble of binary models. + +mod bnn; mod constants; -mod embedded; -mod layers; -mod reader; -mod runtime; #[cfg(test)] mod tests; -mod tokenizer; mod window; -use self::{constants::CLASSES, runtime::Model, window::build_window}; +use self::{bnn::Bnn, constants::CLASSES, window::build_window}; use crate::{Detection, Language}; pub(crate) fn detect(source: &[u8]) -> Detection { let Some(window) = build_window(source) else { return Detection::from_predictions(Vec::new()); }; - let model = Model::get(); - let units = model.tokenize_units(&window); - let logits = model.logits(&units); + let model = Bnn::get(); + let planes = model.encode_planes(window.bytes()); + let logits = model.logits(&planes); detection_from_logits(&logits) } diff --git a/src/model/activation.rs b/src/model/activation.rs deleted file mode 100644 index a0f9ffb..0000000 --- a/src/model/activation.rs +++ /dev/null @@ -1,148 +0,0 @@ -use fearless_simd::{Simd, SimdBase, SimdFloat, f32x4}; - -#[inline(always)] -pub(crate) fn gelu(x: f32) -> f32 { - // Tanh-form GELU (Hendrycks & Gimpel), matching how the student was trained. - // 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x³))) - 0.5 * x * (1.0 + tanh_approx(0.797_884_6 * (x + 0.044_715 * x * x * x))) -} - -/// [7/6] Padé approximation of tanh, accurate to ~1e-4 over the clamped range. -/// Replaces the libm `tanh` call so the GELU step auto-vectorizes (no function call). -#[inline(always)] -fn tanh_approx(x: f32) -> f32 { - let x = x.clamp(-5.0, 5.0); - let x2 = x * x; - let num = x * (135135.0 + x2 * (17325.0 + x2 * (378.0 + x2))); - let den = 135135.0 + x2 * (62370.0 + x2 * (3150.0 + x2 * 28.0)); - num / den -} - -#[inline(always)] -pub(crate) fn gelu_simd(simd: S, x: f32x4) -> f32x4 { - let half = f32x4::splat(simd, 0.5); - let one = f32x4::splat(simd, 1.0); - let c1 = f32x4::splat(simd, 0.797_884_6); - let c2 = f32x4::splat(simd, 0.044_715); - let x2 = x * x; - let x3 = x * x2; - let inner = c1 * (x + c2 * x3); - let t = tanh_approx_simd(simd, inner); - half * x * (one + t) -} - -/// SIMD [7/6] Padé tanh approximation matching `tanh_approx`. -#[inline(always)] -fn tanh_approx_simd(simd: S, x: f32x4) -> f32x4 { - let neg5 = f32x4::splat(simd, -5.0); - let pos5 = f32x4::splat(simd, 5.0); - let x = x.max(neg5).min(pos5); - let x2 = x * x; - let c28 = f32x4::splat(simd, 28.0); - let c378 = f32x4::splat(simd, 378.0); - let c3150 = f32x4::splat(simd, 3150.0); - let c17325 = f32x4::splat(simd, 17325.0); - let c62370 = f32x4::splat(simd, 62370.0); - let c135135 = f32x4::splat(simd, 135135.0); - let num = x * (c135135 + x2 * (c17325 + x2 * (c378 + x2))); - let den = c135135 + x2 * (c62370 + x2 * (c3150 + c28 * x2)); - num / den -} - -#[cfg(test)] -mod tests { - use super::*; - use fearless_simd::{Level, Simd, SimdBase, dispatch, f32x4}; - - #[test] - fn tanh_approx_tracks_reference_over_clamped_range() { - for i in 0..=240 { - let x = -6.0 + i as f32 * 12.0 / 240.0; - let expected = x.clamp(-5.0, 5.0).tanh(); - let actual = tanh_approx(x); - - assert!( - (actual - expected).abs() <= 1.25e-4, - "x={x}, actual={actual}, expected={expected}" - ); - assert!(actual.is_finite(), "x={x}, actual={actual}"); - } - } - - #[test] - fn tanh_approx_is_odd_and_clamps_inputs() { - for x in [0.0, 0.25, 1.0, 3.0, 5.0, 6.0, 32.0] { - let positive = tanh_approx(x); - let negative = tanh_approx(-x); - assert!( - (positive + negative).abs() <= f32::EPSILON, - "x={x}, positive={positive}, negative={negative}" - ); - } - - assert_eq!(tanh_approx(6.0), tanh_approx(5.0)); - assert_eq!(tanh_approx(-6.0), tanh_approx(-5.0)); - } - - #[test] - fn gelu_tracks_tanh_form_reference() { - for i in 0..=120 { - let x = -6.0 + i as f32 * 12.0 / 120.0; - let inner = 0.797_884_6 * (x + 0.044_715 * x * x * x); - let expected = 0.5 * x * (1.0 + inner.tanh()); - let actual = gelu(x); - - assert!( - (actual - expected).abs() <= 2.5e-4, - "x={x}, actual={actual}, expected={expected}" - ); - } - } - - #[test] - fn simd_approximations_match_scalar_lanes() { - let level = Level::new(); - dispatch!(level, simd => { - assert_tanh_approx_simd_matches_scalar(simd); - assert_gelu_simd_matches_scalar(simd); - }); - } - - fn assert_tanh_approx_simd_matches_scalar(simd: S) { - let inputs = [ - -10.0, -6.0, -5.0, -3.0, -1.25, -0.5, 0.0, 0.5, 1.25, 3.0, 5.0, 10.0, - ]; - - let (inputs, []) = inputs.as_chunks::<4>() else { - unreachable!(); - }; - for input in inputs { - let mut actual = [0.0; 4]; - tanh_approx_simd(simd, f32x4::from_slice(simd, input)).store_slice(&mut actual); - - for (&x, &actual) in input.iter().zip(actual.iter()) { - let expected = tanh_approx(x); - assert_eq!(actual, expected, "x={x}"); - } - } - } - - fn assert_gelu_simd_matches_scalar(simd: S) { - let inputs = [ - -10.0, -6.0, -5.0, -3.0, -1.25, -0.5, 0.0, 0.5, 1.25, 3.0, 5.0, 10.0, - ]; - - let (inputs, []) = inputs.as_chunks::<4>() else { - unreachable!(); - }; - for input in inputs { - let mut actual = [0.0; 4]; - gelu_simd(simd, f32x4::from_slice(simd, input)).store_slice(&mut actual); - - for (&x, &actual) in input.iter().zip(actual.iter()) { - let expected = gelu(x); - assert_eq!(actual, expected, "x={x}"); - } - } - } -} diff --git a/src/model/bnn.rs b/src/model/bnn.rs new file mode 100644 index 0000000..4c81bc1 --- /dev/null +++ b/src/model/bnn.rs @@ -0,0 +1,447 @@ +//! Binary CNN inference over the Shannon-coded byte window. +//! +//! The raw (begin + end) byte window is entropy-coded with a canonical +//! Shannon code learned from the training byte distribution, producing two +//! fixed 16384-bit planes (code bits + codeword-boundary markers). Inference +//! is entirely bitwise/integer: XNOR + popcount convolutions (optionally a +//! sum of binary bases combined with fixed-point integer scales), folded +//! integer thresholds, OR-pooling, segmented popcount head counts, and an +//! integer classifier (per-class f32 scale). Multiple models may be stored +//! in one artifact; their logits are summed (ensemble). + +use super::constants::CLASSES; +use std::sync::OnceLock; + +pub(crate) static BNN_BYTES: &[u8] = include_bytes!("../../assets/bnn/source-bnn2.bin"); + +const MAGIC: u32 = 0x324E_4242; // "BBN2" +const VERSION: u32 = 2; +pub(crate) const BITS: usize = 16_384; +const BIT_WORDS: usize = BITS / 64; + +struct Layer { + m: usize, + cin: usize, + k: usize, + stride: usize, + pad: usize, + cout: usize, + /// Per basis: `cout` rows of `(cin * k).div_ceil(64)` words + /// (bit index within a row = tap * cin + ci). + weights: Vec>, + /// Per basis: fixed-point (x4096) per-out-channel scales. + aq: Vec>, + /// 1 => fires when z_q >= thr, 0 => fires when z_q <= thr. + sign: Vec, + thr: Vec, +} + +struct Model { + c: usize, + segs: usize, + stem: Layer, + conv1: Layer, + conv2: Layer, + head_weights: Vec, + head_scale: Vec, + head_bias: Vec, +} + +pub(crate) struct Bnn { + /// Canonical Shannon codes: (code value, bit length) per byte. + codes: [(u32, u32); 256], + models: Vec, +} + +struct Reader<'a> { + bytes: &'a [u8], + cur: usize, +} + +impl<'a> Reader<'a> { + fn u32(&mut self) -> u32 { + let v = u32::from_le_bytes(self.bytes[self.cur..self.cur + 4].try_into().unwrap()); + self.cur += 4; + v + } + + fn u8s(&mut self, n: usize) -> &'a [u8] { + let s = &self.bytes[self.cur..self.cur + n]; + self.cur += n; + s + } + + fn i32s(&mut self, n: usize) -> Vec { + (0..n).map(|_| self.u32() as i32).collect() + } + + fn i16s(&mut self, n: usize) -> Vec { + (0..n) + .map(|_| { + let v = i16::from_le_bytes(self.bytes[self.cur..self.cur + 2].try_into().unwrap()); + self.cur += 2; + v + }) + .collect() + } + + fn f32s(&mut self, n: usize) -> Vec { + (0..n).map(|_| f32::from_bits(self.u32())).collect() + } + + /// Reads `rows` bit-rows of `nbits` (u32-packed, LSB-first) into u64 rows. + fn bit_rows(&mut self, rows: usize, nbits: usize) -> Vec { + let words32 = nbits.div_ceil(32); + let words64 = nbits.div_ceil(64); + let mut out = vec![0u64; rows * words64]; + for r in 0..rows { + for w in 0..words32 { + let v = u64::from(self.u32()); + out[r * words64 + w / 2] |= v << ((w % 2) * 32); + } + } + out + } +} + +/// Rebuild canonical codes from code lengths: codes are assigned in +/// (length, byte) order, left-justified within the max length. +fn canonical_codes(lengths: &[u8]) -> [(u32, u32); 256] { + let mut order: Vec = (0..=255u8).collect(); + order.sort_by_key(|&b| (lengths[b as usize], b)); + let mut codes = [(0u32, 0u32); 256]; + let mut code = 0u32; + let mut prev_len = 0u32; + for &b in &order { + let len = u32::from(lengths[b as usize]); + code <<= len - prev_len; + codes[b as usize] = (code, len); + code += 1; + prev_len = len; + } + codes +} + +fn read_layer(r: &mut Reader) -> Layer { + let m = r.u32() as usize; + let cin = r.u32() as usize; + let k = r.u32() as usize; + // cout is implied by the enclosing model's C; stored thresholds follow. + // Weight rows and aq come per basis. + // The caller patches stride/pad/cout. + let mut layer = Layer { + m, + cin, + k, + stride: 1, + pad: k / 2, + cout: 0, + weights: Vec::with_capacity(m), + aq: Vec::with_capacity(m), + sign: Vec::new(), + thr: Vec::new(), + }; + layer + .weights + .reserve_exact(m.saturating_sub(layer.weights.capacity())); + layer +} + +impl Bnn { + fn load() -> Self { + let mut r = Reader { + bytes: BNN_BYTES, + cur: 0, + }; + assert_eq!(r.u32(), MAGIC, "bad BBN2 magic"); + assert_eq!(r.u32(), VERSION, "bad BBN2 version"); + assert_eq!(r.u32() as usize, BITS); + assert_eq!(r.u32() as usize, CLASSES); + let n_models = r.u32() as usize; + + let lengths: [u8; 256] = r.u8s(256).try_into().unwrap(); + let codes = canonical_codes(&lengths); + + let mut models = Vec::with_capacity(n_models); + for _ in 0..n_models { + let c = r.u32() as usize; + let stem_k = r.u32() as usize; + let stem_stride = r.u32() as usize; + let segs = r.u32() as usize; + let mut layers = Vec::with_capacity(3); + for li in 0..3 { + let mut layer = read_layer(&mut r); + layer.cout = c; + if li == 0 { + assert_eq!(layer.k, stem_k); + layer.stride = stem_stride; + layer.pad = stem_k / 2; + } + for _ in 0..layer.m { + layer.weights.push(r.bit_rows(c, layer.cin * layer.k)); + layer.aq.push(r.i32s(c)); + } + layer.sign = r.u8s(c).to_vec(); + layer.thr = r.i32s(c); + layers.push(layer); + } + let conv2 = layers.pop().unwrap(); + let conv1 = layers.pop().unwrap(); + let stem = layers.pop().unwrap(); + let head_weights = r.i16s(CLASSES * segs * c); + let head_scale = r.f32s(CLASSES); + let head_bias = r.f32s(CLASSES); + models.push(Model { + c, + segs, + stem, + conv1, + conv2, + head_weights, + head_scale, + head_bias, + }); + } + assert_eq!(r.cur, BNN_BYTES.len(), "unexpected BBN2 payload length"); + Self { codes, models } + } + + pub(crate) fn get() -> &'static Self { + static MODEL: OnceLock = OnceLock::new(); + MODEL.get_or_init(Self::load) + } + + /// Shannon-encode window bytes into two fixed-length bitplanes + /// (code bits, codeword-boundary markers), truncated at the last + /// codeword that fits and zero padded. + pub(crate) fn encode_planes(&self, window: &[u8]) -> ([u64; BIT_WORDS], [u64; BIT_WORDS]) { + let mut bits = [0u64; BIT_WORDS]; + let mut bounds = [0u64; BIT_WORDS]; + let mut pos = 0usize; + for &b in window { + let (code, len) = self.codes[b as usize]; + if pos + len as usize > BITS { + break; + } + bounds[pos / 64] |= 1u64 << (pos % 64); + // MSB of the codeword first, matching the trainer. + for i in (0..len).rev() { + if (code >> i) & 1 == 1 { + bits[pos / 64] |= 1u64 << (pos % 64); + } + pos += 1; + } + } + (bits, bounds) + } + + pub(crate) fn logits(&self, planes: &([u64; BIT_WORDS], [u64; BIT_WORDS])) -> [f32; CLASSES] { + let mut logits = [0.0f32; CLASSES]; + for model in &self.models { + let ml = model.logits(planes); + for (acc, l) in logits.iter_mut().zip(ml.iter()) { + *acc += l; + } + } + let n = self.models.len().max(1) as f32; + for logit in &mut logits { + *logit /= n; + } + logits + } +} + +impl Model { + fn logits(&self, planes: &([u64; BIT_WORDS], [u64; BIT_WORDS])) -> [f32; CLASSES] { + let c = self.c; + let cw = c.div_ceil(64); + // Stem over the two input planes -> position-major rows of C bits, + // OR-pooled by 4. + let l0 = (BITS + 2 * self.stem.pad - self.stem.k) / self.stem.stride + 1; + let l1 = (l0 / 4).max(1); + let mut a1 = vec![0u64; l1 * cw]; + for t in 0..l0 { + let start = (t * self.stem.stride) as isize - self.stem.pad as isize; + let (x, mask, n_valid) = stem_window(planes, start, self.stem.k); + let slot = t / 4; + if slot >= l1 { + break; + } + for ch in 0..c { + let mut zq = 0i64; + for j in 0..self.stem.m { + let w = self.stem.weights[j][ch]; + let m = (!(x ^ w) & mask).count_ones() as i64; + let z = 2 * m - n_valid as i64; + zq += i64::from(self.stem.aq[j][ch]) * z; + } + let fire = if self.stem.sign[ch] == 1 { + zq >= i64::from(self.stem.thr[ch]) + } else { + zq <= i64::from(self.stem.thr[ch]) + }; + if fire { + a1[slot * cw + ch / 64] |= 1u64 << (ch % 64); + } + } + } + + let l2 = l1 / 4; + let a2 = conv3(&a1, l1, &self.conv1, 4, c); + let a3 = conv3(&a2, l2, &self.conv2, 1, c); + + // Segmented popcount head. + let segs = self.segs; + let seg = l2 / segs; + let mut counts = vec![0i32; segs * c]; + for t in 0..seg * segs { + let row = &a3[t * cw..(t + 1) * cw]; + let base = (t / seg) * c; + for (wi, &word) in row.iter().enumerate() { + let mut bitsw = word; + while bitsw != 0 { + let b = bitsw.trailing_zeros() as usize; + counts[base + wi * 64 + b] += 1; + bitsw &= bitsw - 1; + } + } + } + + let mut logits = [0.0f32; CLASSES]; + let width = segs * c; + for (k, logit) in logits.iter_mut().enumerate() { + let row = &self.head_weights[k * width..(k + 1) * width]; + let mut acc = 0i64; + for (w, &f) in row.iter().zip(counts.iter()) { + acc += i64::from(*w) * i64::from(f); + } + *logit = acc as f32 * self.head_scale[k] + self.head_bias[k]; + } + logits + } +} + +/// Gather a stem window over the two input planes: bit index tap*2 + plane. +/// Returns (bits, valid mask, n_valid). +fn stem_window( + planes: &([u64; BIT_WORDS], [u64; BIT_WORDS]), + start: isize, + k: usize, +) -> (u64, u64, usize) { + let mut x = 0u64; + let mut mask = 0u64; + let mut n_valid = 0usize; + for tap in 0..k { + let p = start + tap as isize; + if p >= 0 && (p as usize) < BITS { + let p = p as usize; + n_valid += 2; + mask |= 0b11 << (tap * 2); + if planes.0[p / 64] >> (p % 64) & 1 == 1 { + x |= 1u64 << (tap * 2); + } + if planes.1[p / 64] >> (p % 64) & 1 == 1 { + x |= 1u64 << (tap * 2 + 1); + } + } + } + (x, mask, n_valid) +} + +/// k=3 (pad=1) multi-basis binary conv + threshold + OR-pool over +/// position-major rows of `c` bits. Output: len/pool rows. +fn conv3(rows: &[u64], len: usize, layer: &Layer, pool: usize, c: usize) -> Vec { + let cw = c.div_ceil(64); + let kw = 3 * cw; + let out_len = len / pool; + let mut out = vec![0u64; out_len * cw]; + let zero = vec![0u64; cw]; + for t in 0..len { + let taps: [&[u64]; 3] = [ + if t == 0 { + &zero + } else { + &rows[(t - 1) * cw..t * cw] + }, + &rows[t * cw..(t + 1) * cw], + if t + 1 >= len { + &zero + } else { + &rows[(t + 1) * cw..(t + 2) * cw] + }, + ]; + let edge = [t == 0, false, t + 1 >= len]; + let n_valid = (3 - edge.iter().filter(|&&e| e).count()) * c; + let slot = t / pool; + if slot >= out_len { + break; + } + for ch in 0..c { + let mut zq = 0i64; + for j in 0..layer.m { + let w = &layer.weights[j][ch * kw..(ch + 1) * kw]; + let mut m = 0i64; + for tap in 0..3 { + if edge[tap] { + continue; + } + for wi in 0..cw { + m += i64::from((!(taps[tap][wi] ^ w[tap * cw + wi])).count_ones()); + } + } + let z = 2 * m - n_valid as i64; + zq += i64::from(layer.aq[j][ch]) * z; + } + let fire = if layer.sign[ch] == 1 { + zq >= i64::from(layer.thr[ch]) + } else { + zq <= i64::from(layer.thr[ch]) + }; + if fire { + out[slot * cw + ch / 64] |= 1u64 << (ch % 64); + } + } + } + out +} + +#[cfg(test)] +mod parity_tests { + use super::*; + + #[test] + fn python_parity() { + let Ok(fixtures) = std::env::var("BNN_FIXTURES") else { + return; + }; + let raw = std::fs::read_to_string(fixtures).unwrap(); + let model = Bnn::get(); + // fixtures: [{"window_hex": "...", "logits": [f, ...]}, ...] + let mut cases = Vec::new(); + for chunk in raw.split("\"window_hex\": \"").skip(1) { + let hex_end = chunk.find('"').unwrap(); + let hex = &chunk[..hex_end]; + let rest = &chunk[hex_end..]; + let ls = rest.find('[').unwrap() + 1; + let le = rest.find(']').unwrap(); + let logits: Vec = rest[ls..le] + .split(',') + .map(|s| s.trim().parse().unwrap()) + .collect(); + let bytes: Vec = (0..hex.len()) + .step_by(2) + .map(|i| u8::from_str_radix(&hex[i..i + 2], 16).unwrap()) + .collect(); + cases.push((bytes, logits)); + } + assert!(!cases.is_empty()); + for (window, expected) in cases { + let planes = model.encode_planes(&window); + let got = model.logits(&planes); + for (g, e) in got.iter().zip(expected.iter()) { + let tol = 1e-3 * g.abs().max(e.abs()).max(1.0); + assert!((g - e).abs() < tol, "logit mismatch: {g} vs {e}"); + } + } + } +} diff --git a/src/model/constants.rs b/src/model/constants.rs index 8716087..1d86476 100644 --- a/src/model/constants.rs +++ b/src/model/constants.rs @@ -1,43 +1,6 @@ -pub(crate) const MODEL_MAGIC: [u8; 8] = [0x4d, 0x53, 0x51, 0x31, 0x01, 0x00, 0x00, 0x00]; -pub(crate) const SCALE_COUNT: usize = 6; - pub(crate) const MAGIKA_BEG_SIZE: usize = 1_024; pub(crate) const MAGIKA_END_SIZE: usize = 1_024; pub(crate) const MAGIKA_WINDOW_SIZE: usize = MAGIKA_BEG_SIZE + MAGIKA_END_SIZE; pub(crate) const MAGIKA_BLOCK_SIZE: usize = 4_096; -// Wordseq architecture constants. Must match what the model was trained with. -pub(crate) const BINS: usize = 1_024; -pub(crate) const MAX_UNITS: usize = 2_048; -pub(crate) const EMBED: usize = 24; -pub(crate) const CONV0_KERNEL: usize = 7; -pub(crate) const CONV0: usize = 64; -pub(crate) const CONV0_POOL: usize = 4; -pub(crate) const CONV1_KERNEL: usize = 5; -pub(crate) const CONV1: usize = 128; -pub(crate) const CONV1_POOL: usize = 2; -pub(crate) const CONV2_KERNEL: usize = 3; -pub(crate) const CONV2: usize = 128; -pub(crate) const POOLED: usize = CONV2 * 2; // GlobalMax + GlobalAvg -pub(crate) const DENSE: usize = 96; pub(crate) const CLASSES: usize = 48; - -const EMBED_SCRATCH: usize = MAX_UNITS * EMBED; -const POOL0_SCRATCH: usize = (MAX_UNITS / CONV0_POOL) * CONV0; -const POOL1_SCRATCH: usize = (MAX_UNITS / CONV0_POOL / CONV1_POOL) * CONV1; - -const fn max_usize(a: usize, b: usize) -> usize { - if a > b { a } else { b } -} - -pub(crate) const ACTIVATION_SCRATCH: usize = - max_usize(EMBED_SCRATCH + POOL0_SCRATCH, POOL0_SCRATCH + POOL1_SCRATCH); -pub(crate) const CONV_SCRATCH: usize = 4 * CONV2; -pub(crate) const INFERENCE_SCRATCH: usize = ACTIVATION_SCRATCH + CONV_SCRATCH; - -// Tokenizer flag bits. Must match `_PUNCT_FLAG`/etc. in the Python trainer. -pub(crate) const WORD_MASK: u32 = 0x00FF_FFFF; -pub(crate) const PUNCT_FLAG: u32 = 0x1000_0000; -pub(crate) const INDENT_FLAG: u32 = 0x2000_0000; -pub(crate) const NUM_FLAG: u32 = 0x4000_0000; -pub(crate) const BRACKET_FLAG: u32 = 0x5000_0000; diff --git a/src/model/embedded.rs b/src/model/embedded.rs deleted file mode 100644 index ecbe81a..0000000 --- a/src/model/embedded.rs +++ /dev/null @@ -1 +0,0 @@ -pub(crate) static MODEL_BYTES: &[u8] = include_bytes!("../../assets/magika/source-student-q4.bin"); diff --git a/src/model/layers.rs b/src/model/layers.rs deleted file mode 100644 index b7a5e8d..0000000 --- a/src/model/layers.rs +++ /dev/null @@ -1,878 +0,0 @@ -use super::{ - activation::gelu_simd, - constants::{BINS, EMBED}, -}; -use fearless_simd::{Level, Simd, SimdBase, SimdFloat, dispatch, f32x4}; -use std::{ops::Range, sync::OnceLock}; - -/// Detect the best available SIMD level once per process. -fn simd_level() -> Level { - static LEVEL: OnceLock = OnceLock::new(); - *LEVEL.get_or_init(Level::new) -} - -#[inline(always)] -fn as_array_chunks(slice: &[f32]) -> &[[f32; N]] { - let (chunks, remainder) = slice.as_chunks::(); - debug_assert!(remainder.is_empty()); - chunks -} - -#[inline(always)] -fn as_array_chunks_mut(slice: &mut [f32]) -> &mut [[f32; N]] { - let (chunks, remainder) = slice.as_chunks_mut::(); - debug_assert!(remainder.is_empty()); - chunks -} - -/// Sum the K=3 hashed embedding rows for one unit-id into `dst` (length EMBED). -#[inline] -pub(crate) fn embed_position(embedding: &[f32], unit: u32, dst: &mut [f32]) { - let b0 = hash_bin(unit, 0); - let b1 = hash_bin(unit, 1); - let b2 = hash_bin(unit, 2); - let rows = as_array_chunks::(embedding); - let row0 = &rows[b0]; - let row1 = &rows[b1]; - let row2 = &rows[b2]; - for (((d, &v0), &v1), &v2) in dst.iter_mut().zip(row0).zip(row1).zip(row2) { - *d = v0 + v1 + v2; - } -} - -/// K=3 prime-mix hash matching `hash_unit_indices` in the Python trainer. -fn hash_bin(unit: u32, head: usize) -> usize { - const PRIMES: [u32; 4] = [2_654_435_761, 2_246_822_519, 3_266_489_917, 668_265_263]; - let p1 = PRIMES[head % PRIMES.len()]; - let p2 = PRIMES[(head + 1) % PRIMES.len()]; - let mut h = unit.wrapping_mul(p1); - h ^= h >> 13; - h = h.wrapping_mul(p2); - (h as usize) % BINS -} - -/// Conv1d (SAME pad) for a block of `BLOCK` consecutive output positions -/// starting at `t_base`. Accumulates into `accs` (`BLOCK * out_channels` long). -/// -/// `kernel` is `[k][in_c][out_c]` with the inner row contiguous over -/// out_channels. -/// -/// BLOCK=4 + all positions in-bounds → SIMD fast path. When out_channels -/// is a multiple of 16, the group-16 kernel keeps 16 `f32x4` accumulators -/// register-resident across the full (k, in_c) inner loop; otherwise a -/// chunk-inner SIMD path with 4 named acc rows is used. -/// Edges and partial blocks use the plain f32 path in this kernel. -#[inline(always)] -#[allow(clippy::too_many_arguments)] -fn conv1d_block( - simd: S, - input: &[f32], - seq_len: usize, - in_channels: usize, - kernel: &[f32], - kernel_size: usize, - out_channels: usize, - bias: &[f32], - t_base: usize, - accs: &mut [f32], -) { - debug_assert_eq!(accs.len(), BLOCK * out_channels); - debug_assert!(out_channels.is_multiple_of(4)); - let pad = (kernel_size - 1) / 2; - if BLOCK == 4 { - let all_in_bounds = - t_base >= pad && t_base + 4 + (kernel_size - 1).saturating_sub(pad) <= seq_len; - if all_in_bounds { - conv1d_block4_simd_inner( - simd, - input, - in_channels, - kernel, - kernel_size, - out_channels, - bias, - t_base, - pad, - accs, - ); - return; - } - } - // Edge/partial-block path. - for acc in accs.chunks_exact_mut(out_channels) { - acc.copy_from_slice(bias); - } - for k in 0..kernel_size { - let src_t_at_s0 = t_base as isize + k as isize - pad as isize; - let s_lo = if src_t_at_s0 < 0 { - ((-src_t_at_s0) as usize).min(BLOCK) - } else { - 0 - }; - let s_hi_signed = seq_len as isize - src_t_at_s0; - let s_hi = if s_hi_signed > 0 { - (s_hi_signed as usize).min(BLOCK) - } else { - 0 - }; - if s_lo >= s_hi { - continue; - } - let kbase = k * in_channels * out_channels; - let krows = kernel[kbase..kbase + in_channels * out_channels].chunks_exact(out_channels); - for (in_c, krow) in krows.enumerate() { - for (s, acc) in accs - .chunks_exact_mut(out_channels) - .enumerate() - .skip(s_lo) - .take(s_hi - s_lo) - { - let src_t = (src_t_at_s0 + s as isize) as usize; - let x = input[src_t * in_channels + in_c]; - for (a, &w) in acc.iter_mut().zip(krow) { - *a = w.mul_add(x, *a); - } - } - } - } -} - -/// BLOCK=4 SIMD hot path. Out-channels are processed in 16-wide groups so -/// that each group's 4 accumulator rows × 4 chunks (= 16 NEON registers) -/// live in regs across the entire (k, in_c) inner loop. Combined with the -/// chunk-inner inner pattern, kernel and input reads stay cache-sequential -/// while the per-iter accs load/store traffic is eliminated. -#[inline(always)] -#[allow(clippy::too_many_arguments)] -fn conv1d_block4_simd_inner( - simd: S, - input: &[f32], - in_channels: usize, - kernel: &[f32], - kernel_size: usize, - out_channels: usize, - bias: &[f32], - t_base: usize, - pad: usize, - accs: &mut [f32], -) { - const GROUP: usize = 16; - if out_channels.is_multiple_of(GROUP) { - conv1d_block4_group16::( - simd, - input, - in_channels, - kernel, - kernel_size, - out_channels, - bias, - t_base, - pad, - accs, - ); - return; - } - // Chunk-inner SIMD path for out_channels not a multiple of 16. - for acc in accs.chunks_exact_mut(out_channels) { - acc.copy_from_slice(bias); - } - let (a01, a23) = accs.split_at_mut(2 * out_channels); - let (a0, a1) = a01.split_at_mut(out_channels); - let (a2, a3) = a23.split_at_mut(out_channels); - for k in 0..kernel_size { - let base_t = t_base + k - pad; - let row0_off = base_t * in_channels; - let row1_off = (base_t + 1) * in_channels; - let row2_off = (base_t + 2) * in_channels; - let row3_off = (base_t + 3) * in_channels; - let kbase = k * in_channels * out_channels; - let krows = kernel[kbase..kbase + in_channels * out_channels].chunks_exact(out_channels); - for (in_c, krow) in krows.enumerate() { - let xv0 = f32x4::splat(simd, input[row0_off + in_c]); - let xv1 = f32x4::splat(simd, input[row1_off + in_c]); - let xv2 = f32x4::splat(simd, input[row2_off + in_c]); - let xv3 = f32x4::splat(simd, input[row3_off + in_c]); - for ((((kr_c, a0_c), a1_c), a2_c), a3_c) in as_array_chunks::<4>(krow) - .iter() - .zip(as_array_chunks_mut::<4>(a0).iter_mut()) - .zip(as_array_chunks_mut::<4>(a1).iter_mut()) - .zip(as_array_chunks_mut::<4>(a2).iter_mut()) - .zip(as_array_chunks_mut::<4>(a3).iter_mut()) - { - let kr = f32x4::from_slice(simd, kr_c); - let av0 = f32x4::from_slice(simd, a0_c); - let av1 = f32x4::from_slice(simd, a1_c); - let av2 = f32x4::from_slice(simd, a2_c); - let av3 = f32x4::from_slice(simd, a3_c); - kr.mul_add(xv0, av0).store_slice(a0_c); - kr.mul_add(xv1, av1).store_slice(a1_c); - kr.mul_add(xv2, av2).store_slice(a2_c); - kr.mul_add(xv3, av3).store_slice(a3_c); - } - } - } -} - -/// Group-16 SIMD kernel. For each 16-wide out-channel group, holds -/// `4 acc rows × 4 chunks = 16 f32x4` accumulators in NEON registers -/// across the entire (k, in_c) inner loop. -/// -/// Per (k, in_c) iter: 4 kernel f32x4 loads + 4 broadcast x loads + 16 FMAs. -/// All 16 accumulators stay register-resident — no per-iter accs load/store. -#[inline(always)] -#[allow(clippy::too_many_arguments)] -fn conv1d_block4_group16( - simd: S, - input: &[f32], - in_channels: usize, - kernel: &[f32], - kernel_size: usize, - out_channels: usize, - bias: &[f32], - t_base: usize, - pad: usize, - accs: &mut [f32], -) { - let (acc01, acc23) = accs.split_at_mut(2 * out_channels); - let (acc0, acc1) = acc01.split_at_mut(out_channels); - let (acc2, acc3) = acc23.split_at_mut(out_channels); - for (g, ((((bias_group, acc0_group), acc1_group), acc2_group), acc3_group)) in - as_array_chunks::<16>(bias) - .iter() - .zip(as_array_chunks_mut::<16>(acc0).iter_mut()) - .zip(as_array_chunks_mut::<16>(acc1).iter_mut()) - .zip(as_array_chunks_mut::<16>(acc2).iter_mut()) - .zip(as_array_chunks_mut::<16>(acc3).iter_mut()) - .enumerate() - { - let g_off = g * 16; - let bias_chunks = as_array_chunks::<4>(bias_group); - let b0 = f32x4::from_slice(simd, &bias_chunks[0]); - let b1 = f32x4::from_slice(simd, &bias_chunks[1]); - let b2 = f32x4::from_slice(simd, &bias_chunks[2]); - let b3 = f32x4::from_slice(simd, &bias_chunks[3]); - let (mut a0_0, mut a0_1, mut a0_2, mut a0_3) = (b0, b1, b2, b3); - let (mut a1_0, mut a1_1, mut a1_2, mut a1_3) = (b0, b1, b2, b3); - let (mut a2_0, mut a2_1, mut a2_2, mut a2_3) = (b0, b1, b2, b3); - let (mut a3_0, mut a3_1, mut a3_2, mut a3_3) = (b0, b1, b2, b3); - - for k in 0..kernel_size { - let base_t = t_base + k - pad; - let row0_off = base_t * in_channels; - let row1_off = (base_t + 1) * in_channels; - let row2_off = (base_t + 2) * in_channels; - let row3_off = (base_t + 3) * in_channels; - let kbase = k * in_channels * out_channels; - let krows = - kernel[kbase..kbase + in_channels * out_channels].chunks_exact(out_channels); - for (in_c, krow) in krows.enumerate() { - let kernel_chunks = as_array_chunks::<4>(&krow[g_off..g_off + 16]); - let kr0 = f32x4::from_slice(simd, &kernel_chunks[0]); - let kr1 = f32x4::from_slice(simd, &kernel_chunks[1]); - let kr2 = f32x4::from_slice(simd, &kernel_chunks[2]); - let kr3 = f32x4::from_slice(simd, &kernel_chunks[3]); - let xv0 = f32x4::splat(simd, input[row0_off + in_c]); - let xv1 = f32x4::splat(simd, input[row1_off + in_c]); - let xv2 = f32x4::splat(simd, input[row2_off + in_c]); - let xv3 = f32x4::splat(simd, input[row3_off + in_c]); - a0_0 = kr0.mul_add(xv0, a0_0); - a0_1 = kr1.mul_add(xv0, a0_1); - a0_2 = kr2.mul_add(xv0, a0_2); - a0_3 = kr3.mul_add(xv0, a0_3); - a1_0 = kr0.mul_add(xv1, a1_0); - a1_1 = kr1.mul_add(xv1, a1_1); - a1_2 = kr2.mul_add(xv1, a1_2); - a1_3 = kr3.mul_add(xv1, a1_3); - a2_0 = kr0.mul_add(xv2, a2_0); - a2_1 = kr1.mul_add(xv2, a2_1); - a2_2 = kr2.mul_add(xv2, a2_2); - a2_3 = kr3.mul_add(xv2, a2_3); - a3_0 = kr0.mul_add(xv3, a3_0); - a3_1 = kr1.mul_add(xv3, a3_1); - a3_2 = kr2.mul_add(xv3, a3_2); - a3_3 = kr3.mul_add(xv3, a3_3); - } - } - - // Store the 4×4 accumulator tile back into accs[s][g_off..g_off+16]. - let acc0_chunks = as_array_chunks_mut::<4>(acc0_group); - a0_0.store_slice(&mut acc0_chunks[0]); - a0_1.store_slice(&mut acc0_chunks[1]); - a0_2.store_slice(&mut acc0_chunks[2]); - a0_3.store_slice(&mut acc0_chunks[3]); - let acc1_chunks = as_array_chunks_mut::<4>(acc1_group); - a1_0.store_slice(&mut acc1_chunks[0]); - a1_1.store_slice(&mut acc1_chunks[1]); - a1_2.store_slice(&mut acc1_chunks[2]); - a1_3.store_slice(&mut acc1_chunks[3]); - let acc2_chunks = as_array_chunks_mut::<4>(acc2_group); - a2_0.store_slice(&mut acc2_chunks[0]); - a2_1.store_slice(&mut acc2_chunks[1]); - a2_2.store_slice(&mut acc2_chunks[2]); - a2_3.store_slice(&mut acc2_chunks[3]); - let acc3_chunks = as_array_chunks_mut::<4>(acc3_group); - a3_0.store_slice(&mut acc3_chunks[0]); - a3_1.store_slice(&mut acc3_chunks[1]); - a3_2.store_slice(&mut acc3_chunks[2]); - a3_3.store_slice(&mut acc3_chunks[3]); - } -} - -#[derive(Clone, Debug)] -struct RepeatedRows { - range: Range, - extends_right_padding: bool, -} - -impl RepeatedRows { - fn new(start: usize, end: usize, extends_right_padding: bool) -> Self { - Self { - range: start..end, - extends_right_padding, - } - } - - fn empty(at: usize) -> Self { - Self::new(at, at, false) - } - - fn is_empty(&self) -> bool { - self.range.is_empty() - } -} - -#[derive(Clone, Debug)] -pub(crate) struct Tensor<'a> { - data: &'a [f32], - rows: usize, - channels: usize, - repeated_tail: RepeatedRows, -} - -impl<'a> Tensor<'a> { - pub(crate) fn with_repeated_tail( - data: &'a [f32], - rows: usize, - channels: usize, - start: usize, - ) -> Self { - let start = start.min(rows); - let materialized_rows = if start < rows { start + 1 } else { rows }; - debug_assert!(data.len() >= materialized_rows * channels); - debug_assert!(data.len() <= rows * channels); - Self { - data, - rows, - channels, - repeated_tail: RepeatedRows::new(start, rows, true), - } - } - - #[cfg(test)] - pub(crate) fn copy_to_dense(&self, out: &mut [f32]) { - debug_assert_eq!(out.len(), self.rows * self.channels); - for row_index in 0..self.rows { - out[row_index * self.channels..(row_index + 1) * self.channels] - .copy_from_slice(self.row(row_index)); - } - } - - fn row(&self, index: usize) -> &'a [f32] { - debug_assert!(index < self.rows); - let region = &self.repeated_tail; - let index = if region.range.contains(&index) { - region.range.start - } else { - index - }; - row(self.data, self.channels, index) - } -} - -struct ConvSpec<'a> { - seq_len: usize, - in_channels: usize, - kernel: &'a [f32], - kernel_size: usize, - out_channels: usize, - bias: &'a [f32], - pool: usize, -} - -impl ConvSpec<'_> { - fn row_count(&self) -> usize { - self.seq_len / self.pool - } -} - -#[allow(clippy::too_many_arguments)] -pub(crate) fn conv_gelu_maxpool_tensor<'a>( - input: Tensor<'_>, - kernel: &[f32], - kernel_size: usize, - out_channels: usize, - bias: &[f32], - pool: usize, - out: &'a mut [f32], - scratch: &mut [f32], -) -> Tensor<'a> { - let spec = ConvSpec { - seq_len: input.rows, - in_channels: input.channels, - kernel, - kernel_size, - out_channels, - bias, - pool, - }; - let out_region = pooled_repeated_region(&input, &spec); - let row_count = spec.row_count(); - debug_assert_eq!(out.len(), row_count * out_channels); - - if out_region.is_empty() { - conv_gelu_range(&input, &spec, 0..row_count, out, scratch); - return Tensor::with_repeated_tail(out, row_count, out_channels, row_count); - } - - let out_start = out_region.range.start; - let out_end = out_region.range.end; - conv_gelu_range(&input, &spec, 0..out_start, out, scratch); - let value_row = row_range_mut(out, out_channels, out_start..out_start + 1); - const_output_value(&input, &spec, value_row); - conv_gelu_range(&input, &spec, out_end..row_count, out, scratch); - - Tensor { - data: out, - rows: row_count, - channels: out_channels, - repeated_tail: out_region, - } -} - -#[allow(clippy::too_many_arguments)] -pub(crate) fn conv_gelu_global_pool_tensor( - input: Tensor<'_>, - kernel: &[f32], - kernel_size: usize, - out_channels: usize, - bias: &[f32], - out_max: &mut [f32], - out_avg: &mut [f32], - tmp: &mut [f32], - scratch: &mut [f32], -) { - let rows_len = input.rows * out_channels; - let conv = conv_gelu_maxpool_tensor( - input, - kernel, - kernel_size, - out_channels, - bias, - 1, - &mut tmp[..rows_len], - scratch, - ); - global_max_avg_pool(conv, out_max, out_avg); -} - -pub(crate) fn global_max_avg_pool(input: Tensor<'_>, out_max: &mut [f32], out_avg: &mut [f32]) { - debug_assert_eq!(out_max.len(), input.channels); - debug_assert_eq!(out_avg.len(), input.channels); - out_max.fill(f32::NEG_INFINITY); - out_avg.fill(0.0); - if input.rows == 0 { - return; - } - - let region = &input.repeated_tail; - if region.is_empty() { - accumulate_pool_rows(input.data, input.channels, out_max, out_avg); - } else { - accumulate_pool_rows( - row_range(input.data, input.channels, 0..region.range.start), - input.channels, - out_max, - out_avg, - ); - accumulate_const_rows( - input.row(region.range.start), - region.range.len(), - out_max, - out_avg, - ); - accumulate_pool_rows( - row_range(input.data, input.channels, region.range.end..input.rows), - input.channels, - out_max, - out_avg, - ); - } - - let inv = 1.0 / input.rows as f32; - for avg in out_avg { - *avg *= inv; - } -} - -fn pooled_repeated_region(input: &Tensor<'_>, spec: &ConvSpec<'_>) -> RepeatedRows { - let input_region = &input.repeated_tail; - if input_region.is_empty() { - return RepeatedRows::empty(spec.row_count()); - } - - let pad = (spec.kernel_size - 1) / 2; - let right = spec.kernel_size - 1 - pad; - let conv_start = input_region - .range - .start - .saturating_add(pad) - .min(spec.seq_len); - let conv_end = if input_region.extends_right_padding { - spec.seq_len - } else { - input_region - .range - .end - .saturating_sub(right) - .min(spec.seq_len) - }; - - if conv_start >= conv_end { - return RepeatedRows::empty(spec.row_count()); - } - - let start = conv_start.div_ceil(spec.pool); - let end = conv_end / spec.pool; - if start >= end { - RepeatedRows::empty(spec.row_count()) - } else { - RepeatedRows::new(start, end, false) - } -} - -fn conv_gelu_range( - input: &Tensor<'_>, - spec: &ConvSpec<'_>, - rows: Range, - out: &mut [f32], - scratch: &mut [f32], -) { - if rows.is_empty() { - return; - } - let row_start = rows.start; - let range = row_range_mut(out, spec.out_channels, rows.clone()); - if !range_touches_repeated(input, spec, rows) { - conv_gelu_maxpool_range( - input.data, - spec.seq_len, - spec.in_channels, - spec.kernel, - spec.kernel_size, - spec.out_channels, - spec.bias, - spec.pool, - row_start, - range, - scratch, - ); - } else { - conv_gelu_maxpool_tensor_range(input, spec, row_start, range, scratch); - } -} - -fn range_touches_repeated(input: &Tensor<'_>, spec: &ConvSpec<'_>, rows: Range) -> bool { - let region = &input.repeated_tail; - if rows.is_empty() { - return false; - } - if region.is_empty() { - return false; - } - let pad = (spec.kernel_size - 1) / 2; - let right = spec.kernel_size - 1 - pad; - let first_output = rows.start * spec.pool; - let last_output = (rows.end - 1) * spec.pool + (spec.pool - 1); - let src_start = first_output.saturating_sub(pad); - let src_end = (last_output + right + 1).min(input.rows); - src_start < region.range.end && region.range.start < src_end -} - -#[allow(clippy::too_many_arguments)] -fn conv_gelu_maxpool_range( - input: &[f32], - seq_len: usize, - in_channels: usize, - kernel: &[f32], - kernel_size: usize, - out_channels: usize, - bias: &[f32], - pool: usize, - row_start: usize, - out: &mut [f32], - scratch: &mut [f32], -) { - let row_count = out.len() / out_channels; - if row_count == 0 { - return; - } - let level = simd_level(); - dispatch!(level, simd => conv_gelu_maxpool_range_simd( - simd, input, seq_len, in_channels, kernel, kernel_size, - out_channels, bias, pool, row_start, out, scratch, - )); -} - -#[inline(always)] -#[allow(clippy::too_many_arguments)] -fn conv_gelu_maxpool_range_simd( - simd: S, - input: &[f32], - seq_len: usize, - in_channels: usize, - kernel: &[f32], - kernel_size: usize, - out_channels: usize, - bias: &[f32], - pool: usize, - row_start: usize, - out: &mut [f32], - scratch: &mut [f32], -) { - match pool { - 4 => conv_gelu_maxpool_range_run::( - simd, - input, - seq_len, - in_channels, - kernel, - kernel_size, - out_channels, - bias, - row_start, - out, - scratch, - ), - 2 => conv_gelu_maxpool_range_run::( - simd, - input, - seq_len, - in_channels, - kernel, - kernel_size, - out_channels, - bias, - row_start, - out, - scratch, - ), - _ => conv_gelu_maxpool_range_run::( - simd, - input, - seq_len, - in_channels, - kernel, - kernel_size, - out_channels, - bias, - row_start, - out, - scratch, - ), - } -} - -#[inline(always)] -#[allow(clippy::too_many_arguments)] -fn conv_gelu_maxpool_range_run( - simd: S, - input: &[f32], - seq_len: usize, - in_channels: usize, - kernel: &[f32], - kernel_size: usize, - out_channels: usize, - bias: &[f32], - row_start: usize, - out: &mut [f32], - scratch: &mut [f32], -) { - debug_assert!(scratch.len() >= POOL * out_channels); - let accs = &mut scratch[..POOL * out_channels]; - for (row_offset, dst) in out.chunks_exact_mut(out_channels).enumerate() { - let t_base = (row_start + row_offset) * POOL; - conv1d_block::( - simd, - input, - seq_len, - in_channels, - kernel, - kernel_size, - out_channels, - bias, - t_base, - &mut *accs, - ); - pool_gelu_rows(simd, accs, out_channels, dst); - } -} - -fn conv_gelu_maxpool_tensor_range( - input: &Tensor<'_>, - spec: &ConvSpec<'_>, - row_start: usize, - out: &mut [f32], - scratch: &mut [f32], -) { - debug_assert!(scratch.len() >= spec.pool * spec.out_channels); - let accs = &mut scratch[..spec.pool * spec.out_channels]; - for (row_offset, dst) in out.chunks_exact_mut(spec.out_channels).enumerate() { - let t_base = (row_start + row_offset) * spec.pool; - conv1d_tensor_block(input, spec, t_base, spec.pool, accs); - for value in accs.iter_mut() { - *value = super::activation::gelu(*value); - } - pool_rows_scalar(accs, spec.out_channels, dst); - } -} - -fn conv1d_tensor_block( - input: &Tensor<'_>, - spec: &ConvSpec<'_>, - t_base: usize, - block: usize, - accs: &mut [f32], -) { - debug_assert_eq!(accs.len(), block * spec.out_channels); - for acc in accs.chunks_exact_mut(spec.out_channels) { - acc.copy_from_slice(spec.bias); - } - let pad = (spec.kernel_size - 1) / 2; - for k in 0..spec.kernel_size { - let src_t_at_s0 = t_base as isize + k as isize - pad as isize; - let kbase = k * spec.in_channels * spec.out_channels; - let krows = spec.kernel[kbase..kbase + spec.in_channels * spec.out_channels] - .chunks_exact(spec.out_channels); - for (in_c, krow) in krows.enumerate() { - for (s, acc) in accs.chunks_exact_mut(spec.out_channels).enumerate() { - let src_t = src_t_at_s0 + s as isize; - if src_t < 0 || src_t >= input.rows as isize { - continue; - } - let x = input.row(src_t as usize)[in_c]; - for (a, &w) in acc.iter_mut().zip(krow) { - *a = w.mul_add(x, *a); - } - } - } - } -} - -fn pool_rows_scalar(rows: &[f32], channels: usize, dst: &mut [f32]) { - dst.copy_from_slice(&rows[..channels]); - for row in rows[channels..].chunks_exact(channels) { - for (dst, &value) in dst.iter_mut().zip(row) { - *dst = dst.max(value); - } - } -} - -fn const_output_value(input: &Tensor<'_>, spec: &ConvSpec<'_>, out: &mut [f32]) { - let input_region = &input.repeated_tail; - debug_assert!(!input_region.is_empty()); - let input_row = input.row(input_region.range.start); - out.copy_from_slice(spec.bias); - for krow in spec - .kernel - .chunks_exact(spec.in_channels * spec.out_channels) - { - for (&x, weights) in input_row.iter().zip(krow.chunks_exact(spec.out_channels)) { - for (o, &w) in out.iter_mut().zip(weights) { - *o = w.mul_add(x, *o); - } - } - } - for v in out { - *v = super::activation::gelu(*v); - } -} - -fn row(data: &[f32], channels: usize, index: usize) -> &[f32] { - &data[index * channels..(index + 1) * channels] -} - -fn row_range(data: &[f32], channels: usize, rows: Range) -> &[f32] { - &data[rows.start * channels..rows.end * channels] -} - -fn row_range_mut(data: &mut [f32], channels: usize, rows: Range) -> &mut [f32] { - &mut data[rows.start * channels..rows.end * channels] -} - -fn accumulate_pool_rows(rows: &[f32], channels: usize, out_max: &mut [f32], out_avg: &mut [f32]) { - for row in rows.chunks_exact(channels) { - for ((mx, avg), &v) in out_max.iter_mut().zip(out_avg.iter_mut()).zip(row) { - *mx = mx.max(v); - *avg += v; - } - } -} - -fn accumulate_const_rows(value: &[f32], count: usize, out_max: &mut [f32], out_avg: &mut [f32]) { - for ((mx, avg), &v) in out_max.iter_mut().zip(out_avg.iter_mut()).zip(value) { - *mx = mx.max(v); - for _ in 0..count { - *avg += v; - } - } -} - -#[inline(always)] -fn pool_gelu_rows(simd: S, rows: &[f32], channels: usize, dst: &mut [f32]) { - debug_assert_eq!(dst.len(), channels); - let mut rows = rows.chunks_exact(channels); - let first = rows.next().expect("maxpool requires at least one row"); - store_gelu_row(simd, first, dst); - for row in rows { - max_gelu_row(simd, row, dst); - } -} - -#[inline(always)] -fn store_gelu_row(simd: S, row: &[f32], dst: &mut [f32]) { - for (d_c, a_c) in as_array_chunks_mut::<4>(dst) - .iter_mut() - .zip(as_array_chunks::<4>(row)) - { - let v = f32x4::from_slice(simd, a_c); - gelu_simd(simd, v).store_slice(d_c); - } -} - -#[inline(always)] -fn max_gelu_row(simd: S, row: &[f32], dst: &mut [f32]) { - for (d_c, a_c) in as_array_chunks_mut::<4>(dst) - .iter_mut() - .zip(as_array_chunks::<4>(row)) - { - let v = f32x4::from_slice(simd, a_c); - let g = gelu_simd(simd, v); - let dv = f32x4::from_slice(simd, d_c); - g.max(dv).store_slice(d_c); - } -} - -pub(crate) fn dense_forward(input: &[f32], kernel: &[f32], bias: &[f32], out: &mut [f32]) { - let in_len = input.len(); - let out_len = out.len(); - debug_assert_eq!(kernel.len(), in_len * out_len); - out.copy_from_slice(bias); - if out_len == 0 { - return; - } - for (&x, krow) in input.iter().zip(kernel.chunks_exact(out_len)) { - for (o, &w) in out.iter_mut().zip(krow) { - *o = w.mul_add(x, *o); - } - } -} diff --git a/src/model/reader.rs b/src/model/reader.rs deleted file mode 100644 index f12bb70..0000000 --- a/src/model/reader.rs +++ /dev/null @@ -1,52 +0,0 @@ -use super::embedded::MODEL_BYTES; - -pub(crate) fn read_int4_dequant(cursor: &mut usize, count: usize, scale: f32) -> Box<[f32]> { - let bytes = count.div_ceil(2); - let payload = &MODEL_BYTES[*cursor..*cursor + bytes]; - *cursor += bytes; - let mut out = Vec::with_capacity(count); - for &packed in payload { - if out.len() < count { - let lo = (packed & 0x0f) as i8 - 8; - out.push(lo as f32 * scale); - } - if out.len() < count { - let hi = (packed >> 4) as i8 - 8; - out.push(hi as f32 * scale); - } - } - out.into_boxed_slice() -} - -pub(crate) fn read_ternary_dequant(cursor: &mut usize, count: usize, scale: f32) -> Box<[f32]> { - let bytes = count.div_ceil(4); - let payload = &MODEL_BYTES[*cursor..*cursor + bytes]; - *cursor += bytes; - let mut out = Vec::with_capacity(count); - for &packed in payload { - for shift in [0u32, 2, 4, 6] { - if out.len() >= count { - break; - } - let code = (packed >> shift) & 0x03; - let v = match code { - 0 => -scale, - 2 => scale, - _ => 0.0, // 1 → 0; codes outside {0,1,2} treated as 0 too - }; - out.push(v); - } - } - out.into_boxed_slice() -} - -pub(crate) fn read_f32_array(cursor: &mut usize) -> [f32; N] { - let mut out = [0.0; N]; - let (chunks, remainder) = MODEL_BYTES[*cursor..*cursor + N * 4].as_chunks::<4>(); - debug_assert!(remainder.is_empty()); - for (value, bytes) in out.iter_mut().zip(chunks) { - *value = f32::from_le_bytes(*bytes); - } - *cursor += N * 4; - out -} diff --git a/src/model/runtime.rs b/src/model/runtime.rs deleted file mode 100644 index 10d1c7e..0000000 --- a/src/model/runtime.rs +++ /dev/null @@ -1,194 +0,0 @@ -use super::{ - activation::gelu, - constants::*, - embedded::MODEL_BYTES, - layers::{ - Tensor, conv_gelu_global_pool_tensor, conv_gelu_maxpool_tensor, dense_forward, - embed_position, - }, - reader::{read_f32_array, read_int4_dequant, read_ternary_dequant}, - tokenizer::tokenize, - window::TokenWindow, -}; -use std::sync::OnceLock; - -#[derive(Debug)] -pub(crate) struct Model { - /// 1024 × 24 dequantized embedding rows. - pub(crate) embedding: Box<[f32]>, - /// `[k][in_c][out_c]` — inner kernel row is contiguous over out_channels. - conv0_kernel: Box<[f32]>, - conv0_bias: [f32; CONV0], - conv1_kernel: Box<[f32]>, - conv1_bias: [f32; CONV1], - conv2_kernel: Box<[f32]>, - conv2_bias: [f32; CONV2], - /// (POOLED, DENSE) flattened. - dense0_kernel: Box<[f32]>, - dense0_bias: [f32; DENSE], - /// (DENSE, CLASSES) flattened. - pub(crate) output_kernel: Box<[f32]>, - output_bias: [f32; CLASSES], -} - -impl Model { - fn load() -> Self { - debug_assert!(MODEL_BYTES.starts_with(&MODEL_MAGIC), "bad MSQ1 magic"); - let mut cur = MODEL_MAGIC.len(); - let scales = read_f32_array::(&mut cur); - - // q_hash_embedding: weights [(1024, 24)] int4 - let embedding = read_int4_dequant(&mut cur, BINS * EMBED, scales[0]); - - // q_conv_0: weights [(7, 24, 64)] ternary, bias [(64,)] f32 - let conv0_kernel = read_ternary_dequant(&mut cur, CONV0_KERNEL * EMBED * CONV0, scales[1]); - let conv0_bias = read_f32_array::(&mut cur); - - // q_conv_1: weights [(5, 64, 128)] ternary, bias [(128,)] f32 - let conv1_kernel = read_ternary_dequant(&mut cur, CONV1_KERNEL * CONV0 * CONV1, scales[2]); - let conv1_bias = read_f32_array::(&mut cur); - - // q_conv_2: weights [(3, 128, 128)] ternary, bias [(128,)] f32 - let conv2_kernel = read_ternary_dequant(&mut cur, CONV2_KERNEL * CONV1 * CONV2, scales[3]); - let conv2_bias = read_f32_array::(&mut cur); - - // q_dense_0: weights [(256, 96)] ternary, bias [(96,)] f32 - let dense0_kernel = read_ternary_dequant(&mut cur, POOLED * DENSE, scales[4]); - let dense0_bias = read_f32_array::(&mut cur); - - // q_output: weights [(96, 48)] int4, bias [(48,)] f32 - let output_kernel = read_int4_dequant(&mut cur, DENSE * CLASSES, scales[5]); - let output_bias = read_f32_array::(&mut cur); - - debug_assert_eq!(cur, MODEL_BYTES.len(), "unexpected model payload length"); - - Self { - embedding, - conv0_kernel, - conv0_bias, - conv1_kernel, - conv1_bias, - conv2_kernel, - conv2_bias, - dense0_kernel, - dense0_bias, - output_kernel, - output_bias, - } - } - - pub(crate) fn get() -> &'static Self { - static MODEL: OnceLock = OnceLock::new(); - MODEL.get_or_init(Self::load) - } - - /// Run the full forward pass on a unit-id sequence using the fixed padded - /// shape used by the shipped Python evaluator. - pub(crate) fn logits(&self, units: &[i32]) -> [f32; CLASSES] { - let mut scratch = vec![0.0f32; INFERENCE_SCRATCH].into_boxed_slice(); - debug_assert!(scratch.len() >= INFERENCE_SCRATCH); - let t1 = MAX_UNITS / CONV0_POOL; - let t2 = t1 / CONV1_POOL; - let embed_len = MAX_UNITS * EMBED; - let pool0_len = t1 * CONV0; - let pool1_len = t2 * CONV1; - let unit_count = units.len().min(MAX_UNITS); - let (activations, conv_scratch) = scratch.split_at_mut(ACTIVATION_SCRATCH); - - { - let (embed_storage, pool0_storage) = activations.split_at_mut(embed_len); - - // 1) HashEmbedding + GELU. - let (embed_rows, embed_remainder) = embed_storage.as_chunks_mut::(); - debug_assert!(embed_remainder.is_empty()); - for (&id, dst) in units.iter().take(unit_count).zip(embed_rows.iter_mut()) { - debug_assert!(id >= 0); - embed_position(&self.embedding, id as u32, dst); - for v in dst.iter_mut() { - *v = gelu(*v); - } - } - if unit_count < MAX_UNITS { - embed_rows[unit_count].fill(0.0); - } - - // 2) Conv0 + GELU + MaxPool(4). - let pool0 = &mut pool0_storage[..pool0_len]; - let embed_materialized_rows = if unit_count < MAX_UNITS { - unit_count + 1 - } else { - MAX_UNITS - }; - let embed_tensor = Tensor::with_repeated_tail( - &embed_storage[..embed_materialized_rows * EMBED], - MAX_UNITS, - EMBED, - unit_count, - ); - let pool0_tensor = conv_gelu_maxpool_tensor( - embed_tensor, - &self.conv0_kernel, - CONV0_KERNEL, - CONV0, - &self.conv0_bias, - CONV0_POOL, - pool0, - &mut *conv_scratch, - ); - - // 3) Conv1 + GELU + MaxPool(2). - let pool1 = &mut embed_storage[..pool1_len]; - let pool1_tensor = conv_gelu_maxpool_tensor( - pool0_tensor, - &self.conv1_kernel, - CONV1_KERNEL, - CONV1, - &self.conv1_bias, - CONV1_POOL, - pool1, - &mut *conv_scratch, - ); - - // 4) Conv2 + GELU + GlobalMax/AvgPool. - let mut pooled = [0.0f32; POOLED]; - let (max_slice, avg_slice) = pooled.split_at_mut(CONV2); - conv_gelu_global_pool_tensor( - pool1_tensor, - &self.conv2_kernel, - CONV2_KERNEL, - CONV2, - &self.conv2_bias, - max_slice, - avg_slice, - pool0_storage, - &mut *conv_scratch, - ); - - // 5) Dense + GELU. - let mut dense0_out = [0.0f32; DENSE]; - dense_forward( - &pooled, - &self.dense0_kernel, - &self.dense0_bias, - &mut dense0_out, - ); - for v in &mut dense0_out { - *v = gelu(*v); - } - - // 6) Output logits. - let mut logits = [0.0f32; CLASSES]; - dense_forward( - &dense0_out, - &self.output_kernel, - &self.output_bias, - &mut logits, - ); - logits - } - } - - pub(crate) fn tokenize_units(&self, window: &TokenWindow) -> Vec { - tokenize(window) - } -} diff --git a/src/model/tests.rs b/src/model/tests.rs index 0d0ad1d..5ea240f 100644 --- a/src/model/tests.rs +++ b/src/model/tests.rs @@ -1,61 +1,7 @@ -use super::{ - constants::*, - layers::{Tensor, conv_gelu_global_pool_tensor, conv_gelu_maxpool_tensor}, - runtime::Model, - tokenizer::{hash_unit_bytes, tokenize}, - window::build_window, -}; use crate::Language; -use rand::{Rng, SeedableRng, rngs::StdRng}; use std::collections::HashSet; use std::{fs, path::Path}; -#[test] -fn loads_embedded_model() { - let model = Model::get(); - assert_eq!(model.embedding.len(), BINS * EMBED); - assert_eq!(model.output_kernel.len(), DENSE * CLASSES); -} - -#[test] -fn tokenizer_casefolds_and_isolates_brackets() { - let source = b"Foo(foo)\n"; - let window = build_window(source).expect("source should build a model window"); - let units = tokenize(&window); - - assert_eq!(units[0] as u32, hash_unit_bytes(b"foo") & WORD_MASK); - assert!(units.contains(&((BRACKET_FLAG | b'(' as u32) as i32))); - assert!(units.contains(&((BRACKET_FLAG | b')' as u32) as i32))); -} - -#[test] -fn tokenizer_matches_legacy_buffer_model_on_fuzzed_windows() { - let mut rng = StdRng::seed_from_u64(0x544f_4b45_4e49_5a45); - - for case in [ - &b"abc123!x"[..], - b"1.2 .. ...\r\n\tfoo(bar)[baz]{qux}", - b" indented\n next\n", - b"UPPER_lower 000.111 <=> -> ::", - ] { - let window = build_window(case).expect("case should build a model window"); - assert_eq!(tokenize(&window), legacy_tokenize_bytes(window.bytes())); - } - - for _ in 0..1024 { - let len = rng.gen_range(8..6000); - let mut source = Vec::with_capacity(len); - source.extend(std::iter::repeat_n(b' ', rng.gen_range(0..16))); - while source.len() < len { - source.push(rng.gen_range(0..=255)); - } - - if let Some(window) = build_window(&source) { - assert_eq!(tokenize(&window), legacy_tokenize_bytes(window.bytes())); - } - } -} - #[test] fn detects_rust_from_source() { let detection = crate::detect("use std::fmt;\nfn main() { println!(\"hi\"); }"); @@ -256,238 +202,6 @@ fn probabilities_sum_to_one_across_model_languages() { assert!((sum - 1.0).abs() < 1e-5, "{sum}"); } -#[test] -fn runtime_inference_accepts_short_sources() { - let source = "use std::fmt;\nfn main() { println!(\"hi\"); }\n"; - let Some(window) = build_window(source.as_bytes()) else { - panic!("expected source to build a model window"); - }; - let model = Model::get(); - let units = model.tokenize_units(&window); - assert!(units.len() < MAX_UNITS); - - let logits = model.logits(&units); - assert!(logits.iter().all(|logit| logit.is_finite())); -} - -#[test] -fn repeated_tensor_layers_match_full_convolution() { - for seed in 0..32 { - check_repeated_tensor_layers(seed, None); - } - check_repeated_tensor_layers(32, Some(128)); -} - -fn check_repeated_tensor_layers(seed: u64, tail_start: Option) { - let mut rng = StdRng::seed_from_u64(seed); - let seq_len = 128; - let in_channels = 8; - let mid_channels = 16; - let out_channels = 12; - let tail_start = tail_start.unwrap_or_else(|| rng.gen_range(8..seq_len - 8)); - - let mut input = random_f32s(&mut rng, seq_len * in_channels); - for value in &mut input[tail_start * in_channels..] { - *value = 0.0; - } - let input_tensor = Tensor::with_repeated_tail(&input, seq_len, in_channels, tail_start); - - let kernel0 = random_f32s(&mut rng, 5 * in_channels * mid_channels); - let bias0 = random_f32s(&mut rng, mid_channels); - let mut full_pool = vec![0.0; (seq_len / 4) * mid_channels]; - let mut const_pool = vec![0.0; full_pool.len()]; - let mut scratch = vec![0.0; 4 * mid_channels.max(out_channels)]; - let full_input_tensor = Tensor::with_repeated_tail(&input, seq_len, in_channels, seq_len); - conv_gelu_maxpool_tensor( - full_input_tensor, - &kernel0, - 5, - mid_channels, - &bias0, - 4, - &mut full_pool, - &mut scratch, - ); - let pool_tensor = conv_gelu_maxpool_tensor( - input_tensor, - &kernel0, - 5, - mid_channels, - &bias0, - 4, - &mut const_pool, - &mut scratch, - ); - let mut materialized_pool = vec![0.0; full_pool.len()]; - pool_tensor.copy_to_dense(&mut materialized_pool); - assert_f32s_eq(&materialized_pool, &full_pool); - - let kernel1 = random_f32s(&mut rng, 3 * mid_channels * out_channels); - let bias1 = random_f32s(&mut rng, out_channels); - let full_pool_tensor = - Tensor::with_repeated_tail(&full_pool, seq_len / 4, mid_channels, seq_len / 4); - let mut full_max = vec![0.0; out_channels]; - let mut full_avg = vec![0.0; out_channels]; - let mut full_tmp = vec![0.0; (seq_len / 4) * out_channels]; - conv_gelu_global_pool_tensor( - full_pool_tensor, - &kernel1, - 3, - out_channels, - &bias1, - &mut full_max, - &mut full_avg, - &mut full_tmp, - &mut scratch, - ); - - let mut const_max = vec![0.0; out_channels]; - let mut const_avg = vec![0.0; out_channels]; - let mut tmp = vec![0.0; (seq_len / 4) * out_channels]; - conv_gelu_global_pool_tensor( - pool_tensor, - &kernel1, - 3, - out_channels, - &bias1, - &mut const_max, - &mut const_avg, - &mut tmp, - &mut scratch, - ); - assert_f32s_eq(&const_max, &full_max); - assert_f32s_eq(&const_avg, &full_avg); -} - -fn random_f32s(rng: &mut StdRng, len: usize) -> Vec { - (0..len).map(|_| rng.gen_range(-0.25..0.25)).collect() -} - -fn legacy_tokenize_bytes(bytes: &[u8]) -> Vec { - let mut out: Vec = Vec::with_capacity(MAX_UNITS); - let mut word: Vec = Vec::new(); - let mut number: Vec = Vec::new(); - let mut punct: Vec = Vec::new(); - let mut at_line_start = true; - let mut indent_units: u32 = 0; - - for &raw_value in bytes { - let value = raw_value.to_ascii_lowercase(); - let is_letter = value.is_ascii_lowercase() || value == b'_'; - let is_digit = value.is_ascii_digit(); - let is_newline = value == b'\n'; - let is_cr = value == b'\r'; - let is_space = value == b' ' || value == b'\t'; - let is_bracket = matches!(value, b'(' | b')' | b'[' | b']' | b'{' | b'}'); - - if !is_letter { - legacy_flush(&mut word, &mut out, 0); - } - if !(is_digit || value == b'.') { - legacy_flush(&mut number, &mut out, NUM_FLAG); - } - let need_flush_punct = - is_letter || is_digit || is_space || is_newline || is_cr || is_bracket || value == b'.'; - if need_flush_punct { - legacy_flush(&mut punct, &mut out, PUNCT_FLAG); - } - - if out.len() >= MAX_UNITS { - break; - } - - if is_letter { - if at_line_start { - push_legacy_indent(&mut out, indent_units); - } - at_line_start = false; - indent_units = 0; - word.push(value); - continue; - } - if is_digit || value == b'.' { - if value == b'.' && number.is_empty() { - if at_line_start { - push_legacy_indent(&mut out, indent_units); - } - at_line_start = false; - indent_units = 0; - punct.push(value); - continue; - } - if at_line_start { - push_legacy_indent(&mut out, indent_units); - } - at_line_start = false; - indent_units = 0; - number.push(value); - continue; - } - if is_newline { - if at_line_start { - push_legacy_indent(&mut out, indent_units); - } - if out.len() < MAX_UNITS { - out.push(((b'\n' as u32) | PUNCT_FLAG) as i32); - } - at_line_start = true; - indent_units = 0; - continue; - } - if is_cr { - continue; - } - if at_line_start && is_space { - indent_units += if value == b' ' { 1 } else { 4 }; - continue; - } - if at_line_start { - push_legacy_indent(&mut out, indent_units); - } - at_line_start = false; - indent_units = 0; - if is_space { - let space_token = ((b' ' as u32) | PUNCT_FLAG) as i32; - if out.last() != Some(&space_token) && out.len() < MAX_UNITS { - out.push(space_token); - } - continue; - } - if is_bracket { - if out.len() < MAX_UNITS { - out.push(((value as u32) | BRACKET_FLAG) as i32); - } - continue; - } - punct.push(value); - } - - legacy_flush(&mut word, &mut out, 0); - legacy_flush(&mut number, &mut out, NUM_FLAG); - legacy_flush(&mut punct, &mut out, PUNCT_FLAG); - out -} - -fn legacy_flush(buffer: &mut Vec, out: &mut Vec, flag: u32) { - if !buffer.is_empty() && out.len() < MAX_UNITS { - out.push(((hash_unit_bytes(buffer) & WORD_MASK) | flag) as i32); - } - buffer.clear(); -} - -fn push_legacy_indent(out: &mut Vec, indent: u32) { - if indent > 0 && out.len() < MAX_UNITS { - out.push((indent.min(63) | INDENT_FLAG) as i32); - } -} - -fn assert_f32s_eq(actual: &[f32], expected: &[f32]) { - assert_eq!(actual.len(), expected.len()); - for (index, (&actual, &expected)) in actual.iter().zip(expected).enumerate() { - assert_eq!(actual.to_bits(), expected.to_bits(), "index {index}"); - } -} - #[test] fn empty_input_returns_empty_detection() { assert!(crate::detect("").top_languages().next().is_none()); diff --git a/src/model/tokenizer.rs b/src/model/tokenizer.rs deleted file mode 100644 index e243971..0000000 --- a/src/model/tokenizer.rs +++ /dev/null @@ -1,161 +0,0 @@ -use super::{ - constants::{BRACKET_FLAG, INDENT_FLAG, MAX_UNITS, NUM_FLAG, PUNCT_FLAG, WORD_MASK}, - window::TokenWindow, -}; - -pub(crate) fn hash_unit_bytes(bytes: &[u8]) -> u32 { - const PRIME: u32 = 2_654_435_761; - let mut h = 0u32; - for &b in bytes { - h = h.wrapping_mul(PRIME).wrapping_add(b as u32); - } - h -} - -fn push_indent_unit(out: &mut Vec, indent: u32) { - if indent > 0 && out.len() < MAX_UNITS { - out.push((indent.min(63) | INDENT_FLAG) as i32); - } -} - -#[derive(Default, Clone, Copy, PartialEq, Eq)] -enum TokenKind { - #[default] - Empty, - Word, - Number, - Punct, -} - -#[derive(Default)] -struct TokenBuffer { - kind: TokenKind, - bytes: Vec, -} - -impl TokenBuffer { - fn is_number(&self) -> bool { - self.kind == TokenKind::Number - } - - fn flush(&mut self, out: &mut Vec) { - let flag = match self.kind { - TokenKind::Empty => return, - TokenKind::Word => 0, - TokenKind::Number => NUM_FLAG, - TokenKind::Punct => PUNCT_FLAG, - }; - if out.len() < MAX_UNITS { - out.push(((hash_unit_bytes(&self.bytes) & WORD_MASK) | flag) as i32); - } - self.bytes.clear(); - self.kind = TokenKind::Empty; - } - - fn push(&mut self, kind: TokenKind, value: u8, out: &mut Vec) { - if self.kind != kind { - self.flush(out); - self.kind = kind; - } - self.bytes.push(value); - } -} - -/// Production word-unit tokenizer version 3. -/// -/// Case-folds word hashes and emits unambiguous brackets as BRACKET_FLAG tokens. -pub(crate) fn tokenize(window: &TokenWindow) -> Vec { - let bytes = window.bytes(); - let mut out: Vec = Vec::with_capacity(MAX_UNITS); - let mut current = TokenBuffer::default(); - let mut at_line_start = true; - let mut indent_units: u32 = 0; - - for &raw_value in bytes { - let value = raw_value.to_ascii_lowercase(); - let is_letter = value.is_ascii_lowercase() || value == b'_'; - let is_digit = value.is_ascii_digit(); - let is_newline = value == b'\n'; - let is_cr = value == b'\r'; - let is_space = value == b' ' || value == b'\t'; - let is_bracket = matches!(value, b'(' | b')' | b'[' | b']' | b'{' | b'}'); - - if out.len() >= MAX_UNITS { - break; - } - - if is_letter { - if at_line_start { - push_indent_unit(&mut out, indent_units); - } - at_line_start = false; - indent_units = 0; - current.push(TokenKind::Word, value, &mut out); - continue; - } - if is_digit || value == b'.' { - if value == b'.' && !current.is_number() { - if at_line_start { - push_indent_unit(&mut out, indent_units); - } - at_line_start = false; - indent_units = 0; - current.flush(&mut out); - current.push(TokenKind::Punct, value, &mut out); - continue; - } - if at_line_start { - push_indent_unit(&mut out, indent_units); - } - at_line_start = false; - indent_units = 0; - current.push(TokenKind::Number, value, &mut out); - continue; - } - if is_newline { - current.flush(&mut out); - if at_line_start { - push_indent_unit(&mut out, indent_units); - } - if out.len() < MAX_UNITS { - out.push(((b'\n' as u32) | PUNCT_FLAG) as i32); - } - at_line_start = true; - indent_units = 0; - continue; - } - if is_cr { - current.flush(&mut out); - continue; - } - if at_line_start && is_space { - indent_units += if value == b' ' { 1 } else { 4 }; - continue; - } - if at_line_start { - push_indent_unit(&mut out, indent_units); - } - at_line_start = false; - indent_units = 0; - if is_space { - current.flush(&mut out); - let space_token = ((b' ' as u32) | PUNCT_FLAG) as i32; - if out.last() != Some(&space_token) && out.len() < MAX_UNITS { - out.push(space_token); - } - continue; - } - if is_bracket { - current.flush(&mut out); - if out.len() < MAX_UNITS { - out.push(((value as u32) | BRACKET_FLAG) as i32); - } - continue; - } - current.push(TokenKind::Punct, value, &mut out); - } - - current.flush(&mut out); - - out -}