From 08a941a14e0fe281853deee197b3eae1029ab695 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 10 Jun 2026 21:49:51 +0200 Subject: [PATCH] feat: added standalone snapcompact package and model-specific frame shaping - Added a new @oh-my-pi/snapcompact package and redirected compaction call sites to it. - Added provider-aware snapcompact shape resolution for model-specific mixed-frame behavior. - Added optional image detail support by extending ImageContent and passing hints through OpenAI providers. - Added native snapcompact render options, including 5x8/8x8 font loading and palette/geometry controls. --- README.md | 4 + bun.lock | 17 + crates/pi-natives/src/fonts/unscii-8.hex | 255 +++++++ crates/pi-natives/src/snapcompact.rs | 501 ++++++++++++-- docs/compaction.md | 8 +- .../examples/hello-extension/package.json | 1 + .../mini-marketplace/my-plugin/package.json | 1 + docs/skills/examples/safety-hook/package.json | 1 + package.json | 2 + packages/agent/CHANGELOG.md | 7 +- packages/agent/package.json | 1 + packages/agent/src/compaction/compaction.ts | 2 +- packages/agent/src/compaction/index.ts | 1 - .../compaction/prompts/snapcompact-summary.md | 6 +- packages/agent/src/compaction/snapcompact.ts | 316 --------- .../agent/test/snapcompact-frames.test.ts | 44 ++ packages/ai/CHANGELOG.md | 4 + .../ai/src/providers/openai-completions.ts | 2 + .../src/providers/openai-responses-shared.ts | 4 +- packages/ai/src/types.ts | 6 + packages/coding-agent/CHANGELOG.md | 1 + .../extensions/with-deps/package.json | 1 + packages/coding-agent/package.json | 1 + .../coding-agent/src/session/agent-session.ts | 6 +- .../src/session/session-manager.ts | 2 +- .../test/job-poll-displacement.test.ts | 12 +- .../test/streaming-preview-height.test.ts | 17 +- packages/natives/CHANGELOG.md | 2 +- packages/natives/native/index.d.ts | 38 +- packages/snapcompact/.gitignore | 2 + packages/snapcompact/CHANGELOG.md | 10 + packages/snapcompact/package.json | 62 ++ .../snapcompact/research/anthropic_api.py | 74 ++ packages/snapcompact/research/bdf.py | 247 +++++++ .../snapcompact/research/exp01_patchalign.py | 284 ++++++++ .../snapcompact/research/exp02_surprisal.py | 427 ++++++++++++ .../snapcompact/research/exp03_numhard.py | 382 +++++++++++ packages/snapcompact/research/exp04_layout.py | 374 ++++++++++ .../snapcompact/research/exp05_anchors.py | 344 ++++++++++ .../snapcompact/research/exp06_rolecolor.py | 417 ++++++++++++ .../snapcompact/research/exp07_readtax.py | 273 ++++++++ .../snapcompact/research/exp08_foveate.py | 388 +++++++++++ .../snapcompact/research/exp09_cacheappend.py | 365 ++++++++++ .../snapcompact/research/exp10_profiles.py | 358 ++++++++++ .../snapcompact/research/exp11_memhier.py | 301 +++++++++ .../snapcompact/research/exp12_arbitrage.py | 320 +++++++++ .../snapcompact/research/exp13_extractive.py | 255 +++++++ .../snapcompact/research/exp14_bestgpt.py | 475 +++++++++++++ .../snapcompact/research/exp15_bestgemini.py | 478 +++++++++++++ .../snapcompact/research/exp16_bestfable.py | 476 +++++++++++++ .../snapcompact/research/exp17_bestopus.py | 419 ++++++++++++ .../snapcompact/research/exp18_bestkimi.py | 439 ++++++++++++ .../snapcompact/research/exp19_bestglm.py | 477 +++++++++++++ packages/snapcompact/research/exp20_8x8u.py | 362 ++++++++++ .../snapcompact/research/exp21_braille.py | 276 ++++++++ packages/snapcompact/research/exp22_ttf6pt.py | 243 +++++++ packages/snapcompact/research/final.py | 331 +++++++++ packages/snapcompact/research/mono.py | 154 +++++ .../research/prompts/exp02-qa-image.md | 8 + .../research/prompts/exp04-qa-image.md | 8 + .../research/prompts/exp05-qa-image-ctl.md | 10 + .../research/prompts/exp05-qa-image.md | 13 + .../research/prompts/exp06-prov-image.md | 6 + .../research/prompts/exp06-qa-image-tag.md | 6 + .../research/prompts/exp06-qa-image.md | 6 + .../research/prompts/exp07-answer-bands.md | 4 + .../research/prompts/exp07-locate.md | 5 + .../research/prompts/exp07-qa-image.md | 7 + .../research/prompts/exp08-archive-eager.md | 10 + .../research/prompts/exp08-archive-phrase.md | 9 + .../research/prompts/exp08-archive.md | 7 + .../research/prompts/exp08-zoom.md | 6 + .../research/prompts/exp09-frame.md | 3 + .../research/prompts/exp09-page.md | 1 + .../snapcompact/research/prompts/exp09-qa.md | 6 + .../research/prompts/exp11-qa-hier.md | 12 + .../research/prompts/exp13-extract.md | 10 + .../research/prompts/exp19-qa-doc.md | 8 + .../research/prompts/exp21-qa-braille.md | 11 + .../research/prompts/qa-image-cols.md | 6 + .../research/prompts/qa-image-multi.md | 6 + .../snapcompact/research/prompts/qa-image.md | 6 + .../research/prompts/qa-remote-compact.md | 8 + .../snapcompact/research/prompts/qa-text.md | 10 + .../research/prompts/session-frame.md | 5 + packages/snapcompact/research/providers.py | 261 +++++++ packages/snapcompact/research/run.py | 353 ++++++++++ .../snapcompact_3d_activation_html.py | 192 ++++++ .../research/snapcompact_3d_activation_viz.py | 230 +++++++ .../research/snapcompact_activation_probe.py | 225 ++++++ .../snapcompact_blackbox_occlusion.py | 261 +++++++ .../research/snapcompact_blog_viz.py | 231 +++++++ .../snapcompact_carrier_convergence.py | 212 ++++++ .../research/snapcompact_convergence_3d.py | 187 +++++ .../snapcompact_convergence_extras.py | 207 ++++++ .../research/snapcompact_convergence_viz.py | 219 ++++++ .../snapcompact_lockon_anatomy_viz.py | 288 ++++++++ .../research/snapcompact_logit_lens_dump.py | 143 ++++ .../research/snapcompact_logit_lens_viz.py | 169 +++++ .../research/snapcompact_materialize_sweep.py | 226 +++++++ .../research/snapcompact_materialize_viz.py | 171 +++++ .../research/snapcompact_pricing_viz.py | 137 ++++ .../snapcompact_qwen_control_intervention.py | 363 ++++++++++ .../snapcompact_qwen_spotlight_viz.py | 278 ++++++++ .../research/snapcompact_r2_chord.py | 265 ++++++++ .../research/snapcompact_r2_crystal.py | 338 ++++++++++ .../research/snapcompact_r2_filmstrip.py | 256 +++++++ .../research/snapcompact_r2_hero.py | 497 ++++++++++++++ .../research/snapcompact_r2_metro.py | 348 ++++++++++ .../research/snapcompact_tensor_heatmap.py | 354 ++++++++++ .../research/snapcompact_text_image_3d_viz.py | 211 ++++++ .../snapcompact_text_image_compare.py | 489 ++++++++++++++ .../research/snapcompact_token_entry_dump.py | 158 +++++ .../research/snapcompact_token_entry_viz.py | 225 ++++++ .../research/snapcompact_viz_atlas.py | 377 +++++++++++ .../research/snapcompact_viz_circuit.py | 427 ++++++++++++ .../research/snapcompact_viz_city.py | 237 +++++++ .../research/snapcompact_viz_explainer.py | 444 ++++++++++++ .../research/snapcompact_viz_glass_stack.py | 255 +++++++ .../research/snapcompact_viz_glyph_matrix.py | 397 +++++++++++ .../research/snapcompact_viz_radial.py | 219 ++++++ .../research/snapcompact_viz_token_grid.py | 396 +++++++++++ .../research/snapcompact_viz_volume.py | 235 +++++++ .../research/snapcompact_viz_waterfall.py | 301 +++++++++ packages/snapcompact/research/squad.py | 120 ++++ packages/snapcompact/src/index.ts | 1 + .../src/prompts/file-operations.md | 10 + .../src/prompts/snapcompact-summary.md | 17 + packages/snapcompact/src/snapcompact.ts | 638 ++++++++++++++++++ .../test/snapcompact.test.ts | 191 ++++-- packages/snapcompact/tsconfig.json | 7 + packages/snapcompact/tsconfig.publish.json | 25 + .../typescript-edit-benchmark/package.json | 1 + packages/utils/test/prompt.test.ts | 2 +- python/robomp/web/package.json | 1 + scripts/ci-release-publish.ts | 1 + scripts/install-tests/run-ci.sh | 6 +- 137 files changed, 21786 insertions(+), 489 deletions(-) create mode 100644 crates/pi-natives/src/fonts/unscii-8.hex delete mode 100644 packages/agent/src/compaction/snapcompact.ts create mode 100644 packages/agent/test/snapcompact-frames.test.ts create mode 100644 packages/snapcompact/.gitignore create mode 100644 packages/snapcompact/CHANGELOG.md create mode 100644 packages/snapcompact/package.json create mode 100644 packages/snapcompact/research/anthropic_api.py create mode 100644 packages/snapcompact/research/bdf.py create mode 100644 packages/snapcompact/research/exp01_patchalign.py create mode 100644 packages/snapcompact/research/exp02_surprisal.py create mode 100644 packages/snapcompact/research/exp03_numhard.py create mode 100644 packages/snapcompact/research/exp04_layout.py create mode 100644 packages/snapcompact/research/exp05_anchors.py create mode 100644 packages/snapcompact/research/exp06_rolecolor.py create mode 100644 packages/snapcompact/research/exp07_readtax.py create mode 100644 packages/snapcompact/research/exp08_foveate.py create mode 100644 packages/snapcompact/research/exp09_cacheappend.py create mode 100644 packages/snapcompact/research/exp10_profiles.py create mode 100644 packages/snapcompact/research/exp11_memhier.py create mode 100644 packages/snapcompact/research/exp12_arbitrage.py create mode 100644 packages/snapcompact/research/exp13_extractive.py create mode 100644 packages/snapcompact/research/exp14_bestgpt.py create mode 100644 packages/snapcompact/research/exp15_bestgemini.py create mode 100644 packages/snapcompact/research/exp16_bestfable.py create mode 100644 packages/snapcompact/research/exp17_bestopus.py create mode 100644 packages/snapcompact/research/exp18_bestkimi.py create mode 100644 packages/snapcompact/research/exp19_bestglm.py create mode 100644 packages/snapcompact/research/exp20_8x8u.py create mode 100644 packages/snapcompact/research/exp21_braille.py create mode 100644 packages/snapcompact/research/exp22_ttf6pt.py create mode 100644 packages/snapcompact/research/final.py create mode 100644 packages/snapcompact/research/mono.py create mode 100644 packages/snapcompact/research/prompts/exp02-qa-image.md create mode 100644 packages/snapcompact/research/prompts/exp04-qa-image.md create mode 100644 packages/snapcompact/research/prompts/exp05-qa-image-ctl.md create mode 100644 packages/snapcompact/research/prompts/exp05-qa-image.md create mode 100644 packages/snapcompact/research/prompts/exp06-prov-image.md create mode 100644 packages/snapcompact/research/prompts/exp06-qa-image-tag.md create mode 100644 packages/snapcompact/research/prompts/exp06-qa-image.md create mode 100644 packages/snapcompact/research/prompts/exp07-answer-bands.md create mode 100644 packages/snapcompact/research/prompts/exp07-locate.md create mode 100644 packages/snapcompact/research/prompts/exp07-qa-image.md create mode 100644 packages/snapcompact/research/prompts/exp08-archive-eager.md create mode 100644 packages/snapcompact/research/prompts/exp08-archive-phrase.md create mode 100644 packages/snapcompact/research/prompts/exp08-archive.md create mode 100644 packages/snapcompact/research/prompts/exp08-zoom.md create mode 100644 packages/snapcompact/research/prompts/exp09-frame.md create mode 100644 packages/snapcompact/research/prompts/exp09-page.md create mode 100644 packages/snapcompact/research/prompts/exp09-qa.md create mode 100644 packages/snapcompact/research/prompts/exp11-qa-hier.md create mode 100644 packages/snapcompact/research/prompts/exp13-extract.md create mode 100644 packages/snapcompact/research/prompts/exp19-qa-doc.md create mode 100644 packages/snapcompact/research/prompts/exp21-qa-braille.md create mode 100644 packages/snapcompact/research/prompts/qa-image-cols.md create mode 100644 packages/snapcompact/research/prompts/qa-image-multi.md create mode 100644 packages/snapcompact/research/prompts/qa-image.md create mode 100644 packages/snapcompact/research/prompts/qa-remote-compact.md create mode 100644 packages/snapcompact/research/prompts/qa-text.md create mode 100644 packages/snapcompact/research/prompts/session-frame.md create mode 100644 packages/snapcompact/research/providers.py create mode 100644 packages/snapcompact/research/run.py create mode 100644 packages/snapcompact/research/snapcompact_3d_activation_html.py create mode 100644 packages/snapcompact/research/snapcompact_3d_activation_viz.py create mode 100644 packages/snapcompact/research/snapcompact_activation_probe.py create mode 100644 packages/snapcompact/research/snapcompact_blackbox_occlusion.py create mode 100644 packages/snapcompact/research/snapcompact_blog_viz.py create mode 100644 packages/snapcompact/research/snapcompact_carrier_convergence.py create mode 100644 packages/snapcompact/research/snapcompact_convergence_3d.py create mode 100644 packages/snapcompact/research/snapcompact_convergence_extras.py create mode 100644 packages/snapcompact/research/snapcompact_convergence_viz.py create mode 100644 packages/snapcompact/research/snapcompact_lockon_anatomy_viz.py create mode 100644 packages/snapcompact/research/snapcompact_logit_lens_dump.py create mode 100644 packages/snapcompact/research/snapcompact_logit_lens_viz.py create mode 100644 packages/snapcompact/research/snapcompact_materialize_sweep.py create mode 100644 packages/snapcompact/research/snapcompact_materialize_viz.py create mode 100644 packages/snapcompact/research/snapcompact_pricing_viz.py create mode 100644 packages/snapcompact/research/snapcompact_qwen_control_intervention.py create mode 100644 packages/snapcompact/research/snapcompact_qwen_spotlight_viz.py create mode 100755 packages/snapcompact/research/snapcompact_r2_chord.py create mode 100755 packages/snapcompact/research/snapcompact_r2_crystal.py create mode 100755 packages/snapcompact/research/snapcompact_r2_filmstrip.py create mode 100755 packages/snapcompact/research/snapcompact_r2_hero.py create mode 100755 packages/snapcompact/research/snapcompact_r2_metro.py create mode 100644 packages/snapcompact/research/snapcompact_tensor_heatmap.py create mode 100644 packages/snapcompact/research/snapcompact_text_image_3d_viz.py create mode 100644 packages/snapcompact/research/snapcompact_text_image_compare.py create mode 100644 packages/snapcompact/research/snapcompact_token_entry_dump.py create mode 100644 packages/snapcompact/research/snapcompact_token_entry_viz.py create mode 100644 packages/snapcompact/research/snapcompact_viz_atlas.py create mode 100644 packages/snapcompact/research/snapcompact_viz_circuit.py create mode 100644 packages/snapcompact/research/snapcompact_viz_city.py create mode 100644 packages/snapcompact/research/snapcompact_viz_explainer.py create mode 100644 packages/snapcompact/research/snapcompact_viz_glass_stack.py create mode 100644 packages/snapcompact/research/snapcompact_viz_glyph_matrix.py create mode 100644 packages/snapcompact/research/snapcompact_viz_radial.py create mode 100644 packages/snapcompact/research/snapcompact_viz_token_grid.py create mode 100644 packages/snapcompact/research/snapcompact_viz_volume.py create mode 100755 packages/snapcompact/research/snapcompact_viz_waterfall.py create mode 100644 packages/snapcompact/research/squad.py create mode 100644 packages/snapcompact/src/index.ts create mode 100644 packages/snapcompact/src/prompts/file-operations.md create mode 100644 packages/snapcompact/src/prompts/snapcompact-summary.md create mode 100644 packages/snapcompact/src/snapcompact.ts rename packages/{agent => snapcompact}/test/snapcompact.test.ts (64%) create mode 100644 packages/snapcompact/tsconfig.json create mode 100644 packages/snapcompact/tsconfig.publish.json diff --git a/README.md b/README.md index a9fbe3395..ac6ec2743 100644 --- a/README.md +++ b/README.md @@ -482,12 +482,16 @@ For architecture and contribution guidelines, see [packages/coding-agent/DEVELOP | Package | Description | | --------------------------------------------------------- | -------------------------------------------------------------------------- | | **[@oh-my-pi/pi-ai](packages/ai)** | Multi-provider LLM client with streaming and model/provider integration | +| **[@oh-my-pi/pi-catalog](packages/catalog)** | Model catalog: bundled model database, provider descriptors, and identity | | **[@oh-my-pi/pi-agent-core](packages/agent)** | Agent runtime with tool calling and state management | | **[@oh-my-pi/pi-coding-agent](packages/coding-agent)** | Interactive coding agent CLI and SDK | | **[@oh-my-pi/pi-tui](packages/tui)** | Terminal UI library with differential rendering | | **[@oh-my-pi/pi-natives](packages/natives)** | N-API bindings for grep, shell, image, text, syntax highlighting, and more | | **[@oh-my-pi/omp-stats](packages/stats)** | Local observability dashboard for AI usage statistics | | **[@oh-my-pi/pi-utils](packages/utils)** | Shared utilities (logging, streams, dirs/env/process helpers) | +| **[@oh-my-pi/hashline](packages/hashline)** | Line-anchored patch language and applier behind the `edit` tool | +| **[@oh-my-pi/pi-mnemopi](packages/mnemopi)** | Local SQLite memory engine for Oh My Pi agents | +| **[@oh-my-pi/pi-snapcompact](packages/snapcompact)** | SQuAD eval suite for snapcompact bitmap-frame context compression | | **[@oh-my-pi/swarm-extension](packages/swarm-extension)** | Swarm orchestration extension package | ### Rust Crates diff --git a/bun.lock b/bun.lock index d097dee50..5fe030003 100644 --- a/bun.lock +++ b/bun.lock @@ -21,6 +21,7 @@ "@oh-my-pi/pi-catalog": "catalog:", "@oh-my-pi/pi-natives": "catalog:", "@oh-my-pi/pi-utils": "catalog:", + "@oh-my-pi/snapcompact": "catalog:", "@opentelemetry/api": "catalog:", }, "devDependencies": { @@ -76,6 +77,7 @@ "@oh-my-pi/pi-natives": "catalog:", "@oh-my-pi/pi-tui": "catalog:", "@oh-my-pi/pi-utils": "catalog:", + "@oh-my-pi/snapcompact": "catalog:", "@opentelemetry/api": "catalog:", "@opentelemetry/context-async-hooks": "catalog:", "@opentelemetry/exporter-trace-otlp-proto": "catalog:", @@ -141,6 +143,18 @@ "@types/bun": "catalog:", }, }, + "packages/snapcompact": { + "name": "@oh-my-pi/snapcompact", + "version": "15.10.12", + "dependencies": { + "@oh-my-pi/pi-ai": "catalog:", + "@oh-my-pi/pi-natives": "catalog:", + "@oh-my-pi/pi-utils": "catalog:", + }, + "devDependencies": { + "@types/bun": "catalog:", + }, + }, "packages/stats": { "name": "@oh-my-pi/omp-stats", "version": "15.10.12", @@ -276,6 +290,7 @@ "@oh-my-pi/pi-natives": "15.10.12", "@oh-my-pi/pi-tui": "15.10.12", "@oh-my-pi/pi-utils": "15.10.12", + "@oh-my-pi/snapcompact": "15.10.12", "@opentelemetry/api": "^1.9.1", "@opentelemetry/context-async-hooks": "^2.7.1", "@opentelemetry/exporter-trace-otlp-proto": "^0.218.0", @@ -672,6 +687,8 @@ "@oh-my-pi/pi-utils": ["@oh-my-pi/pi-utils@workspace:packages/utils"], + "@oh-my-pi/snapcompact": ["@oh-my-pi/snapcompact@workspace:packages/snapcompact"], + "@oh-my-pi/swarm-extension": ["@oh-my-pi/swarm-extension@workspace:packages/swarm-extension"], "@oh-my-pi/typescript-edit-benchmark": ["@oh-my-pi/typescript-edit-benchmark@workspace:packages/typescript-edit-benchmark"], diff --git a/crates/pi-natives/src/fonts/unscii-8.hex b/crates/pi-natives/src/fonts/unscii-8.hex new file mode 100644 index 000000000..cc21a2b06 --- /dev/null +++ b/crates/pi-natives/src/fonts/unscii-8.hex @@ -0,0 +1,255 @@ +00001:E080EA2AEE0A0A00 +00002:E080EA2AE40A0A00 +00003:E080CA8AE40A0A00 +00004:E080CE84E4040400 +00005:E080CE8AEA0E0400 +00006:E0A0EAAAAC0A0A00 +00007:C0A0C8A8C8080E00 +00008:C0A0CEA8CE020E00 +00009:A0A0EEA4A4040400 +0000A:80808E88EC080800 +0000B:A0A0AEA444040400 +0000C:E080CE888C080800 +0000D:E0808E8AEE0C0A00 +0000E:E080EE2AEA0A0E00 +0000F:E080EE24E4040E00 +00010:C0A0A8A8C8080E00 +00011:C0A0A4ACC4040E00 +00012:C0A0AEA2CE080E00 +00013:C0A0AEA2C6020E00 +00014:C0A0AAAACE020200 +00015:E0A0AAAAAC0A0A00 +00016:E080EA2AEE040400 +00017:E080CC8AEC0A0C00 +00018:E0808E8AEA0A0A00 +00019:E080CA8EEA0A0A00 +0001A:3C66663018001800 +0001B:E080CE88E8080E00 +0001C:E080CE888E020E00 +0001D:E080AEA8EE020E00 +0001E:E0A0EEC8AE020E00 +0001F:A0A0AEA8EE020E00 +00020:0000000000000000 +00021:1818181818001800 +00022:6666660000000000 +00023:6C6CFE6CFE6C6C00 +00024:183E603C067C1800 +00025:00C6CC183066C600 +00026:386C3876DCCC7600 +00027:1818300000000000 +00028:0C18303030180C00 +00029:30180C0C0C183000 +0002A:00663CFF3C660000 +0002B:0018187E18180000 +0002C:0000000000181830 +0002D:0000007E00000000 +0002E:0000000000181800 +0002F:03060C183060C000 +00030:3C666E7666663C00 +00031:1838181818187E00 +00032:3C660C1830607E00 +00033:3C66061C06663C00 +00034:1C3C6CCCFE0C0C00 +00035:7E607C0606663C00 +00036:1C30607C66663C00 +00037:7E06060C18181800 +00038:3C66663C66663C00 +00039:3C66663E060C3800 +0003A:0018180000181800 +0003B:0018180000181830 +0003C:0C18306030180C00 +0003D:00007E007E000000 +0003E:6030180C18306000 +0003F:3C66060C18001800 +00040:7CC6DEDEDEC07C00 +00041:183C66667E666600 +00042:7C66667C66667C00 +00043:3C66606060663C00 +00044:786C6666666C7800 +00045:7E60607C60607E00 +00046:7E60607C60606000 +00047:3C66606E66663E00 +00048:6666667E66666600 +00049:7E18181818187E00 +0004A:0606060606663C00 +0004B:C6CCD8F0D8CCC600 +0004C:6060606060607E00 +0004D:C6EEFED6C6C6C600 +0004E:C6E6F6DECEC6C600 +0004F:3C66666666663C00 +00050:7C66667C60606000 +00051:3C666666666C3600 +00052:7C66667C6C666600 +00053:3C66603C06663C00 +00054:7E18181818181800 +00055:6666666666663C00 +00056:66666666663C1800 +00057:C6C6C6D6FEEEC600 +00058:C3663C183C66C300 +00059:C3663C1818181800 +0005A:7E060C1830607E00 +0005B:3C30303030303C00 +0005C:C06030180C060300 +0005D:3C0C0C0C0C0C3C00 +0005E:10386CC600000000 +0005F:00000000000000FF +00060:180C060000000000 +00061:00003C063E663E00 +00062:60607C6666667C00 +00063:00003C6060603C00 +00064:06063E6666663E00 +00065:00003C667E603C00 +00066:1C307C3030303000 +00067:00003E66663E067C +00068:60607C6666666600 +00069:1800381818181E00 +0006A:0C000C0C0C0C0C78 +0006B:6060666C786C6600 +0006C:3818181818181E00 +0006D:0000CCFED6D6C600 +0006E:00007C6666666600 +0006F:00003C6666663C00 +00070:00007C66667C6060 +00071:00003E66663E0606 +00072:00007C6660606000 +00073:00003E603C067C00 +00074:30307E3030301E00 +00075:0000666666663E00 +00076:00006666663C1800 +00077:0000C6C6D67C6C00 +00078:0000C66C386CC600 +00079:00006666663E063C +0007A:00007E0C18307E00 +0007B:0E18187018180E00 +0007C:1818181818181800 +0007D:7018180E18187000 +0007E:76DC000000000000 +0007F:C0A0AEA4C4040400 +00080:E0A0EEAAEA0A0E00 +00081:E0A0E4ACE4040E00 +00082:E0A0EEA2EE080E00 +00083:E0A0EEA2EE020E00 +00084:E0404E4AEA0A0A00 +00085:E0A0A8A8A8080E00 +00086:E080EE28EE020E00 +00087:E080CE88EE020E00 +00088:A0A0EEA8AE020E00 +00089:A0A0E2A2A20A0E00 +0008A:A0A0AEA84E020E00 +0008B:E0A0EC8A8A0A0C00 +0008C:E0A0EA8A8A0A0E00 +0008D:C0A0CEA4A4040E00 +0008E:E080EE22EE080E00 +0008F:E080EE22E6020E00 +00090:C0A0AEA8C8080E00 +00091:E0A0E48C84040E00 +00092:E0A0EE828E080E00 +00093:E080EE28EC080E00 +00094:E0808E88E8080E00 +00095:A0E0EAAAAE0E0A00 +00096:E080EE2AEE080800 +00097:E080CE8AEE080800 +00098:E0A0EE2AEE0A0E00 +00099:E0A0EE2AEE020E00 +0009A:E0A0EE2AEE0A0A00 +0009B:E0808E88EE020E00 +0009C:E080EE24E4040400 +0009D:E0A0AEA8EE020E00 +0009E:E0A0EA8E8E0A0A00 +0009F:E0A0EEAAAE080800 +000A0:0000000000000000 +000A1:1800181818181800 +000A2:00187ED8D87E1800 +000A3:386C60F06066FC00 +000A4:663C663C66000000 +000A5:C3663C183C181800 +000A6:1818180018181800 +000A7:3C603C66663C063C +000A8:6600000000000000 +000A9:7E819DB1B19D817E +000AA:3C6C6C3E007E0000 +000AB:003366CC66330000 +000AC:007E060600000000 +000AD:0000003C00000000 +000AE:7E81B9A5B9A5817E +000AF:7E00000000000000 +000B0:3C663C0000000000 +000B1:18187E1818007E00 +000B2:7018306078000000 +000B3:780C180C78000000 +000B4:0C18300000000000 +000B5:00006666667C60C0 +000B6:3E7A7A3A1A1A1A00 +000B7:0000001800000000 +000B8:0000000000000818 +000B9:3070303030000000 +000BA:386C6C38007C0000 +000BB:00CC663366CC0000 +000BC:40C64C5A366ACF02 +000BD:40C64C5E3366CC0F +000BE:C0266C3AF66ACF02 +000BF:1800183060663C00 +000C0:70003C667E666600 +000C1:0E003C667E666600 +000C2:1866003C667E6600 +000C3:76DC003C667E6600 +000C4:66003C667E666600 +000C5:1818003C667E6600 +000C6:3F6CCCFECCCCCF00 +000C7:3C66606060663C18 +000C8:7000FEC0F8C0FE00 +000C9:0E00FEC0F8C0FE00 +000CA:186600FEF0C0FE00 +000CB:6600FEC0F8C0FE00 +000CC:70007E1818187E00 +000CD:0E007E1818187E00 +000CE:1866007E18187E00 +000CF:66007E1818187E00 +000D0:786C66F6666C7800 +000D1:76DC00C6F6DEC600 +000D2:70007CC6C6C67C00 +000D3:0E007CC6C6C67C00 +000D4:1866007CC6C67C00 +000D5:76DC007CC6C67C00 +000D6:66007CC6C6C67C00 +000D7:00C66C386CC60000 +000D8:3E666E7E76667C00 +000D9:7000C6C6C6C67C00 +000DA:0E00C6C6C6C67C00 +000DB:186600C6C6C67C00 +000DC:6600C6C6C6C67C00 +000DD:0E0066663C181800 +000DE:C0C0FCC6FCC0C000 +000DF:3C66666C66666C00 +000E0:70003C063E663E00 +000E1:0E003C063E663E00 +000E2:1866003E66C67E00 +000E3:76DC003E66C67E00 +000E4:66003C063E663E00 +000E5:1818003E66C67E00 +000E6:00007E1B7FD87700 +000E7:00003C6060603C18 +000E8:70003C667E603C00 +000E9:0E003C667E603C00 +000EA:1866003C7E603C00 +000EB:66003C667E603C00 +000EC:7000381818183C00 +000ED:0E00381818183C00 +000EE:1866003818183C00 +000EF:6600381818183C00 +000F0:0C3E0C7CCCCC7800 +000F1:76DC007C66666600 +000F2:70003C6666663C00 +000F3:0E003C6666663C00 +000F4:1866003C66663C00 +000F5:76DC003C66663C00 +000F6:66003C6666663C00 +000F7:1818007E00181800 +000F8:00027CCED6E67C80 +000F9:7000666666663E00 +000FA:0E00666666663E00 +000FB:1866006666663E00 +000FC:6600666666663E00 +000FD:0E006666663E063C +000FE:60607C66667C6060 +000FF:66006666663E063C diff --git a/crates/pi-natives/src/snapcompact.rs b/crates/pi-natives/src/snapcompact.rs index e1da346a4..3805a8b08 100644 --- a/crates/pi-natives/src/snapcompact.rs +++ b/crates/pi-natives/src/snapcompact.rs @@ -1,46 +1,61 @@ //! Snapcompact frame rendering. //! -//! Rasterizes pre-normalized conversation text onto a square 8-bit indexed -//! bitmap using the bundled public-domain X.org `5x8` BDF font, with glyph -//! ink cycling through six hues at sentence boundaries, then encodes the -//! bitmap as an indexed PNG. +//! Rasterizes pre-normalized conversation text onto a square bitmap using one +//! of the bundled public-domain pixel fonts, then encodes it as PNG: //! -//! Text normalization, frame chunking, and archive management live in -//! `packages/agent/src/compaction/snapcompact.ts`; this module is only the -//! hot `text -> PNG bytes` path. +//! - `5x8` — X.org BDF font (legacy shape). +//! - `8x8` — unscii-8 hex font (Latin-1 subset), the square cell that won the +//! snapcompact `SQuAD` evals. +//! +//! Shape controls, all eval-validated in `packages/snapcompact`: +//! +//! - **variant** — `sent` cycles glyph ink through six hues at sentence +//! boundaries; `bw` prints plain black ink (best for Anthropic readers). +//! - **lineRepeat** — prints every text line N times; copies after the first +//! sit on a pale highlight band. Redundancy coding: two looks per glyph at +//! half the density ("8x8r" shapes). +//! - **cellWidth/cellHeight** — target cell size. When it differs from the +//! font's natural cell, glyphs are rasterized at native size and the canvas +//! is Lanczos3-resampled to the target (anisotropic stretch, e.g. the +//! OpenAI-optimal "6x6u" shape), producing an anti-aliased RGB frame. +//! +//! Text normalization, frame chunking, provider shape selection, and archive +//! management live in `packages/agent/src/compaction/snapcompact.ts`; this +//! module is only the hot `text -> PNG bytes` path. -use std::{borrow::Cow, collections::HashMap, sync::LazyLock}; +use std::{borrow::Cow, collections::HashMap, f32::consts::PI, sync::LazyLock}; use napi::bindgen_prelude::*; use napi_derive::napi; -/// Horizontal advance per glyph cell in pixels. -const GLYPH_ADVANCE_X: usize = 5; -/// Vertical pitch per text row in pixels. -const GLYPH_PITCH_Y: usize = 8; - /// Upper bound on the frame edge: a hard stop against absurd allocations /// (`size * size` pixel buffer), far above the 2576px production frame. const MAX_FRAME_SIZE: u32 = 16384; -/// Sentence ink palette: index 0 is the white background, 1-6 are the six -/// dark hues from the eval renderer (HLS l=0.22 s=0.95, h ∈ {0, .08, .3, -/// .5, .62, .78}), pre-baked to sRGB. -const PALETTE: [[u8; 3]; 7] = [ +/// Indexed palette: 0 is the white background, 1-6 are the six dark sentence +/// hues from the eval renderer (HLS l=0.22 s=0.95, h ∈ {0, .08, .3, .5, .62, +/// .78}), 7 is plain black ink (`bw` variant), 8 is the pale highlight band +/// behind repeated line copies. +const PALETTE: [[u8; 3]; 9] = [ [255, 255, 255], - [109, 2, 2], // red - [109, 53, 2], // amber - [24, 109, 2], // green - [2, 109, 109], // teal - [2, 32, 109], // blue - [75, 2, 109], // violet + [109, 2, 2], // red + [109, 53, 2], // amber + [24, 109, 2], // green + [2, 109, 109], // teal + [2, 32, 109], // blue + [75, 2, 109], // violet + [0, 0, 0], // bw ink + [255, 247, 194], // repeat highlight band ]; -const INK_COLORS: usize = PALETTE.len() - 1; +const INK_COLORS: usize = 6; +const INK_BLACK: u8 = 7; +const BG_REPEAT: u8 = 8; -static FONT: LazyLock = LazyLock::new(|| parse_bdf(include_str!("fonts/5x8.bdf"))); +static FONT_5X8: LazyLock = LazyLock::new(|| parse_bdf(include_str!("fonts/5x8.bdf"), 5, 8)); +static FONT_8X8: LazyLock = LazyLock::new(|| parse_hex(include_str!("fonts/unscii-8.hex"))); struct Glyph { - /// Glyph width in pixels (≤ 8 for this font). + /// Glyph width in pixels (≤ 8 for the bundled fonts). w: u8, /// Glyph height in pixels. h: i32, @@ -54,9 +69,13 @@ struct Font { /// Glyphs keyed by Unicode code point (ASCII + Latin-1 coverage). glyphs: HashMap, ascent: i32, + /// Natural cell advance (x) in pixels. + cell_w: usize, + /// Natural cell pitch (y) in pixels. + cell_h: usize, } -fn parse_bdf(text: &str) -> Font { +fn parse_bdf(text: &str, cell_w: usize, cell_h: usize) -> Font { let mut glyphs = HashMap::new(); let mut ascent = 0i32; let mut enc = -1i64; @@ -91,27 +110,88 @@ fn parse_bdf(text: &str) -> Font { } } } - Font { glyphs, ascent } + Font { glyphs, ascent, cell_w, cell_h } } -/// Rasterize `text` onto a `size` x `size` palette-indexed bitmap, row-major -/// with no word wrap. Ink color advances at sentence boundaries (terminator -/// in `.!?` followed by a space). Characters beyond the grid capacity are -/// ignored; code points missing from the font leave their cell blank. -fn render_bitmap(text: &str, size: usize, font: &Font) -> Vec { - let cols = size / GLYPH_ADVANCE_X; - let grid_rows = size / GLYPH_PITCH_Y; - let capacity = cols * grid_rows; - let mut pixels = vec![0u8; size * size]; // 0 = white background +/// Parse a unifont-style `.hex` font (`CODEPOINT:16-hex-digit bitmap`, one +/// byte per row of an 8x8 glyph). Baseline sits at row 7 (`ascent` 7 with a +/// one-pixel descender row), matching the eval renderer. +fn parse_hex(text: &str) -> Font { + let mut glyphs = HashMap::new(); + for line in text.lines() { + let Some((cp, bits)) = line.split_once(':') else { + continue; + }; + let Ok(enc) = u32::from_str_radix(cp.trim(), 16) else { + continue; + }; + let bits = bits.trim(); + if bits.len() != 16 { + continue; + } + let rows: Vec = (0..8) + .map(|i| u8::from_str_radix(&bits[i * 2..i * 2 + 2], 16).unwrap_or(0)) + .collect(); + glyphs.insert(enc, Glyph { w: 8, h: 8, xoff: 0, yoff: -1, rows }); + } + Font { glyphs, ascent: 7, cell_w: 8, cell_h: 8 } +} + +fn resolve_font(name: &str) -> Option<&'static Font> { + match name { + "5x8" => Some(&FONT_5X8), + "8x8" => Some(&FONT_8X8), + _ => None, + } +} + +/// Frame grid geometry shared with the TypeScript caller. +struct Grid { + cols: usize, + rows: usize, + repeat: usize, +} + +/// Rasterize `text` onto a `width` x `height` palette-indexed bitmap at the +/// font's natural cell size, row-major with no word wrap. Each text line is +/// printed `grid.repeat` times; copies after the first sit on the highlight +/// band. Ink cycles through six hues at sentence boundaries (terminator in +/// `.!?` followed by a space) unless `black_ink` pins it to black. Characters +/// beyond `cols * rows` are ignored; code points missing from the font leave +/// their cell blank. +fn render_bitmap( + text: &str, + width: usize, + height: usize, + font: &Font, + grid: &Grid, + black_ink: bool, +) -> Vec { + let mut pixels = vec![0u8; width * height]; // 0 = white background + let capacity = grid.cols * grid.rows; if capacity == 0 { return pixels; } + if grid.repeat > 1 { + for row in 0..grid.rows { + for copy in 1..grid.repeat { + let band_top = (row * grid.repeat + copy) * font.cell_h; + for y in band_top..(band_top + font.cell_h).min(height) { + pixels[y * width..y * width + width].fill(BG_REPEAT); + } + } + } + } let codes: Vec = text.chars().map(|ch| ch as u32).collect(); let count = codes.len().min(capacity); let mut sentence = 0usize; for i in 0..count { let code = codes[i]; - let ink = (1 + sentence % INK_COLORS) as u8; + let ink = if black_ink { + INK_BLACK + } else { + (1 + sentence % INK_COLORS) as u8 + }; if matches!(code, 0x2e | 0x21 | 0x3f) && codes.get(i + 1) == Some(&0x20) { sentence += 1; } @@ -121,24 +201,27 @@ fn render_bitmap(text: &str, size: usize, font: &Font) -> Vec { if glyph.rows.is_empty() { continue; } - let row = i / cols; - let col = i - row * cols; - let top = (row * GLYPH_PITCH_Y) as i32 + font.ascent - glyph.h - glyph.yoff; - let left = (col * GLYPH_ADVANCE_X) as i32 + glyph.xoff; - for (r, &bits) in glyph.rows.iter().enumerate() { - if bits == 0 { - continue; - } - let y = top + r as i32; - if y < 0 || y >= size as i32 { - continue; - } - let row_base = y as usize * size; - for b in 0..glyph.w { - if bits & (0x80u8 >> b) != 0 { - let x = left + i32::from(b); - if x >= 0 && (x as usize) < size { - pixels[row_base + x as usize] = ink; + let row = i / grid.cols; + let col = i - row * grid.cols; + let left = (col * font.cell_w) as i32 + glyph.xoff; + for copy in 0..grid.repeat { + let cell_top = ((row * grid.repeat + copy) * font.cell_h) as i32; + let top = cell_top + font.ascent - glyph.h - glyph.yoff; + for (r, &bits) in glyph.rows.iter().enumerate() { + if bits == 0 { + continue; + } + let y = top + r as i32; + if y < 0 || y >= height as i32 { + continue; + } + let row_base = y as usize * width; + for b in 0..glyph.w { + if bits & (0x80u8 >> b) != 0 { + let x = left + i32::from(b); + if x >= 0 && (x as usize) < width { + pixels[row_base + x as usize] = ink; + } } } } @@ -147,8 +230,89 @@ fn render_bitmap(text: &str, size: usize, font: &Font) -> Vec { pixels } +// ============================================================================ +// Lanczos3 resampling (stretch shapes) +// ============================================================================ + +fn lanczos3(x: f32) -> f32 { + let x = x.abs(); + if x < 1e-6 { + return 1.0; + } + if x >= 3.0 { + return 0.0; + } + let pix = PI * x; + (pix.sin() / pix) * ((pix / 3.0).sin() / (pix / 3.0)) +} + +/// Per-output-pixel kernel contributions for one axis, PIL-convention +/// (`center = (i + 0.5) * scale`, kernel stretched by `max(scale, 1)`, +/// weights normalized). +fn contributions(src_len: usize, dst_len: usize) -> Vec<(usize, Vec)> { + let scale = src_len as f32 / dst_len as f32; + let filt_scale = scale.max(1.0); + let support = 3.0 * filt_scale; + let mut out = Vec::with_capacity(dst_len); + for i in 0..dst_len { + let center = (i as f32 + 0.5) * scale; + let begin = ((center - support) as isize).max(0) as usize; + let end = ((center + support).ceil() as usize).min(src_len); + let mut weights = Vec::with_capacity(end - begin); + let mut total = 0.0f32; + for x in begin..end { + let w = lanczos3((x as f32 + 0.5 - center) / filt_scale); + weights.push(w); + total += w; + } + if total != 0.0 { + for w in &mut weights { + *w /= total; + } + } + out.push((begin, weights)); + } + out +} + +/// Separable Lanczos3 resize of an interleaved RGB f32 buffer. +fn resize_rgb(src: &[f32], sw: usize, sh: usize, dw: usize, dh: usize) -> Vec { + let horiz = contributions(sw, dw); + let mut tmp = vec![0f32; dw * sh * 3]; + for y in 0..sh { + let src_row = &src[y * sw * 3..(y + 1) * sw * 3]; + let dst_row = &mut tmp[y * dw * 3..(y + 1) * dw * 3]; + for (x, (begin, weights)) in horiz.iter().enumerate() { + let mut acc = [0f32; 3]; + for (k, &w) in weights.iter().enumerate() { + let s = (begin + k) * 3; + acc[0] = src_row[s].mul_add(w, acc[0]); + acc[1] = src_row[s + 1].mul_add(w, acc[1]); + acc[2] = src_row[s + 2].mul_add(w, acc[2]); + } + dst_row[x * 3..x * 3 + 3].copy_from_slice(&acc); + } + } + let vert = contributions(sh, dh); + let mut out = vec![0f32; dw * dh * 3]; + for (y, (begin, weights)) in vert.iter().enumerate() { + let dst_row = &mut out[y * dw * 3..(y + 1) * dw * 3]; + for (k, &w) in weights.iter().enumerate() { + let src_row = &tmp[(begin + k) * dw * 3..(begin + k + 1) * dw * 3]; + for (d, &s) in dst_row.iter_mut().zip(src_row) { + *d = s.mul_add(w, *d); + } + } + } + out +} + +// ============================================================================ +// PNG encoding +// ============================================================================ + /// Pack one-byte-per-pixel palette indices into 4-bit PNG scanline data -/// (two pixels per byte, high nibble first). With only 7 palette entries, +/// (two pixels per byte, high nibble first). With only 9 palette entries, /// 4-bit depth halves the pre-deflate stream vs 8-bit. fn pack_nibbles(pixels: &[u8], size: usize) -> Vec { let row_bytes = size.div_ceil(2); @@ -196,49 +360,236 @@ fn encode_indexed_png( Ok(out) } -/// Render one snapcompact frame: print pre-normalized text onto a -/// `size` x `size` 4-bit indexed bitmap and encode it as a PNG. +/// Encode an interleaved RGB8 buffer as PNG. Stretched frames are +/// continuous-tone, so adaptive filtering (the `Balanced` default) helps. +fn encode_rgb_png(pixels: &[u8], size: usize, compression: png::Compression) -> Result> { + let mut out = Vec::new(); + let mut encoder = png::Encoder::new(&mut out, size as u32, size as u32); + encoder.set_color(png::ColorType::Rgb); + encoder.set_depth(png::BitDepth::Eight); + encoder.set_compression(compression); + let mut writer = encoder + .write_header() + .map_err(|err| Error::from_reason(format!("Failed to write PNG header: {err}")))?; + writer + .write_image_data(pixels) + .map_err(|err| Error::from_reason(format!("Failed to write PNG data: {err}")))?; + writer + .finish() + .map_err(|err| Error::from_reason(format!("Failed to finish PNG stream: {err}")))?; + Ok(out) +} + +// ============================================================================ +// Entry point +// ============================================================================ + +/// Shape options for one snapcompact frame. +#[napi(object)] +#[derive(Default)] +pub struct SnapcompactRenderOptions { + /// Frame edge in pixels. + pub size: u32, + /// Bundled font: `"5x8"` (X.org BDF) or `"8x8"` (unscii-8). Default `"5x8"`. + pub font: Option, + /// Target cell advance in pixels. Differing from the font's natural cell + /// triggers the Lanczos stretch path. Default: font natural width. + pub cell_width: Option, + /// Target cell pitch in pixels. Default: font natural height. + pub cell_height: Option, + /// Ink variant: `"sent"` (six-hue sentence cycling) or `"bw"` (black). + /// Default `"sent"`. + pub variant: Option, + /// Print each text line this many times; copies after the first sit on a + /// pale highlight band. Default 1. + pub line_repeat: Option, +} + +/// Render one snapcompact frame: print pre-normalized text onto a square +/// bitmap and encode it as PNG. /// -/// The glyph grid holds `floor(size/5) * floor(size/8)` characters; input -/// beyond that is ignored (the caller chunks text to capacity). Returns the -/// PNG bytes. +/// The glyph grid holds `floor(size/cellWidth) * +/// floor(size/cellHeight/lineRepeat)` characters; input beyond that is ignored +/// (the caller chunks text to capacity). Native-cell shapes encode as 4-bit +/// indexed PNG; stretched shapes (target cell != font cell) encode as RGB. +/// Returns the PNG bytes. #[napi] -pub fn render_snapcompact_png(text: String, size: u32) -> Result { +pub fn render_snapcompact_png( + text: String, + options: SnapcompactRenderOptions, +) -> Result { + let size = options.size; if size == 0 || size > MAX_FRAME_SIZE { return Err(Error::from_reason(format!( "Invalid frame size {size}: expected 1..={MAX_FRAME_SIZE}" ))); } - let pixels = render_bitmap(&text, size as usize, &FONT); - Ok(encode_indexed_png(&pixels, size as usize, png::Compression::Balanced)?.into()) + let font_name = options.font.as_deref().unwrap_or("5x8"); + let font = resolve_font(font_name).ok_or_else(|| { + Error::from_reason(format!( + "Unknown snapcompact font {font_name:?}: expected \"5x8\" or \"8x8\"" + )) + })?; + let black_ink = match options.variant.as_deref().unwrap_or("sent") { + "sent" => false, + "bw" => true, + other => { + return Err(Error::from_reason(format!( + "Unknown snapcompact variant {other:?}: expected \"sent\" or \"bw\"" + ))); + }, + }; + let target_w = options.cell_width.unwrap_or(font.cell_w as u32).max(1) as usize; + let target_h = options.cell_height.unwrap_or(font.cell_h as u32).max(1) as usize; + let repeat = options.line_repeat.unwrap_or(1).max(1) as usize; + let size = size as usize; + let grid = Grid { cols: size / target_w, rows: size / target_h / repeat, repeat }; + if grid.cols == 0 || grid.rows == 0 { + return Err(Error::from_reason(format!( + "Frame size {size} cannot fit a {target_w}x{target_h} cell grid (repeat {repeat})" + ))); + } + + if (target_w, target_h) == (font.cell_w, font.cell_h) { + // Native cell: rasterize straight onto the frame, indexed. + let pixels = render_bitmap(&text, size, size, font, &grid, black_ink); + return Ok(encode_indexed_png(&pixels, size, png::Compression::Balanced)?.into()); + } + + // Stretch shape: rasterize at the font's natural cell on a tight canvas, + // Lanczos3-resample to the target cell, paste onto the white frame. + let src_w = grid.cols * font.cell_w; + let src_h = grid.rows * grid.repeat * font.cell_h; + let dst_w = grid.cols * target_w; + let dst_h = grid.rows * grid.repeat * target_h; + let indexed = render_bitmap(&text, src_w, src_h, font, &grid, black_ink); + let mut rgb = vec![0f32; src_w * src_h * 3]; + for (dst, &idx) in rgb.chunks_exact_mut(3).zip(&indexed) { + let [r, g, b] = PALETTE[idx as usize]; + dst[0] = f32::from(r); + dst[1] = f32::from(g); + dst[2] = f32::from(b); + } + let resized = resize_rgb(&rgb, src_w, src_h, dst_w, dst_h); + let mut frame = vec![255u8; size * size * 3]; + for y in 0..dst_h.min(size) { + let src_row = &resized[y * dst_w * 3..(y + 1) * dst_w * 3]; + let dst_row = &mut frame[y * size * 3..]; + for (d, &s) in dst_row[..dst_w.min(size) * 3].iter_mut().zip(src_row) { + *d = s.round().clamp(0.0, 255.0) as u8; + } + } + Ok(encode_rgb_png(&frame, size, png::Compression::Balanced)?.into()) } #[cfg(test)] mod tests { use super::*; + fn opts(size: u32) -> SnapcompactRenderOptions { + SnapcompactRenderOptions { size, ..Default::default() } + } + #[test] - fn font_parses_ascii_coverage() { - let font = &*FONT; - assert_eq!(font.ascent, 7); - // Every printable ASCII char must have a glyph. - for cp in 0x20u32..0x7f { - assert!(font.glyphs.contains_key(&cp), "missing glyph for U+{cp:04X}"); + fn fonts_parse_ascii_coverage() { + for (font, ascent) in [(&*FONT_5X8, 7), (&*FONT_8X8, 7)] { + assert_eq!(font.ascent, ascent); + // Every printable ASCII char must have a glyph. + for cp in 0x20u32..0x7f { + assert!(font.glyphs.contains_key(&cp), "missing glyph for U+{cp:04X}"); + } } } #[test] fn bitmap_inks_sentences_and_caps_capacity() { - let font = &*FONT; - // 40px -> 8 cols x 5 rows = 40 cells. - let pixels = render_bitmap("Hi. Ok.", 40, font); + // 40px -> 8 cols x 5 rows = 40 cells (5x8 font). + let grid = Grid { cols: 8, rows: 5, repeat: 1 }; + let pixels = render_bitmap("Hi. Ok.", 40, 40, &FONT_5X8, &grid, false); let inks: Vec = pixels.iter().copied().filter(|&p| p != 0).collect(); assert!(inks.contains(&1), "first sentence should use ink 1"); assert!(inks.contains(&2), "second sentence should use ink 2"); assert!(!inks.contains(&3), "no third sentence ink expected"); // Overflow input renders without panicking and stays in-bounds. - let overflow = render_bitmap(&"x".repeat(100), 40, font); + let overflow = render_bitmap(&"x".repeat(100), 40, 40, &FONT_5X8, &grid, false); assert_eq!(overflow.len(), 40 * 40); } + + #[test] + fn bw_variant_prints_black_only() { + let grid = Grid { cols: 8, rows: 8, repeat: 1 }; + let pixels = render_bitmap("Hi. Ok.", 64, 64, &FONT_8X8, &grid, true); + let inks: Vec = pixels.iter().copied().filter(|&p| p != 0).collect(); + assert!(!inks.is_empty()); + assert!(inks.iter().all(|&p| p == INK_BLACK), "bw must ink only black"); + } + + #[test] + fn line_repeat_duplicates_rows_on_highlight_bands() { + // 64px, 8x8 font, repeat 2 -> 8 cols x 4 unique rows. + let grid = Grid { cols: 8, rows: 4, repeat: 2 }; + let pixels = render_bitmap("ABCDEFGH", 64, 64, &FONT_8X8, &grid, true); + // Copy band (rows 8..16) carries the highlight background. + assert!(pixels[9 * 64..10 * 64].contains(&BG_REPEAT), "duplicate band must be highlighted"); + // Identical glyph ink in both copies: compare full 8-row bands modulo + // background. + for y in 0..8 { + for x in 0..64 { + let a = pixels[y * 64 + x]; + let b = pixels[(y + 8) * 64 + x]; + assert_eq!(a == INK_BLACK, b == INK_BLACK, "copy ink mismatch at ({x},{y})"); + } + } + } + + #[test] + fn render_native_is_indexed_and_stretch_is_rgb() { + let native = render_snapcompact_png("Hello world. Again.".into(), SnapcompactRenderOptions { + size: 128, + font: Some("8x8".into()), + variant: Some("bw".into()), + line_repeat: Some(2), + ..Default::default() + }) + .unwrap(); + // PNG color type lives at byte 25 of the IHDR: 3 = indexed. + assert_eq!(native[25], 3); + + let stretched = + render_snapcompact_png("Hello world. Again.".into(), SnapcompactRenderOptions { + size: 128, + font: Some("8x8".into()), + cell_width: Some(6), + cell_height: Some(6), + ..Default::default() + }) + .unwrap(); + // 2 = truecolor RGB. + assert_eq!(stretched[25], 2); + // Stretched output must contain anti-aliased (non-extreme) pixels. + let legacy = render_snapcompact_png("Hi. Ok.".into(), opts(40)).unwrap(); + assert_eq!(legacy[25], 3, "default shape stays the legacy 5x8 indexed path"); + } + + #[test] + fn rejects_bad_shapes() { + assert!(render_snapcompact_png("x".into(), opts(0)).is_err()); + assert!( + render_snapcompact_png("x".into(), SnapcompactRenderOptions { + size: 64, + font: Some("9x9".into()), + ..Default::default() + }) + .is_err() + ); + assert!( + render_snapcompact_png("x".into(), SnapcompactRenderOptions { + size: 64, + variant: Some("zebra".into()), + ..Default::default() + }) + .is_err() + ); + } } diff --git a/docs/compaction.md b/docs/compaction.md index 789e36aa1..46af16287 100644 --- a/docs/compaction.md +++ b/docs/compaction.md @@ -10,7 +10,7 @@ Both are persisted as session entries and converted back into user-context messa ## Key implementation files - `packages/agent/src/compaction/compaction.ts` (context-full summarization and handoff generation) -- `packages/agent/src/compaction/snapcompact.ts` (snapcompact strategy: history archived as dense bitmap images) +- `packages/snapcompact/src/snapcompact.ts` (snapcompact strategy: history archived as dense bitmap images) - `packages/agent/src/compaction/branch-summarization.ts` - `packages/agent/src/compaction/pruning.ts` - `packages/agent/src/compaction/utils.ts` @@ -129,13 +129,13 @@ The automatic paths are intentionally different: ### Snapcompact strategy -`compaction.strategy: "snapcompact"` replaces the LLM summarization call with a local, deterministic archival pass (`snapcompactCompact` in `packages/agent/src/compaction/snapcompact.ts`): +`compaction.strategy: "snapcompact"` replaces the LLM summarization call with a local, deterministic archival pass (`snapcompactCompact` from `@oh-my-pi/snapcompact`): -- The discarded history is serialized (same `serializeConversation` path the summarizer uses), whitespace-collapsed, and printed onto square PNG frames using the bundled public-domain X.org `5x8` pixel font — 2576px frames carrying ~165k characters each. Glyph ink cycles through six hues at sentence boundaries to help the model keep its place. +- The discarded history is serialized, whitespace-collapsed, and printed onto provider-aware square PNG frames using bundled public-domain pixel fonts. Anthropic-family and unknown APIs use repeated black `8x8` cells, Google uses repeated sentence-colored `8x8` cells, and OpenAI uses dense stretched `6x6` cells with `detail: "original"`. - Frames persist under `CompactionEntry.preserveData.snapcompact` and are re-attached to the `compactionSummary` message as image blocks on every context rebuild; the entry's `summary` is a deterministic reading guide (grid geometry, role tags, truncation notes) plus the usual file-operation lists. - Later compactions carry earlier frames forward. Beyond an 8-frame budget the archive fades from the middle out: the earliest frame (session head — the original request, or the filmed summary of older history) is pinned, and the oldest *unpinned* frames are evicted, so head and tail both survive. If the previous compaction was text-based, its summary is printed at the head of the frame archive as `[Summary of earlier history]`. - No model, API key, or network is involved, so snapcompact is also safe for overflow recovery. It requires a vision-capable current model (`model.input` includes `"image"`); otherwise the run falls back to context-full and emits a warning notice (auto and manual paths). Manual `/compact` honors the strategy unless custom instructions are given (those imply a directed LLM summary). -- Rationale: a vision model reads ~50 chars per image token after provider downscaling, ~7x cheaper than raw text at near-parity recall (SQuAD eval: F1 0.878 vs 0.899). +- Rationale: the shape table comes from the snapcompact 200k-token evals in `packages/snapcompact`, where bitmap frames preserved QA recall at lower billed-token cost than raw text for vision-capable models. ### Display transcript diff --git a/docs/skills/examples/hello-extension/package.json b/docs/skills/examples/hello-extension/package.json index 1c00ef2be..5027b45a4 100644 --- a/docs/skills/examples/hello-extension/package.json +++ b/docs/skills/examples/hello-extension/package.json @@ -2,6 +2,7 @@ "name": "hello-extension", "version": "1.0.0", "description": "Minimal oh-my-pi extension example", + "homepage": "https://omp.sh", "omp": { "extensions": ["./index.ts"] } diff --git a/docs/skills/examples/mini-marketplace/my-plugin/package.json b/docs/skills/examples/mini-marketplace/my-plugin/package.json index e971f0913..7c4636884 100644 --- a/docs/skills/examples/mini-marketplace/my-plugin/package.json +++ b/docs/skills/examples/mini-marketplace/my-plugin/package.json @@ -1,6 +1,7 @@ { "name": "my-plugin", "version": "0.1.0", + "homepage": "https://omp.sh", "omp": { "extensions": ["./index.ts"] } diff --git a/docs/skills/examples/safety-hook/package.json b/docs/skills/examples/safety-hook/package.json index 8c6ab8271..b8b14fd1c 100644 --- a/docs/skills/examples/safety-hook/package.json +++ b/docs/skills/examples/safety-hook/package.json @@ -2,6 +2,7 @@ "name": "safety-hook", "version": "1.0.0", "description": "oh-my-pi extension example: block rm -rf / via tool_call hook", + "homepage": "https://omp.sh", "omp": { "extensions": ["./index.ts"] } diff --git a/package.json b/package.json index 173cd305c..12da8abf7 100644 --- a/package.json +++ b/package.json @@ -1,5 +1,6 @@ { "name": "omp-monorepo", + "homepage": "https://omp.sh", "private": true, "type": "module", "packageManager": "bun@1.3.14", @@ -30,6 +31,7 @@ "@oh-my-pi/pi-natives": "15.10.12", "@oh-my-pi/pi-tui": "15.10.12", "@oh-my-pi/pi-utils": "15.10.12", + "@oh-my-pi/snapcompact": "15.10.12", "@opentelemetry/api": "^1.9.1", "@opentelemetry/context-async-hooks": "^2.7.1", "@opentelemetry/exporter-trace-otlp-proto": "^0.218.0", diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 607bfc6ee..9fe70f2b0 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -1,9 +1,9 @@ # Changelog ## [Unreleased] - ### Breaking Changes +- Removed `compaction/index.ts` re-export of snapcompact helpers, so snapcompact utilities are no longer available from the agent compaction barrel and should be imported from `@oh-my-pi/snapcompact` - Removed the `convertToLlm` alias export from `compaction/messages` — it duplicated `defaultConvertToLlm` under a second name. Import `defaultConvertToLlm` (array form) or the new `convertMessageToLlm` (single-message form) instead ### Added @@ -11,7 +11,8 @@ - Added `convertMessageToLlm()`: the single-message core transformer behind `defaultConvertToLlm()`. Embedders with app-specific message roles should handle their own roles and delegate every core role (`user`/`developer`/`assistant`/`toolResult`/`custom`/`hookMessage`/`branchSummary`/`compactionSummary`) to it instead of duplicating the conversion — a duplicated `compactionSummary` case is how snapcompact frames once silently dropped off provider requests - Added `pruneSupersededToolResults()` and the opt-in `PruneConfig.supersedeKey` hook so harnesses can prune stale tool results superseded by a newer read of the same file; superseded results are pruned ahead of age-based victims during overflow pruning and replaced with a `[Superseded by a newer read of this file]` placeholder. Without the new config, `pruneToolOutputs()` behavior is unchanged. - Added `readToolSupersedeKey()` implementing the read-tool path/selector grammar (selector-free reads supersede range reads of the same file; URL-scheme paths exempt). Pruning honors prompt-cache economics: per-turn prunes only fire when the post-candidate suffix is small or the cache is cold (idle gap). -- Added the `snapcompact` compaction strategy (`snapcompactCompact()` in `compaction/snapcompact.ts`): instead of an LLM summary, discarded history is printed onto dense 2576px PNG frames with the public-domain X.org `5x8` pixel font (ink cycles per sentence) and re-attached to the compaction summary message as image blocks. Fully local — no model call; ~7x cheaper than raw text at near-parity recall. `CompactionSummaryMessage` gains an optional `images` field, `estimateTokens()` charges per attached frame, and frames persist under `preserveData.snapcompact` with an 8-frame budget that evicts middle-out: the session-head frame is pinned, the oldest unpinned frames drop first, so head and tail both survive. Rasterization and PNG encoding run in native code (`renderSnapcompactPng()` from `@oh-my-pi/pi-natives`), emitting 4-bit indexed PNGs +- Added the `snapcompact` compaction strategy via `@oh-my-pi/snapcompact`: instead of an LLM summary, discarded history is printed onto dense bitmap frames and re-attached to the compaction summary message as image blocks. `CompactionSummaryMessage` gains an optional `images` field, `estimateTokens()` charges per attached frame, and frames persist under `preserveData.snapcompact` with an 8-frame middle-out eviction budget. +- Snapcompact frames are now rendered in a provider-aware shape (`SNAPCOMPACT_SHAPES` + `resolveSnapcompactShape(api)`), following the snapcompact 200k-token monolithic evals: Anthropic-family and unknown APIs get `8x8r-bw` (unscii-8 square cells, black ink, every line printed twice with the copy on a pale highlight band — read at F1 parity with raw text at ~2x lower cost and the most refusal-robust), Google gets `8x8r-sent` (sentence-hue ink, ~2.9x cheaper), and OpenAI gets `6x6u-sent` (unscii Lanczos-stretched to 6x6 cells — OpenAI bills a flat ~2.9k tokens per image, so frame count is the only cost lever) with `detail: "original"` on the frame images. `snapcompactCompact()` accepts `model`/`shape` options, frames persist their shape metadata, mixed-shape archives (provider switches, legacy 5x8 frames) are flagged in the reading instructions, and `snapcompactGeometry()`/`renderSnapcompactFrame()` now take a shape ### Fixed @@ -680,4 +681,4 @@ Initial release under @oh-my-pi scope. See previous releases at [badlogic/pi-mon - `Agent` constructor now has all options optional (empty options use defaults). -- `queueMessage()` is now synchronous (no longer returns a Promise). +- `queueMessage()` is now synchronous (no longer returns a Promise). \ No newline at end of file diff --git a/packages/agent/package.json b/packages/agent/package.json index bd035f3ba..6037d08f3 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -39,6 +39,7 @@ "@oh-my-pi/pi-catalog": "catalog:", "@oh-my-pi/pi-natives": "catalog:", "@oh-my-pi/pi-utils": "catalog:", + "@oh-my-pi/snapcompact": "catalog:", "@opentelemetry/api": "catalog:" }, "devDependencies": { diff --git a/packages/agent/src/compaction/compaction.ts b/packages/agent/src/compaction/compaction.ts index 565f801dc..09b6cf3f0 100644 --- a/packages/agent/src/compaction/compaction.ts +++ b/packages/agent/src/compaction/compaction.ts @@ -18,6 +18,7 @@ import { import { clampThinkingLevelForModel } from "@oh-my-pi/pi-catalog/model-thinking"; import { countTokens } from "@oh-my-pi/pi-natives"; import { logger, prompt } from "@oh-my-pi/pi-utils"; +import { SNAPCOMPACT_FRAME_TOKEN_ESTIMATE } from "@oh-my-pi/snapcompact"; import { type AgentTelemetry, instrumentedCompleteSimple } from "../telemetry"; import { ThinkingLevel } from "../thinking"; import type { AgentMessage } from "../types"; @@ -37,7 +38,6 @@ import compactionSummaryPrompt from "./prompts/compaction-summary.md" with { typ import compactionTurnPrefixPrompt from "./prompts/compaction-turn-prefix.md" with { type: "text" }; import compactionUpdateSummaryPrompt from "./prompts/compaction-update-summary.md" with { type: "text" }; import handoffDocumentPrompt from "./prompts/handoff-document.md" with { type: "text" }; -import { SNAPCOMPACT_FRAME_TOKEN_ESTIMATE } from "./snapcompact"; import { computeFileLists, diff --git a/packages/agent/src/compaction/index.ts b/packages/agent/src/compaction/index.ts index 1425b2976..401215724 100644 --- a/packages/agent/src/compaction/index.ts +++ b/packages/agent/src/compaction/index.ts @@ -10,5 +10,4 @@ export * from "./messages"; export * from "./openai"; export * from "./pruning"; export * from "./shake"; -export * from "./snapcompact"; export * from "./utils"; diff --git a/packages/agent/src/compaction/prompts/snapcompact-summary.md b/packages/agent/src/compaction/prompts/snapcompact-summary.md index 3e25c2dcb..0d8d22323 100644 --- a/packages/agent/src/compaction/prompts/snapcompact-summary.md +++ b/packages/agent/src/compaction/prompts/snapcompact-summary.md @@ -1,6 +1,10 @@ Prior conversation history has been archived verbatim onto {{frameCount}} snapcompact frame{{#if multipleFrames}}s{{/if}} — the bitmap image{{#if multipleFrames}}s{{/if}} attached below{{#if multipleFrames}}, ordered oldest to newest{{/if}}. -Reading a frame: monospace 5x8 pixel font on a white background, {{cols}} characters per row, {{rows}} rows per frame; read left to right, top to bottom. Text flows continuously with no word wrap, so words may break across row ends. Whitespace runs (including newlines) were collapsed to single spaces. Ink color cycles through six colors, advancing at sentence boundaries — a color change marks a new sentence. Roles are tagged inline as [User]:, [Assistant]:, [Assistant thinking]:, [Assistant tool calls]:, and [Tool result]:. +Reading a frame: monospace {{fontCell}} pixel font on a white background, {{cols}} characters per row, {{rows}} text rows per frame; read left to right, top to bottom. Text flows continuously with no word wrap, so words may break across row ends. Whitespace runs (including newlines) were collapsed to single spaces. {{#if sentenceInk}}Ink color cycles through six colors, advancing at sentence boundaries — a color change marks a new sentence.{{else}}Glyphs are plain black ink.{{/if}}{{#if lineRepeated}} Every text line is printed twice in a row — first on the white background, then repeated on a pale yellow band. The copies are identical: read each line once and use the duplicate only to double-check hard glyphs.{{/if}} Roles are tagged inline as [User]:, [Assistant]:, [Assistant thinking]:, [Assistant tool calls]:, and [Tool result]:. +{{#if mixedShapes}} + +Older frames may use a different font, grid, or ink coloring than described above; the reading order is always the same (left to right, top to bottom, oldest frame first). +{{/if}} {{#if includedPreviousSummary}} The earliest frame begins with "[Summary of earlier history]" — a condensed digest of context that predates the archived conversation. diff --git a/packages/agent/src/compaction/snapcompact.ts b/packages/agent/src/compaction/snapcompact.ts deleted file mode 100644 index 5236a26eb..000000000 --- a/packages/agent/src/compaction/snapcompact.ts +++ /dev/null @@ -1,316 +0,0 @@ -/** - * Snapcompact compaction: archive conversation history as dense bitmap images. - * - * Instead of asking an LLM to summarize discarded history, the serialized - * conversation is rendered into square PNG frames using the X.org `5x8` - * pixel font (public domain) — one character per 5x8 cell, row-major, glyph - * ink cycling through six hues at sentence boundaries. Vision models read - * the frames back directly, like an archivist at a snapcompact frame reader. - * - * Validated by the imageee SQuAD eval (`img-5x8-sent`, 2576px frames): - * ~0.88 F1 recall vs ~0.90 for raw text, at roughly 7x fewer input tokens. - * The provider downscales frames to its image cap (1568px for Anthropic), - * so one frame costs ~3.3k tokens while carrying ~165k characters. - * - * The whole pass is local and deterministic — no LLM call, no API key, no - * latency beyond rendering. Rasterization and PNG encoding happen in native - * code (`renderSnapcompactPng` in `crates/pi-natives/src/snapcompact.rs`). - * Frames persist in the compaction entry's `preserveData` and are - * re-attached to the compaction summary message on every context rebuild. - */ - -import type { ImageContent } from "@oh-my-pi/pi-ai"; -import { renderSnapcompactPng } from "@oh-my-pi/pi-natives"; -import { prompt } from "@oh-my-pi/pi-utils"; -import type { CompactionDetails, CompactionPreparation, CompactionResult } from "./compaction"; -import { type ConvertToLlm, defaultConvertToLlm } from "./messages"; -import { withOpenAiRemoteCompactionPreserveData } from "./openai"; -import snapcompactSummaryPrompt from "./prompts/snapcompact-summary.md" with { type: "text" }; -import { computeFileLists, serializeConversation, upsertFileOperations } from "./utils"; - -// ============================================================================ -// Constants -// ============================================================================ - -/** Frame edge in pixels. 2576px is the eval-validated sweet spot: the provider - * downscale to 1568px anti-aliases the 1px glyph strokes instead of shearing - * them, which reads *better* than rendering at 1568 directly. */ -export const SNAPCOMPACT_FRAME_SIZE = 2576; - -/** Glyph cell geometry of the bundled `5x8` BDF font. */ -const GLYPH_ADVANCE_X = 5; -const GLYPH_PITCH_Y = 8; - -/** Maximum frames carried on a compaction entry. Oldest frames are dropped - * first once the budget is exceeded (mirrors how iterative text summaries - * fade the oldest detail). 8 frames ≈ 26k image tokens ≈ 1.3M chars. */ -export const SNAPCOMPACT_MAX_FRAMES = 8; - -/** Token cost estimate per frame. Frames render at ≥1568px, so providers bill - * the downscaled long-edge cap: 1568*1568/750 ≈ 3,278 tokens (Anthropic). */ -export const SNAPCOMPACT_FRAME_TOKEN_ESTIMATE = 3300; - -/** Key under `CompactionEntry.preserveData` holding the frame archive. */ -export const SNAPCOMPACT_PRESERVE_KEY = "snapcompact"; - -// ============================================================================ -// Types -// ============================================================================ - -/** One developed snapcompact frame: a base64 PNG plus its reading geometry. */ -export interface SnapcompactFrame { - /** Base64-encoded PNG. */ - data: string; - mimeType: string; - /** Characters per row in the frame grid. */ - cols: number; - /** Rows in the frame grid. */ - rows: number; - /** Characters actually printed onto this frame. */ - chars: number; -} - -/** Frame archive persisted under `preserveData[SNAPCOMPACT_PRESERVE_KEY]`. */ -export interface SnapcompactArchive { - /** Frames ordered oldest to newest. */ - frames: SnapcompactFrame[]; - /** Characters currently readable across all frames. */ - totalChars: number; - /** Characters dropped so far to respect the frame budget. */ - truncatedChars: number; -} - -export interface SnapcompactGeometry { - cols: number; - rows: number; - /** Characters that fit one frame (cols * rows). */ - capacity: number; -} - -export interface SnapcompactOptions { - /** App-level message transformer (same contract as `SummaryOptions.convertToLlm`). */ - convertToLlm?: ConvertToLlm; - /** Frame edge in pixels. Defaults to {@link SNAPCOMPACT_FRAME_SIZE}. */ - frameSize?: number; - /** Frame budget. Defaults to {@link SNAPCOMPACT_MAX_FRAMES}. */ - maxFrames?: number; -} - -/** Result of rendering one frame, before base64 packing. */ -export interface RenderedFrame { - png: Uint8Array; - cols: number; - rows: number; - /** Characters printed (input may be shorter than capacity). */ - chars: number; -} - -// ============================================================================ -// Text normalization -// ============================================================================ - -/** Folds for common non-Latin-1 characters the 5x8 subset cannot draw. */ -const CHAR_FOLD: Record = { - "\u2018": "'", - "\u2019": "'", - "\u201a": "'", - "\u201b": "'", - "\u201c": '"', - "\u201d": '"', - "\u201e": '"', - "\u2013": "-", - "\u2014": "-", - "\u2015": "-", - "\u2212": "-", - "\u2026": "...", - "\u2022": "*", - "\u25cf": "*", - "\u25a0": "*", - "\u25aa": "*", - "\u2190": "<-", - "\u2192": "->", - "\u21d2": "=>", - "\u2713": "v", - "\u2714": "v", - "\u2717": "x", - "\u2718": "x", -}; - -/** - * Prepare text for printing: collapse whitespace runs (incl. newlines) to - * single spaces — the eval's "paragraph breaks collapsed to spaces" format — - * then fold everything outside the font's ASCII + Latin-1 coverage to ASCII - * approximations (`?` as the last resort). - */ -export function normalizeForSnapcompact(text: string): string { - const collapsed = text.replace(/\s+/g, " ").trim(); - let out = ""; - for (const ch of collapsed) { - const cp = ch.codePointAt(0) as number; - if (cp < 0x7f || (cp >= 0xa0 && cp <= 0xff)) { - out += ch; - continue; - } - const fold = CHAR_FOLD[ch]; - if (fold !== undefined) { - out += fold; - } else if (cp >= 0x2500 && cp <= 0x257f) { - // Box drawing: keep table skeletons legible. - out += cp === 0x2502 || cp === 0x2503 ? "|" : cp === 0x2500 || cp === 0x2501 ? "-" : "+"; - } else { - out += "?"; - } - } - return out; -} - -// ============================================================================ -// Rendering -// ============================================================================ -export function snapcompactGeometry(size: number = SNAPCOMPACT_FRAME_SIZE): SnapcompactGeometry { - const cols = Math.floor(size / GLYPH_ADVANCE_X); - const rows = Math.floor(size / GLYPH_PITCH_Y); - return { cols, rows, capacity: cols * rows }; -} - -/** Render one snapcompact frame from already-normalized text. */ -export function renderSnapcompactFrame(text: string, size: number = SNAPCOMPACT_FRAME_SIZE): RenderedFrame { - const { cols, rows, capacity } = snapcompactGeometry(size); - const chars = Math.min(text.length, capacity); - return { png: renderSnapcompactPng(text, size), cols, rows, chars }; -} - -// ============================================================================ -// Archive helpers -// ============================================================================ - -/** Validate and extract a persisted frame archive from `preserveData`. */ -export function getPreservedSnapcompactArchive( - preserveData: Record | undefined, -): SnapcompactArchive | undefined { - const candidate = preserveData?.[SNAPCOMPACT_PRESERVE_KEY]; - if (!candidate || typeof candidate !== "object") return undefined; - const archive = candidate as SnapcompactArchive; - if (!Array.isArray(archive.frames)) return undefined; - const frames = archive.frames.filter( - frame => - !!frame && - typeof frame.data === "string" && - frame.data.length > 0 && - typeof frame.mimeType === "string" && - typeof frame.cols === "number" && - typeof frame.rows === "number" && - typeof frame.chars === "number", - ); - if (frames.length === 0) return undefined; - return { - frames, - totalChars: typeof archive.totalChars === "number" ? archive.totalChars : 0, - truncatedChars: typeof archive.truncatedChars === "number" ? archive.truncatedChars : 0, - }; -} - -/** Convert archive frames into LLM image blocks (oldest first). */ -export function snapcompactImages(archive: SnapcompactArchive): ImageContent[] { - return archive.frames.map(frame => ({ type: "image", data: frame.data, mimeType: frame.mimeType })); -} - -// ============================================================================ -// Compaction entry point -// ============================================================================ - -/** - * Run a snapcompact compaction over prepared messages. Fully local: serializes - * the discarded history, prints it onto PNG frames, merges previously - * archived frames (oldest dropped beyond the budget), and produces a - * deterministic summary explaining how to read the frames. - * - * If the previous compaction was text-based, its summary is printed at the - * head of the frame archive as `[Summary of earlier history]` so no continuity is lost. - */ -export async function snapcompactCompact( - preparation: CompactionPreparation, - options?: SnapcompactOptions, -): Promise { - const { firstKeptEntryId, tokensBefore, previousSummary, previousPreserveData, fileOps } = preparation; - if (!firstKeptEntryId) { - throw new Error("First kept entry has no ID - session may need migration"); - } - const frameSize = options?.frameSize ?? SNAPCOMPACT_FRAME_SIZE; - const maxFrames = Math.max(1, options?.maxFrames ?? SNAPCOMPACT_MAX_FRAMES); - const geometry = snapcompactGeometry(frameSize); - - const messages = preparation.messagesToSummarize.concat(preparation.turnPrefixMessages); - const llmMessages = (options?.convertToLlm ?? defaultConvertToLlm)(messages); - let archiveText = normalizeForSnapcompact(serializeConversation(llmMessages)); - - const previousArchive = getPreservedSnapcompactArchive(previousPreserveData); - const includedPreviousSummary = !previousArchive && !!previousSummary; - if (includedPreviousSummary && previousSummary) { - const head = `[Summary of earlier history] ${normalizeForSnapcompact(previousSummary)}`; - archiveText = archiveText.length > 0 ? `${head} [Recent conversation] ${archiveText}` : head; - } - - let truncatedChars = previousArchive?.truncatedChars ?? 0; - - const newFrames: SnapcompactFrame[] = []; - for (let offset = 0; offset < archiveText.length; offset += geometry.capacity) { - const chunk = archiveText.slice(offset, offset + geometry.capacity); - const rendered = renderSnapcompactFrame(chunk, frameSize); - newFrames.push({ - data: Buffer.from(rendered.png).toBase64(), - mimeType: "image/png", - cols: rendered.cols, - rows: rendered.rows, - chars: rendered.chars, - }); - // Keep the event loop responsive between native render passes. - await Bun.sleep(0); - } - - const frames = [...(previousArchive?.frames ?? []), ...newFrames]; - if (frames.length > maxFrames) { - // Pin the earliest frame: it anchors the session head (the original - // request, or the filmed summary of even older history) the way the - // LLM-summary strategies keep the original goal alive across rounds. - // Eviction removes the oldest *unpinned* frames, so the archive fades - // from the middle out — head and tail survive. With a budget of one - // frame the pin is moot; keep the newest frame instead. - const evictStart = maxFrames >= 2 ? 1 : 0; - const dropped = frames.splice(evictStart, frames.length - maxFrames); - for (const frame of dropped) truncatedChars += frame.chars; - } - const totalChars = frames.reduce((sum, frame) => sum + frame.chars, 0); - - let summary: string; - if (frames.length === 0) { - summary = "No prior history."; - } else { - summary = prompt.render(snapcompactSummaryPrompt, { - frameCount: frames.length, - multipleFrames: frames.length > 1, - cols: geometry.cols, - rows: geometry.rows, - totalChars, - truncatedChars, - includedPreviousSummary, - }); - } - const { readFiles, modifiedFiles } = computeFileLists(fileOps); - summary = upsertFileOperations(summary, readFiles, modifiedFiles); - - // A snapcompact pass replaces any provider-side replacement history; strip the - // OpenAI remote-compaction payload like the default summarizer path does. - // OpenAI remote-compaction payload like the default summarizer path does. - const basePreserve = withOpenAiRemoteCompactionPreserveData(previousPreserveData, undefined) ?? {}; - const archive: SnapcompactArchive = { frames, totalChars, truncatedChars }; - - return { - summary, - shortSummary: `Archived ${totalChars.toLocaleString()} chars of history onto ${frames.length} snapcompact frame${frames.length === 1 ? "" : "s"}`, - firstKeptEntryId, - tokensBefore, - details: { readFiles, modifiedFiles } as CompactionDetails, - preserveData: { ...basePreserve, [SNAPCOMPACT_PRESERVE_KEY]: archive }, - }; -} diff --git a/packages/agent/test/snapcompact-frames.test.ts b/packages/agent/test/snapcompact-frames.test.ts new file mode 100644 index 000000000..1e24f74a0 --- /dev/null +++ b/packages/agent/test/snapcompact-frames.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, it } from "bun:test"; +import type { ImageContent } from "@oh-my-pi/pi-ai"; +import { SNAPCOMPACT_FRAME_TOKEN_ESTIMATE } from "@oh-my-pi/snapcompact"; +import { estimateTokens } from "../src/compaction/compaction"; +import { createCompactionSummaryMessage, defaultConvertToLlm } from "../src/compaction/messages"; + +describe("compaction summary message with snapcompact frames", () => { + const images: ImageContent[] = [ + { type: "image", data: "ZmFrZQ==", mimeType: "image/png" }, + { type: "image", data: "ZmFrZTI=", mimeType: "image/png" }, + ]; + + it("estimateTokens charges per attached frame", () => { + const bare = createCompactionSummaryMessage("summary text", 1000, new Date().toISOString()); + const withFrames = createCompactionSummaryMessage( + "summary text", + 1000, + new Date().toISOString(), + undefined, + undefined, + images, + ); + expect(estimateTokens(withFrames) - estimateTokens(bare)).toBe(2 * SNAPCOMPACT_FRAME_TOKEN_ESTIMATE); + }); + + it("defaultConvertToLlm appends frames as image blocks after the summary text", () => { + const message = createCompactionSummaryMessage( + "the snapcompact archive", + 1000, + new Date().toISOString(), + undefined, + undefined, + images, + ); + const [converted] = defaultConvertToLlm([message]); + expect(converted.role).toBe("user"); + const content = converted.content as Array<{ type: string; text?: string; data?: string }>; + expect(content.length).toBe(3); + expect(content[0].type).toBe("text"); + expect(content[0].text).toContain("the snapcompact archive"); + expect(content[1]).toEqual(images[0]); + expect(content[2]).toEqual(images[1]); + }); +}); diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index af15979fe..cebe1dad9 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added optional `ImageContent.detail` (`"auto" | "low" | "high" | "original"`): an OpenAI resolution hint forwarded by the `openai-responses` serializers (default stays `auto`) and by `openai-completions` for the values Chat Completions supports. `"original"` preserves native resolution — required for snapcompact frames, whose pixel-font glyphs do not survive the default downscale. Providers without a detail knob ignore the field. + ## [15.10.12] - 2026-06-10 ### Added diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index 0a1d83eb3..1b3401a94 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/packages/ai/src/providers/openai-completions.ts @@ -1651,6 +1651,8 @@ export function convertMessages( type: "image_url", image_url: { url: `data:${item.mimeType};base64,${item.data}`, + // Chat Completions has no "original"; omit it (provider default). + ...(item.detail && item.detail !== "original" ? { detail: item.detail } : {}), }, } satisfies ChatCompletionContentPartImage); } else { diff --git a/packages/ai/src/providers/openai-responses-shared.ts b/packages/ai/src/providers/openai-responses-shared.ts index 9e28d60e0..5ddf64b81 100644 --- a/packages/ai/src/providers/openai-responses-shared.ts +++ b/packages/ai/src/providers/openai-responses-shared.ts @@ -289,7 +289,7 @@ export function convertResponsesInputContent( for (const item of imageBlocks) { normalizedContent.push({ type: "input_image", - detail: "auto", + detail: item.detail ?? "auto", image_url: `data:${item.mimeType};base64,${item.data}`, } satisfies ResponseInputImage); } @@ -448,7 +448,7 @@ export function appendResponsesToolResultMessages( if (block.type === "image") { contentParts.push({ type: "input_image", - detail: "auto", + detail: block.detail ?? "auto", image_url: `data:${block.mimeType};base64,${block.data}`, } satisfies ResponseInputImage); } diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 1aa43478d..6d64e09c2 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -409,6 +409,12 @@ export interface ImageContent { type: "image"; data: string; // base64 encoded image data mimeType: string; // e.g., "image/jpeg", "image/png" + /** + * OpenAI-only resolution hint. `"original"` preserves native resolution + * (required for snapcompact frames, whose glyphs do not survive the + * default `auto` downscale). Providers without a detail knob ignore it. + */ + detail?: "auto" | "low" | "high" | "original"; } export interface ToolCall { diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 6fcc94773..44ce4e906 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -12,6 +12,7 @@ ### Added - Added pre-TUI startup input capture so users could type while interactive sessions initialize and keep their draft while the application loads +- Snapcompact compaction now passes the session model so frames render in the provider-optimal shape (unscii `8x8r-bw` for Anthropic-family/unknown APIs, `8x8r-sent` for Google, Lanczos-stretched `6x6u-sent` with `detail: "original"` for OpenAI), per the snapcompact 200k-token evals - Added queued submission replay so Enter presses made before startup completion are submitted automatically once interactive mode begins - Added per-turn supersede pruning of stale `read` results: when a file is re-read, older copies of the same path/selector are pruned from context at cache-favorable moments (small suffix, idle gap, or alongside overflow pruning). Gated by the new `compaction.supersedeReads` setting (default on) - Added soft request budgets for task subagents (explore/quick_task 40, others 90, configurable via `task.softRequestBudget`, 0 disables): crossing the budget injects a one-time wrap-up steer into the child; crossing 1.5× aborts the run gracefully diff --git a/packages/coding-agent/examples/extensions/with-deps/package.json b/packages/coding-agent/examples/extensions/with-deps/package.json index 347be16b5..06dc2cb08 100644 --- a/packages/coding-agent/examples/extensions/with-deps/package.json +++ b/packages/coding-agent/examples/extensions/with-deps/package.json @@ -1,6 +1,7 @@ { "name": "pi-extension-with-deps", "version": "1.0.0", + "homepage": "https://omp.sh", "type": "module", "omp": { "extensions": [ diff --git a/packages/coding-agent/package.json b/packages/coding-agent/package.json index 9dd841227..dbd5acefe 100644 --- a/packages/coding-agent/package.json +++ b/packages/coding-agent/package.json @@ -56,6 +56,7 @@ "@oh-my-pi/pi-natives": "catalog:", "@oh-my-pi/pi-tui": "catalog:", "@oh-my-pi/pi-utils": "catalog:", + "@oh-my-pi/snapcompact": "catalog:", "@opentelemetry/api": "catalog:", "@opentelemetry/context-async-hooks": "catalog:", "@opentelemetry/exporter-trace-otlp-proto": "catalog:", diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index d0746f35a..51f82b5b5 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -55,7 +55,6 @@ import { type ShakeRegion, type SummaryOptions, shouldCompact, - snapcompactCompact, } from "@oh-my-pi/pi-agent-core/compaction"; import { DEFAULT_PRUNE_CONFIG, @@ -107,6 +106,7 @@ import { relativePathWithinRoot, Snowflake, } from "@oh-my-pi/pi-utils"; +import { snapcompactCompact } from "@oh-my-pi/snapcompact"; import { type AsyncJob, type AsyncJobDeliveryState, AsyncJobManager } from "../async"; import { classifyDifficulty } from "../auto-thinking/classifier"; import { reset as resetCapabilities } from "../capability"; @@ -6323,7 +6323,7 @@ export class AgentSession { details = compactionPrep.details; preserveData = compactionPrep.preserveData; } else if (snapcompactReady) { - const snapcompactResult = await snapcompactCompact(preparation, { convertToLlm }); + const snapcompactResult = await snapcompactCompact(preparation, { convertToLlm, model: this.model }); summary = snapcompactResult.summary; shortSummary = snapcompactResult.shortSummary; firstKeptEntryId = snapcompactResult.firstKeptEntryId; @@ -7825,7 +7825,7 @@ export class AgentSession { } else if (action === "snapcompact") { // Local, deterministic: render discarded history onto PNG frames. // No model candidates, no API key, no retry loop. - const snapcompactResult = await snapcompactCompact(preparation, { convertToLlm }); + const snapcompactResult = await snapcompactCompact(preparation, { convertToLlm, model: this.model }); summary = snapcompactResult.summary; shortSummary = snapcompactResult.shortSummary; firstKeptEntryId = snapcompactResult.firstKeptEntryId; diff --git a/packages/coding-agent/src/session/session-manager.ts b/packages/coding-agent/src/session/session-manager.ts index bc0c23c40..2c8f49c86 100644 --- a/packages/coding-agent/src/session/session-manager.ts +++ b/packages/coding-agent/src/session/session-manager.ts @@ -2,7 +2,6 @@ import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; -import { getPreservedSnapcompactArchive, snapcompactImages } from "@oh-my-pi/pi-agent-core/compaction"; import type { ImageContent, Message, @@ -28,6 +27,7 @@ import { Snowflake, toError, } from "@oh-my-pi/pi-utils"; +import { getPreservedSnapcompactArchive, snapcompactImages } from "@oh-my-pi/snapcompact"; import { ArtifactManager } from "./artifacts"; import { type BlobPutOptions, diff --git a/packages/coding-agent/test/job-poll-displacement.test.ts b/packages/coding-agent/test/job-poll-displacement.test.ts index 4ce22f38d..81fa7d08d 100644 --- a/packages/coding-agent/test/job-poll-displacement.test.ts +++ b/packages/coding-agent/test/job-poll-displacement.test.ts @@ -97,7 +97,10 @@ describe("job waiting-poll block lifecycle", () => { }); it("never marks non-job tools displaceable", () => { - const component = trackComponent(created, new ToolExecutionComponent("bash", { command: "ls" }, {}, undefined, uiStub)); + const component = trackComponent( + created, + new ToolExecutionComponent("bash", { command: "ls" }, {}, undefined, uiStub), + ); component.updateResult(pollResult(["running"]), false); expect(component.isDisplaceableBlock()).toBe(false); }); @@ -145,7 +148,12 @@ describe("EventController displaces consecutive waiting polls", () => { } async function runPoll(controller: EventController, children: Component[], toolCallId: string) { - await controller.handleEvent({ type: "tool_execution_start", toolCallId, toolName: "job", args: { poll: ["j0"] } }); + await controller.handleEvent({ + type: "tool_execution_start", + toolCallId, + toolName: "job", + args: { poll: ["j0"] }, + }); const component = children[children.length - 1] as ToolExecutionComponent; trackComponent(created, component); await controller.handleEvent({ diff --git a/packages/coding-agent/test/streaming-preview-height.test.ts b/packages/coding-agent/test/streaming-preview-height.test.ts index acf37be34..fdf2476c9 100644 --- a/packages/coding-agent/test/streaming-preview-height.test.ts +++ b/packages/coding-agent/test/streaming-preview-height.test.ts @@ -372,12 +372,12 @@ describe("streaming tool call preview height (bounded across renderers)", () => } }, 30_000); - test("task pending preview stays bounded with a long multiline assignment", () => { + test("task pending preview keeps the full assignment brief", () => { // CONTRACT CHANGE with the single-spawn task rework: the old uncapped - // multi-task `context` rendering is gone with the field. The pending - // preview now intentionally bounds the assignment (first line + a - // "more lines" marker when collapsed; 12 lines when expanded), like - // bash/ssh, so a long assignment can no longer strand the block top. + // multi-task `context` rendering is gone with the field. The assignment + // brief is the durable record of what the subagent was asked to do, so + // the pending preview renders it in full (like eval code) instead of + // windowing it like bash/ssh command previews. const longLines = Array.from({ length: 80 }, (_, i) => `line-${i}`); const { lines, text } = renderPending("task", { agent: "task", @@ -386,12 +386,11 @@ describe("streaming tool call preview height (bounded across renderers)", () => assignment: longLines.join("\n"), }); - expect(lines.length, "task preview should stay bounded").toBeLessThan(20); + expect(lines.length, "task assignment brief should not be capped").toBeGreaterThan(80); expect(text).toContain("preview"); expect(text).toContain("line-0"); - expect(text).not.toContain("line-40"); - expect(text).not.toContain("line-79"); - expect(text, "task preview should advertise truncation").toMatch(/more lines/); + expect(text).toContain("line-40"); + expect(text).toContain("line-79"); }); test("eval pending preview preserves full code (never collapsed)", () => { diff --git a/packages/natives/CHANGELOG.md b/packages/natives/CHANGELOG.md index 1184a6341..23e0aaac0 100644 --- a/packages/natives/CHANGELOG.md +++ b/packages/natives/CHANGELOG.md @@ -4,7 +4,7 @@ ### Added -- Added `renderSnapcompactPng(text, size)`: rasterizes pre-normalized text onto a square 4-bit indexed PNG using the bundled public-domain X.org `5x8` BDF font, with glyph ink cycling through six hues at sentence boundaries. Replaces the JS rasterizer/PNG writer previously in `@oh-my-pi/pi-agent-core`; the font ships inside the crate (`crates/pi-natives/src/fonts/5x8.bdf`). +- Added `renderSnapcompactPng(text, options)`: rasterizes pre-normalized text onto a square PNG in an eval-validated snapcompact shape. Options select the bundled font (`5x8` X.org BDF or `8x8` unscii-8, both public domain, shipped in `crates/pi-natives/src/fonts/`), the ink variant (`sent` six-hue sentence cycling or `bw` black), line repetition (each text line printed N times, copies on a pale highlight band), and a target cell size — cells differing from the font's natural cell render via Lanczos3 stretch into an anti-aliased RGB frame (e.g. the OpenAI-optimal 6x6 unscii shape); native-cell shapes encode as 4-bit indexed PNG. Replaces the JS rasterizer/PNG writer previously in `@oh-my-pi/pi-agent-core`. ## [15.10.12] - 2026-06-10 diff --git a/packages/natives/native/index.d.ts b/packages/natives/native/index.d.ts index a61406901..bbeeb5364 100644 --- a/packages/natives/native/index.d.ts +++ b/packages/natives/native/index.d.ts @@ -1285,14 +1285,15 @@ export interface PtyStartOptions { export declare function readImageFromClipboard(): Promise /** - * Render one snapcompact frame: print pre-normalized text onto a - * `size` x `size` 4-bit indexed bitmap and encode it as a PNG. + * Render one snapcompact frame: print pre-normalized text onto a square + * bitmap and encode it as PNG. * - * The glyph grid holds `floor(size/5) * floor(size/8)` characters; input - * beyond that is ignored (the caller chunks text to capacity). Returns the - * PNG bytes. + * The glyph grid holds `floor(size/cellWidth) * floor(size/cellHeight/lineRepeat)` + * characters; input beyond that is ignored (the caller chunks text to + * capacity). Native-cell shapes encode as 4-bit indexed PNG; stretched + * shapes (target cell != font cell) encode as RGB. Returns the PNG bytes. */ -export declare function renderSnapcompactPng(text: string, size: number): Uint8Array +export declare function renderSnapcompactPng(text: string, options: SnapcompactRenderOptions): Uint8Array /** * Search content for a pattern (one-shot, compiles pattern each time). @@ -1423,6 +1424,31 @@ export interface SliceResult { */ export declare function sliceWithWidth(line: string, startCol: number, length: number, strict: boolean | undefined | null, tabWidth: number): SliceResult +/** Shape options for one snapcompact frame. */ +export interface SnapcompactRenderOptions { + /** Frame edge in pixels. */ + size: number + /** Bundled font: `"5x8"` (X.org BDF) or `"8x8"` (unscii-8). Default `"5x8"`. */ + font?: string + /** + * Target cell advance in pixels. Differing from the font's natural cell + * triggers the Lanczos stretch path. Default: font natural width. + */ + cellWidth?: number + /** Target cell pitch in pixels. Default: font natural height. */ + cellHeight?: number + /** + * Ink variant: `"sent"` (six-hue sentence cycling) or `"bw"` (black). + * Default `"sent"`. + */ + variant?: string + /** + * Print each text line this many times; copies after the first sit on a + * pale highlight band. Default 1. + */ + lineRepeat?: number +} + export declare function summarizeCode(options: SummaryOptions): SummaryResult export interface SummaryOptions { diff --git a/packages/snapcompact/.gitignore b/packages/snapcompact/.gitignore new file mode 100644 index 000000000..839945bad --- /dev/null +++ b/packages/snapcompact/.gitignore @@ -0,0 +1,2 @@ +.cache/ +results/ diff --git a/packages/snapcompact/CHANGELOG.md b/packages/snapcompact/CHANGELOG.md new file mode 100644 index 000000000..d62403267 --- /dev/null +++ b/packages/snapcompact/CHANGELOG.md @@ -0,0 +1,10 @@ +# Changelog + +## [Unreleased] +### Added + +- Added provider-specific snapcompact frame-shape presets and shape helpers (`SNAPCOMPACT_SHAPES`, `resolveSnapcompactShape`, `isSnapcompactShape`) so callers can consistently select validated image-frame geometry for archive renders +- Added `file-operations.md` and `snapcompact-summary.md` prompts to preserve file-read/write context and frame metadata in the compaction prompt flow +- Added a full `packages/snapcompact/research` experiment and visualization suite for running snapcompact SQuAD studies, provider probes, and activation-style analyses +- Added package-level TypeScript exports and publication config so consumers can import `@oh-my-pi/snapcompact` with typed access to snapcompact APIs +- Published `@oh-my-pi/snapcompact` as the reusable snapcompact compaction package, including bitmap-frame rendering helpers, archive helpers, and the local `snapcompactCompact()` strategy. diff --git a/packages/snapcompact/package.json b/packages/snapcompact/package.json new file mode 100644 index 000000000..bc6bab478 --- /dev/null +++ b/packages/snapcompact/package.json @@ -0,0 +1,62 @@ +{ + "type": "module", + "name": "@oh-my-pi/snapcompact", + "version": "15.10.12", + "description": "Bitmap-frame context compression for vision-capable LLMs", + "homepage": "https://omp.sh", + "author": "Can Boluk", + "license": "MIT", + "repository": { + "type": "git", + "url": "git+https://github.com/can1357/oh-my-pi.git", + "directory": "packages/snapcompact" + }, + "bugs": { + "url": "https://github.com/can1357/oh-my-pi/issues" + }, + "keywords": [ + "context-compression", + "vision", + "compaction", + "llm" + ], + "main": "./src/index.ts", + "types": "./src/index.ts", + "scripts": { + "check": "biome check . && bun run check:types", + "check:types": "tsgo -p tsconfig.json --noEmit", + "lint": "biome lint .", + "test": "bun test --parallel", + "fix": "biome check --write --unsafe .", + "fmt": "biome format --write ." + }, + "dependencies": { + "@oh-my-pi/pi-ai": "catalog:", + "@oh-my-pi/pi-natives": "catalog:", + "@oh-my-pi/pi-utils": "catalog:" + }, + "devDependencies": { + "@types/bun": "catalog:" + }, + "engines": { + "bun": ">=1.3.14" + }, + "files": [ + "src", + "CHANGELOG.md" + ], + "exports": { + ".": { + "types": "./src/index.ts", + "import": "./src/index.ts" + }, + "./snapcompact": { + "types": "./src/snapcompact.ts", + "import": "./src/snapcompact.ts" + }, + "./*": { + "types": "./src/*.ts", + "import": "./src/*.ts" + } + } +} diff --git a/packages/snapcompact/research/anthropic_api.py b/packages/snapcompact/research/anthropic_api.py new file mode 100644 index 000000000..d6a94fa6e --- /dev/null +++ b/packages/snapcompact/research/anthropic_api.py @@ -0,0 +1,74 @@ +"""Minimal Anthropic Messages API client. Key from ~/.env, no SDK.""" + +import base64 +import json +import time +import urllib.error +import urllib.request +from pathlib import Path + +API_URL = "https://api.anthropic.com/v1/messages" +API_VERSION = "2023-06-01" + + +def load_api_key(env_path: str = "~/.env") -> str: + """Last ANTHROPIC_API_KEY assignment wins (mirrors shell sourcing).""" + key = None + for line in Path(env_path).expanduser().read_text().splitlines(): + line = line.strip() + if line.startswith("export "): + line = line[len("export ") :] + if line.startswith("ANTHROPIC_API_KEY="): + key = line.split("=", 1)[1].strip().strip("'\"") + if not key: + raise SystemExit(f"no ANTHROPIC_API_KEY in {env_path}") + return key + + +def image_block(png_path: Path) -> dict: + data = base64.b64encode(png_path.read_bytes()).decode() + return {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": data}} + + +def complete( + api_key: str, + model: str, + messages: list[dict], + system: str | None = None, + max_tokens: int = 8192, + effort: str | None = None, + retries: int = 4, +) -> tuple[str, dict, str]: + """Returns (joined text content, usage dict, stop_reason). + + effort: adaptive-thinking effort (low|medium|high|xhigh|max); None = provider default. + """ + body: dict = {"model": model, "max_tokens": max_tokens, "messages": messages} + if system: + body["system"] = system + if effort: + body["output_config"] = {"effort": effort} + payload = json.dumps(body).encode() + headers = { + "content-type": "application/json", + "x-api-key": api_key, + "anthropic-version": API_VERSION, + } + if effort: + headers["anthropic-beta"] = "effort-2025-11-24" + req = urllib.request.Request(API_URL, data=payload, headers=headers) + for attempt in range(retries + 1): + try: + with urllib.request.urlopen(req, timeout=300) as resp: + out = json.load(resp) + text = "".join(b.get("text", "") for b in out["content"] if b.get("type") == "text") + return text, out.get("usage", {}), out.get("stop_reason", "") + except urllib.error.HTTPError as err: + detail = err.read().decode(errors="replace")[:500] + if err.code in (429, 500, 502, 503, 529) and attempt < retries: + wait = 2.0 * 2**attempt + print(f" HTTP {err.code}, retrying in {wait:.0f}s: {detail[:120]}") + time.sleep(wait) + continue + raise SystemExit(f"API error {err.code}: {detail}") from err + raise AssertionError("unreachable") diff --git a/packages/snapcompact/research/bdf.py b/packages/snapcompact/research/bdf.py new file mode 100644 index 000000000..15003636d --- /dev/null +++ b/packages/snapcompact/research/bdf.py @@ -0,0 +1,247 @@ +"""BDF/HEX pixel-font parsing and dense text-to-image rendering.""" + +import colorsys +import urllib.request +from dataclasses import dataclass +from pathlib import Path + +from PIL import Image + +XORG_RAW = "https://gitlab.freedesktop.org/xorg/font/misc-misc/-/raw/master/{name}.bdf" +TOM_THUMB = "https://robey.lag.net/downloads/tom-thumb.bdf" +UNSCII_HEX = "https://raw.githubusercontent.com/viznut/unscii/master/fontfiles/{name}.hex" + + +@dataclass(frozen=True) +class FontCfg: + """One density configuration: a BDF font drawn on an adv x pitch cell grid.""" + + name: str # condition label, e.g. "6x10" + source: str # bdf file stem or "tom-thumb" + adv: int # x advance per character cell, px + pitch: int # y advance per row, px + ascent: int | None = None # override; default from FONT_ASCENT + native: tuple[int, int] | None = None # rasterize at this cell size, then resize (stretch) to adv x pitch + repeat: int = 1 # render each text line this many times (copy 0 plain, later copies bg-highlighted) + + +def ensure_font(cfg: FontCfg, cache: Path) -> Path: + hexfont = cfg.source.startswith("unscii") + path = cache / f"{cfg.source}.{'hex' if hexfont else 'bdf'}" + if not path.exists(): + if hexfont: + url = UNSCII_HEX.format(name=cfg.source) + else: + url = TOM_THUMB if cfg.source == "tom-thumb" else XORG_RAW.format(name=cfg.source) + urllib.request.urlretrieve(url, path) + return path + + +def parse_bdf(path: Path) -> tuple[dict[int, dict], int]: + """Returns ({codepoint: {bbx, rows}}, font_ascent).""" + glyphs: dict[int, dict] = {} + ascent = 0 + cur: dict = {} + lines = path.read_text().splitlines() + i = 0 + while i < len(lines): + ln = lines[i] + if ln.startswith("FONT_ASCENT"): + ascent = int(ln.split()[1]) + elif ln.startswith("STARTCHAR"): + cur = {"rows": []} + elif ln.startswith("ENCODING"): + cur["enc"] = int(ln.split()[1]) + elif ln.startswith("BBX"): + cur["bbx"] = tuple(map(int, ln.split()[1:5])) + elif ln.startswith("BITMAP"): + i += 1 + while not lines[i].startswith("ENDCHAR"): + cur["rows"].append(int(lines[i], 16)) + i += 1 + glyphs[cur["enc"]] = cur + i += 1 + return glyphs, ascent + + +def parse_hex(path: Path) -> tuple[dict[int, dict], int]: + """Unifont-style .hex (unscii-8: 8x8, one byte per row). Baseline at row 7.""" + glyphs: dict[int, dict] = {} + for line in path.read_text().splitlines(): + cp, _, bits = line.partition(":") + data = bytes.fromhex(bits.strip()) + if len(data) == 8: + glyphs[int(cp, 16)] = {"bbx": (8, 8, 0, -1), "rows": list(data)} + return glyphs, 7 + + +def load_font(cfg: FontCfg, cache: Path) -> tuple[dict[int, dict], int]: + path = ensure_font(cfg, cache) + return parse_hex(path) if path.suffix == ".hex" else parse_bdf(path) + + +# 6 hues; dark variant for glyphs, pale variant for the row's background band. +_HUES = [0.0, 0.08, 0.3, 0.5, 0.62, 0.78] +_DARK = [tuple(int(c * 255) for c in colorsys.hls_to_rgb(h, 0.22, 0.95)) for h in _HUES] +_PALE = [tuple(int(c * 255) for c in colorsys.hls_to_rgb(h, 0.94, 0.6)) for h in _HUES] +_BRIGHT = [tuple(int(c * 255) for c in colorsys.hls_to_rgb(h, 0.70, 0.95)) for h in _HUES] + +VARIANTS = ("color", "zebra", "bw", "sent", "dark", "dark-sent", "dim", "sent-dim", "dark-sent-dim") +_BLACK = (0, 0, 0) +_WHITE = (255, 255, 255) +_GRAY = (232, 232, 232) +_DIMMED = (176, 176, 176) +_DIMMED_DARK = (104, 104, 104) +_REP_LIGHT = (255, 247, 194) # pale yellow highlight for repeated-line copies +_REP_DARK = (44, 44, 24) + +# High-frequency function words a reader can reconstruct from context; the dim +# variants render them in light gray so content words carry the contrast. +_STOPWORDS = frozenset( + "the a an and or of to in on at as is are was were be been by for with that this it its from had has have not but " + "he she his her they their them which also who whom when where while will would could should there then than " + "into over under about after before between during each such these those some most more other only same so".split() +) + + +def _row_palette(variant: str, row: int) -> tuple[tuple[int, int, int], tuple[int, int, int]]: + """(background, default glyph) colors for a row under the given render variant.""" + if variant == "color": + return _PALE[row % 6], _DARK[row % 6] + if variant == "zebra": + return (_WHITE if row % 2 == 0 else _GRAY), _BLACK + if variant in ("bw", "sent", "dim", "sent-dim"): + return _WHITE, _BLACK + if variant in ("dark", "dark-sent", "dark-sent-dim"): + return _BLACK, _WHITE + raise ValueError(f"unknown variant: {variant}") + + +def _sentence_indices(text: str) -> list[int]: + """Running sentence index per character (boundary: terminator + space).""" + out = [0] * len(text) + idx = 0 + for i, ch in enumerate(text): + out[i] = idx + if ch in ".!?" and i + 1 < len(text) and text[i + 1] == " ": + idx += 1 + return out + + +def _stopword_mask(text: str) -> list[bool]: + """True for every character of a word in _STOPWORDS.""" + mask = [False] * len(text) + i = 0 + while i < len(text): + if text[i].isalpha(): + j = i + while j < len(text) and text[j].isalpha(): + j += 1 + if text[i:j].lower() in _STOPWORDS: + for k in range(i, j): + mask[k] = True + i = j + else: + i += 1 + return mask + + +def capacity(cfg: FontCfg, size: int = 1568, columns: int = 1) -> tuple[int, int, int]: + """(cols per line, rows, chars) that fit a size x size image with `columns` newspaper columns.""" + rows = size // cfg.pitch // cfg.repeat + gutter = 2 * cfg.adv if columns > 1 else 0 + cols = (size - (columns - 1) * gutter) // columns // cfg.adv + return cols, rows, columns * cols * rows + + +def render( + text: str, cfg: FontCfg, cache: Path, size: int = 1568, variant: str = "color", columns: int = 1 +) -> Image.Image: + """Fill a size x size grid with `text`; styling per `variant`. + + Layout: full-width row-major when columns == 1; otherwise newspaper flow + (fill the leftmost column top-to-bottom, then the next), columns separated + by a 2-cell gutter with a hairline rule. + + When cfg.native is set, glyphs are rasterized at the native cell size and + the whole canvas is Lanczos-resized to the adv x pitch target (anisotropic + stretch: anti-aliased, no cropping or ink fusion). + + Variants: + color per-row hue cycle on pale row bands + zebra black text, alternating white/gray bands + bw black on white + sent white background, glyph hue cycles per sentence + dark white text on black + dark-sent bright sentence hues on black + dim black on white, stopwords dimmed gray + sent-dim sentence hues, stopwords dimmed gray + dark-sent-dim bright sentence hues on black, stopwords dimmed + """ + glyphs, font_ascent = load_font(cfg, cache) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, cap = capacity(cfg, size, columns) + text = text[:cap] + sent_idx = _sentence_indices(text) if variant in ("sent", "dark-sent", "sent-dim", "dark-sent-dim") else None + dim_mask = _stopword_mask(text) if variant in ("dim", "sent-dim", "dark-sent-dim") else None + sent_palette = _BRIGHT if variant in ("dark-sent", "dark-sent-dim") else _DARK + dark_bg = variant in ("dark", "dark-sent", "dark-sent-dim") + base_color = _BLACK if dark_bg else _WHITE + if cfg.native is not None: + aw, ph = cfg.native + gutter = 2 * aw if columns > 1 else 0 + span = cols * aw + gutter + canvas_w, canvas_h = columns * span - gutter, rows * cfg.repeat * ph + else: + aw, ph = cfg.adv, cfg.pitch + gutter = 2 * aw if columns > 1 else 0 + span = cols * aw + gutter + canvas_w = canvas_h = size + img = Image.new("RGB", (canvas_w, canvas_h), base_color) + px = img.load() + for row in range(rows): + bg, row_fg = _row_palette(variant, row) + for copy in range(cfg.repeat): + y0 = (row * cfg.repeat + copy) * ph + cbg = bg if copy == 0 else (_REP_DARK if dark_bg else _REP_LIGHT) + for y in range(y0, min(y0 + ph, canvas_h)): + for x in range(canvas_w): + px[x, y] = cbg + for blk in range(columns): + for col in range(cols): + i = (blk * rows + row) * cols + col + if i >= len(text): + break + glyph = glyphs.get(ord(text[i])) + if glyph is None: + continue + fg = row_fg + if sent_idx is not None: + fg = sent_palette[sent_idx[i] % 6] + if dim_mask is not None and dim_mask[i]: + fg = _DIMMED_DARK if dark_bg else _DIMMED + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < canvas_h: + continue + for b in range(w): + if bits & (shift >> b): + x = blk * span + col * aw + xoff + b + if 0 <= x < canvas_w: + px[x, y] = fg + rule = (96, 96, 96) if dark_bg else (204, 204, 204) + for blk in range(1, columns): + x = blk * span - gutter // 2 + if 0 <= x < canvas_w: + for y in range(canvas_h): + px[x, y] = rule + if cfg.native is not None: + img = img.resize((canvas_w * cfg.adv // aw, canvas_h * cfg.pitch // ph), Image.LANCZOS) + if img.size != (size, size): + out = Image.new("RGB", (size, size), base_color) + out.paste(img, (0, 0)) + img = out + return img diff --git a/packages/snapcompact/research/exp01_patchalign.py b/packages/snapcompact/research/exp01_patchalign.py new file mode 100644 index 000000000..29431073c --- /dev/null +++ b/packages/snapcompact/research/exp01_patchalign.py @@ -0,0 +1,284 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp01: ViT patch-alignment hypothesis. + +Vision encoders patch at fixed pixel grids (14/16 px, often 2x2-merged; Gemini +tiles at 768 px = 48*16). Current fonts (pitch 10/12/13) straddle patch +boundaries so glyph rows smear across visual tokens. Test patch-aligned cell +grids against same-glyph-budget misaligned controls: + + aligned: img-7x14-sent (native X.Org 7x14), + img-8x16-sent (native Spleen 8x16; X.Org misc-misc has no 8x16), + img-8on16-sent (8x13 glyphs on an 8x16 cell: identical glyphs to + the control, ONLY the pitch changes), + img-6on7x14-sent (6x12 glyphs on a 7x14 cell) + controls: img-7x13-sent, img-8x13-sent (same glyph width, pitch 13) + +Render-size probe: 1568 (baseline size) vs 1536 = 2*768 (exact Gemini tile +multiple; integer 2x downsample keeps 16 px pitch on an 8 px grid; also 3*512 +for OpenAI tile schemes) for the 16 px-pitch fonts. + +Usage: uv run exp01_patchalign.py # full grid + uv run exp01_patchalign.py --report # re-aggregate from cache +""" + +import argparse +import csv +import json +import sys +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import FontCfg, capacity, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODELS = { + "gpt-5.5": (2.0, 16.0), + "google/gemini-3.5-flash": (0.6, 4.0), +} +FONTS = { + "7x14": FontCfg("7x14", "7x14", 7, 14), # aligned: pitch = 14 px patch + "8x16": FontCfg("8x16", "spleen-8x16", 8, 16), # aligned: native 16 px font (Spleen) + "8on16": FontCfg("8on16", "8x13", 8, 16), # aligned: 8x13 glyphs, pitch 16 cell + "6on7x14": FontCfg("6on7x14", "6x12", 7, 14), # aligned: 6x12 glyphs, 7x14 cell + "7x13": FontCfg("7x13", "7x13", 7, 13), # control for 7x14 (same glyph budget, pitch 13) + "8x13": FontCfg("8x13", "8x13", 8, 13), # control for 8x16/8on16 +} +SPLEEN_URL = "https://raw.githubusercontent.com/fcambus/spleen/master/spleen-8x16.bdf" +# (condition, render size) cells. 1536 only for 16 px-pitch fonts (14 does not +# divide 768 or 1536). +GRID = ( + ("img-7x14-sent", 1568), + ("img-7x13-sent", 1568), + ("img-8x16-sent", 1568), + ("img-8on16-sent", 1568), + ("img-8x13-sent", 1568), + ("img-6on7x14-sent", 1568), + ("img-8x16-sent", 1536), + ("img-8on16-sent", 1536), +) +LENGTHS = (50, 150) + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def parse_img_condition(name: str) -> tuple[str, str]: + _, font, variant = name.split("-", 2) + return font, variant + + +def ensure_spleen() -> None: + """bdf.ensure_font only knows X.Org/tom-thumb URLs; stage Spleen 8x16 ourselves.""" + path = CACHE / "spleen-8x16.bdf" + if path.exists() and path.stat().st_size > 0: + return + tmp = path.with_suffix(".tmp.bdf") + urllib.request.urlretrieve(SPLEEN_URL, tmp) + tmp.replace(path) + + +def run_cell_chunk(model: str, cond: str, size: int, start: int, end: int, ctx: dict) -> list[dict]: + """One (model, condition, size, chunk) unit: render carrier, QA, score.""" + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + + font, variant = parse_img_condition(cond) + png = CACHE / f"exp01-{font}-{variant}-{sha8(chunk_text, str(size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render(chunk_text, FONTS[font], CACHE, size, variant).save(tmp) + tmp.replace(png) + cols, rows, _ = capacity(FONTS[font], size) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("qa-image.md").format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, "exp01-qa", {"messages": messages, "size": size, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "size": size, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--report", action="store_true", help="reprint from cache only (re-runs cells; all should hit cache)") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + ensure_spleen() + out_dir = RESULTS / "exp01-patchalign" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for cond, size in GRID: + budget = capacity(FONTS[parse_img_condition(cond)[0]], size)[2] + for start in range(0, len(flow), budget): + tasks.append((model, cond, size, start, min(start + budget, len(flow)), ctx)) + print(f"grid: {len(models)} models x {len(lengths)} lengths x {len(GRID)} cells = {len(tasks)} chunk tasks") + + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell_chunk, m, c, sz, s, e, ctx) for m, c, sz, s, e, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + if done % 10 == 0: + print(f" {done}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond, size in GRID: + sub = [ + r for r in records + if r["model"] == model and r["length"] == length and r["cond"] == cond and r["size"] == size + ] + if not sub: + continue + cells.append( + { + "model": model, + "length": length, + "condition": f"{cond}@{size}", + **aggregate(sub, *MODELS[model]), + } + ) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for length in lengths: + print(f"\n== {length} passages == (F1 +-se / n / $)") + hdr = f"{'condition':<22}" + "".join(f"{m:>34}" for m in models) + print(hdr + "\n" + "-" * len(hdr)) + for cond, size in GRID: + label = f"{cond}@{size}" + row = f"{label:<22}" + for model in models: + cell = next( + (c for c in cells if c["model"] == model and c["length"] == length and c["condition"] == label), + None, + ) + row += ( + f"{cell['f1']:>10.3f} +-{cell['f1_se']:.3f} {cell['n']:>4} {cell['cost_usd']:>7.3f}" + if cell + else f"{'-':>34}" + ) + print(row) + print(f"\ndataset -> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp02_surprisal.py b/packages/snapcompact/research/exp02_surprisal.py new file mode 100644 index 000000000..63fe857f6 --- /dev/null +++ b/packages/snapcompact/research/exp02_surprisal.py @@ -0,0 +1,427 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "wordfreq"] +# /// +"""exp02: surprisal-weighted glyph contrast + stopword-removal/disemvowel density. + +Conditions (all 6x10 font, 1568px, vs baseline img-6x10-sent): + img-6x10-surp full text; glyph gray level = unigram surprisal bucket + (wordfreq zipf: ultra-common -> light gray, rare -> black) + img-6x10-sent-surp sentence hues; lightness scaled by surprisal bucket + img-6x10-disemv stopwords dropped + lowercase non-entity words disemvoweled, + rendered with plain `sent` hues; chunks sized by *transformed* + capacity so each page carries more original chars (fewer pages) + +Usage: uv run exp02_surprisal.py [--report] [--fresh] +Keys: OPENAI_API_KEY + OPENROUTER_API_KEY from ~/.env. +""" + +import argparse +import colorsys +import csv +import json +import re +import sys +from concurrent.futures import ThreadPoolExecutor +from functools import lru_cache +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _sentence_indices, capacity, ensure_font, parse_bdf, render # noqa: E402 +from PIL import Image # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 +from wordfreq import zipf_frequency # noqa: E402 + +MODELS = {"gpt-5.5": (2.0, 16.0), "google/gemini-3.5-flash": (0.6, 4.0)} +LENGTHS = (50, 150) +CONDITIONS = ("img-6x10-surp", "img-6x10-sent-surp", "img-6x10-disemv") +FONT = FONTS["6x10"] + +# Baseline cells from results/optimal-*/matrix.csv (img-6x10-sent): F1, se, cost$. +BASELINE = { + ("gpt-5.5", 50): (0.850, 0.051, 0.068), + ("gpt-5.5", 150): (0.822, 0.029, 0.245), + ("google/gemini-3.5-flash", 50): (0.984, 0.012, 0.018), + ("google/gemini-3.5-flash", 150): (0.805, 0.035, 0.097), +} + +SURP_NOTE = ( + "Glyph darkness encodes word informativeness: very common words are printed in lighter gray, " + "rarer / more informative words in darker ink. All words are spelled out in full." +) +DISEMV_NOTE = ( + "To fit more text, very common function words (the, of, and, is, ...) were removed, and other " + 'common lowercase words are abbreviated by stripping their interior vowels (e.g. "qck brwn fx jmpd" ' + 'means "quick brown fox jumped"; "gvrnmnt" means "government"). Proper nouns, capitalized words, ' + "numbers, dates, and words adjacent to numbers are kept verbatim. Mentally reconstruct the original " + "wording; ALWAYS write your answers in normal, fully spelled English." +) + +# --- text transform: stopword drop + disemvowel ------------------------------ + +# bdf._STOPWORDS minus "not" (removal flips meaning; dimming it was harmless). +_STOP = frozenset( + "the a an and or of to in on at as is are was were be been by for with that this it its from had has have but " + "he she his her they their them which also who whom when where while will would could should there then than " + "into over under about after before between during each such these those some most more other only same so".split() +) +# Spelled-out numbers/ordinals/units stay verbatim: golds are full of them. +_NUMWORDS = frozenset( + "zero one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen sixteen " + "seventeen eighteen nineteen twenty thirty forty fifty sixty seventy eighty ninety hundred thousand million " + "billion trillion first second third fourth fifth sixth seventh eighth ninth tenth half quarter percent".split() +) +_VOWELS = frozenset("aeiou") +_TOKEN_RE = re.compile(r"([^A-Za-z]*)([a-z]+)([^A-Za-z]*)") +_WORD_RE = re.compile(r"[A-Za-z]+") + + +def _has_digit(s: str) -> bool: + return any(c.isdigit() for c in s) + + +def _disemvowel(w: str) -> str: + if len(w) <= 2: + return w + return w[0] + "".join(c for c in w[1:-1] if c not in _VOWELS) + w[-1] + + +def transform(text: str) -> str: + """Drop lowercase stopwords, disemvowel other lowercase words. + + Never touches: capitalized words, tokens containing digits or apostrophes, + tokens adjacent to a digit-bearing token, spelled-out numbers/units. + """ + toks = text.split() + out: list[str] = [] + for i, tok in enumerate(toks): + m = _TOKEN_RE.fullmatch(tok) + if not m: # uppercase, digits, apostrophes, hyphens-with-letters: verbatim + out.append(tok) + continue + pre, core, suf = m.groups() + if ( + _has_digit(pre) + or _has_digit(suf) + or (i > 0 and _has_digit(toks[i - 1])) + or (i + 1 < len(toks) and _has_digit(toks[i + 1])) + or core in _NUMWORDS + ): + out.append(tok) + continue + if core in _STOP: + if pre + suf: # keep punctuation of dropped words + out.append(pre + suf) + continue + out.append(pre + _disemvowel(core) + suf) + return " ".join(out) + + +# --- surprisal-weighted rendering -------------------------------------------- + +# 4 buckets, near-black (rare) .. light gray (ultra-common). Lightest matches +# the old _DIMMED (176) readability point. +_GRAYS = [(185, 185, 185), (135, 135, 135), (75, 75, 75), (0, 0, 0)] +_LIGHT = [0.72, 0.55, 0.40, 0.22] # lightness per bucket for sent hues +_HUES = [0.0, 0.08, 0.3, 0.5, 0.62, 0.78] +_SENT_SURP = [ + [tuple(int(c * 255) for c in colorsys.hls_to_rgb(h, l, 0.95)) for l in _LIGHT] for h in _HUES +] +_WHITE = (255, 255, 255) + + +@lru_cache(maxsize=65536) +def _zipf_bucket(word: str) -> int: + """0 = ultra-common (lightest) .. 3 = rare/unknown (black).""" + z = zipf_frequency(word, "en") + if z >= 6.0: + return 0 + if z >= 5.0: + return 1 + if z >= 4.0: + return 2 + return 3 + + +def _shade_indices(text: str) -> list[int]: + """Per-char surprisal bucket; non-letters (digits, punctuation) stay black.""" + out = [3] * len(text) + for m in _WORD_RE.finditer(text): + b = _zipf_bucket(m.group().lower()) + if b != 3: + for k in range(m.start(), m.end()): + out[k] = b + return out + + +def render_surp(text: str, cfg, cache: Path, size: int = 1568, sent_hues: bool = False) -> Image.Image: + """bdf.render() with the boolean dim_mask generalized to surprisal buckets.""" + glyphs, font_ascent = parse_bdf(ensure_font(cfg, cache)) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, cap = capacity(cfg, size) + text = text[:cap] + shades = _shade_indices(text) + sent_idx = _sentence_indices(text) if sent_hues else None + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for row in range(rows): + y0 = row * cfg.pitch + for col in range(cols): + i = row * cols + col + if i >= len(text): + break + glyph = glyphs.get(ord(text[i])) + if glyph is None: + continue + b = shades[i] + fg = _SENT_SURP[sent_idx[i] % 6][b] if sent_idx is not None else _GRAYS[b] + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for bcol in range(w): + if bits & (shift >> bcol): + x = col * cfg.adv + xoff + bcol + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- chunk planning ----------------------------------------------------------- + + +def plan_chunks(flow: str, cond: str, size: int) -> list[tuple[int, int]]: + cap = capacity(FONT, size)[2] + if cond != "img-6x10-disemv": + return [(s, min(s + cap, len(flow))) for s in range(0, len(flow), cap)] + # Greedy: max original span whose *transformed* text fits a page; word-snapped. + chunks: list[tuple[int, int]] = [] + s, n = 0, len(flow) + while s < n: + if len(transform(flow[s:n])) <= cap: + chunks.append((s, n)) + break + lo, hi = s + cap, n # transform never lengthens => s+cap always fits + while lo < hi: + mid = (lo + hi + 1) // 2 + if len(transform(flow[s:mid])) <= cap: + lo = mid + else: + hi = mid - 1 + e = lo + sp = flow.rfind(" ", s, e) # don't split a word across pages + if sp > s: + e = sp + chunks.append((s, e)) + s = e + (1 if e < n and flow[e] == " " else 0) + return chunks + + +# --- harness (mirrors final.py) ----------------------------------------------- + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def build_png(cond: str, render_text: str, size: int) -> Path: + png = CACHE / f"exp02-{cond}-{sha8(render_text, str(size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + if cond == "img-6x10-surp": + img = render_surp(render_text, FONT, CACHE, size, sent_hues=False) + elif cond == "img-6x10-sent-surp": + img = render_surp(render_text, FONT, CACHE, size, sent_hues=True) + else: # disemv: text transform is the variable; keep baseline sent hues + img = render(render_text, FONT, CACHE, size, "sent") + img.save(tmp) + tmp.replace(png) + return png + + +def run_chunk(model: str, cond: str, start: int, end: int, png: Path, render_chars: int, ctx: dict) -> list[dict]: + args, keys = ctx["args"], ctx["keys"] + questions = squad.sample_chunk_questions(ctx["paras"], ctx["offsets"], start, end, args.qpc, args.seed) + if not questions: + return [] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + cols, rows, _ = capacity(FONT, args.size) + extra = DISEMV_NOTE if cond == "img-6x10-disemv" else SURP_NOTE + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("exp02-qa-image.md").format(cols=cols, rows=rows, extra=extra)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, "exp02-qa", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + records[0]["chunk_orig_chars"] = end - start + records[0]["chunk_render_chars"] = render_chars + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + pages = sum(1 for r in records if "chunk_orig_chars" in r) + orig_chars = sum(r.get("chunk_orig_chars", 0) for r in records) + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + "pages": pages, + "orig_chars": orig_chars, + "orig_chars_per_page": round(orig_chars / pages) if pages else 0, + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--report", action="store_true", help="reprint from cache only") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp02-surprisal" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + + keys = {} + if not args.report: + keys["openai"] = load_env_key("OPENAI_API_KEY", args.env) + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for cond in conditions: + chunks = plan_chunks(flow, cond, args.size) + for start, end in chunks: + orig = flow[start:end] + render_text = transform(orig) if cond == "img-6x10-disemv" else orig + png = build_png(cond, render_text, args.size) # pre-render: no tmp races in pool + if cond == "img-6x10-disemv": + print( + f" len={length} disemv chunk [{start},{end}): {end - start} orig -> " + f"{len(render_text)} rendered chars (x{(end - start) / len(render_text):.2f})" + ) + for model in models: + tasks.append((model, cond, start, end, png, len(render_text), ctx)) + print(f"grid: {len(models)} models x {len(lengths)} lengths x {len(conditions)} conditions = {len(tasks)} chunk tasks") + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_chunk, *t) for t in tasks] + for i, fut in enumerate(futures): + records.extend(fut.result()) + print(f" {i + 1}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + print(f"\n{'model':<26}{'len':>5}{'condition':<22}{'n':>4}{'EM':>7}{'F1':>7}{'se':>7}{'cost$':>8}{'dF1 vs base':>13}") + for c in cells: + base = BASELINE.get((c["model"], c["length"])) + d = f"{c['f1'] - base[0]:+.3f}" if base else "-" + print( + f"{c['model']:<26}{c['length']:>5} {c['condition']:<20}{c['n']:>4}{c['em']:>7.3f}" + f"{c['f1']:>7.3f}{c['f1_se']:>7.3f}{c['cost_usd']:>8.3f}{d:>13}" + ) + print(f"\ndataset -> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp03_numhard.py b/packages/snapcompact/research/exp03_numhard.py new file mode 100644 index 000000000..115fbb96c --- /dev/null +++ b/packages/snapcompact/research/exp03_numhard.py @@ -0,0 +1,382 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp03_numhard: harden digit glyphs in optical-compaction renders. + +Extractive QA on rendered pages dies on 0/O, 1/l, rn/m confusions, and SQuAD +answers skew heavily toward numbers/dates. This experiment re-renders the +img-6x10 baseline with a hardening mask over digits (plus number/date +punctuation directly adjacent to a digit: "1,000", "3.5%", "1914-18"): + + numbold double-strike: every masked glyph pixel painted at x and x+1 + (6x10 digit ink spans cols 0-4 of the 6px cell, so x+1 stays + in-cell -- no bleed into the neighbor) + numred masked glyphs painted pure dark red, overriding sent/bw color + +Conditions: img-6x10-sent-numbold, img-6x10-sent-numred, img-6x10-bw-numred. +Methodology matches the baseline grid (seed 42, qpc 30, size 1568, +max_tokens 32768, effort None) so question sets are identical and the +numeric-gold subset can be compared apples-to-apples against the baseline +records in results/optimal-gpt55 / results/optimal-gemini. + +Usage: uv run exp03_numhard.py [--report] +""" + +import argparse +import csv +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DARK, _row_palette, _sentence_indices, capacity, ensure_font, parse_bdf # noqa: E402 +from final import MODELS, aggregate, cached # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp03" +SLUG = "numhard" +FONT = FONTS["6x10"] +CONDITIONS = ("img-6x10-sent-numbold", "img-6x10-sent-numred", "img-6x10-bw-numred") +RUN_MODELS = ("gpt-5.5", "google/gemini-3.5-flash") +BASELINE_COND = "img-6x10-sent" +BASELINE_RECORDS = { + "gpt-5.5": RESULTS / "optimal-gpt55" / "records.jsonl", + "google/gemini-3.5-flash": RESULTS / "optimal-gemini" / "records.jsonl", +} +RED = (220, 0, 0) +_WHITE = (255, 255, 255) +_NUM_PUNCT = set(".,:/%-\u2013$") + + +def number_mask(text: str) -> list[bool]: + """True for digits and number/date punctuation directly adjacent to a digit.""" + mask = [False] * len(text) + for i, ch in enumerate(text): + if ch.isdigit(): + mask[i] = True + elif ch in _NUM_PUNCT: + if (i > 0 and text[i - 1].isdigit()) or (i + 1 < len(text) and text[i + 1].isdigit()): + mask[i] = True + return mask + + +def render_hard(text: str, cfg, cache: Path, size: int, variant: str, hard: str) -> Image.Image: + """Copy of bdf.render() restricted to white-bg variants (sent/bw), with a + digit-hardening pass: `numbold` double-strikes masked glyphs, `numred` + recolors them pure red.""" + glyphs, font_ascent = parse_bdf(ensure_font(cfg, cache)) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, cap = capacity(cfg, size) + text = text[:cap] + sent_idx = _sentence_indices(text) if variant == "sent" else None + num_mask = number_mask(text) + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for row in range(rows): + bg, fg_default = _row_palette(variant, row) + y0 = row * cfg.pitch + for y in range(y0, min(y0 + cfg.pitch, size)): + for x in range(size): + px[x, y] = bg + for col in range(cols): + i = row * cols + col + if i >= len(text): + break + glyph = glyphs.get(ord(text[i])) + if glyph is None: + continue + fg = _DARK[sent_idx[i] % 6] if sent_idx is not None else fg_default + hardened = num_mask[i] + if hardened and hard == "numred": + fg = RED + bold = hardened and hard == "numbold" + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = col * cfg.adv + xoff + b + if 0 <= x < size: + px[x, y] = fg + if bold and x + 1 < size: + px[x + 1, y] = fg + return img + + +def parse_cond(cond: str) -> tuple[str, str]: + """img-6x10-- -> (base, hard).""" + parts = cond.split("-") + return parts[2], parts[3] + + +def chunk_png(chunk_text: str, size: int, base: str, hard: str) -> Path: + png = CACHE / f"{EXP}-img-6x10-{base}-{hard}-{sha8(chunk_text, str(size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render_hard(chunk_text, FONT, CACHE, size, base, hard).save(tmp) + tmp.replace(png) # atomic; cache dir is shared across agents + return png + + +def run_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + args, flow = ctx["args"], ctx["flow"] + questions = squad.sample_chunk_questions(ctx["paras"], ctx["offsets"], start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + base, hard = parse_cond(cond) + png = chunk_png(chunk_text, args.size, base, hard) + cols, rows, _ = capacity(FONT, args.size) + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("qa-image.md").format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, f"{EXP}-qa", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(ctx["keys"], model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +# --- analysis --- + + +def is_numeric_gold(golds: list[str]) -> bool: + return any(any(c.isdigit() for c in g) for g in golds) + + +def f1_stats(records: list[dict]) -> dict: + n = len(records) + if n == 0: + return {"n": 0, "em": 0.0, "f1": 0.0, "f1_se": 0.0} + f1s = [r["f1"] for r in records] + mean = sum(f1s) / n + se = (sum((x - mean) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + return {"n": n, "em": sum(r["em"] for r in records) / n, "f1": mean, "f1_se": se} + + +def load_baseline(model: str, lengths: list[int]) -> list[dict]: + out = [] + with BASELINE_RECORDS[model].open() as fh: + for line in fh: + r = json.loads(line) + if r["cond"] == BASELINE_COND and r["length"] in lengths: + out.append(r) + return out + + +def numeric_subset_cells(records: list[dict], models: list[str], lengths: list[int], conditions: list[str]) -> list[dict]: + """Per (model, length): baseline vs each condition, restricted to numeric-gold + questions present in BOTH runs (matched by question text).""" + cells = [] + for model in models: + base = load_baseline(model, lengths) + for length in lengths: + base_num = {r["q"]: r for r in base if r["length"] == length and is_numeric_gold(r["golds"])} + for cond in conditions: + mine = [ + r for r in records + if r["model"] == model and r["length"] == length and r["cond"] == cond + and is_numeric_gold(r["golds"]) and r["q"] in base_num + ] + if not mine: + continue + base_match = [base_num[r["q"]] for r in mine] + cells.append( + { + "model": model, + "length": length, + "condition": cond, + **f1_stats(mine), + "baseline_f1": f1_stats(base_match)["f1"], + "baseline_se": f1_stats(base_match)["f1_se"], + } + ) + return cells + + +def save_sample(records_ctx_flow: str, size: int, base: str, hard: str, out_dir: Path) -> Path: + """Crop a digit-dense region from the first chunk's PNG, 4x nearest upscale.""" + cols, rows, cap = capacity(FONT, size) + text = records_ctx_flow[:cap] + mask = number_mask(text) + # densest 60-char window + best_i, best_n = 0, -1 + win = 60 + run = sum(mask[:win]) + for i in range(len(text) - win): + if run > best_n: + best_n, best_i = run, i + run += mask[i + win] - mask[i] + row, col = best_i // cols, best_i % cols + png = chunk_png(records_ctx_flow[:cap], size, base, hard) + img = Image.open(png) + x0 = max(0, min(col, cols - win) * FONT.adv) + y0 = max(0, (row - 1) * FONT.pitch) + crop = img.crop((x0, y0, min(x0 + win * FONT.adv, size), min(y0 + 4 * FONT.pitch, size))) + crop = crop.resize((crop.width * 4, crop.height * 4), Image.NEAREST) + out = out_dir / f"sample-{base}-{hard}.png" + crop.save(out) + return out + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(RUN_MODELS)) + ap.add_argument("--lengths", default="50,150") + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--report", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / f"{EXP}-{SLUG}" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + + keys = {} + if not args.report: + keys["openai"] = load_env_key("OPENAI_API_KEY", args.env) + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + budget = capacity(FONT, args.size)[2] + tasks = [] + flows: dict[int, str] = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + flows[length] = flow + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for cond in conditions: + for start in range(0, len(flow), budget): + tasks.append((model, cond, start, min(start + budget, len(flow)), ctx)) + print(f"{EXP}: {len(models)} models x {len(lengths)} lengths x {len(conditions)} conditions = {len(tasks)} chunk tasks") + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_chunk, m, c, s, e, ctx) for m, c, s, e, ctx in tasks] + for k, fut in enumerate(futures): + records.extend(fut.result()) + print(f" {k + 1}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + # overall cells (same shape as final.py matrix) + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if sub: + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + num_cells = numeric_subset_cells(records, models, lengths, conditions) + + # baseline numeric subset stats (per model/length, from the SAME matched questions) + base_overall = {} + for model in models: + base = load_baseline(model, lengths) + for length in lengths: + sub = [r for r in base if r["length"] == length] + base_overall[(model, length)] = f1_stats(sub) + + (out_dir / "summary.json").write_text( + json.dumps( + { + "args": vars(args), + "cells": cells, + "numeric_subset": num_cells, + "baseline_overall": {f"{m}|{l}": v for (m, l), v in base_overall.items()}, + }, + indent=1, + ) + ) + + samples = [] + for cond in conditions: + base, hard = parse_cond(cond) + samples.append(str(save_sample(flows[lengths[0]], args.size, base, hard, out_dir))) + + print("\n== overall ==") + print(f"{'model':<24}{'len':>5}{'condition':<28}{'n':>4}{'EM':>7}{'F1':>7}{'se':>7}{'cost$':>8}{'base F1':>9}{'d':>7}") + for c in cells: + b = base_overall[(c["model"], c["length"])] + print( + f"{c['model']:<24}{c['length']:>5}{c['condition']:<28}{c['n']:>4}{c['em']:>7.3f}{c['f1']:>7.3f}" + f"{c['f1_se']:>7.3f}{c['cost_usd']:>8.3f}{b['f1']:>9.3f}{c['f1'] - b['f1']:>+7.3f}" + ) + print("\n== numeric-gold subset (matched questions vs img-6x10-sent baseline) ==") + print(f"{'model':<24}{'len':>5}{'condition':<28}{'n':>4}{'F1':>7}{'se':>7}{'base F1':>9}{'base se':>8}{'d':>7}") + for c in num_cells: + print( + f"{c['model']:<24}{c['length']:>5}{c['condition']:<28}{c['n']:>4}{c['f1']:>7.3f}{c['f1_se']:>7.3f}" + f"{c['baseline_f1']:>9.3f}{c['baseline_se']:>8.3f}{c['f1'] - c['baseline_f1']:>+7.3f}" + ) + print(f"\nresults -> {out_dir}/ samples: {', '.join(samples)}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp04_layout.py b/packages/snapcompact/research/exp04_layout.py new file mode 100644 index 000000000..0d8a1d861 --- /dev/null +++ b/packages/snapcompact/research/exp04_layout.py @@ -0,0 +1,374 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp04: document-style (in-distribution) layout vs row-major grid. + +Hypothesis: VLMs are pretrained on documents, not 261-col row-major char +grids. A two-column newspaper page (word-wrap, paragraph breaks, headings) +costs capacity (gutter, blank lines, ragged right) but may read better per +token. Conditions: img-6x10-doc (near-black, document-plain) and +img-6x10-doc-sent (sentence-hue glyphs). Same 6x10 font/page size as the +img-6x10-sent baseline; chunking is paragraph-aligned page packing. +""" + +import argparse +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DARK, capacity, parse_bdf, ensure_font # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 +from final import cached # noqa: E402 + +MODELS = {"gpt-5.5": (2.0, 16.0), "google/gemini-3.5-flash": (0.6, 4.0)} +LENGTHS = (50, 150) +CONDITIONS = ("img-6x10-doc", "img-6x10-doc-sent") +FONT = FONTS["6x10"] +GUTTER = 3 # char cells between columns +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) +_INK = (24, 24, 24) # near-black body text, like a printed page + + +# --- document layout ------------------------------------------------------- + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: # pathological; never hit on SQuAD prose + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Article title changes become headings (heading is repeated at the top of + a page even when the article continues from the previous page, since each + page is read in isolation). Paragraphs are separated by one blank line. + """ + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +# --- renderer (glyph loop copied from bdf.render, two-column layout) ------- + + +def _sentence_colors(lines: list[dict]) -> list[list[tuple[int, int, int]]]: + """Per-line per-char glyph color cycling hue per sentence across the page.""" + joined = "\n".join(ln["text"] for ln in lines) + idx, out_idx = 0, [] + for i, ch in enumerate(joined): + out_idx.append(idx) + if ch in ".!?" and i + 1 < len(joined) and joined[i + 1] in " \n": + idx += 1 + colors, pos = [], 0 + for ln in lines: + n = len(ln["text"]) + colors.append([_DARK[out_idx[pos + k] % 6] for k in range(n)]) + pos += n + 1 # the joining newline + return colors + + +def render_doc(lines: list[dict], size: int, variant: str, cache: Path) -> Image.Image: + """Two-column page: left column rows top-to-bottom, then right column.""" + glyphs, font_ascent = parse_bdf(ensure_font(FONT, cache)) + ascent = FONT.ascent if FONT.ascent is not None else font_ascent + cols, rows, _ = capacity(FONT, size) + col_w = (cols - GUTTER) // 2 + sent_colors = _sentence_colors(lines) if variant == "doc-sent" else None + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * FONT.adv + y0 = row * FONT.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + if ln["kind"] == "heading": + fg = _BLACK + elif sent_colors is not None: + fg = sent_colors[li][ci] + else: + fg = _INK + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * FONT.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runner ----------------------------------------------------------------- + + +def run_page(model: str, cond: str, page: tuple[int, int], ctx: dict) -> list[dict]: + args, paras, offsets, keys = ctx["args"], ctx["paras"], ctx["offsets"], ctx["keys"] + i, j = page + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + variant = cond.removeprefix("img-6x10-") + lines = ctx["lines"][page] + page_key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"exp04-{variant}-{page_key}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render_doc(lines, args.size, variant, CACHE).save(tmp) + tmp.replace(png) + cols, rows, _ = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("exp04-qa-image.md").format(col_w=col_w, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, "exp04-qa", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true", help="render pages + capacity stats, no API") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp04-layout" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + + keys = {} + if not args.render_only: + keys["openai"] = load_env_key("OPENAI_API_KEY", args.env) + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + cols, rows, grid_cap = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + max_lines = 2 * rows + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + capacity_stats = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + pages = pack_pages(paras, col_w, max_lines) + page_lines = {pg: layout_page(paras[pg[0] : pg[1]], col_w) for pg in pages} + page_chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + capacity_stats[length] = { + "pages": len(pages), + "chars_per_page": page_chars, + "mean_chars_page": round(sum(page_chars) / len(pages)), + "grid_chars_page": grid_cap, + "corpus_chars": len(flow), + "grid_pages": -(-len(flow) // grid_cap), + } + ctx = { + "args": args, "paras": paras, "offsets": offsets, "keys": keys, + "length": length, "lines": page_lines, + } + for model in models: + for cond in conditions: + for pg in pages: + tasks.append((model, cond, pg, ctx)) + + print(f"layout: {cols} cols -> 2 x {col_w} + gutter {GUTTER}; {max_lines} line slots/page") + for length, st in capacity_stats.items(): + print( + f" len {length}: {st['pages']} doc pages (mean {st['mean_chars_page']} chars/page; " + f"grid {st['grid_chars_page']} chars/page -> {st['grid_pages']} pages)" + ) + if args.render_only: + for length in lengths: + paras = all_paras[:length] + pages = pack_pages(paras, col_w, max_lines) + for cond in conditions: + variant = cond.removeprefix("img-6x10-") + i, j = pages[0] + lines = layout_page(paras[i:j], col_w) + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"exp04-{variant}-{key}.png" + tmp = png.with_suffix(".tmp.png") + render_doc(lines, args.size, variant, CACHE).save(tmp) + tmp.replace(png) + print(f" sample: {png}") + return + + print(f"grid: {len(tasks)} page tasks") + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_page, m, c, pg, ctx) for m, c, pg, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} pages", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text( + json.dumps({"args": vars(args), "capacity": capacity_stats, "cells": cells}, indent=1) + ) + import csv + + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"{c['model']:<26} len {c['length']:<4} {c['condition']:<20} " + f"n={c['n']:<4} EM {c['em']:.3f} F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f}" + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp05_anchors.py b/packages/snapcompact/research/exp05_anchors.py new file mode 100644 index 000000000..2cea6a75f --- /dev/null +++ b/packages/snapcompact/research/exp05_anchors.py @@ -0,0 +1,344 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp05: margin row-number ruler ("anchors") on img-6x10-sent. + +Hypothesis: a hex-editor-style row ruler makes the image addressable, so the +model can localize answers instead of transcribing the whole bitmap — cutting +the reasoning/output "read tax" and possibly improving F1. + +Conditions (both chunked by the ruler content capacity, 39936 chars, so the +two conditions share identical question sets): + img-6x10-sent-ruler 5-col blue row ruler + anti-transcription prompt + asking for `answer | row≈N` + img-6x10-sent-noruler baseline render (261 cols) + anti-transcription + prompt only (separates ruler effect from prompt effect) + +Diagnostic: claimed row vs true row of the gold answer (char offset // cols). + +Usage: uv run exp05_anchors.py [--report] [--lengths 50,150] [--models ...] +""" + +import argparse +import csv +import json +import re +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _BLACK, _DARK, _WHITE, FontCfg, _row_palette, _sentence_indices, ensure_font, parse_bdf, render # noqa: E402 +from final import MODELS, cached # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +FONT = FONTS["6x10"] +SIZE = 1568 +MARGIN_COLS = 5 # 4 digit cells + 1 gap cell +RULER_STEP = 5 +RULER_FG = (90, 120, 215) # medium blue: legible, never competes with content +COND_RULER = "img-6x10-sent-ruler" +COND_CTL = "img-6x10-sent-noruler" + +# ctx.md baseline (img-6x10-sent, seed 42, qpc 30): (f1, se, cost) +BASELINE = { + ("gpt-5.5", 50): (0.850, 0.051, 0.068), + ("gpt-5.5", 150): (0.822, 0.029, 0.245), + ("google/gemini-3.5-flash", 50): (0.984, 0.012, 0.018), + ("google/gemini-3.5-flash", 150): (0.805, 0.035, 0.097), +} + + +def ruler_capacity(cfg: FontCfg, size: int = SIZE) -> tuple[int, int, int]: + """(content_cols, rows, chars) once MARGIN_COLS are reserved for the ruler.""" + cols, rows = size // cfg.adv - MARGIN_COLS, size // cfg.pitch + return cols, rows, cols * rows + + +def render_ruler(text: str, cfg: FontCfg, cache: Path, size: int = SIZE, variant: str = "sent") -> Image.Image: + """bdf.render() with a left row-number ruler every RULER_STEP rows. + + Content glyphs are shifted right by MARGIN_COLS cells; row indices are + drawn 0-based, right-aligned in the first MARGIN_COLS-1 cells, in RULER_FG. + """ + glyphs, font_ascent = parse_bdf(ensure_font(cfg, cache)) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, cap = ruler_capacity(cfg, size) + text = text[:cap] + sent_idx = _sentence_indices(text) if variant in ("sent", "dark-sent", "sent-dim") else None + sent_palette = _DARK + img = Image.new("RGB", (size, size), _BLACK if variant in ("dark", "dark-sent") else _WHITE) + px = img.load() + + def draw_glyph(ch: str, cell_col: int, y0: int, fg: tuple[int, int, int]) -> None: + glyph = glyphs.get(ord(ch)) + if glyph is None: + return + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = cell_col * cfg.adv + xoff + b + if 0 <= x < size: + px[x, y] = fg + + for row in range(rows): + bg, fg = _row_palette(variant, row) + y0 = row * cfg.pitch + for y in range(y0, min(y0 + cfg.pitch, size)): + for x in range(size): + px[x, y] = bg + if row % RULER_STEP == 0: + label = str(row) + for j, ch in enumerate(label): + draw_glyph(ch, MARGIN_COLS - 1 - len(label) + j, y0, RULER_FG) + for col in range(cols): + i = row * cols + col + if i >= len(text): + break + if sent_idx is not None: + fg = sent_palette[sent_idx[i] % 6] + draw_glyph(text[i], MARGIN_COLS + col, y0, fg) + return img + + +_ROW_CLAIM = re.compile(r"\|\s*rows?\s*[≈~=:]*\s*(\d+)", re.IGNORECASE) + + +def parse_answers_rows(text: str, n: int) -> tuple[list[str], list[int | None]]: + """Numbered answers, with optional `| row≈N` suffixes stripped into rows.""" + answers, rows = [""] * n, [None] * n + for line in text.splitlines(): + m = re.match(r"\s*(\d+)[.):]\s*(.*\S)?\s*$", line) + if not (m and m.group(2)): + continue + idx = int(m.group(1)) - 1 + if not (0 <= idx < n) or answers[idx]: + continue + body = m.group(2).strip() + rm = _ROW_CLAIM.search(body) + if rm: + rows[idx] = int(rm.group(1)) + body = body[: rm.start()].strip() + answers[idx] = body + return answers, rows + + +def gold_row(chunk_text: str, q: dict, chunk_len: int, content_cols: int) -> int | None: + """True row of the gold answer in the rendered grid (0-based).""" + approx = max(0, int(q["pos_rel"] * chunk_len) - 10) + for g in q["golds"]: + i = chunk_text.find(g, approx) + if i < 0: + i = chunk_text.find(g) + if i < 0: + i = chunk_text.lower().find(g.lower()) + if i >= 0: + return i // content_cols + return None + + +def _ensure_png(png: Path, make) -> None: + """Render-once with atomic publish; unique tmp avoids cross-thread races.""" + if png.exists() and png.stat().st_size > 0: + return + tmp = png.with_suffix(f".{uuid.uuid4().hex}.tmp.png") + make().save(tmp) + tmp.replace(png) + + +def run_cell_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + + if cond == COND_RULER: + cols, rows, _ = ruler_capacity(FONT, args.size) + png = CACHE / f"exp05-ruler-{sha8(chunk_text, str(args.size))}.png" + _ensure_png(png, lambda: render_ruler(chunk_text, FONT, CACHE, args.size, "sent")) + last_label = (rows - 1) // RULER_STEP * RULER_STEP + prompt = load_prompt("exp05-qa-image.md").format(cols=cols, rows=rows, last_label=last_label) + else: # control: baseline render, anti-transcription prompt only + cols, rows = args.size // FONT.adv, args.size // FONT.pitch + png = CACHE / f"exp05-ctl-{sha8(chunk_text, str(args.size))}.png" + _ensure_png(png, lambda: render(chunk_text, FONT, CACHE, args.size, "sent")) + prompt = load_prompt("exp05-qa-image-ctl.md").format(cols=cols, rows=rows) + + messages = [ + { + "role": "user", + "content": [{"text": prompt}, {"image_path": png}, {"text": q_block}], + } + ] + qa = cached( + model, f"exp05-qa-{cond}", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers, claimed_rows = parse_answers_rows(qa["text"], len(questions)) + records = [] + for q, a, crow in zip(questions, answers, claimed_rows): + trow = gold_row(chunk_text, q, end - start, cols) if cond == COND_RULER else None + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + "claimed_row": crow, + "true_row": trow, + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + loc = [(r["claimed_row"], r["true_row"]) for r in records if r["claimed_row"] is not None and r["true_row"] is not None] + row_stats = {} + if loc: + errs = [abs(c - t) for c, t in loc] + row_stats = { + "row_n": len(loc), + "row_claimed_frac": round(sum(r["claimed_row"] is not None for r in records) / n, 3), + "row_mae": round(sum(errs) / len(errs), 2), + "row_within2": round(sum(e <= 2 for e in errs) / len(errs), 3), + "row_within5": round(sum(e <= 5 for e in errs) / len(errs), 3), + } + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + **row_stats, + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default="gpt-5.5,google/gemini-3.5-flash") + ap.add_argument("--lengths", default="50,150") + ap.add_argument("--conditions", default=f"{COND_RULER},{COND_CTL}") + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=SIZE) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--report", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp05-anchors" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + + keys = {} + if not args.report: + keys["openai"] = load_env_key("OPENAI_API_KEY", args.env) + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + budget = ruler_capacity(FONT, args.size)[2] # both conds: identical question sets + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for cond in conditions: + for start in range(0, len(flow), budget): + tasks.append((model, cond, start, min(start + budget, len(flow)), ctx)) + print(f"grid: {len(tasks)} chunk tasks, chunk budget {budget} chars") + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell_chunk, m, c, s, e, ctx) for m, c, s, e, ctx in tasks] + for i, fut in enumerate(futures): + records.extend(fut.result()) + print(f" {i + 1}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cell = {"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])} + base = BASELINE.get((model, length)) + if base: + cell["base_f1"] = base[0] + cell["d_f1"] = round(cell["f1"] - base[0], 3) + cell["base_cost"] = base[2] + cell["d_cost"] = round(cell["cost_usd"] - base[2], 4) + cells.append(cell) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + fieldnames = sorted({k for c in cells for k in c}, key=lambda k: (k not in ("model", "length", "condition"), k)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"{c['model']:>26} L{c['length']:<4}{c['condition']:<24} n={c['n']:<4} em={c['em']:.3f} f1={c['f1']:.3f}±{c['f1_se']:.3f}" + f" cost=${c['cost_usd']:.3f} out={c['tok_out']} reas={c['tok_reasoning']}" + + (f" rowMAE={c['row_mae']} w5={c['row_within5']}" if "row_mae" in c else "") + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp06_rolecolor.py b/packages/snapcompact/research/exp06_rolecolor.py new file mode 100644 index 000000000..e1c81be2c --- /dev/null +++ b/packages/snapcompact/research/exp06_rolecolor.py @@ -0,0 +1,417 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp06 rolecolor: hue = message role, as zero-char metadata in optical compaction. + +Synthetic transcript: each SQuAD passage gets a role (user/assistant/tool), +round-robin within shuffled triples (seed 42). Three conditions, same passages, +same questions: + + img-6x10-role plain text, glyph hue per role (blue=user, green=assistant, + red=tool) -- metadata at zero character cost + img-6x10-tagbw black-on-white, inline "[user] "/"[asst] "/"[tool] " tag + before each passage -- text-equivalent control, ~7 chars/passage + img-6x10-nometa baseline sent variant, no role metadata -- provenance floor + (content F1 for this condition == img-6x10-sent baseline table) + +Two evals per chunk image: (a) standard SQuAD QA (content F1; not run for +nometa), (b) provenance QA: "which role's message contains the answer?", +scored as plain accuracy against the assigned role of the source passage. + +Chunking: greedy consecutive passages such that the TAGGED text fits the +6x10 capacity (40716 chars), so all three conditions share identical chunks +and question sets. Boundaries therefore shift slightly vs the baseline run +(which chunks the plain flow at exactly 40716); question sets overlap heavily +but are not char-identical to the baseline cells. + +Saturation-decay-by-recency was considered and deliberately omitted: no eval +question tests recency, and desaturating old passages risks destroying the +hue signal the provenance task measures. + +Run from the snapcompact dir: uv run exp06_rolecolor.py +""" + +import argparse +import colorsys +import csv +import json +import random +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, ensure_font, parse_bdf, render # noqa: E402 +from final import cached # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODELS = {"gpt-5.5": (2.0, 16.0), "google/gemini-3.5-flash": (0.6, 4.0)} +LENGTHS = (50, 150) +CONDITIONS = ("img-6x10-role", "img-6x10-tagbw", "img-6x10-nometa") +FONT = FONTS["6x10"] + +ROLES = ("user", "assistant", "tool") +TAGS = {"user": "user", "assistant": "asst", "tool": "tool"} # all "[xxxx] " = 7 chars +ROLE_HUES = {"user": 0.62, "assistant": 0.33, "tool": 0.02} +ROLE_RGB = {r: tuple(int(c * 255) for c in colorsys.hls_to_rgb(h, 0.27, 0.90)) for r, h in ROLE_HUES.items()} +_WHITE = (255, 255, 255) + +ENCODING = { + "img-6x10-role": ( + "Glyph color encodes the author role: dark blue = user, dark green = assistant, dark red = tool. " + "A message boundary is where the glyph color changes." + ), + "img-6x10-tagbw": ( + "Each message is preceded by a bracketed role tag rendered in the text: [user], [asst], or [tool]." + ), + "img-6x10-nometa": ( + "The rendering does NOT visually indicate roles; glyph colors only cycle per sentence and carry " + "no role information. Use your best guess." + ), +} +QA_PROMPT = {"img-6x10-role": "exp06-qa-image.md", "img-6x10-tagbw": "exp06-qa-image-tag.md"} + + +def assign_roles(n: int, seed: int) -> list[str]: + """Round-robin role assignment: each consecutive triple of passages contains + every role exactly once, triple-internal order shuffled deterministically.""" + rng = random.Random(seed) + roles: list[str] = [] + while len(roles) < n: + triple = list(ROLES) + rng.shuffle(triple) + roles.extend(triple) + return roles[:n] + + +def build_chunks(paras: list[dict], budget: int) -> list[tuple[int, int]]: + """Greedy consecutive passage ranges [a, b) whose TAGGED rendering fits budget.""" + chunks, cur, cur_len = [], 0, 0 + for i, p in enumerate(paras): + add = 7 + len(p["ctx"]) + 1 # "[xxxx] " + ctx + " " + if cur_len + add > budget and i > cur: + chunks.append((cur, i)) + cur, cur_len = i, 0 + cur_len += add + chunks.append((cur, len(paras))) + return chunks + + +def sample_questions(paras: list[dict], offsets: list[int], start: int, end: int, n: int, seed: int) -> list[dict]: + """squad.sample_chunk_questions with the source passage index recorded (same rng sequence).""" + rng = random.Random(seed * 1_000_003 + start) + eligible = [i for i in range(len(offsets)) if offsets[i] >= start and offsets[i] + len(paras[i]["ctx"]) <= end] + if not eligible: + return [] + n = min(n, len(eligible)) + step = len(eligible) / n + picked = [] + for k in range(n): + pi = eligible[int(k * step)] + qa = rng.choice(paras[pi]["qas"]) + picked.append( + { + "q": " ".join(qa["question"].split()), + "golds": sorted({a["text"] for a in qa["answers"]}), + "pos_rel": (offsets[pi] - start) / (end - start), + "pi": pi, + } + ) + return picked + + +def render_role(text: str, colors: list[tuple[int, int, int]], size: int) -> Image.Image: + """bdf.render() copy, simplified: white background, per-character glyph color.""" + glyphs, font_ascent = parse_bdf(ensure_font(FONT, CACHE)) + ascent = FONT.ascent if FONT.ascent is not None else font_ascent + cols, rows, cap = capacity(FONT, size) + text = text[:cap] + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for row in range(rows): + y0 = row * FONT.pitch + for col in range(cols): + i = row * cols + col + if i >= len(text): + break + glyph = glyphs.get(ord(text[i])) + if glyph is None: + continue + fg = colors[i] + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = col * FONT.adv + xoff + b + if 0 <= x < size: + px[x, y] = fg + return img + + +def chunk_carriers(paras: list[dict], roles: list[str], a: int, b: int) -> dict: + """Plain text + per-char role colors, and tagged text, for passages [a, b).""" + plain_parts, colors, tagged_parts = [], [], [] + for i in range(a, b): + seg = paras[i]["ctx"] + " " + plain_parts.append(seg) + colors.extend([ROLE_RGB[roles[i]]] * len(seg)) + tagged_parts.append(f"[{TAGS[roles[i]]}] {seg}") + return {"plain": "".join(plain_parts), "colors": colors, "tagged": "".join(tagged_parts)} + + +def atomic_png(png: Path, make) -> Path: + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + make().save(tmp) + tmp.replace(png) + return png + + +def build_image(cond: str, car: dict, size: int) -> Path: + if cond == "img-6x10-role": + png = CACHE / f"exp06-role-{sha8(car['plain'], str(size))}.png" + return atomic_png(png, lambda: render_role(car["plain"], car["colors"], size)) + if cond == "img-6x10-tagbw": + png = CACHE / f"exp06-tagbw-{sha8(car['tagged'], str(size))}.png" + return atomic_png(png, lambda: render(car["tagged"], FONT, CACHE, size, "bw")) + png = CACHE / f"exp06-nometa-{sha8(car['plain'], str(size))}.png" + return atomic_png(png, lambda: render(car["plain"], FONT, CACHE, size, "sent")) + + +def norm_role(answer: str) -> str: + a = answer.lower().strip(" \t.[]()\"'`*") + if "assist" in a or a == "asst": + return "assistant" + if "user" in a or "human" in a: + return "user" + if "tool" in a or "function" in a: + return "tool" + return a + + +def run_cell(model: str, cond: str, length: int, ci: int, chunk: dict, args, keys) -> list[dict]: + """One (model, cond, chunk): content QA (role/tagbw only) + provenance QA.""" + questions, car = chunk["questions"], chunk["car"] + if not questions: + return [] + png = build_image(cond, car, args.size) + cols, rows, _ = capacity(FONT, args.size) + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + usage_rows: list[tuple[str, dict]] = [] + answers = [""] * len(questions) + + if cond in QA_PROMPT: + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt(QA_PROMPT[cond]).format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, "exp06-qa", {"cond": cond, "messages": messages}, + lambda: dict( + zip(("text", "usage", "stop"), llm_complete(keys, model, messages, max_tokens=args.max_tokens)) + ), + args.fresh, + ) + usage_rows.append(("qa", qa["usage"])) + answers = squad.parse_numbered(qa["text"], len(questions)) + + prov_messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("exp06-prov-image.md").format(cols=cols, rows=rows, encoding=ENCODING[cond])}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + prov = cached( + model, "exp06-prov", {"cond": cond, "messages": prov_messages}, + lambda: dict( + zip(("text", "usage", "stop"), llm_complete(keys, model, prov_messages, max_tokens=args.max_tokens)) + ), + args.fresh, + ) + usage_rows.append(("prov", prov["usage"])) + prov_answers = squad.parse_numbered(prov["text"], len(questions)) + + records = [] + for q, a, pa in zip(questions, answers, prov_answers): + gold_role = chunk["roles"][q["pi"]] + scored = cond in QA_PROMPT + records.append( + { + "model": model, + "length": length, + "cond": cond, + "chunk": ci, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]) if scored else None, + "f1": squad.f1(a, q["golds"]) if scored else None, + "abstained": "unreadable" in a.lower() if scored else None, + "prov_answer": pa, + "prov_gold": gold_role, + "prov_correct": float(norm_role(pa) == gold_role), + } + ) + records[0]["usage"] = [{"phase": p, **u} for p, u in usage_rows] + return records + + +def phase_cost(records: list[dict], phase: str, price_in: float, price_out: float) -> tuple[dict, float]: + us = [u for r in records if "usage" in r for u in r["usage"] if u["phase"] == phase] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + tok["out"] / 1e6 * price_out + return tok, cost + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records if r["f1"] is not None] + if f1s: + mean_f1 = sum(f1s) / len(f1s) + se = (sum((x - mean_f1) ** 2 for x in f1s) / (len(f1s) * (len(f1s) - 1))) ** 0.5 if len(f1s) > 1 else 0.0 + em = sum(r["em"] for r in records if r["em"] is not None) / len(f1s) + abstained = sum(r["abstained"] for r in records if r["abstained"] is not None) + else: + mean_f1 = se = em = None + abstained = None + pacc = sum(r["prov_correct"] for r in records) / n + pse = (pacc * (1 - pacc) / n) ** 0.5 + qa_tok, qa_cost = phase_cost(records, "qa", price_in, price_out) + _, prov_cost = phase_cost(records, "prov", price_in, price_out) + return { + "n": n, + "em": em, + "f1": mean_f1, + "f1_se": se, + "abstained": abstained, + "prov_acc": round(pacc, 4), + "prov_se": round(pse, 4), + **{f"tok_{k}": v for k, v in qa_tok.items()}, + "qa_cost_usd": round(qa_cost, 4), + "prov_cost_usd": round(prov_cost, 4), + "cost_usd": round(qa_cost + prov_cost, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp06-rolecolor" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + + budget = capacity(FONT, args.size)[2] + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + roles = assign_roles(length, args.seed) + flow, offsets = squad.build_flow(paras) + for ci, (a, b) in enumerate(build_chunks(paras, budget)): + start, end = offsets[a], offsets[b - 1] + len(paras[b - 1]["ctx"]) + chunk = { + "questions": sample_questions(paras, offsets, start, end, args.qpc, args.seed), + "car": chunk_carriers(paras, roles, a, b), + "roles": roles, + } + for model in models: + for cond in conditions: + tasks.append((model, cond, length, ci, chunk)) + print(f"grid: {len(models)} models x {len(lengths)} lengths x {len(conditions)} conditions = {len(tasks)} cells") + + records: list[dict] = [] + failed = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell, m, c, ln, ci, ch, args, keys) for m, c, ln, ci, ch in tasks] + for done, (fut, t) in enumerate(zip(futures, tasks), 1): + try: + records.extend(fut.result()) + except Exception as err: # noqa: BLE001 -- partial results still get written; rerun resumes from cache + failed += 1 + print(f" FAIL {t[0]} {t[1]} len={t[2]} chunk={t[3]}: {type(err).__name__}: {err}", flush=True) + print(f" {done}/{len(tasks)} cells", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if sub: + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for length in lengths: + print(f"\n== {length} passages == (content F1 / prov acc / $total)") + hdr = f"{'condition':<18}" + "".join(f"{m:>26}" for m in models) + print(hdr + "\n" + "-" * len(hdr)) + for cond in conditions: + row = f"{cond:<18}" + for model in models: + cell = next( + (c for c in cells if c["model"] == model and c["length"] == length and c["condition"] == cond), + None, + ) + if cell: + f1 = f"{cell['f1']:.3f}" if cell["f1"] is not None else " - " + row += f"{f1:>10} {cell['prov_acc']:>6.3f} {cell['cost_usd']:>8.3f}" + else: + row += f"{'-':>26}" + print(row) + print(f"\ndataset -> {out_dir}/records.jsonl, matrix.csv, summary.json") + if failed: + raise SystemExit(f"{failed} cells failed -- rerun to resume from cache") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp07_readtax.py b/packages/snapcompact/research/exp07_readtax.py new file mode 100644 index 000000000..1a69432fb --- /dev/null +++ b/packages/snapcompact/research/exp07_readtax.py @@ -0,0 +1,273 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp07 read tax: kill reasoning/output-token inflation on image conditions. + +Image-condition cost is dominated by reasoning tokens — models transcribe the +whole bitmap in CoT before answering. Carrier is the baseline winner +(img-6x10-sent, 1568px) unchanged; only the QA protocol varies: + + baseline - exact baseline protocol, re-measured for a latency + reference (baseline matrices have no wall-clock data) + effort-low - qa-image.md, effort="low" + effort-minimal - qa-image.md, effort="minimal" (gpt-5.5 only, if accepted) + no-transcribe - prompts/exp07-qa-image.md: explicit "do not transcribe, + locate the region per question, read only that region" + locate-then-answer - two turns: (1) effort=low, output only a row-band guess + per question; (2) same conversation + "read only those + bands, answer". Combined cost/latency tracked. + +Wall-clock latency is measured around every llm_complete and stored inside the +response cache entry, so resumes keep real timings. +""" + +import argparse +import csv +import json +import statistics +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from final import cached # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODELS = {"gpt-5.5": (2.0, 16.0), "google/gemini-3.5-flash": (0.6, 4.0)} +LENGTHS = (50, 150) +CONDITIONS = ( + "baseline", "effort-low", "effort-minimal", "no-transcribe", + "locate-then-answer", "locate-low", "locate-none", +) +FONT, VARIANT = "6x10", "sent" + + +def ensure_png(chunk_text: str, size: int) -> Path: + """Reuse the baseline cache PNG when present (identical render), else write exp07-prefixed.""" + h = sha8(chunk_text, str(size)) + base = CACHE / f"img-{FONT}-{VARIANT}-{h}.png" + if base.exists() and base.stat().st_size > 0: + return base + png = CACHE / f"exp07-img-{FONT}-{VARIANT}-{h}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render(chunk_text, FONTS[FONT], CACHE, size, VARIANT).save(tmp) + tmp.replace(png) + return png + + +def timed_call(keys: dict, model: str, messages: list[dict], max_tokens: int, effort: str | None) -> dict: + t0 = time.monotonic() + text, usage, stop = llm_complete(keys, model, messages, max_tokens=max_tokens, effort=effort) + return {"text": text, "usage": usage, "stop": stop, "latency_s": round(time.monotonic() - t0, 2)} + + +def probe_min_effort(keys: dict) -> str | None: + """Lowest reasoning effort gpt-5.5 accepts: try "minimal", fall back to "none".""" + for effort in ("minimal", "none"): + try: + llm_complete( + keys, "gpt-5.5", + [{"role": "user", "content": [{"text": "Reply with the single word OK."}]}], + max_tokens=64, effort=effort, + ) + return effort + except SystemExit as err: + print(f"effort={effort} rejected by gpt-5.5: {str(err)[:160]}") + return None + + +def run_cell_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + png = ensure_png(chunk_text, args.size) + cols, rows, _ = capacity(FONTS[FONT], args.size) + usage_rows: list[tuple[str, dict]] = [] + + def qa_messages(prompt_file: str) -> list[dict]: + return [ + { + "role": "user", + "content": [ + {"text": load_prompt(prompt_file).format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + + if cond.startswith("locate"): + turn2_effort = {"locate-then-answer": None, "locate-low": "low", "locate-none": "none"}[cond] + locate_msgs = qa_messages("exp07-locate.md") + locate = cached( + model, "exp07-locate", {"messages": locate_msgs, "effort": "low"}, + lambda: timed_call(keys, model, locate_msgs, args.max_tokens, "low"), + args.fresh, + ) + usage_rows.append(("locate", {**locate["usage"], "latency_s": locate.get("latency_s", 0)})) + answer_msgs = locate_msgs + [ + {"role": "assistant", "content": [{"text": locate["text"]}]}, + {"role": "user", "content": [{"text": load_prompt("exp07-answer-bands.md")}]}, + ] + qa = cached( + model, "exp07-qa", {"cond": cond, "messages": answer_msgs, "effort": turn2_effort}, + lambda: timed_call(keys, model, answer_msgs, args.max_tokens, turn2_effort), + args.fresh, + ) + else: + prompt_file = "exp07-qa-image.md" if cond == "no-transcribe" else "qa-image.md" + effort = cond.removeprefix("effort-") if cond.startswith("effort-") else None + messages = qa_messages(prompt_file) + qa = cached( + model, "exp07-qa", {"cond": cond, "messages": messages, "effort": effort}, + lambda: timed_call(keys, model, messages, args.max_tokens, effort), + args.fresh, + ) + usage_rows.append(("qa", {**qa["usage"], "latency_s": qa.get("latency_s", 0)})) + + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": p, **u} for p, u in usage_rows] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + # Per-chunk latency = sum over phases (locate + qa for the two-turn protocol). + chunk_lat = [sum(u.get("latency_s", 0) for u in r["usage"]) for r in records if "usage" in r] + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "latency_p50_s": round(statistics.median(chunk_lat), 1) if chunk_lat else 0.0, + "latency_max_s": round(max(chunk_lat), 1) if chunk_lat else 0.0, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + ap.add_argument("--out", default="exp07-readtax") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / args.out + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + min_effort = probe_min_effort(keys) if "effort-minimal" in conditions and "gpt-5.5" in models else None + print(f"lowest gpt-5.5 effort: {min_effort or 'unavailable -> condition skipped'}") + if "effort-minimal" in conditions: + conditions = [f"effort-{min_effort}" if c == "effort-minimal" and min_effort else c for c in conditions] + conditions = [c for c in conditions if c != "effort-minimal"] + + budget = capacity(FONTS[FONT], args.size)[2] + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for cond in conditions: + if cond in (f"effort-{min_effort}", "locate-none") and model != "gpt-5.5": + continue + if cond == "locate-none" and min_effort != "none": + continue + for start in range(0, len(flow), budget): + tasks.append((model, cond, start, min(start + budget, len(flow)), ctx)) + print(f"grid: {len(tasks)} chunk tasks") + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell_chunk, m, c, s, e, ctx) for m, c, s, e, ctx in tasks] + for i, fut in enumerate(futures, 1): + records.extend(fut.result()) + print(f" {i}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"{c['model']:>24} L{c['length']:<4}{c['condition']:<20}" + f"F1 {c['f1']:.3f}±{c['f1_se']:.3f} reas {c['tok_reasoning']:>6} out {c['tok_out']:>6}" + f" p50 {c['latency_p50_s']:>6.1f}s ${c['cost_usd']:.4f}" + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp08_foveate.py b/packages/snapcompact/research/exp08_foveate.py new file mode 100644 index 000000000..235765e79 --- /dev/null +++ b/packages/snapcompact/research/exp08_foveate.py @@ -0,0 +1,388 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp08: foveated two-tier reading — aggressive 5x8 archive + on-demand zoom. + +Turn 1: a 5x8 (313 cols x 196 rows = 61348 chars/page, 1.5x denser than the +6x10 winner) archive image + all questions; the model answers what it can and +replies `ZOOM rows A-B` where the region is too small. Turn 2: the union of +requested row bands is sliced from the chunk text (row r covers chars +[(r-1)*cols, r*cols)) and re-rendered at a comfortable 8x13 font as zoom +image(s); the conversation continues with the zoom images + the pending +questions. Answers are merged; F1/cost vs the img-6x10-sent baseline. + +Run from the snapcompact dir: uv run exp08_foveate.py +""" + +import argparse +import csv +import json +import re +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODELS = { + "gpt-5.5": (2.0, 16.0), + "google/gemini-3.5-flash": (0.6, 4.0), +} +LENGTHS = (50, 150) +CONDITIONS = ("fov-5x8-bw", "fov-5x8-sent") +ARCHIVE_FONT = "5x8" +ZOOM_FONT = "8x13" +ZOOM_SIZES = (520, 784, 1040, 1568) # smallest square that fits the band wins +# fov = conservative prompt (zoom only when too small), rows addressing, tight pad +# fov2 = eager prompt (zoom unless fully certain), rows addressing, wide pad +# fov3 = eager prompt, phrase addressing (model quotes partially-read anchor words; +# harness fuzzy-locates them in the chunk and zooms that row band) +PROTO = { + "fov": ("exp08-archive.md", 2, "rows"), + "fov2": ("exp08-archive-eager.md", 12, "rows"), + "fov3": ("exp08-archive-phrase.md", 12, "phrase"), +} + +# img-6x10-sent baseline (results/optimal-gpt55 + optimal-gemini): f1, se, cost$ +BASELINE = { + ("gpt-5.5", 50): (0.850, 0.0508, 0.068), + ("gpt-5.5", 150): (0.8218, 0.0290, 0.2452), + ("google/gemini-3.5-flash", 50): (0.9841, 0.0119, 0.0181), + ("google/gemini-3.5-flash", 150): (0.8046, 0.0349, 0.097), +} + +_ZOOM_RANGE = re.compile(r"(?i)\bzoom\b[^\d]*(\d+)\s*(?:[-\u2013\u2014]|to\b)\s*(\d+)") +_ZOOM_SINGLE = re.compile(r"(?i)\bzoom\b[^\d]*(\d+)") +_ZOOM_PHRASE = re.compile(r"(?i)\bzoom\b\s*[\"\u201c']+(.+?)[\"\u201d']*\s*$") + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def atomic_png(img, path: Path) -> None: + tmp = path.with_name(f"{path.stem}.{uuid.uuid4().hex[:8]}.tmp.png") + img.save(tmp) + tmp.replace(path) + + +def parse_zoom(answer: str) -> tuple[int, int] | None: + """`ZOOM rows A-B` (or single row) -> (A, B); None if not a zoom request.""" + m = _ZOOM_RANGE.search(answer) + if m: + a, b = int(m.group(1)), int(m.group(2)) + return (a, b) if a <= b else (b, a) + m = _ZOOM_SINGLE.search(answer) + if m: + r = int(m.group(1)) + return (r, r) + return None + + +def locate_phrase(chunk: str, phrase: str) -> tuple[int, int] | None: + """Best-effort char span of a (possibly misread) anchor phrase in the chunk.""" + lower = re.sub(r"[^a-z0-9]", " ", chunk.lower()) # length-preserving normalize + p = re.sub(r"[^a-z0-9]", " ", phrase.lower()).split() + if not p: + return None + exact = " ".join(p) + i = lower.find(exact) + if i >= 0: + return i, i + len(exact) + words = [(m.start(), m.group()) for m in re.finditer(r"\S+", lower)] + pset = set(p) + k = max(2 * len(p), 8) + best_score, best_pos = 0, None + for s in range(len(words)): + score = sum(1 for _, w in words[s : s + k] if w in pset) + if score > best_score: + best_score, best_pos = score, words[s][0] + if best_pos is not None and best_score >= max(2, (len(p) + 1) // 2): + return best_pos, best_pos + len(exact) + return None + + +def merge_bands(bands: list[tuple[int, int]], max_row: int, pad: int) -> list[tuple[int, int]]: + """Pad by `pad`, clamp to [1, max_row], merge overlapping/adjacent bands.""" + padded = sorted((max(1, a - pad), min(max_row, b + pad)) for a, b in bands) + merged: list[tuple[int, int]] = [] + for a, b in padded: + if merged and a <= merged[-1][1] + 1: + merged[-1] = (merged[-1][0], max(merged[-1][1], b)) + else: + merged.append((a, b)) + return merged + + +def zoom_renders(chunk_text: str, bands: list[tuple[int, int]], arch_cols: int) -> list[tuple[tuple[int, int], Path]]: + """Slice each band's rows from the chunk and render at ZOOM_FONT; oversized bands split.""" + zcfg = FONTS[ZOOM_FONT] + max_rows = capacity(zcfg, ZOOM_SIZES[-1])[2] // arch_cols # archive rows per zoom page + out = [] + for a, b in bands: + pieces = [(s, min(s + max_rows - 1, b)) for s in range(a, b + 1, max_rows)] + for pa, pb in pieces: + txt = chunk_text[(pa - 1) * arch_cols : pb * arch_cols] + if not txt.strip(): + continue + size = next((s for s in ZOOM_SIZES if capacity(zcfg, s)[2] >= len(txt)), ZOOM_SIZES[-1]) + png = CACHE / f"exp08-zoom-{ZOOM_FONT}-{size}-{sha8(txt)}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_png(render(txt, zcfg, CACHE, size, "bw"), png) + out.append(((pa, pb), png)) + return out + + +def run_cell_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + """One (model, condition, chunk): archive QA turn, optional zoom turn, merge, score.""" + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + proto, _, variant = cond.split("-", 2) + prompt_file, pad, mode = PROTO[proto] + cfg = FONTS[ARCHIVE_FONT] + cols, rows, _ = capacity(cfg, args.size) + + png = CACHE / f"exp08-arch-{ARCHIVE_FONT}-{variant}-{sha8(chunk_text, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_png(render(chunk_text, cfg, CACHE, args.size, variant), png) + + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt(prompt_file).format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa1 = cached( + model, "exp08-qa1", {"messages": messages, "effort": args.effort}, + lambda: dict(zip(("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort))), + args.fresh, + ) + usage_rows = [("qa1", qa1["usage"])] + answers1 = squad.parse_numbered(qa1["text"], len(questions)) + zoom_req: list[tuple[int, int] | None] = [] + anchors: list[str | None] = [] + for a in answers1: + anchor = None + if mode == "phrase": + m = _ZOOM_PHRASE.search(a) + text = m.group(1) if m else None + if text is None and re.search(r"(?i)\bzoom\b", a) and not parse_zoom(a): + text = re.sub(r"(?i)^.*?\bzoom\b[:\s]*", "", a).strip("\"'\u201c\u201d ") + if text and len(text.split()) >= 2: + anchor = text + span = locate_phrase(chunk_text, anchor) + zoom_req.append((span[0] // cols + 1, span[1] // cols + 1) if span else None) + elif re.search(r"(?i)\bzoom\b", a): + zoom_req.append(parse_zoom(a)) # rows fallback + else: + zoom_req.append(None) + else: + zoom_req.append(parse_zoom(a)) + anchors.append(anchor) + requested = [i for i, a in enumerate(answers1) if re.search(r"(?i)\bzoom\b", a)] + pending = [i for i, z in enumerate(zoom_req) if z is not None] + final = list(answers1) + for i in requested: + if zoom_req[i] is None: + final[i] = "UNREADABLE" # zoom requested but band unresolvable + + if pending: + bands = merge_bands([zoom_req[i] for i in pending], rows, pad) + zooms = zoom_renders(chunk_text, bands, cols) + z_content: list[dict] = [{"text": load_prompt("exp08-zoom.md")}] + for (a, b), zpng in zooms: + z_content.append({"text": f"Zoom of archive rows {a}-{b}:"}) + z_content.append({"image_path": zpng}) + z_content.append({"text": "\n".join(f"{i + 1}. {questions[i]['q']}" for i in pending)}) + messages2 = messages + [ + {"role": "assistant", "content": [{"text": qa1["text"]}]}, + {"role": "user", "content": z_content}, + ] + qa2 = cached( + model, "exp08-qa2", {"messages": messages2, "effort": args.effort}, + lambda: dict(zip(("text", "usage", "stop"), + llm_complete(keys, model, messages2, max_tokens=args.max_tokens, effort=args.effort))), + args.fresh, + ) + usage_rows.append(("qa2", qa2["usage"])) + answers2 = squad.parse_numbered(qa2["text"], len(questions)) + for i in pending: + final[i] = answers2[i] or "UNREADABLE" + + records = [] + for i, (q, a) in enumerate(zip(questions, final)): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "answer_turn1": answers1[i], + "zoomed": i in requested, + "zoom_band": list(zoom_req[i]) if zoom_req[i] else None, + "anchor": anchors[i], + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": p, **u} for p, u in usage_rows] + return records + + +def _phase_cost(us: list[dict], price_in: float, price_out: float) -> float: + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r")} + return (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + tok["out"] / 1e6 * price_out + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + zoomed = sum(r["zoomed"] for r in records) + zoom_chunks = sum(1 for r in records if "usage" in r and any(u["phase"] == "qa2" for u in r["usage"])) + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + "zoom_q": zoomed, + "zoom_rate": round(zoomed / n, 4), + "no_zoom_pct": round(100 * (n - zoomed) / n, 1), + "zoom_chunks": zoom_chunks, + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + "cost_zoom_usd": round(_phase_cost([u for u in us if u["phase"] == "qa2"], price_in, price_out), 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + ap.add_argument("--out", default="exp08-foveate") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / args.out + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + + budget = capacity(FONTS[ARCHIVE_FONT], args.size)[2] + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for cond in conditions: + for start in range(0, len(flow), budget): + tasks.append((model, cond, start, min(start + budget, len(flow)), ctx)) + print(f"grid: {len(models)} models x {len(lengths)} lengths x {len(conditions)} conditions = {len(tasks)} chunk tasks") + + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell_chunk, m, c, s, e, ctx) for m, c, s, e, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cell = {"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])} + base = BASELINE.get((model, length)) + if base: + cell["base_f1"] = base[0] + cell["base_cost_usd"] = base[2] + cell["d_f1"] = round(cell["f1"] - base[0], 4) + cell["d_cost_usd"] = round(cell["cost_usd"] - base[2], 4) + cells.append(cell) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"{c['model']:>26} L{c['length']:<4}{c['condition']:<14} n={c['n']:<4} f1={c['f1']:.3f}±{c['f1_se']:.3f} " + f"zoom={c['zoom_rate']:.0%} ${c['cost_usd']:.3f} (zoom ${c['cost_zoom_usd']:.3f}) " + f"vs base f1={c.get('base_f1', float('nan')):.3f} ${c.get('base_cost_usd', float('nan')):.3f}" + ) + print(f"\ndataset -> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() + + diff --git a/packages/snapcompact/research/exp09_cacheappend.py b/packages/snapcompact/research/exp09_cacheappend.py new file mode 100644 index 000000000..1d67070e0 --- /dev/null +++ b/packages/snapcompact/research/exp09_cacheappend.py @@ -0,0 +1,365 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp09: append-only optical pages vs rewrite-compaction — prompt-cache economics. + +Simulates a growing session over K sequential chunks of the SQuAD flow (length 150 +-> 3 pages of img-6x10-sent @ 1568). At each step k the context is pages 1..k and a +QA turn samples questions over ALL pages so far (seed 42). + +Regimes: + A append-optical : prefix = fixed frame + k byte-identical PNG pages (rendered + once, reused), QA message last. Prefix grows append-only -> + provider prompt cache should re-bill old pages at 0.1x. + Zero LLM calls on the write path. + B rewrite-compact: each step re-summarizes the whole history text with + agent compaction-summary.md (one fresh LLM call per step = + write-path cost), QA over the fresh summary. The summary + rewrite invalidates any prompt-cache prefix. + +Cache probe: the same multi-image prefix is sent twice back-to-back (disk cache +bypassed via distinct probe-call payloads) and cache_r is read on both calls — +do image input tokens actually get prefix-cache hits on OpenAI Responses and +OpenRouter/Gemini? + +Outputs: results/exp09-cacheappend/{records.jsonl, steps.csv, matrix.csv, summary.json} +""" + +import argparse +import csv +import hashlib +import io +import json +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from final import ACK, cached, session_frame # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, FONTS, agent_prompt, load_prompt, sha8 # noqa: E402 + +MODELS = {"gpt-5.5": (2.0, 16.0), "google/gemini-3.5-flash": (0.6, 4.0)} +FONT = "6x10" +VARIANT = "sent" +LENGTH = 150 +MAX_STEPS = 4 +PROBE_TAIL = "Reply with exactly the word OK and nothing else." + + +def call(keys: dict, model: str, messages: list[dict], system: str | None = None, max_tokens: int = 32768) -> dict: + t0 = time.monotonic() + text, usage, stop = llm_complete(keys, model, messages, system=system, max_tokens=max_tokens) + return {"text": text, "usage": usage, "stop": stop, "secs": round(time.monotonic() - t0, 2)} + + +def usd(u: dict, p_in: float, p_out: float) -> float: + return (u.get("in", 0) + 0.1 * u.get("cache_r", 0)) / 1e6 * p_in + u.get("out", 0) / 1e6 * p_out + + +def usd_nocache(u: dict, p_in: float, p_out: float) -> float: + return (u.get("in", 0) + u.get("cache_r", 0)) / 1e6 * p_in + u.get("out", 0) / 1e6 * p_out + + +def render_pages(flow: str, size: int) -> tuple[list[tuple[int, int, Path]], dict]: + """Render each page once; byte-identical files reused across steps. Returns pages + determinism info.""" + budget = capacity(FONTS[FONT], size)[2] + pages = [] + for i, start in enumerate(range(0, len(flow), budget)): + chunk = flow[start : start + budget] + png = CACHE / f"exp09-page{i + 1}-{sha8(chunk, str(size), VARIANT)}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render(chunk, FONTS[FONT], CACHE, size, VARIANT).save(tmp) + tmp.replace(png) + pages.append((start, min(start + budget, len(flow)), png)) + # Write-path determinism of regime A: render page 1 twice in-memory, compare bytes. + chunk1 = flow[:budget] + digests = [] + for _ in range(2): + buf = io.BytesIO() + render(chunk1, FONTS[FONT], CACHE, size, VARIANT).save(buf, format="PNG") + digests.append(hashlib.sha256(buf.getvalue()).hexdigest()) + det = {"render_sha256": digests, "deterministic": digests[0] == digests[1]} + return pages, det + + +def prefix_messages(k: int, pages: list, cols: int, rows: int) -> list[dict]: + """Append-only prefix: frame + pages 1..k, each ACKed. Byte-stable across steps.""" + msgs = [ + {"role": "user", "content": [{"text": load_prompt("exp09-frame.md").format(cols=cols, rows=rows)}, {"image_path": pages[0][2]}]}, + {"role": "assistant", "content": [{"text": ACK}]}, + ] + for i in range(1, k): + msgs.append({"role": "user", "content": [{"text": load_prompt("exp09-page.md").format(page=i + 1)}, {"image_path": pages[i][2]}]}) + msgs.append({"role": "assistant", "content": [{"text": ACK}]}) + return msgs + + +def step_questions(paras: list, offsets: list, end: int, qpc: int, seed: int) -> tuple[list[dict], str]: + qs = squad.sample_chunk_questions(paras, offsets, 0, end, qpc, seed) + return qs, "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(qs)) + + +def score_records(model: str, regime: str, step: int, questions: list[dict], text: str) -> list[dict]: + answers = squad.parse_numbered(text, len(questions)) + return [ + { + "model": model, + "length": LENGTH, + "cond": regime, + "step": step, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + for q, a in zip(questions, answers) + ] + + +def common_prefix_len(a: str, b: str) -> int: + n = 0 + for x, y in zip(a, b): + if x != y: + break + n += 1 + return n + + +def run_model(model: str, ctx: dict) -> dict: + args, keys, flow, paras, offsets, pages = ctx["args"], ctx["keys"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["pages"] + p_in, p_out = MODELS[model] + cols, rows, _ = capacity(FONTS[FONT], args.size) + K = len(pages) + records: list[dict] = [] + steps: list[dict] = [] + + # --- Regime A: append-optical (sequential; step k's prefix warms step k+1's cache) --- + cum_a = 0.0 + for k in range(1, K + 1): + end = pages[k - 1][1] + questions, q_block = step_questions(paras, offsets, end, args.qpc, args.seed) + msgs = prefix_messages(k, pages, cols, rows) + [ + {"role": "user", "content": [{"text": load_prompt("exp09-qa.md").format(questions=q_block)}]} + ] + qa = cached( + model, "exp09-A-qa", {"step": k, "messages": msgs}, + lambda: call(keys, model, msgs, max_tokens=args.max_tokens), args.fresh, + ) + recs = score_records(model, "append-optical", k, questions, qa["text"]) + recs[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + records += recs + cost = usd(qa["usage"], p_in, p_out) + cum_a += cost + u = qa["usage"] + steps.append( + { + "model": model, "regime": "append-optical", "step": k, "n": len(recs), + "f1": round(sum(r["f1"] for r in recs) / len(recs), 3), + "write_in": 0, "write_out": 0, "write_secs": 0.0, "write_cost": 0.0, + "qa_in": u["in"], "qa_cache_r": u["cache_r"], "qa_out": u["out"], + "qa_reasoning": u.get("reasoning", 0), "qa_secs": qa["secs"], + "step_cost": round(cost, 4), "step_cost_nocache": round(usd_nocache(u, p_in, p_out), 4), + "cum_cost": round(cum_a, 4), + } + ) + print(f" {model} A step {k}: in={u['in']} cache_r={u['cache_r']} out={u['out']} f1={steps[-1]['f1']}", flush=True) + + # --- Cache probe: identical multi-image prefix twice in a row (disk cache bypassed via call index) --- + probe_msgs = prefix_messages(K, pages, cols, rows) + [{"role": "user", "content": [{"text": PROBE_TAIL}]}] + probe = [] + for i in (1, 2, 3): + r = cached( + model, "exp09-probe", {"call": i, "messages": probe_msgs}, + lambda: call(keys, model, probe_msgs, max_tokens=args.max_tokens), args.fresh, + ) + probe.append({"call": i, **r["usage"], "secs": r["secs"]}) + print(f" {model} probe call {i}: in={r['usage']['in']} cache_r={r['usage']['cache_r']}", flush=True) + + # --- Regime B: rewrite-compact (fresh summary each step = write path) --- + cum_b = 0.0 + summaries: list[str] = [] + for k in range(1, K + 1): + end = pages[k - 1][1] + text_k = flow[:end] + questions, q_block = step_questions(paras, offsets, end, args.qpc, args.seed) + sm = cached( + model, "exp09-B-sum", {"step": k, "chunk": sha8(text_k)}, + lambda: call( + keys, model, + session_frame(text_k) + [{"role": "user", "content": [{"text": agent_prompt("compaction-summary.md")}]}], + system=agent_prompt("summarization-system.md"), max_tokens=args.max_tokens, + ), + args.fresh, + ) + summaries.append(sm["text"]) + qa_msgs = [ + {"role": "user", "content": [{"text": load_prompt("qa-text.md").format(context=sm["text"])}, {"text": q_block}]} + ] + qa = cached( + model, "exp09-B-qa", {"step": k, "summary": sm["text"], "q": q_block}, + lambda: call(keys, model, qa_msgs, max_tokens=args.max_tokens), args.fresh, + ) + recs = score_records(model, "rewrite-compact", k, questions, qa["text"]) + recs[0]["usage"] = [{"phase": "summarize", **sm["usage"]}, {"phase": "qa", **qa["usage"]}] + records += recs + w_cost = usd(sm["usage"], p_in, p_out) + q_cost = usd(qa["usage"], p_in, p_out) + cum_b += w_cost + q_cost + su, qu = sm["usage"], qa["usage"] + steps.append( + { + "model": model, "regime": "rewrite-compact", "step": k, "n": len(recs), + "f1": round(sum(r["f1"] for r in recs) / len(recs), 3), + "write_in": su["in"] + su["cache_r"], "write_out": su["out"], + "write_secs": sm["secs"], "write_cost": round(w_cost, 4), + "qa_in": qu["in"], "qa_cache_r": qu["cache_r"], "qa_out": qu["out"], + "qa_reasoning": qu.get("reasoning", 0), "qa_secs": qa["secs"], + "step_cost": round(w_cost + q_cost, 4), + "step_cost_nocache": round(usd_nocache(su, p_in, p_out) + usd_nocache(qu, p_in, p_out), 4), + "cum_cost": round(cum_b, 4), + } + ) + print(f" {model} B step {k}: write {su['in']}+{su['cache_r']}c->{su['out']} ({sm['secs']}s) f1={steps[-1]['f1']}", flush=True) + + # Write-path determinism of regime B: re-run the step-1 summarize with identical payload (fresh key). + det = cached( + model, "exp09-B-sum-det", {"step": 1, "chunk": sha8(flow[: pages[0][1]])}, + lambda: call( + keys, model, + session_frame(flow[: pages[0][1]]) + [{"role": "user", "content": [{"text": agent_prompt("compaction-summary.md")}]}], + system=agent_prompt("summarization-system.md"), max_tokens=args.max_tokens, + ), + args.fresh, + ) + b_det = { + "identical": det["text"] == summaries[0], + "common_prefix_chars": common_prefix_len(det["text"], summaries[0]), + "len_a": len(summaries[0]), "len_b": len(det["text"]), + } + # Cross-step summary prefix stability (the thing the prompt cache would need). + step_stability = [ + {"steps": f"{k}->{k + 1}", "common_prefix_chars": common_prefix_len(summaries[k - 1], summaries[k]), + "len_prev": len(summaries[k - 1]), "len_next": len(summaries[k])} + for k in range(1, K) + ] + return {"records": records, "steps": steps, "probe": probe, "b_determinism": b_det, "b_step_stability": step_stability} + + +def aggregate(records: list[dict], p_in: float, p_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 0.1 * tok["cache_r"]) / 1e6 * p_in + cost_out = tok["out"] / 1e6 * p_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--qpc", type=int, default=10) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp09-cacheappend" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + + paras = squad.load_paragraphs(CACHE)[:LENGTH] + flow, offsets = squad.build_flow(paras) + pages, a_det = render_pages(flow, args.size) + pages = pages[:MAX_STEPS] + print(f"flow {len(flow)} chars -> {len(pages)} pages (K={len(pages)} steps); render deterministic: {a_det['deterministic']}") + + ctx = {"args": args, "keys": keys, "flow": flow, "paras": paras, "offsets": offsets, "pages": pages} + with ThreadPoolExecutor(min(2, len(models))) as pool: + results = dict(zip(models, pool.map(lambda m: run_model(m, ctx), models))) + + records = [r for m in models for r in results[m]["records"]] + steps = [s for m in models for s in results[m]["steps"]] + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + with (out_dir / "steps.csv").open("w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=list(steps[0].keys())) + w.writeheader() + w.writerows(steps) + + cells = [] + for model in models: + for cond in ("append-optical", "rewrite-compact"): + sub = [r for r in records if r["model"] == model and r["cond"] == cond] + final_step = max(r["step"] for r in sub) + fin = [r for r in sub if r["step"] == final_step] + cell = {"model": model, "length": LENGTH, "condition": cond, **aggregate(sub, *MODELS[model])} + cell["final_step_f1"] = round(sum(r["f1"] for r in fin) / len(fin), 3) + cell["final_step_em"] = round(sum(r["em"] for r in fin) / len(fin), 3) + cells.append(cell) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + w.writeheader() + w.writerows(cells) + + summary = { + "args": vars(args), + "pages": [{"start": s, "end": e, "png": p.name} for s, e, p in pages], + "render_determinism": a_det, + "cells": cells, + "steps": steps, + "probe": {m: results[m]["probe"] for m in models}, + "b_determinism": {m: results[m]["b_determinism"] for m in models}, + "b_step_stability": {m: results[m]["b_step_stability"] for m in models}, + } + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1, default=str)) + + print("\n== per-step (qa_in / qa_cache_r / write_cost / step_cost / cum_cost / f1) ==") + for s in steps: + print( + f"{s['model']:<26} {s['regime']:<16} k={s['step']} in={s['qa_in']:>6} cache_r={s['qa_cache_r']:>6} " + f"write=${s['write_cost']:.4f} step=${s['step_cost']:.4f} cum=${s['cum_cost']:.4f} f1={s['f1']:.3f}" + ) + print("\n== cache probe (same multi-image prefix twice) ==") + for m in models: + for p in results[m]["probe"]: + print(f"{m:<26} call {p['call']}: in={p['in']:>6} cache_r={p['cache_r']:>6} secs={p['secs']}") + print(f"\ndataset -> {out_dir}/records.jsonl, steps.csv, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp10_profiles.py b/packages/snapcompact/research/exp10_profiles.py new file mode 100644 index 000000000..127edcd5f --- /dev/null +++ b/packages/snapcompact/research/exp10_profiles.py @@ -0,0 +1,358 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp10: per-model optical-profile calibration sweep for gpt-5.5 and gemini-3.5-flash. + +Sibling models each preferred a different (font, variant) density point +(fable: 6x12-dim, opus: 8x13-bw, kimi: 8x13-sent-dim, glm: 8x13-dark-sent), but +gpt-5.5 / gemini-3.5-flash were only ever measured at 6x10-sent (+ 5x8-bw / +6x9-sent-dim). This script runs the calibration sweep an "optical profile" +catalog entry would ship with: + + phase A (screen, length 150): sibling-winner combos + a variant probe at 8x13 + + 6x12-{dim,sent} + phase B (ladder, length 150): density ladder {8x13..5x8} at each model's best + variant from phase A + phase C (cross, length 150): each model evaluated on the *other* model's + optimal cell (profile transferability) + phase D (confirm): each model's top cell at lengths 50 and 250 + +Methodology matches the optimal-* baselines exactly (seed 42, qpc 30, size +1568, max_tokens 32768, effort None) so identical cells hit the shared +.cache/qa/ response cache for free. Cache policy: reads try the canonical +"qa" tag first (free reuse of optimal-run payloads), new responses are written +under the exp10-namespaced tag. + +Usage: uv run exp10_profiles.py # full self-driving sweep + uv run exp10_profiles.py --report # re-aggregate from cache only +""" + +import argparse +import json +import os +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp10" +OUT_DIR = RESULTS / f"{EXP}-profiles" +MODELS = { # ($/M in, $/M out); cached reads bill 0.1x input + "gpt-5.5": (2.0, 16.0), + "google/gemini-3.5-flash": (0.6, 4.0), +} +BASELINE = { # img-6x10-sent from results/optimal-{gpt55,gemini}/matrix.csv + ("gpt-5.5", 50): (0.850, 0.051, 0.068), + ("gpt-5.5", 150): (0.822, 0.029, 0.245), + ("gpt-5.5", 250): (0.822, 0.026, 0.380), + ("google/gemini-3.5-flash", 50): (0.984, 0.012, 0.018), + ("google/gemini-3.5-flash", 150): (0.805, 0.035, 0.097), + ("google/gemini-3.5-flash", 250): (0.755, 0.033, 0.147), +} + +# Phase A screening cells (length 150). Sibling winners + variant probe at +# 8x13 + the 6x12 bridge. img-6x10-sent is the baseline -- not re-run. +SCREEN = ( + "img-6x12-dim", # fable's winner + "img-8x13-bw", # opus's winner + "img-8x13-sent-dim", # kimi's winner + "img-8x13-dark-sent", # glm's winner + "img-8x13-sent", + "img-8x13-dim", + "img-6x12-sent", +) +LADDER_FONTS = ("8x13", "7x13", "6x12", "6x10", "6x9", "5x8") + + +def parse_img_condition(name: str) -> tuple[str, str]: + _, font, variant = name.split("-", 2) + return font, variant + + +def cached(model: str, payload: object, fn, fresh: bool) -> dict: + """Like final.cached(), but dual-key: read canonical "qa" tag first (free + reuse of the optimal-run cache), then our exp10 tag; write under exp10.""" + blob = json.dumps(payload, sort_keys=True, default=str) + canon = QA_CACHE / f"{sha8(model, 'qa', blob)}.json" + mine = QA_CACHE / f"{sha8(model, f'{EXP}-qa', blob)}.json" + if not fresh: + for path in (canon, mine): + if path.exists(): + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {sha8(model, f'{EXP}-qa', blob)}") + else: + tmp = mine.with_suffix(f".{os.getpid()}.tmp") + tmp.write_text(json.dumps(out)) + tmp.replace(mine) + return out + + +def render_png(chunk_text: str, font: str, variant: str, size: int) -> Path: + """Canonical final.py naming so QA payloads (which embed the path) match + the shared cache; render is deterministic, tmp-then-replace is atomic.""" + png = CACHE / f"img-{font}-{variant}-{sha8(chunk_text, str(size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") + render(chunk_text, FONTS[font], CACHE, size, variant).save(tmp) + tmp.replace(png) + return png + + +def run_cell_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + """One (model, condition, chunk): render carrier image, QA, score. + Copied from final.run_cell_chunk, image conditions only.""" + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + font, variant = parse_img_condition(cond) + png = render_png(chunk_text, font, variant, args.size) + cols, rows, _ = capacity(FONTS[font], args.size) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("qa-image.md").format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, {"messages": messages, "extra": None, "effort": None}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=None), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +class Runner: + def __init__(self, args, keys): + self.args = args + self.keys = keys + self.records: list[dict] = [] + self.done: set[tuple[str, int, str]] = set() + self.ctxs: dict[int, dict] = {} + self.all_paras = squad.load_paragraphs(CACHE) + + def ctx(self, length: int) -> dict: + if length not in self.ctxs: + paras = self.all_paras[:length] + flow, offsets = squad.build_flow(paras) + self.ctxs[length] = { + "args": self.args, "flow": flow, "paras": paras, + "offsets": offsets, "keys": self.keys, "length": length, + } + return self.ctxs[length] + + def run(self, cells: list[tuple[str, int, str]], label: str) -> None: + cells = [c for c in cells if c not in self.done] + self.done.update(cells) + tasks = [] + for model, length, cond in cells: + ctx = self.ctx(length) + flow = ctx["flow"] + budget = capacity(FONTS[parse_img_condition(cond)[0]], self.args.size)[2] + for start in range(0, len(flow), budget): + tasks.append((model, cond, start, min(start + budget, len(flow)), ctx)) + if not tasks: + return + print(f"[{label}] {len(cells)} cells -> {len(tasks)} chunk tasks") + with ThreadPoolExecutor(self.args.workers) as pool: + futures = [pool.submit(run_cell_chunk, *t) for t in tasks] + for fut in futures: + self.records.extend(fut.result()) + + def cell(self, model: str, length: int, cond: str) -> dict | None: + sub = [r for r in self.records if r["model"] == model and r["length"] == length and r["cond"] == cond] + return {"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])} if sub else None + + def cells_for(self, model: str, length: int) -> list[dict]: + conds = sorted({r["cond"] for r in self.records if r["model"] == model and r["length"] == length}) + return [c for cond in conds if (c := self.cell(model, length, cond))] + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + runner = Runner(args, keys) + models = list(MODELS) + + # -- phase A: screen at length 150 -------------------------------------- + runner.run([(m, 150, c) for m in models for c in SCREEN], "A screen") + + # -- phase B: density ladder at each model's best variant --------------- + ladder = [(m, 150, "img-6x10-sent") for m in models] # baseline; free via shared cache + best_variant = {} + for m in models: + top = max(runner.cells_for(m, 150), key=lambda c: c["f1"]) + v = parse_img_condition(top["condition"])[1] + best_variant[m] = v + ladder += [(m, 150, f"img-{f}-{v}") for f in LADDER_FONTS] + runner.run(ladder, "B ladder") + + # -- phase C: cross-profile transfer (each model on the other's optimum) - + top150 = {m: max(runner.cells_for(m, 150), key=lambda c: c["f1"]) for m in models} + cross = [(other, 150, top150[m]["condition"]) for m in models for other in models if other != m] + runner.run(cross, "C cross") + top150 = {m: max(runner.cells_for(m, 150), key=lambda c: c["f1"]) for m in models} + + # -- phase D: confirm top cell at lengths 50 and 250 --------------------- + runner.run([(m, ln, top150[m]["condition"]) for m in models for ln in (50, 250)], "D confirm") + + # -- outputs -------------------------------------------------------------- + with (OUT_DIR / "records.jsonl").open("w") as fh: + for r in runner.records: + fh.write(json.dumps(r) + "\n") + cells = [] + for m in models: + for ln in (50, 150, 250): + cells.extend(runner.cells_for(m, ln)) + cells.sort(key=lambda c: (c["model"], c["length"], -c["f1"])) + import csv + + with (OUT_DIR / "matrix.csv").open("w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + w.writeheader() + w.writerows(cells) + + profiles = {} + for m in models: + top = top150[m] + font, variant = parse_img_condition(top["condition"]) + cols, rows, chars = capacity(FONTS[font], args.size) + confirm = {ln: runner.cell(m, ln, top["condition"]) for ln in (50, 250)} + other = next(o for o in models if o != m) + transfer = runner.cell(other, 150, top["condition"]) + profiles[m] = { + "model": m, + "optical_profile": { + "font": font, + "variant": variant, + "px": args.size, + "cols": cols, + "rows": rows, + "chars_per_page": chars, + "prompt": "qa-image.md", + }, + "expected": { + "f1_at_150": round(top["f1"], 4), + "f1_se_at_150": round(top["f1_se"], 4), + "cost_usd_at_150": top["cost_usd"], + "confirm": { + str(ln): {"f1": round(c["f1"], 4), "se": round(c["f1_se"], 4), "cost_usd": c["cost_usd"]} + for ln, c in confirm.items() if c + }, + }, + "baseline_img_6x10_sent_f1_at_150": BASELINE[(m, 150)][0], + "transfer_f1_on_other_model_at_150": round(transfer["f1"], 4) if transfer else None, + } + (OUT_DIR / "profiles.json").write_text(json.dumps(profiles, indent=1)) + (OUT_DIR / "summary.json").write_text( + json.dumps({"args": vars(args), "best_variant": best_variant, "cells": cells}, indent=1) + ) + + # -- console report ------------------------------------------------------- + spend = 0.0 + for m in models: + print(f"\n== {m} (length 150 screening, sorted by F1) ==") + base_f1, base_se, base_cost = BASELINE[(m, 150)] + print(f"{'condition':<22}{'n':>5}{'EM':>7}{'F1':>7}{'se':>7}{'abst':>6}{'cost$':>8}{'dF1':>8}") + for c in sorted(runner.cells_for(m, 150), key=lambda c: -c["f1"]): + spend += c["cost_usd"] + print( + f"{c['condition']:<22}{c['n']:>5}{c['em']:>7.3f}{c['f1']:>7.3f}{c['f1_se']:>7.3f}" + f"{c['abstained']:>6}{c['cost_usd']:>8.3f}{c['f1'] - base_f1:>+8.3f}" + ) + print(f"{'img-6x10-sent [base]':<22}{'':>5}{'':>7}{base_f1:>7.3f}{base_se:>7.3f}{'':>6}{base_cost:>8.3f}{0:>+8.3f}") + for ln in (50, 250): + c = runner.cell(m, ln, top150[m]["condition"]) + if c: + spend += c["cost_usd"] + b = BASELINE[(m, ln)] + print( + f"confirm@{ln}: {c['condition']} F1={c['f1']:.3f}+-{c['f1_se']:.3f} cost=${c['cost_usd']:.3f}" + f" (baseline {b[0]:.3f}+-{b[1]:.3f} ${b[2]:.3f})" + ) + print(f"\ntotal cell cost (incl. cache-free cells): ${spend:.2f}") + print(f"-> {OUT_DIR}/records.jsonl, matrix.csv, summary.json, profiles.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp11_memhier.py b/packages/snapcompact/research/exp11_memhier.py new file mode 100644 index 000000000..a3badacc3 --- /dev/null +++ b/packages/snapcompact/research/exp11_memhier.py @@ -0,0 +1,301 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp11 memhier: memory hierarchy instead of either/or compression. + +Split the passage flow by age into thirds (StreamingLLM/H2O intuition, client-side): + L3 (oldest) -> narrative compaction summary (agent compaction-summary prompt) + L2 (middle) -> optical pages (img-6x10-sent, standard bdf render) + L1 (newest) -> verbatim text +One QA context = [L3 summary] + [L2 images] + [L1 text] + questions, framed by +prompts/exp11-qa-hier.md. `hier-appendix` additionally attaches the L3 text as a +dense 5x8-sent optical appendix (summary for gist + image for exact lookup). + +Question sets are identical to the uniform text baseline (same seed/qpc/40716-char +chunk grid), so cells are directly comparable. Each record carries a *global* +pos_rel and its tier, enabling the pos_rel-tercile breakdown. +""" + +import argparse +import csv +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from final import MODELS, aggregate, cached, session_frame # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, TEXT_CHUNK, agent_prompt, load_prompt, sha8 # noqa: E402 + +L2_FONT, L2_VAR = "6x10", "sent" +APX_FONT, APX_VAR = "5x8", "sent" +CONDITIONS = ("hier", "hier-appendix") +APPENDIX_NOTE = ( + ", plus {n_apx} dense bitmap appendix image(s) of the SAME oldest text " + "(monospace pixel font, {acols} characters per row, {arows} rows; use the appendix " + "for exact lookups the summary lacks)" +) + + +def tier_bounds(offsets: list[int], flow_len: int) -> tuple[int, int]: + """Passage-start offsets nearest to 1/3 and 2/3 of the flow (no passage straddles a tier).""" + b1 = min(offsets, key=lambda o: abs(o - flow_len / 3)) + b2 = min(offsets, key=lambda o: abs(o - 2 * flow_len / 3)) + return b1, b2 + + +def render_pages(text: str, font: str, var: str, size: int) -> list[Path]: + cap = capacity(FONTS[font], size)[2] + pages = [] + for s in range(0, len(text), cap): + seg = text[s : s + cap] + png = CACHE / f"exp11-img-{font}-{var}-{sha8(seg, str(size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render(seg, FONTS[font], CACHE, size, var).save(tmp) + tmp.replace(png) + pages.append(png) + return pages + + +def gen_summary(model: str, keys: dict, l3_text: str, max_tokens: int, fresh: bool) -> dict: + return cached( + model, "exp11-summary", {"chunk": l3_text}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete( + keys, model, + session_frame(l3_text) + + [{"role": "user", "content": [{"text": agent_prompt("compaction-summary.md")}]}], + system=agent_prompt("summarization-system.md"), + max_tokens=max_tokens, + ), + ) + ), + fresh, + ) + + +def context_blocks(cond: str, summary: str, l2_pages: list[Path], apx_pages: list[Path], l1_text: str, size: int) -> list[dict]: + cols, rows, _ = capacity(FONTS[L2_FONT], size) + apx_note = "" + if cond == "hier-appendix": + acols, arows, _ = capacity(FONTS[APX_FONT], size) + apx_note = APPENDIX_NOTE.format(n_apx=len(apx_pages), acols=acols, arows=arows) + frame = load_prompt("exp11-qa-hier.md").format( + appendix_note=apx_note, n_pages=len(l2_pages), cols=cols, rows=rows + ) + blocks: list[dict] = [{"text": frame}, {"text": f"TIER 3 — SUMMARY OF OLDEST THIRD:\n\n{summary}"}] + if cond == "hier-appendix": + for i, p in enumerate(apx_pages): + blocks.append({"text": f"TIER 3 appendix image {i + 1}/{len(apx_pages)} (same oldest text as dense bitmap):"}) + blocks.append({"image_path": p}) + for i, p in enumerate(l2_pages): + blocks.append({"text": f"TIER 2 page {i + 1}/{len(l2_pages)} (middle third as bitmap):"}) + blocks.append({"image_path": p}) + blocks.append({"text": f"TIER 1 — VERBATIM NEWEST THIRD:\n\n\n{l1_text}\n"}) + return blocks + + +def run_chunk(model: str, cond: str, start: int, end: int, cell: dict) -> list[dict]: + """One QA call: shared hierarchical context + this chunk's question batch.""" + args, keys, flow = cell["args"], cell["keys"], cell["flow"] + questions = squad.sample_chunk_questions(cell["paras"], cell["offsets"], start, end, args.qpc, args.seed) + if not questions: + return [] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": cell["blocks"][cond] + [{"text": f"QUESTIONS:\n{q_block}"}], + } + ] + qa = cached( + model, "exp11-qa", {"cond": cond, "length": cell["length"], "messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + b1, b2 = cell["bounds"] + records = [] + for q, a in zip(questions, answers): + pos_abs = start + q["pos_rel"] * (end - start) + tier = "L1" if pos_abs >= b2 else ("L2" if pos_abs >= b1 else "L3") + records.append( + { + "model": model, + "length": cell["length"], + "cond": cond, + "chunk": start, + "pos_rel": round(pos_abs / len(flow), 4), + "tier": tier, + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def tier_stats(records: list[dict]) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean = sum(f1s) / n + se = (sum((x - mean) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + return { + "n": n, + "em": round(sum(r["em"] for r in records) / n, 4), + "f1": round(mean, 4), + "f1_se": round(se, 4), + "abstained": sum(r["abstained"] for r in records), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default="gpt-5.5,google/gemini-3.5-flash") + ap.add_argument("--lengths", default="150,250") + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + ap.add_argument("--out", default="exp11-memhier") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / args.out + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + + all_paras = squad.load_paragraphs(CACHE) + cells: dict[tuple[str, int], dict] = {} + summary_usage: dict[tuple[str, int], dict] = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + b1, b2 = tier_bounds(offsets, len(flow)) + l3, l2, l1 = flow[:b1], flow[b1:b2], flow[b2:] + l2_pages = render_pages(l2, L2_FONT, L2_VAR, args.size) + apx_pages = render_pages(l3, APX_FONT, APX_VAR, args.size) if "hier-appendix" in conditions else [] + print( + f"length {length}: flow={len(flow)} chars, tiers L3={len(l3)} L2={len(l2)} L1={len(l1)}, " + f"l2_pages={len(l2_pages)} apx_pages={len(apx_pages)}" + ) + for model in models: + summ = gen_summary(model, keys, l3, args.max_tokens, args.fresh) + if summ.get("stop") == "max_tokens": + raise SystemExit(f"summary truncated for {model} length {length}; raise --max-tokens") + summary_usage[(model, length)] = summ["usage"] + print(f" summary[{model}]: {len(summ['text'])} chars") + cells[(model, length)] = { + "args": args, + "keys": keys, + "flow": flow, + "paras": paras, + "offsets": offsets, + "length": length, + "bounds": (b1, b2), + "blocks": { + cond: context_blocks(cond, summ["text"], l2_pages, apx_pages, l1, args.size) + for cond in conditions + }, + } + + tasks = [] + for (model, length), cell in cells.items(): + for cond in conditions: + for start in range(0, len(cell["flow"]), TEXT_CHUNK): + tasks.append((model, cond, start, min(start + TEXT_CHUNK, len(cell["flow"])), cell)) + print(f"grid: {len(tasks)} QA tasks") + + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_chunk, *t) for t in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} tasks", flush=True) + + # Charge the (cached, shared) summarization once per cell — each condition is a + # standalone strategy that would need its own summary. + charged: set[tuple[str, int, str]] = set() + for r in records: + key = (r["model"], r["length"], r["cond"]) + if key not in charged and "usage" in r: + r["usage"].append({"phase": "summarize", **summary_usage[(r["model"], r["length"])]}) + charged.add(key) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cell_rows, tier_rows = [], [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cell_rows.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + for tier in ("L3", "L2", "L1"): + tsub = [r for r in sub if r["tier"] == tier] + if tsub: + tier_rows.append({"model": model, "length": length, "condition": cond, "tier": tier, **tier_stats(tsub)}) + + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cell_rows, "tiers": tier_rows}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=list(cell_rows[0].keys())) + w.writeheader() + w.writerows(cell_rows) + with (out_dir / "terciles.csv").open("w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=list(tier_rows[0].keys())) + w.writeheader() + w.writerows(tier_rows) + + print("\n== cells ==") + for c in cell_rows: + print( + f"{c['model']:<24} {c['length']:>4} {c['condition']:<14} n={c['n']:<4} EM={c['em']:.3f} " + f"F1={c['f1']:.3f} ±{c['f1_se']:.3f} abst={c['abstained']:<3} ${c['cost_usd']:.3f}" + ) + print("\n== pos_rel terciles (tier = where the answer lives) ==") + for t in tier_rows: + print( + f"{t['model']:<24} {t['length']:>4} {t['condition']:<14} {t['tier']} n={t['n']:<3} " + f"EM={t['em']:.3f} F1={t['f1']:.3f} ±{t['f1_se']:.3f} abst={t['abstained']}" + ) + print(f"\nresults -> {out_dir}/records.jsonl, matrix.csv, terciles.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp12_arbitrage.py b/packages/snapcompact/research/exp12_arbitrage.py new file mode 100644 index 000000000..51f382789 --- /dev/null +++ b/packages/snapcompact/research/exp12_arbitrage.py @@ -0,0 +1,320 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp12: context-window arbitrage measurement. + +(a) chars per input token: optical pages vs plain text (mined from optimal-* runs + + targeted probes at image sizes 1568/1024/768); +(b) TPM accounting: do image tokens dodge text-token throttles? (rate-limit headers + captured around text-heavy vs image-heavy requests); +(c) risk: break-even repricing multiple at which the arbitrage dies. + +Measurement + writeup, no F1 chase. Outputs results/exp12-arbitrage/ +{measurements.json, probes.json, report.md}. +""" + +import argparse +import base64 +import json +import sys +import time +import urllib.error +import urllib.request +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from providers import load_env_key # noqa: E402 +from run import CACHE, FONTS, RESULTS, TEXT_CHUNK, sha8 # noqa: E402 + +OUT = RESULTS / "exp12-arbitrage" +SIZES = (1568, 1024, 768) +PROBE_MODELS = ("gpt-5.5", "google/gemini-3.5-flash") +PRICES = {"gpt-5.5": (2.0, 16.0), "google/gemini-3.5-flash": (0.6, 4.0)} +INSTR = "Reply with exactly: OK" + +# ---------------------------------------------------------------- part A: mine + +MINE_DIRS = ("optimal-combined", "optimal-gpt55", "optimal-gemini", "optimal-fable", + "optimal-opus", "optimal-kimi", "optimal-glm") + + +def cond_budget(cond: str) -> int | None: + """chars per chunk for a condition; None when chars/token is undefined (summaries).""" + if cond == "text": + return TEXT_CHUNK + if cond.startswith("img-"): + font = cond.split("-")[1] + return capacity(FONTS[font], 1568)[2] + return None # compact / handoff carry a generated summary, not the raw chars + + +def mine() -> tuple[list[dict], dict]: + """Per (model, cond): sum carrier chars and total qa input tokens over chunks. + + Also returns per-chunk (chars, tok) detail for the carrier estimation in derive(). + """ + flows = {} + paras = squad.load_paragraphs(CACHE) + for length in (50, 150, 250): + flows[length] = squad.build_flow(paras[:length])[0] + + seen: set[tuple] = set() + agg: dict[tuple[str, str], dict] = {} + detail: dict[str, dict[str, list]] = {} + for d in MINE_DIRS: + path = RESULTS / d / "records.jsonl" + if not path.exists(): + continue + for line in path.read_text().splitlines(): + r = json.loads(line) + if "usage" not in r: + continue + budget = cond_budget(r["cond"]) + if budget is None or r["length"] not in flows: + continue + key = (r["model"], r["cond"], r["length"], r["chunk"]) + if key in seen: + continue # combined is a merge of the per-model dirs + seen.add(key) + qa = next((u for u in r["usage"] if u["phase"] == "qa"), None) + if qa is None: + continue + tok = qa["in"] + qa["cache_r"] + qa["cache_w"] + chars = min(r["chunk"] + budget, len(flows[r["length"]])) - r["chunk"] + cell = agg.setdefault((r["model"], r["cond"]), {"chars": 0, "tok_in": 0, "chunks": 0}) + cell["chars"] += chars + cell["tok_in"] += tok + cell["chunks"] += 1 + detail.setdefault(r["model"], {}).setdefault(r["cond"], []).append((chars, tok)) + rows = [] + for (model, cond), c in sorted(agg.items()): + rows.append({ + "model": model, "cond": cond, "chunks": c["chunks"], "chars": c["chars"], + "tok_in_total": c["tok_in"], + "chars_per_tok": round(c["chars"] / c["tok_in"], 3), + }) + return rows, detail + + +# ------------------------------------------------------------- part B: probes +# Own POST so we can read rate-limit headers (providers._post discards them). + +OPENAI_URL = "https://api.openai.com/v1/responses" +OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions" +RL_PREFIXES = ("x-ratelimit", "ratelimit", "retry-after") + + +def post_h(url: str, body: dict, headers: dict, retries: int = 4) -> tuple[dict, dict]: + payload = json.dumps(body).encode() + req = urllib.request.Request(url, data=payload, headers={"content-type": "application/json", **headers}) + for attempt in range(retries + 1): + try: + with urllib.request.urlopen(req, timeout=600) as resp: + rl = {k.lower(): v for k, v in resp.headers.items() if k.lower().startswith(RL_PREFIXES)} + return json.load(resp), rl + except urllib.error.HTTPError as err: + detail = err.read().decode(errors="replace")[:300] + if err.code in (408, 429, 500, 502, 503, 529) and attempt < retries: + wait = 2.0 * 2**attempt + print(f" HTTP {err.code}, retrying in {wait:.0f}s: {detail[:120]}") + time.sleep(wait) + continue + raise SystemExit(f"API error {err.code} ({url}): {detail}") from err + raise AssertionError("unreachable") + + +def png_b64(path: Path) -> str: + return base64.b64encode(path.read_bytes()).decode() + + +def probe_call(model: str, keys: dict, blocks: list[dict]) -> tuple[dict, dict]: + """One user message of blocks -> (normalized usage, rate-limit headers).""" + if model.startswith("gpt-"): + content = [] + for b in blocks: + if "text" in b: + content.append({"type": "input_text", "text": b["text"]}) + else: + content.append({"type": "input_image", + "image_url": f"data:image/png;base64,{png_b64(b['image_path'])}", + "detail": "original"}) + body = {"model": model, "input": [{"role": "user", "content": content}], + "max_output_tokens": 512, "store": False} + out, rl = post_h(OPENAI_URL, body, {"authorization": f"Bearer {keys['openai']}"}) + u = out.get("usage", {}) + cached = (u.get("input_tokens_details") or {}).get("cached_tokens", 0) + usage = {"in": u.get("input_tokens", 0), "cached": cached, "out": u.get("output_tokens", 0)} + return usage, rl + content = [] + for b in blocks: + if "text" in b: + content.append({"type": "text", "text": b["text"]}) + else: + content.append({"type": "image_url", + "image_url": {"url": f"data:image/png;base64,{png_b64(b['image_path'])}"}}) + body = {"model": model, "messages": [{"role": "user", "content": content}], "max_tokens": 512} + out, rl = post_h(OPENROUTER_URL, body, {"authorization": f"Bearer {keys['openrouter']}"}) + u = out.get("usage", {}) + usage = {"in": u.get("prompt_tokens", 0), + "cached": (u.get("prompt_tokens_details") or {}).get("cached_tokens", 0), + "out": u.get("completion_tokens", 0)} + return usage, rl + + +def probe_pngs(flow: str) -> dict[int, Path]: + pngs = {} + for size in SIZES: + cols, rows, cap = capacity(FONTS["6x10"], size) + text = flow[:cap] + png = CACHE / f"exp12-6x10-sent-{size}-{sha8(text, str(size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render(text, FONTS["6x10"], CACHE, size, "sent").save(tmp) + tmp.replace(png) + pngs[size] = png + return pngs + + +def run_probes(keys: dict, flow: str) -> dict: + """Per model: instr-only, full-page text, and 6x10 pages at SIZES. + + Sequence is deliberate (small, text-heavy, small, image-heavy ...) so the + rate-limit header trail brackets each heavy request with a light one. + """ + pngs = probe_pngs(flow) + page_text = flow[:TEXT_CHUNK] + out: dict = {"page_chars": len(page_text), "models": {}} + for model in PROBE_MODELS: + steps = [("overhead-1", [{"text": INSTR}]), + ("text-page", [{"text": INSTR}, {"text": page_text}]), + ("overhead-2", [{"text": INSTR}])] + steps += [(f"img-{s}", [{"text": INSTR}, {"image_path": pngs[s]}]) for s in SIZES] + rows = [] + for name, blocks in steps: + usage, rl = probe_call(model, keys, blocks) + row = {"step": name, "usage": usage, "ratelimit": rl, "t": time.time()} + rows.append(row) + print(f" {model:>24} {name:<11} in={usage['in']:>6} (cached={usage['cached']}) " + f"out={usage['out']:>5} rl-remaining-tokens={rl.get('x-ratelimit-remaining-tokens', '-')}") + out["models"][model] = rows + return out + + +# --------------------------------------------------------- part C: derivation + + +def estimate_carriers(detail: dict, per_model_probed: dict) -> dict: + """Carrier-only chars/token for the non-probed models, from mined per-chunk data. + + Per-chunk total = carrier + QA overhead (prompt + question block). Overhead is + estimated as the mean (chunk_total - probe-measured carrier) over the two probed + models -- the question blocks are identical across models, tokenizers differ by + only a few %. Validation: applying the same estimate back to the probed models + reproduces their probe-measured chars/text-token within ~2%. + """ + overheads = [] + for model, d in per_model_probed.items(): + carrier = d["images"][1568]["image_tokens"] + overheads += [tok - carrier for _, tok in detail[model]["img-6x10-sent"]] + overhead = sum(overheads) / len(overheads) + page = capacity(FONTS["6x10"], 1568)[2] + est = {} + for model, conds in detail.items(): + if "img-6x10-sent" not in conds or "text" not in conds: + continue + img_rows, text_rows = conds["img-6x10-sent"], conds["text"] + img_tok = sum(t for _, t in img_rows) / len(img_rows) - overhead + text_chars = sum(c for c, _ in text_rows) + text_tok = sum(t for _, t in text_rows) - overhead * len(text_rows) + cpt_img, cpt_text = page / img_tok, text_chars / text_tok + est[model] = { + "est_image_tokens_per_page": round(img_tok), + "est_chars_per_img_tok": round(cpt_img, 3), + "est_chars_per_text_tok": round(cpt_text, 3), + "est_window_stretch": round(cpt_img / cpt_text, 3), + "probed": model in per_model_probed, + } + return {"qa_overhead_tokens_est": round(overhead, 1), "models": est} + + +def derive(mined: list[dict], detail: dict, probes: dict) -> dict: + page_chars = probes["page_chars"] + caps = {s: capacity(FONTS["6x10"], s)[2] for s in SIZES} + per_model = {} + for model, rows in probes["models"].items(): + by = {r["step"]: r["usage"] for r in rows} + overhead = min(by["overhead-1"]["in"], by["overhead-2"]["in"]) + text_tok = by["text-page"]["in"] - overhead + img = {} + for s in SIZES: + itok = by[f"img-{s}"]["in"] - overhead + img[s] = {"image_tokens": itok, "page_chars": caps[s], + "chars_per_img_tok": round(caps[s] / itok, 3), + "tok_per_megapixel": round(itok / (s * s / 1e6), 1)} + cpt_text = page_chars / text_tok + cpt_img = img[1568]["chars_per_img_tok"] + stretch = cpt_img / cpt_text + p_in = PRICES[model][0] + per_model[model] = { + "overhead_tokens": overhead, + "text_tokens_per_page": text_tok, + "chars_per_text_tok": round(cpt_text, 3), + "images": img, + "window_stretch_6x10_1568": round(stretch, 3), + "chars_in_200k_window": {"text": int(200_000 * cpt_text), "img_6x10_1568": int(200_000 * cpt_img)}, + "breakeven_img_token_multiple": round(stretch, 3), + "input_cost_per_mchar": {"text": round(p_in / cpt_text, 4), "img_6x10_1568": round(p_in / cpt_img, 4)}, + } + return {"mined": mined, "probes": probes, "derived": per_model, + "carrier_estimates": estimate_carriers(detail, per_model)} + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--fresh", action="store_true", help="re-run API probes even if probes.json exists") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + OUT.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + mined, detail = mine() + print(f"mined {len(mined)} (model, cond) cells from {', '.join(MINE_DIRS)}") + for r in mined: + print(f" {r['model']:>24} {r['cond']:<18} chunks={r['chunks']:>2} chars={r['chars']:>7} " + f"tok={r['tok_in_total']:>7} chars/tok={r['chars_per_tok']:>7.3f}") + + probes_path = OUT / "probes.json" + if probes_path.exists() and not args.fresh: + probes = json.loads(probes_path.read_text()) + print("reusing probes.json (pass --fresh to re-run)") + else: + keys = {"openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env)} + # 150 paragraphs -> flow ~90k chars, so every probe page (incl. 1568px / 40716 + # chars) is completely full; image token cost is content-independent anyway + # (verified: identical tok/megapixel at three different fill ratios). + flow = squad.build_flow(squad.load_paragraphs(CACHE)[:150])[0] + probes = run_probes(keys, flow) + tmp = probes_path.with_suffix(".tmp.json") + tmp.write_text(json.dumps(probes, indent=1)) + tmp.replace(probes_path) + + measurements = derive(mined, detail, probes) + tmp = (OUT / "measurements.json").with_suffix(".tmp.json") + tmp.write_text(json.dumps(measurements, indent=1)) + tmp.replace(OUT / "measurements.json") + print(f"\nwrote {OUT}/measurements.json") + for model, d in measurements["derived"].items(): + print(f"{model}: text {d['chars_per_text_tok']} c/t | img-1568 " + f"{d['images'][1568 if 1568 in d['images'] else '1568']['chars_per_img_tok']} c/t | " + f"stretch {d['window_stretch_6x10_1568']}x | breakeven {d['breakeven_img_token_multiple']}x") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp13_extractive.py b/packages/snapcompact/research/exp13_extractive.py new file mode 100644 index 000000000..d23db7221 --- /dev/null +++ b/packages/snapcompact/research/exp13_extractive.py @@ -0,0 +1,255 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp13: extractive compaction baseline ("copy the load-bearing sentences verbatim"). + +Fixes the strawman in the text-summary family: the existing `compact`/`handoff` +baselines ask for a *narrative* summary, which is hostile to extractive QA +(gemini abstains UNREADABLE on it, F1=0). Here the compaction prompt instead +asks for verbatim spans up to a character budget sized to match the optical +carrier (~2000 tokens ~= 8000 chars per 40716-char chunk; img-6x10 page costs +~1664 input tok on gemini / ~3396 on gpt-5.5). + +Pipeline per chunk (mirrors final.py's `compact` branch): + session_frame(chunk) + exp13-extract.md -> cached extraction (tag exp13-extract) + qa-text.md(context=extraction) + questions -> cached QA (tag exp13-qa) +Extraction usage is counted in the cell cost, like the summarize phase. + +Extra column vs final.py: gold_survival = fraction of questions whose gold +answer literally survives in the extraction (normalized string containment) — +the recall ceiling of the method, separable from QA ability. + +Run from the snapcompact dir: uv run exp13_extractive.py +""" + +import argparse +import csv +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from squad import _normalize # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, TEXT_CHUNK, load_prompt, sha8 # noqa: E402 + +MODELS = { + "gpt-5.5": (2.0, 16.0), + "google/gemini-3.5-flash": (0.6, 4.0), +} +COND = "extract" +ACK = "Noted. I have read the passages and will keep them in mind." + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def session_frame(chunk_text: str) -> list[dict]: + return [ + {"role": "user", "content": [{"text": load_prompt("session-frame.md").format(context=chunk_text)}]}, + {"role": "assistant", "content": [{"text": ACK}]}, + ] + + +def gold_survives(golds: list[str], extraction_norm: str) -> bool: + return any(_normalize(g) in extraction_norm for g in golds) + + +def run_cell_chunk(model: str, start: int, end: int, ctx: dict) -> list[dict]: + """One (model, chunk) unit: extract verbatim spans, QA over the extraction, score.""" + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + usage_rows: list[tuple[str, dict]] = [] + + extract_prompt = load_prompt("exp13-extract.md").format(budget=args.budget) + gen = cached( + model, "exp13-extract", {"chunk": chunk_text, "budget": args.budget, "effort": args.extract_effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete( + keys, model, + session_frame(chunk_text) + [{"role": "user", "content": [{"text": extract_prompt}]}], + max_tokens=args.extract_max_tokens, + effort=args.extract_effort, + ), + ) + ), + args.fresh, + ) + usage_rows.append(("extract", gen["usage"])) + extraction = gen["text"] + extraction_norm = _normalize(extraction) + + messages = [ + { + "role": "user", + "content": [{"text": load_prompt("qa-text.md").format(context=extraction)}, {"text": q_block}], + } + ] + qa = cached( + model, "exp13-qa", {"messages": messages}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens), + ) + ), + args.fresh, + ) + usage_rows.append(("qa", qa["usage"])) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": COND, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + "gold_survived": gold_survives(q["golds"], extraction_norm), + } + ) + records[0]["usage"] = [{"phase": p, **u} for p, u in usage_rows] + records[0]["extraction_chars"] = len(extraction) + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + "gold_survival": sum(r["gold_survived"] for r in records) / n, + "extraction_chars": sum(r.get("extraction_chars", 0) for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default="50,150,250") + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--budget", type=int, default=8000, help="max extraction chars per chunk") + ap.add_argument("--max-tokens", type=int, default=32768, help="QA max tokens") + ap.add_argument("--extract-max-tokens", type=int, default=16384, help="extraction max tokens (budget+slack)") + ap.add_argument( + "--extract-effort", default="low", + help="reasoning effort for the extraction call only; verbatim copying needs no deliberation " + "(default-effort gemini burns ~16k reasoning tokens verifying quotes and truncates)", + ) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--env", default="~/.env") + ap.add_argument("--out", default="exp13-extractive") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / args.out + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + for m in models: + if m not in MODELS: + raise SystemExit(f"unknown model {m}") + + keys = { + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for start in range(0, len(flow), TEXT_CHUNK): + tasks.append((model, start, min(start + TEXT_CHUNK, len(flow)), ctx)) + print(f"grid: {len(models)} models x {len(lengths)} lengths x 1 condition = {len(tasks)} chunk tasks") + + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell_chunk, m, s, e, ctx) for m, s, e, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + sub = [r for r in records if r["model"] == model and r["length"] == length] + if not sub: + continue + cells.append({"model": model, "length": length, "condition": COND, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"{c['model']:>24} len={c['length']:<4} f1={c['f1']:.3f}±{c['f1_se']:.3f} em={c['em']:.3f} " + f"survival={c['gold_survival']:.3f} abst={c['abstained']}/{c['n']} cost=${c['cost_usd']:.4f}" + ) + print(f"\ndataset -> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() + + diff --git a/packages/snapcompact/research/exp14_bestgpt.py b/packages/snapcompact/research/exp14_bestgpt.py new file mode 100644 index 000000000..9ed35b88d --- /dev/null +++ b/packages/snapcompact/research/exp14_bestgpt.py @@ -0,0 +1,475 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp14: best-of-round-1 combination for gpt-5.5. + +Combines the validated levers from round 1: + - 8x13 glyphs on a patch-aligned 8x16 cell (exp01: 8on16-sent .918@150) + - two-column document layout (exp04: +F1 / -cost / -read-tax at 6x10) + - per-model variant: gpt-5.5 prefers bw (exp10), sent is the runner-up + +Conditions (gpt-5.5 only, 1568px): + img-doc-8on16-bw doc layout, near-black ink (the combination) + img-doc-8on16-sent doc layout, sentence-hue glyphs (variant probe) + img-8on16-bw plain grid, missing round-1 cell (8on16 ran only as sent) + +Phased: screen all three at length 150, confirm the winner at 50/250, +optional effort=none probe at 50. Records merge across runs (cells keyed by +model/length/condition/effort are replaced when re-run, kept otherwise). + +Usage: + uv run exp14_bestgpt.py --render-only # capacity + sample PNGs + uv run exp14_bestgpt.py # screen @150 (default cells) + uv run exp14_bestgpt.py --cells img-doc-8on16-bw@50,img-doc-8on16-bw@250 + uv run exp14_bestgpt.py --cells img-doc-8on16-bw@50 --effort none + uv run exp14_bestgpt.py --report # re-aggregate, no API +""" + +import argparse +import csv +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DARK, FontCfg, capacity, ensure_font, parse_bdf, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp14" +OUT_DIR = RESULTS / f"{EXP}-bestgpt" +MODEL = "gpt-5.5" +PRICE_IN, PRICE_OUT = 2.0, 16.0 +FONT = FontCfg("8on16", "8x13", 8, 16) # exp01 winner: 8x13 glyphs, 16px patch-aligned pitch +GUTTER = 3 # char cells between doc columns (as exp04) +SCREEN_CELLS = "img-doc-8on16-bw@150,img-doc-8on16-sent@150,img-8on16-bw@150" +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) +_INK = (24, 24, 24) # exp04 body ink + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +# --- document layout (ported from exp04, parameterized for FONT) ------------ + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: # pathological; never hit on SQuAD prose + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Article title changes become headings (repeated at the top of a page even + when the article continues, since each page is read in isolation). + Paragraphs are separated by one blank line. + """ + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def _sentence_colors(lines: list[dict]) -> list[list[tuple[int, int, int]]]: + """Per-line per-char glyph color cycling hue per sentence across the page.""" + joined = "\n".join(ln["text"] for ln in lines) + idx, out_idx = 0, [] + for i, ch in enumerate(joined): + out_idx.append(idx) + if ch in ".!?" and i + 1 < len(joined) and joined[i + 1] in " \n": + idx += 1 + colors, pos = [], 0 + for ln in lines: + n = len(ln["text"]) + colors.append([_DARK[out_idx[pos + k] % 6] for k in range(n)]) + pos += n + 1 # the joining newline + return colors + + +def render_doc(lines: list[dict], size: int, variant: str, cache: Path) -> Image.Image: + """Two-column page: left column rows top-to-bottom, then right column.""" + glyphs, font_ascent = parse_bdf(ensure_font(FONT, cache)) + ascent = FONT.ascent if FONT.ascent is not None else font_ascent + cols, rows, _ = capacity(FONT, size) + col_w = (cols - GUTTER) // 2 + sent_colors = _sentence_colors(lines) if variant == "sent" else None + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * FONT.adv + y0 = row * FONT.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + if ln["kind"] == "heading": + fg = _BLACK + elif sent_colors is not None: + fg = sent_colors[li][ci] + else: + fg = _INK + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * FONT.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runner ----------------------------------------------------------------- + + +def atomic_save(img: Image.Image, png: Path) -> None: + tmp = png.with_suffix(".tmp.png") + img.save(tmp) + tmp.replace(png) + + +def qa_call(messages: list[dict], questions: list[dict], length: int, cond: str, + start: int, ctx: dict) -> list[dict]: + """One QA call + scoring; shared by doc and grid paths.""" + args, keys = ctx["args"], ctx["keys"] + qa = cached( + MODEL, f"{EXP}-qa", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, MODEL, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": MODEL, + "length": length, + "cond": cond, + "effort": args.effort, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def run_doc_page(cond: str, length: int, page: tuple[int, int], ctx: dict) -> list[dict]: + args, paras, offsets = ctx["args"], ctx["paras"], ctx["offsets"] + i, j = page + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + variant = cond.removeprefix("img-doc-8on16-") + lines = ctx["lines"][page] + page_key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"{EXP}-doc-{variant}-{page_key}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render_doc(lines, args.size, variant, CACHE), png) + cols, rows, _ = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("exp04-qa-image.md").format(col_w=col_w, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + return qa_call(messages, questions, length, cond, start, ctx) + + +def run_grid_chunk(cond: str, length: int, start: int, end: int, ctx: dict) -> list[dict]: + args, flow, paras, offsets = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + variant = cond.removeprefix("img-8on16-") + png = CACHE / f"{EXP}-8on16-{variant}-{sha8(chunk_text, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render(chunk_text, FONT, CACHE, args.size, variant), png) + cols, rows, _ = capacity(FONT, args.size) + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("qa-image.md").format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + return qa_call(messages, questions, length, cond, start, ctx) + + +def aggregate(records: list[dict]) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * PRICE_IN + cost_out = tok["out"] / 1e6 * PRICE_OUT + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def cell_label(cond: str, effort: str | None) -> str: + return f"{cond}+eff-{effort}" if effort else cond + + +def write_outputs(records: list[dict], capacity_stats: dict, args_dict: dict) -> list[dict]: + with (OUT_DIR / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + cell_keys = sorted({(r["length"], r["cond"], r.get("effort")) for r in records}, + key=lambda k: (k[0], k[1], k[2] or "")) + cells = [] + for length, cond, effort in cell_keys: + sub = [r for r in records if r["length"] == length and r["cond"] == cond and r.get("effort") == effort] + cells.append({"model": MODEL, "length": length, "condition": cell_label(cond, effort), **aggregate(sub)}) + (OUT_DIR / "summary.json").write_text( + json.dumps({"args": args_dict, "capacity": capacity_stats, "cells": cells}, indent=1) + ) + with (OUT_DIR / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + return cells + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--cells", default=SCREEN_CELLS, help="comma list of cond@length") + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true") + ap.add_argument("--report", action="store_true", help="re-aggregate existing records, no API") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + cols, rows, grid_cap = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + max_lines = 2 * rows + print(f"8on16 @ {args.size}px: grid {cols}x{rows} = {grid_cap} chars; " + f"doc 2 x {col_w} cols + gutter {GUTTER}, {max_lines} line slots") + + rec_path = OUT_DIR / "records.jsonl" + existing: list[dict] = [] + if rec_path.exists(): + existing = [json.loads(ln) for ln in rec_path.read_text().splitlines() if ln.strip()] + cap_path = OUT_DIR / "capacity.json" + capacity_stats: dict = json.loads(cap_path.read_text()) if cap_path.exists() else {} + + if args.report: + cells = write_outputs(existing, capacity_stats, vars(args)) + for c in cells: + print(f"len {c['length']:<4} {c['condition']:<28} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f} " + f"out={c['tok_out']} rsn={c['tok_reasoning']}") + return + + cell_specs = [] + for spec in args.cells.split(","): + spec = spec.strip() + if not spec: + continue + cond, _, ln = spec.partition("@") + cell_specs.append((cond, int(ln))) + lengths = sorted({ln for _, ln in cell_specs}) + + keys = {} + if not args.render_only: + keys["openai"] = load_env_key("OPENAI_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + pages = pack_pages(paras, col_w, max_lines) + page_lines = {pg: layout_page(paras[pg[0] : pg[1]], col_w) for pg in pages} + page_chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + capacity_stats[str(length)] = { + "doc_pages": len(pages), + "mean_chars_page": round(sum(page_chars) / len(pages)), + "min_chars_page": min(page_chars), + "max_chars_page": max(page_chars), + "grid_chars_page": grid_cap, + "corpus_chars": len(flow), + "grid_pages": -(-len(flow) // grid_cap), + } + st = capacity_stats[str(length)] + print(f" len {length}: {st['doc_pages']} doc pages (mean {st['mean_chars_page']} chars, " + f"{round(100 * st['mean_chars_page'] / grid_cap)}% of grid {grid_cap}); " + f"grid {st['grid_pages']} pages; corpus {st['corpus_chars']}") + ctx = {"args": args, "paras": paras, "flow": flow, "offsets": offsets, "keys": keys, "lines": page_lines} + for cond, ln in cell_specs: + if ln != length: + continue + if cond.startswith("img-doc-"): + for pg in pages: + tasks.append(("doc", cond, length, pg, ctx)) + else: + for start in range(0, len(flow), grid_cap): + tasks.append(("grid", cond, length, (start, min(start + grid_cap, len(flow))), ctx)) + + cap_path.write_text(json.dumps(capacity_stats, indent=1)) + + if args.render_only: + for kind, cond, length, unit, ctx in tasks: + if unit[0] != 0 and (kind == "grid" or unit != list(ctx["lines"])[0]): + continue + if kind == "doc": + variant = cond.removeprefix("img-doc-8on16-") + i, j = unit + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in ctx["paras"][i:j]]), str(args.size)) + png = CACHE / f"{EXP}-doc-{variant}-{key}.png" + atomic_save(render_doc(ctx["lines"][unit], args.size, variant, CACHE), png) + else: + variant = cond.removeprefix("img-8on16-") + chunk_text = ctx["flow"][unit[0] : unit[1]] + png = CACHE / f"{EXP}-8on16-{variant}-{sha8(chunk_text, str(args.size))}.png" + atomic_save(render(chunk_text, FONT, CACHE, args.size, variant), png) + print(f" sample: {png}") + return + + print(f"{len(tasks)} page/chunk tasks") + new_records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [] + for kind, cond, length, unit, ctx in tasks: + if kind == "doc": + futures.append(pool.submit(run_doc_page, cond, length, unit, ctx)) + else: + futures.append(pool.submit(run_grid_chunk, cond, length, unit[0], unit[1], ctx)) + for fut in futures: + new_records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)}", flush=True) + + # merge: drop existing records for cells just re-run, keep everything else + rerun = {(ln, cond, args.effort) for cond, ln in cell_specs} + kept = [r for r in existing if (r["length"], r["cond"], r.get("effort")) not in rerun] + records = kept + new_records + + cells = write_outputs(records, capacity_stats, vars(args)) + for c in cells: + print(f"len {c['length']:<4} {c['condition']:<28} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f} " + f"out={c['tok_out']} rsn={c['tok_reasoning']}") + print(f"\n-> {OUT_DIR}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp15_bestgemini.py b/packages/snapcompact/research/exp15_bestgemini.py new file mode 100644 index 000000000..1119b32e3 --- /dev/null +++ b/packages/snapcompact/research/exp15_bestgemini.py @@ -0,0 +1,478 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp15: combined best optical profile for google/gemini-3.5-flash. + +Combines the round-1 validated levers: + - 8x13 glyphs on a patch-aligned 8x16 cell (exp01 `8on16`) + - two-column document layout with headings (exp04), parameterized for font/pitch + - gemini's winning variant `sent-dim` (exp10), plus `sent` runner-up + +Phase A (screen @150): img-doc-8on16-sent-dim, img-doc-8on16-sent, and the +missing round-1 grid cell img-8on16-sent-dim. +Phase B (confirm): screening winner at lengths 50 and 250. + +Bar (findings.md best known for gemini): .984@50 / .915@150 / .909@250. +Text ceiling: .989 / .898 / .918. +""" + +import argparse +import csv +import json +import os +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DARK, _DIMMED, _stopword_mask, FontCfg, capacity, parse_bdf, ensure_font, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp15" +OUT_DIR = RESULTS / f"{EXP}-bestgemini" +MODEL = "google/gemini-3.5-flash" +PRICE = (0.6, 4.0) # $/M in, out +FONT = FontCfg("8on16", "8x13", 8, 16) # 8x13 glyphs, ViT-patch-aligned 16px pitch +GUTTER = 3 # char cells between doc columns (exp04) +SCREEN_CONDS = ("img-doc-8on16-sent-dim", "img-doc-8on16-sent", "img-8on16-sent-dim") +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) + +# Best-known cells from local findings (for the printed delta column only). +BEST_KNOWN = {50: (0.984, 0.012), 150: (0.915, 0.019), 250: (0.909, 0.016)} +TEXT_CEILING = {50: 0.989, 150: 0.898, 250: 0.918} + + +def cached(model: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, exp15-qa, payload). Truncations not cached.""" + key = sha8(model, f"{EXP}-qa", json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {key}") + else: + tmp = path.with_suffix(f".{os.getpid()}.tmp") + tmp.write_text(json.dumps(out)) + tmp.replace(path) + return out + + +# --- document layout (ported from exp04, parameterized font) ---------------- + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Headings repeat at the top of a page when an article continues, since + each page is read in isolation.""" + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def _doc_colors(lines: list[dict], variant: str) -> list[list[tuple[int, int, int]]]: + """Per-line per-char glyph colors: sentence hue cycle, optionally with the + stopword dim mask composed on top (sent-dim = exp04 sent + bdf dim).""" + joined = "\n".join(ln["text"] for ln in lines) + sidx, idx = [], 0 + for i, ch in enumerate(joined): + sidx.append(idx) + if ch in ".!?" and i + 1 < len(joined) and joined[i + 1] in " \n": + idx += 1 + dim = _stopword_mask(joined) if variant == "sent-dim" else None + colors, pos = [], 0 + for ln in lines: + n = len(ln["text"]) + colors.append( + [ + _DIMMED if dim is not None and dim[pos + k] else _DARK[sidx[pos + k] % 6] + for k in range(n) + ] + ) + pos += n + 1 # the joining newline + return colors + + +def render_doc(lines: list[dict], size: int, variant: str, cache: Path) -> Image.Image: + """Two-column page at FONT: left column top-to-bottom, then right column.""" + glyphs, font_ascent = parse_bdf(ensure_font(FONT, cache)) + ascent = FONT.ascent if FONT.ascent is not None else font_ascent + cols, rows, _ = capacity(FONT, size) + col_w = (cols - GUTTER) // 2 + colors = _doc_colors(lines, variant) + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * FONT.adv + y0 = row * FONT.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + fg = _BLACK if ln["kind"] == "heading" else colors[li][ci] + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * FONT.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runners ----------------------------------------------------------------- + + +def _qa_call(model: str, messages: list[dict], questions: list[dict], ctx: dict, cond: str, start: int) -> list[dict]: + args, keys = ctx["args"], ctx["keys"] + qa = cached( + model, + {"messages": messages, "extra": None, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def run_doc_page(model: str, cond: str, page: tuple[int, int], ctx: dict) -> list[dict]: + args, paras, offsets = ctx["args"], ctx["paras"], ctx["offsets"] + i, j = page + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + variant = cond.removeprefix("img-doc-8on16-") + lines = ctx["lines"][page] + page_key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"{EXP}-doc-8on16-{variant}-{page_key}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") + render_doc(lines, args.size, variant, CACHE).save(tmp) + tmp.replace(png) + cols, rows, _ = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("exp04-qa-image.md").format(col_w=col_w, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + return _qa_call(model, messages, questions, ctx, cond, start) + + +def run_grid_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + args, flow, paras, offsets = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + variant = cond.removeprefix("img-8on16-") + png = CACHE / f"{EXP}-grid-8on16-{variant}-{sha8(chunk_text, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") + render(chunk_text, FONT, CACHE, args.size, variant).save(tmp) + tmp.replace(png) + cols, rows, _ = capacity(FONT, args.size) + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("qa-image.md").format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + return _qa_call(model, messages, questions, ctx, cond, start) + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + "calls": len(us), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +class Runner: + def __init__(self, args, keys): + self.args = args + self.keys = keys + self.records: list[dict] = [] + self.done: set[tuple[int, str]] = set() + self.ctxs: dict[int, dict] = {} + self.all_paras = squad.load_paragraphs(CACHE) + self.capacity_stats: dict = {} + + def ctx(self, length: int) -> dict: + if length not in self.ctxs: + paras = self.all_paras[:length] + flow, offsets = squad.build_flow(paras) + cols, rows, grid_cap = capacity(FONT, self.args.size) + col_w = (cols - GUTTER) // 2 + pages = pack_pages(paras, col_w, 2 * rows) + page_lines = {pg: layout_page(paras[pg[0] : pg[1]], col_w) for pg in pages} + page_chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + self.capacity_stats[length] = { + "doc_pages": len(pages), + "doc_mean_chars_page": round(sum(page_chars) / len(pages)), + "doc_chars_page": page_chars, + "grid_chars_page": grid_cap, + "corpus_chars": len(flow), + "grid_pages": -(-len(flow) // grid_cap), + } + self.ctxs[length] = { + "args": self.args, "flow": flow, "paras": paras, "offsets": offsets, + "keys": self.keys, "length": length, "pages": pages, "lines": page_lines, + } + return self.ctxs[length] + + def run(self, cells: list[tuple[int, str]], label: str) -> None: + cells = [c for c in cells if c not in self.done] + self.done.update(cells) + tasks = [] + for length, cond in cells: + ctx = self.ctx(length) + if cond.startswith("img-doc-"): + for pg in ctx["pages"]: + tasks.append((run_doc_page, (MODEL, cond, pg, ctx))) + else: + grid_cap = capacity(FONT, self.args.size)[2] + flow = ctx["flow"] + for start in range(0, len(flow), grid_cap): + tasks.append((run_grid_chunk, (MODEL, cond, start, min(start + grid_cap, len(flow)), ctx))) + if not tasks: + return + print(f"[{label}] {len(cells)} cells -> {len(tasks)} page/chunk tasks") + with ThreadPoolExecutor(self.args.workers) as pool: + futures = [pool.submit(fn, *t) for fn, t in tasks] + for k, fut in enumerate(futures): + self.records.extend(fut.result()) + print(f" {k + 1}/{len(tasks)}", flush=True) + + def cell(self, length: int, cond: str) -> dict | None: + sub = [r for r in self.records if r["length"] == length and r["cond"] == cond] + return {"model": MODEL, "length": length, "condition": cond, **aggregate(sub, *PRICE)} if sub else None + + def all_cells(self) -> list[dict]: + keys = sorted({(r["length"], r["cond"]) for r in self.records}) + return [c for ln, cond in keys if (c := self.cell(ln, cond))] + + +def print_cells(cells: list[dict]) -> None: + for c in cells: + best, best_se = BEST_KNOWN.get(c["length"], (None, None)) + extra = "" + if best is not None: + comb = (c["f1_se"] ** 2 + best_se**2) ** 0.5 + d = c["f1"] - best + extra = f" vs best {best:.3f}: {d:+.3f} ({d / comb:+.1f}se)" + if c["f1"] >= TEXT_CEILING[c["length"]] - 1e-9: + extra += " >= TEXT CEILING" + print( + f" len {c['length']:<4} {c['condition']:<24} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} out {c['tok_out']} reas {c['tok_reasoning']} " + f"${c['cost_usd']:.3f}{extra}" + ) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true") + ap.add_argument("--screen-only", action="store_true", help="skip the 50/250 confirmation phase") + ap.add_argument("--confirm-conds", default=None, help="comma list; default = screening F1 winner") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(parents=True, exist_ok=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + keys = {} + if not args.render_only: + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + cols, rows, grid_cap = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + print(f"font 8on16: {cols} cols x {rows} rows; grid cap {grid_cap}; doc 2 x {col_w} + gutter {GUTTER}, {2 * rows} line slots") + + runner = Runner(args, keys) + + if args.render_only: + ctx = runner.ctx(150) + pg = ctx["pages"][0] + for cond in SCREEN_CONDS: + if cond.startswith("img-doc-"): + variant = cond.removeprefix("img-doc-8on16-") + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in ctx["paras"][pg[0] : pg[1]]]), str(args.size)) + png = CACHE / f"{EXP}-doc-8on16-{variant}-{key}.png" + img = render_doc(ctx["lines"][pg], args.size, variant, CACHE) + else: + variant = cond.removeprefix("img-8on16-") + chunk = ctx["flow"][:grid_cap] + png = CACHE / f"{EXP}-grid-8on16-{variant}-{sha8(chunk, str(args.size))}.png" + img = render(chunk, FONT, CACHE, args.size, variant) + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") + img.save(tmp) + tmp.replace(png) + print(f" sample: {png}") + for length, st in runner.capacity_stats.items(): + print(f" len {length}: {st['doc_pages']} doc pages, mean {st['doc_mean_chars_page']} chars/page " + f"(grid {st['grid_chars_page']} -> {st['grid_pages']} pages)") + return + + # Phase A: screen at 150 + runner.run([(150, c) for c in SCREEN_CONDS], "screen@150") + screen_cells = [c for c in runner.all_cells() if c["length"] == 150] + print_cells(screen_cells) + winner = max(screen_cells, key=lambda c: (c["f1"], -c["cost_usd"]))["condition"] + print(f"screen winner: {winner}") + + # Phase B: confirm winner at 50 and 250 + if not args.screen_only: + confirm = [w.strip() for w in args.confirm_conds.split(",")] if args.confirm_conds else [winner] + runner.run([(ln, c) for ln in (50, 250) for c in confirm], "confirm@50/250") + + cells = runner.all_cells() + with (OUT_DIR / "records.jsonl").open("w") as fh: + for r in runner.records: + fh.write(json.dumps(r) + "\n") + (OUT_DIR / "summary.json").write_text( + json.dumps( + {"args": vars(args), "model": MODEL, "capacity": runner.capacity_stats, + "screen_winner": winner, "cells": cells}, + indent=1, + ) + ) + with (OUT_DIR / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=[k for k in cells[0].keys() if k != "doc_chars_page"]) + writer.writeheader() + writer.writerows(cells) + + print_cells(cells) + print(f"\n-> {OUT_DIR}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp16_bestfable.py b/packages/snapcompact/research/exp16_bestfable.py new file mode 100644 index 000000000..8045f588c --- /dev/null +++ b/packages/snapcompact/research/exp16_bestfable.py @@ -0,0 +1,476 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp16: best optical profile for claude-fable-5. + +Round-1 levers (patch-aligned pitch 16, document layout, per-model variant) +were validated on gpt-5.5/gemini only. Fable's known winner is img-6x12-dim +(smaller glyphs than the 8x13 the others prefer, plain dim variant). Test +whether the round-1 levers transfer: + + img-8on16-dim 8x13 glyphs on an 8x16 cell (pitch-16 alignment, big glyphs) + img-6on7x14-dim 6x12 glyphs on a 7x14 cell (alignment, fable's glyph size) + doc-6x12-dim two-column newspaper layout at fable's winner font + doc-8on16-dim layout + alignment combined (run if either lever shows) + +Screen at length 150, confirm the winner at 50/250. Baselines (do not re-run): +img-6x12-dim F1 .956/.911/.923 at $0.132/$0.437/$0.724; text ceiling +.956/.904/.920 at $0.144/$0.498/$0.734. + +Usage: uv run exp16_bestfable.py --render-only # sample PNGs + capacity + uv run exp16_bestfable.py --lengths 150 # screen + uv run exp16_bestfable.py --lengths 50,150,250 --conditions ... # confirm + (re-runs hit .cache/qa/, so the final full invocation rebuilds the + combined records/matrix for free) +""" + +import argparse +import csv +import json +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DIMMED, _stopword_mask, FontCfg, capacity, ensure_font, parse_bdf, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODEL = "claude-fable-5" +PRICE = (10.0, 50.0) # $/M in, out +FONTS = { + "8on16": FontCfg("8on16", "8x13", 8, 16), # 8x13 glyphs, patch-aligned 16 px pitch + "6on7x14": FontCfg("6on7x14", "6x12", 7, 14), # 6x12 glyphs, 7x14 patch-aligned cell + "6x12": FontCfg("6x12", "6x12", 6, 12), # fable's round-0 winner font +} +CONDITIONS = ("img-8on16-dim", "img-6on7x14-dim", "doc-6x12-dim", "doc-8on16-dim") +LENGTHS = (150,) +GUTTER = 3 # char cells between doc columns +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) + +# img-6x12-dim per length: (f1, se, cost); text ceiling: (f1, se, cost). +BASE_IMG = {50: (0.9556, 0.0348, 0.132), 150: (0.9113, 0.0244, 0.437), 250: (0.9233, 0.0163, 0.724)} +BASE_TEXT = {50: (0.9556, 0.0348, 0.144), 150: (0.9043, 0.0216, 0.498), 250: (0.9197, 0.0184, 0.734)} + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def parse_condition(name: str) -> tuple[str, str, str]: + """'img-8on16-dim' -> (kind, font, variant).""" + kind, font, variant = name.split("-", 2) + return kind, font, variant + + +def atomic_save(img: Image.Image, png: Path) -> None: + tmp = png.with_suffix(f".{uuid.uuid4().hex[:8]}.tmp.png") + img.save(tmp) + tmp.replace(png) + + +# --- document layout (ported from exp04_layout.py, parameterized font) ------ + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: # pathological; never hit on SQuAD prose + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Article title changes become headings (repeated at the top of a page even + when the article continues, since each page is read in isolation). + """ + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def _dim_masks(lines: list[dict]) -> list[list[bool]]: + """Per-line stopword mask, computed over the joined page text.""" + joined = "\n".join(ln["text"] for ln in lines) + mask = _stopword_mask(joined) + out, pos = [], 0 + for ln in lines: + n = len(ln["text"]) + out.append(mask[pos : pos + n]) + pos += n + 1 # the joining newline + return out + + +def render_doc(lines: list[dict], cfg: FontCfg, size: int, cache: Path) -> Image.Image: + """Two-column dim page: black content words, gray stopwords, double-strike headings.""" + glyphs, font_ascent = parse_bdf(ensure_font(cfg, cache)) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, _ = capacity(cfg, size) + col_w = (cols - GUTTER) // 2 + masks = _dim_masks(lines) + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * cfg.adv + y0 = row * cfg.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + if ln["kind"] == "heading": + fg = _BLACK + else: + fg = _DIMMED if masks[li][ci] else _BLACK + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * cfg.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runners ----------------------------------------------------------------- + + +def qa_call(cond: str, messages: list[dict], ctx: dict) -> dict: + args, keys = ctx["args"], ctx["keys"] + return cached( + MODEL, f"exp16-qa-{cond}", {"messages": messages, "size": args.size, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, MODEL, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + + +def score(questions: list[dict], qa: dict, cond: str, start: int, ctx: dict) -> list[dict]: + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": MODEL, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def run_grid_chunk(cond: str, start: int, end: int, ctx: dict) -> list[dict]: + """Row-major grid cell: chunk the flow by capacity, one QA call per chunk.""" + args, flow, paras, offsets = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + _, font, variant = parse_condition(cond) + png = CACHE / f"exp16-{font}-{variant}-{sha8(chunk_text, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render(chunk_text, FONTS[font], CACHE, args.size, variant), png) + cols, rows, _ = capacity(FONTS[font], args.size) + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("qa-image.md").format(cols=cols, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + return score(questions, qa_call(cond, messages, ctx), cond, start, ctx) + + +def run_doc_page(cond: str, page: tuple[int, int], ctx: dict) -> list[dict]: + """Document cell: paragraph-aligned page, two-column dim render.""" + args, paras, offsets = ctx["args"], ctx["paras"], ctx["offsets"] + i, j = page + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + _, font, _ = parse_condition(cond) + cfg = FONTS[font] + lines = ctx["lines"][cond][page] + page_key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"exp16-{cond}-{page_key}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render_doc(lines, cfg, args.size, CACHE), png) + cols, rows, _ = capacity(cfg, args.size) + col_w = (cols - GUTTER) // 2 + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": load_prompt("exp04-qa-image.md").format(col_w=col_w, rows=rows)}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + return score(questions, qa_call(cond, messages, ctx), cond, start, ctx) + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true", help="render sample pages + capacity stats, no API") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp16-bestfable" + out_dir.mkdir(parents=True, exist_ok=True) + + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + + keys = {} + if not args.render_only: + keys["anthropic"] = load_env_key("ANTHROPIC_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + capacity_stats: dict[str, dict] = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + page_lines: dict[str, dict] = {} + doc_pages: dict[str, list[tuple[int, int]]] = {} + for cond in conditions: + kind, font, _ = parse_condition(cond) + cfg = FONTS[font] + cols, rows, grid_cap = capacity(cfg, args.size) + if kind == "doc": + col_w = (cols - GUTTER) // 2 + pages = pack_pages(paras, col_w, 2 * rows) + doc_pages[cond] = pages + page_lines[cond] = {pg: layout_page(paras[pg[0] : pg[1]], col_w) for pg in pages} + page_chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + capacity_stats[f"{cond}@{length}"] = { + "pages": len(pages), + "mean_chars_page": round(sum(page_chars) / len(pages)), + "grid_chars_page": grid_cap, + "corpus_chars": len(flow), + } + else: + capacity_stats[f"{cond}@{length}"] = { + "pages": -(-len(flow) // grid_cap), + "mean_chars_page": grid_cap, + "grid_chars_page": grid_cap, + "corpus_chars": len(flow), + } + ctx = { + "args": args, "flow": flow, "paras": paras, "offsets": offsets, + "keys": keys, "length": length, "lines": page_lines, + } + for cond in conditions: + kind, font, _ = parse_condition(cond) + if kind == "doc": + for pg in doc_pages[cond]: + tasks.append(("doc", cond, pg, ctx)) + else: + budget = capacity(FONTS[font], args.size)[2] + for start in range(0, len(flow), budget): + tasks.append(("img", cond, (start, min(start + budget, len(flow))), ctx)) + + for key, st in sorted(capacity_stats.items()): + print( + f" {key}: {st['pages']} pages, mean {st['mean_chars_page']} chars/page " + f"(grid cap {st['grid_chars_page']}; corpus {st['corpus_chars']})" + ) + + if args.render_only: + for length in lengths: + paras = all_paras[:length] + for cond in conditions: + kind, font, _ = parse_condition(cond) + cfg = FONTS[font] + if kind == "doc": + cols, rows, _ = capacity(cfg, args.size) + col_w = (cols - GUTTER) // 2 + i, j = pack_pages(paras, col_w, 2 * rows)[0] + lines = layout_page(paras[i:j], col_w) + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"exp16-{cond}-{key}.png" + atomic_save(render_doc(lines, cfg, args.size, CACHE), png) + else: + flow, _ = squad.build_flow(paras) + cap = capacity(cfg, args.size)[2] + chunk_text = flow[:cap] + _, _, variant = parse_condition(cond) + png = CACHE / f"exp16-{font}-{variant}-{sha8(chunk_text, str(args.size))}.png" + atomic_save(render(chunk_text, cfg, CACHE, args.size, variant), png) + print(f" sample: {png}") + return + + print(f"grid: {len(tasks)} tasks") + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [ + pool.submit(run_doc_page, cond, span, ctx) + if kind == "doc" + else pool.submit(run_grid_chunk, cond, span[0], span[1], ctx) + for kind, cond, span, ctx in tasks + ] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": MODEL, "length": length, "condition": cond, **aggregate(sub, *PRICE)}) + (out_dir / "summary.json").write_text( + json.dumps({"args": vars(args), "capacity": capacity_stats, "cells": cells}, indent=1) + ) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + bi, bt = BASE_IMG.get(c["length"]), BASE_TEXT.get(c["length"]) + comb_se = (c["f1_se"] ** 2 + bi[1] ** 2) ** 0.5 if bi else 0.0 + d_img = f"vs 6x12-dim {c['f1'] - bi[0]:+.3f} ({(c['f1'] - bi[0]) / comb_se:+.1f}se)" if bi else "" + d_txt = f" vs text {c['f1'] - bt[0]:+.3f}" if bt else "" + flag = " ** beats text ceiling" if bt and c["f1"] > bt[0] else "" + print( + f"{MODEL} len {c['length']:<4} {c['condition']:<18} n={c['n']:<4} " + f"EM {c['em']:.3f} F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f} " + f"{d_img}{d_txt}{flag}" + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp17_bestopus.py b/packages/snapcompact/research/exp17_bestopus.py new file mode 100644 index 000000000..01ec418fd --- /dev/null +++ b/packages/snapcompact/research/exp17_bestopus.py @@ -0,0 +1,419 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp17: claude-opus-4-8 model-owner run — do round-1 levers transfer to Anthropic? + +Round-1 (gpt-5.5/gemini) validated: (1) patch-aligned pitch 16 (8x13 glyphs on +an 8x16 cell), (2) two-column document layout, (3) per-model variant. Opus's +round-0 winner is img-8x13-bw; opus emits ~0 reasoning tokens on image QA so +any win here must come from F1, not output-token savings. + +Conditions (variant anchor = bw): + img-8on16-bw grid, 8x13 glyphs on 16 px pitch (alignment only) + img-doc-8on16-bw two-column doc layout at 8on16 (alignment + layout) + img-doc-8x13-bw two-column doc layout at 8x13 (layout only) + +Protocol: screen all three at length 150; only the winner goes to 50/250. +Records accumulate across invocations (re-run cells replace their old rows), +so the final matrix.csv holds the union of screen + confirm runs. + +Usage: uv run exp17_bestopus.py --render-only # capacity + sample PNGs + uv run exp17_bestopus.py # screen @150 + uv run exp17_bestopus.py --conditions img-X --lengths 50,250 # confirm + uv run exp17_bestopus.py --report # reprint from cache +""" + +import argparse +import csv +import json +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import FontCfg, capacity, ensure_font, parse_bdf, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODEL = "claude-opus-4-8" +PRICES = (15.0, 75.0) # $/M in, $/M out +FONTS = { + "8x13": FontCfg("8x13", "8x13", 8, 13), # round-0 opus winner font + "8on16": FontCfg("8on16", "8x13", 8, 16), # same glyphs, patch-aligned 16 px pitch +} +CONDITIONS = ("img-8on16-bw", "img-doc-8on16-bw", "img-doc-8x13-bw") +LENGTHS = (150,) # screening default; confirm via --lengths 50,250 +GUTTER = 3 # char cells between doc columns +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) +_INK = (24, 24, 24) # near-black body text, like a printed page + +# claude-opus-4-8 baselines, results/optimal-combined/matrix.csv (qpc 30, seed 42): +BASELINE = {50: (0.9626, 0.0258, 0.143), 150: (0.8937, 0.0223, 0.380), 250: (0.8708, 0.0196, 0.559)} +TEXT_CEIL = {50: (0.9278, 0.195), 150: (0.9112, 0.637), 250: (0.9268, 0.938)} + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def parse_cond(name: str) -> tuple[str, FontCfg]: + """'img-8on16-bw' -> ('grid', cfg); 'img-doc-8on16-bw' -> ('doc', cfg).""" + parts = name.split("-") + if parts[1] == "doc": + return "doc", FONTS[parts[2]] + return "grid", FONTS[parts[1]] + + +# --- document layout (ported from exp04_layout.py, parameterized font) ------ + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: # pathological; never hit on SQuAD prose + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Headings repeat at the top of a page when an article continues, since + each page is read in isolation. + """ + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def render_doc(lines: list[dict], cfg: FontCfg, size: int) -> Image.Image: + """Two-column bw page: left column rows top-to-bottom, then right column.""" + glyphs, font_ascent = parse_bdf(ensure_font(cfg, CACHE)) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, _ = capacity(cfg, size) + col_w = (cols - GUTTER) // 2 + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * cfg.adv + y0 = row * cfg.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + fg = _BLACK if ln["kind"] == "heading" else _INK + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * cfg.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runner ----------------------------------------------------------------- + + +def render_unit_png(cond: str, unit: tuple[int, int], ctx: dict) -> Path: + """Render (or reuse) the PNG for one chunk/page; atomic tmp-then-replace.""" + args, paras = ctx["args"], ctx["paras"] + kind, cfg = parse_cond(cond) + if kind == "grid": + start, end = unit + chunk_text = ctx["flow"][start:end] + png = CACHE / f"exp17-{cond}-{sha8(chunk_text, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{uuid.uuid4().hex[:8]}.tmp.png") + render(chunk_text, cfg, CACHE, args.size, "bw").save(tmp) + tmp.replace(png) + else: + i, j = unit + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"exp17-{cond}-{key}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{uuid.uuid4().hex[:8]}.tmp.png") + render_doc(ctx["lines"][(cond, unit)], cfg, args.size).save(tmp) + tmp.replace(png) + return png + + +def run_unit(cond: str, unit: tuple[int, int], ctx: dict) -> list[dict]: + """One (condition, chunk-or-page) cell unit: render carrier, QA, score.""" + args, paras, offsets, keys = ctx["args"], ctx["paras"], ctx["offsets"], ctx["keys"] + kind, cfg = parse_cond(cond) + if kind == "grid": + start, end = unit + else: + i, j = unit + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + png = render_unit_png(cond, unit, ctx) + cols, rows, _ = capacity(cfg, args.size) + if kind == "grid": + preamble = load_prompt("qa-image.md").format(cols=cols, rows=rows) + else: + preamble = load_prompt("exp04-qa-image.md").format(col_w=(cols - GUTTER) // 2, rows=rows) + messages = [ + { + "role": "user", + "content": [ + {"text": preamble}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + MODEL, "exp17-qa", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, MODEL, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": MODEL, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true", help="capacity stats + first-page PNGs, no API") + ap.add_argument("--report", action="store_true", help="reprint matrix from accumulated records only") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp17-bestopus" + out_dir.mkdir(parents=True, exist_ok=True) + records_path = out_dir / "records.jsonl" + + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + + all_paras = squad.load_paragraphs(CACHE) + records: list[dict] = [] + + if not args.report: + keys = {} if args.render_only else {"anthropic": load_env_key("ANTHROPIC_API_KEY", args.env)} + tasks = [] + capacity_stats = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = { + "args": args, "flow": flow, "paras": paras, "offsets": offsets, + "keys": keys, "length": length, "lines": {}, + } + for cond in conditions: + kind, cfg = parse_cond(cond) + cols, rows, grid_cap = capacity(cfg, args.size) + if kind == "grid": + units = [(s, min(s + grid_cap, len(flow))) for s in range(0, len(flow), grid_cap)] + chars = [e - s for s, e in units] + else: + col_w = (cols - GUTTER) // 2 + pages = pack_pages(paras, col_w, 2 * rows) + for pg in pages: + ctx["lines"][(cond, pg)] = layout_page(paras[pg[0] : pg[1]], col_w) + units = pages + chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + capacity_stats[f"{cond}@{length}"] = { + "pages": len(units), + "mean_chars_page": round(sum(chars) / len(units)), + "grid_chars_page": grid_cap, + "corpus_chars": len(flow), + } + tasks.extend((cond, u, ctx) for u in units) + + for k, st in capacity_stats.items(): + print( + f" {k}: {st['pages']} pages, mean {st['mean_chars_page']} chars/page " + f"(grid cap {st['grid_chars_page']}; corpus {st['corpus_chars']})" + ) + + if args.render_only: + seen = set() + for cond, unit, ctx in tasks: + if cond in seen: + continue + seen.add(cond) + print(f" sample: {render_unit_png(cond, unit, ctx)}") + return + + print(f"grid: {len(tasks)} unit tasks ({len(conditions)} conds x {lengths})") + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_unit, c, u, ctx) for c, u, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} units", flush=True) + + # Merge: rows for cells just run replace any prior rows for those cells. + ran_cells = {(length, cond) for length in lengths for cond in conditions} + old = [] + if records_path.exists(): + with records_path.open() as fh: + old = [json.loads(ln) for ln in fh if ln.strip()] + records = [r for r in old if (r["length"], r["cond"]) not in ran_cells] + records + with records_path.open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + else: + with records_path.open() as fh: + records = [json.loads(ln) for ln in fh if ln.strip()] + + cells = [] + for length in sorted({r["length"] for r in records}): + for cond in sorted({r["cond"] for r in records if r["length"] == length}): + sub = [r for r in records if r["length"] == length and r["cond"] == cond] + cells.append({"model": MODEL, "length": length, "condition": cond, **aggregate(sub, *PRICES)}) + (out_dir / "summary.json").write_text( + json.dumps({"args": vars(args), "baseline_img_8x13_bw": BASELINE, "cells": cells}, indent=1) + ) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + print(f"\n{'len':<5}{'condition':<22}{'n':>5}{'EM':>7}{'F1':>7}{'+-se':>7}{'$':>8}{'d/se vs 8x13-bw':>17}") + for c in cells: + b_f1, b_se, b_cost = BASELINE[c["length"]] + dse = (c["f1"] - b_f1) / ((c["f1_se"] ** 2 + b_se**2) ** 0.5 or 1) + t_f1, _ = TEXT_CEIL[c["length"]] + flag = " > text ceiling" if c["f1"] > t_f1 else "" + print( + f"{c['length']:<5}{c['condition']:<22}{c['n']:>5}{c['em']:>7.3f}{c['f1']:>7.3f}" + f"{c['f1_se']:>7.3f}{c['cost_usd']:>8.3f}{dse:>+17.2f}{flag}" + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp18_bestkimi.py b/packages/snapcompact/research/exp18_bestkimi.py new file mode 100644 index 000000000..60c2f0601 --- /dev/null +++ b/packages/snapcompact/research/exp18_bestkimi.py @@ -0,0 +1,439 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp18: best optical profile for moonshotai/kimi-k2.6. + +Round-1 levers (patch-aligned pitch 16, two-column doc layout, per-model +variant) were validated on gpt-5.5/gemini only. Kimi's round-0 winner is +img-8x13-sent-dim (beats text at 150) but it pays the worst read tax in the +fleet (~95% of image-cell cost is output tokens; 120k+ out at length 250). + +Screen at length 150, anchored on sent-dim: + img-8on16-sent-dim grid, 8x13 glyphs on an 8x16 cell (alignment only) + img-doc8on16-sent-dim two-column doc layout at 8on16 (alignment + layout) + img-doc8x13-sent-dim two-column doc layout at pitch 13 (layout only) +Confirm the winner at 50 and 250. + +Usage: uv run exp18_bestkimi.py # screening (length 150) + uv run exp18_bestkimi.py --lengths 50,250 --conditions img-... # confirm + uv run exp18_bestkimi.py --render-only # sample pages, no API + uv run exp18_bestkimi.py --report --lengths 50,150,250 # re-aggregate from cache +""" + +import argparse +import csv +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DARK, _DIMMED, FontCfg, _stopword_mask, capacity, ensure_font, parse_bdf, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODEL = "moonshotai/kimi-k2.6" +PRICE_IN, PRICE_OUT = 0.68, 3.41 +FONTS = { + "8on16": FontCfg("8on16", "8x13", 8, 16), # patch-aligned: 8x13 glyphs, 16 px pitch + "8x13": FontCfg("8x13", "8x13", 8, 13), # kimi's round-0 winner pitch +} +# cond -> (kind, font key, variant). All anchored on sent-dim (kimi's winner). +CONDITIONS = { + "img-8on16-sent-dim": ("grid", "8on16", "sent-dim"), + "img-doc8on16-sent-dim": ("doc", "8on16", "sent-dim"), + "img-doc8x13-sent-dim": ("doc", "8x13", "sent-dim"), +} +GUTTER = 3 # char cells between doc columns +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncations are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens": + return hit + out = fn() + if out.get("stop") == "max_tokens": + print(f" WARN truncated, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +# --- document layout (ported from exp04, parameterized for font/pitch) ----- + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Title changes become uppercase double-strike headings; the heading is + repeated at the top of a page when an article continues, since each page + is read in isolation. One blank line between paragraphs. + """ + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def _sentence_indices_doc(lines: list[dict]) -> list[list[int]]: + """Per-line per-char sentence index, cycling across the page (newline counts as boundary space).""" + joined = "\n".join(ln["text"] for ln in lines) + idx, run = 0, [] + for i, ch in enumerate(joined): + run.append(idx) + if ch in ".!?" and i + 1 < len(joined) and joined[i + 1] in " \n": + idx += 1 + out, pos = [], 0 + for ln in lines: + n = len(ln["text"]) + out.append(run[pos : pos + n]) + pos += n + 1 # the joining newline + return out + + +def render_doc(lines: list[dict], cfg: FontCfg, size: int, cache: Path) -> Image.Image: + """Two-column sent-dim page: left column top-to-bottom, then right. + + Body glyph color = sentence hue, overridden to light gray for stopwords + (same composition as bdf.render's sent-dim). Headings: black double-strike. + """ + glyphs, font_ascent = parse_bdf(ensure_font(cfg, cache)) + ascent = cfg.ascent if cfg.ascent is not None else font_ascent + cols, rows, _ = capacity(cfg, size) + col_w = (cols - GUTTER) // 2 + sent_idx = _sentence_indices_doc(lines) + dim_masks = [_stopword_mask(ln["text"]) for ln in lines] + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * cfg.adv + y0 = row * cfg.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + if ln["kind"] == "heading": + fg = _BLACK + elif dim_masks[li][ci]: + fg = _DIMMED + else: + fg = _DARK[sent_idx[li][ci] % 6] + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * cfg.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runner ----------------------------------------------------------------- + + +def doc_png(cond: str, paras: list[dict], lines: list[dict], cfg: FontCfg, size: int) -> Path: + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras]), str(size)) + png = CACHE / f"exp18-{cond}-{key}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render_doc(lines, cfg, size, CACHE).save(tmp) + tmp.replace(png) + return png + + +def run_unit(cond: str, unit: dict, ctx: dict) -> list[dict]: + """One (condition, page/chunk) unit: render carrier, QA, score.""" + args, paras, offsets, keys = ctx["args"], ctx["paras"], ctx["offsets"], ctx["keys"] + start, end = unit["start"], unit["end"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + kind, font, variant = CONDITIONS[cond] + cfg = FONTS[font] + cols, rows, _ = capacity(cfg, args.size) + if kind == "grid": + chunk_text = ctx["flow"][start:end] + png = CACHE / f"exp18-{cond}-{sha8(chunk_text, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(".tmp.png") + render(chunk_text, cfg, CACHE, args.size, variant).save(tmp) + tmp.replace(png) + prompt = load_prompt("qa-image.md").format(cols=cols, rows=rows) + else: + i, j = unit["page"] + png = doc_png(cond, paras[i:j], unit["lines"], cfg, args.size) + col_w = (cols - GUTTER) // 2 + prompt = load_prompt("exp04-qa-image.md").format(col_w=col_w, rows=rows) + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [{"text": prompt}, {"image_path": png}, {"text": q_block}], + } + ] + qa = cached( + MODEL, "exp18-qa", {"messages": messages, "size": args.size, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, MODEL, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": MODEL, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict]) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * PRICE_IN + cost_out = tok["out"] / 1e6 * PRICE_OUT + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--lengths", default="150") + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true", help="render first page per cond + capacity stats, no API") + ap.add_argument("--report", action="store_true", help="re-aggregate (all units should hit cache)") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp18-bestkimi" + out_dir.mkdir(parents=True, exist_ok=True) + + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + for c in conditions: + if c not in CONDITIONS: + sys.exit(f"unknown condition: {c}") + + keys = {} + if not args.render_only: + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + tasks: list[tuple[str, dict, dict]] = [] + capacity_stats: dict = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + capacity_stats[length] = {"corpus_chars": len(flow), "conds": {}} + for cond in conditions: + kind, font, _ = CONDITIONS[cond] + cfg = FONTS[font] + cols, rows, grid_cap = capacity(cfg, args.size) + if kind == "grid": + units = [ + {"start": s, "end": min(s + grid_cap, len(flow))} for s in range(0, len(flow), grid_cap) + ] + chars = [u["end"] - u["start"] for u in units] + else: + col_w = (cols - GUTTER) // 2 + pages = pack_pages(paras, col_w, 2 * rows) + units = [] + for i, j in pages: + units.append( + { + "start": offsets[i], + "end": offsets[j - 1] + len(paras[j - 1]["ctx"]), + "page": (i, j), + "lines": layout_page(paras[i:j], col_w), + } + ) + chars = [u["end"] - u["start"] for u in units] + capacity_stats[length]["conds"][cond] = { + "pages": len(units), + "mean_chars_page": round(sum(chars) / len(units)), + "grid_chars_page": grid_cap, + } + for u in units: + tasks.append((cond, u, ctx)) + + for length, st in capacity_stats.items(): + print(f"len {length}: corpus {st['corpus_chars']} chars") + for cond, cs in st["conds"].items(): + print(f" {cond:<24} {cs['pages']} pages, mean {cs['mean_chars_page']} chars/page (grid cap {cs['grid_chars_page']})") + + if args.render_only: + for cond, u, ctx in tasks: + if u["start"] != 0: + continue + kind, font, variant = CONDITIONS[cond] + cfg = FONTS[font] + if kind == "grid": + chunk_text = ctx["flow"][u["start"] : u["end"]] + png = CACHE / f"exp18-{cond}-{sha8(chunk_text, str(args.size))}.png" + tmp = png.with_suffix(".tmp.png") + render(chunk_text, cfg, CACHE, args.size, variant).save(tmp) + tmp.replace(png) + else: + i, j = u["page"] + png = doc_png(cond, ctx["paras"][i:j], u["lines"], cfg, args.size) + print(f" sample: {png}") + return + + print(f"grid: {len(tasks)} unit tasks on {MODEL}") + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_unit, c, u, ctx) for c, u, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} units", flush=True) + + # Merge with any prior records (confirm runs extend the screening set). + rec_path = out_dir / "records.jsonl" + old: list[dict] = [] + if rec_path.exists(): + ran = {(r["length"], r["cond"]) for r in records} + for line in rec_path.read_text().splitlines(): + r = json.loads(line) + if (r["length"], r["cond"]) not in ran: + old.append(r) + records = old + records + tmp = rec_path.with_suffix(".tmp.jsonl") + with tmp.open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + tmp.replace(rec_path) + + cells = [] + for length in sorted({r["length"] for r in records}): + for cond in CONDITIONS: + sub = [r for r in records if r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": MODEL, "length": length, "condition": cond, **aggregate(sub)}) + (out_dir / "summary.json").write_text( + json.dumps({"args": vars(args), "capacity": capacity_stats, "cells": cells}, indent=1) + ) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"len {c['length']:<4} {c['condition']:<24} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} out {c['tok_out']:>7} (reas {c['tok_reasoning']}) ${c['cost_usd']:.4f}" + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp19_bestglm.py b/packages/snapcompact/research/exp19_bestglm.py new file mode 100644 index 000000000..b8ef651ba --- /dev/null +++ b/packages/snapcompact/research/exp19_bestglm.py @@ -0,0 +1,477 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp19: best optical profile for z-ai/glm-4.6v (weakest optical reader in the fleet). + +glm uniquely needed dark mode (bright hues on black) and its round-0 winner was +img-8x13-dark-sent (.789/.753/.638 vs text .944/.904/.880). Round-1 levers +(patch-aligned pitch 16, two-column doc layout) were only validated on +gpt-5.5/gemini. Screen at length 150: + + img-8on16-dark-sent 8x13 glyphs on an 8x16 cell (patch-aligned), dark + img-doc-8on16-dark-sent doc layout + pitch 16 + dark palette (headings in + bright white double-strike) + img-doc-8x13-dark-sent layout-only control at the baseline pitch 13 + img-doc-8on16-sent light-mode probe: does doc structure remove the + need for dark mode? + +then confirm the screening winner at lengths 50 and 250. Methodology matches +the baselines: seed 42, qpc 30, size 1568, max_tokens 32768, effort default. + +Usage: uv run exp19_bestglm.py # screen @150 + uv run exp19_bestglm.py --lengths 50,250 --conditions X # confirm winner + uv run exp19_bestglm.py --render-only # sample PNGs, no API + +Repeated invocations merge into results/exp19-bestglm/ (records for the +(model, length, condition) cells being run are replaced; others kept). +""" + +import argparse +import csv +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _BRIGHT, _DARK, FontCfg, capacity, ensure_font, parse_bdf, render # noqa: E402 +from final import cached # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +MODELS = {"z-ai/glm-4.6v": (0.30, 0.90)} +LENGTHS = (150,) +FONTS = { + "8on16": FontCfg("8on16", "8x13", 8, 16), # patch-aligned padded cell (exp01 pattern) + "8x13": FontCfg("8x13", "8x13", 8, 13), # baseline pitch +} +# condition -> (kind, font key, palette variant) +CONDITIONS = { + "img-8on16-dark-sent": ("grid", "8on16", "dark-sent"), + "img-doc-8on16-dark-sent": ("doc", "8on16", "dark-sent"), + "img-doc-8x13-dark-sent": ("doc", "8x13", "dark-sent"), + "img-doc-8on16-sent": ("doc", "8on16", "sent"), +} +GUTTER = 3 # char cells between doc columns +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) + + +# --- document layout (ported from exp04, parameterized for font/pitch) ------ + + +def wrap(text: str, width: int) -> list[str]: + """Greedy word-wrap, no mid-word breaks (hard split only for width+ words).""" + lines: list[str] = [] + cur = "" + for word in text.split(): + while len(word) > width: # pathological; never hit on SQuAD prose + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + """Typeset paragraphs into lines: [{kind: heading|body|blank, text}]. + + Article title changes become headings (repeated at the top of a page even + when an article continues, since each page is read in isolation). + Paragraphs are separated by one blank line. + """ + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + """Greedy paragraph-aligned packing: [(i, j)] para ranges, one per page.""" + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def _sentence_colors(lines: list[dict], palette: list) -> list[list[tuple[int, int, int]]]: + """Per-line per-char glyph color cycling hue per sentence across the page.""" + joined = "\n".join(ln["text"] for ln in lines) + idx, out_idx = 0, [] + for i, ch in enumerate(joined): + out_idx.append(idx) + if ch in ".!?" and i + 1 < len(joined) and joined[i + 1] in " \n": + idx += 1 + colors, pos = [], 0 + for ln in lines: + n = len(ln["text"]) + colors.append([palette[out_idx[pos + k] % 6] for k in range(n)]) + pos += n + 1 # the joining newline + return colors + + +def render_doc(lines: list[dict], font: FontCfg, size: int, variant: str, cache: Path) -> Image.Image: + """Two-column page: left column rows top-to-bottom, then right column. + + variant "dark-sent": black page, body glyphs in bright sentence hues, + headings bright white double-strike. variant "sent": white page, body in + dark sentence hues, headings black double-strike. The page background + covers the full image, so a padded cell (pitch > glyph height) is dark + edge-to-edge in dark mode. + """ + glyphs, font_ascent = parse_bdf(ensure_font(font, cache)) + ascent = font.ascent if font.ascent is not None else font_ascent + cols, rows, _ = capacity(font, size) + col_w = (cols - GUTTER) // 2 + dark = variant == "dark-sent" + bg, heading_fg = (_BLACK, _WHITE) if dark else (_WHITE, _BLACK) + sent_colors = _sentence_colors(lines, _BRIGHT if dark else _DARK) + img = Image.new("RGB", (size, size), bg) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break # overflow guard; pack_pages should prevent this + x_origin = column * (col_w + GUTTER) * font.adv + y0 = row * font.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + fg = heading_fg if ln["kind"] == "heading" else sent_colors[li][ci] + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * font.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runner ------------------------------------------------------------------ + + +def save_png(png: Path, img_fn) -> None: + if png.exists() and png.stat().st_size > 0: + return + tmp = png.with_suffix(".tmp.png") + img_fn().save(tmp) + tmp.replace(png) + + +def run_grid_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + """One row-major-grid chunk: render via bdf.render, QA, score.""" + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + _, font_key, variant = CONDITIONS[cond] + font = FONTS[font_key] + chunk_text = flow[start:end] + png = CACHE / f"exp19-{font_key}-{variant}-{sha8(chunk_text, str(args.size))}.png" + save_png(png, lambda: render(chunk_text, font, CACHE, args.size, variant)) + cols, rows, _ = capacity(font, args.size) + prompt = load_prompt("qa-image.md").format(cols=cols, rows=rows) + return qa_and_score(model, cond, prompt, png, questions, start, ctx) + + +def run_doc_page(model: str, cond: str, page: tuple[int, int], ctx: dict) -> list[dict]: + """One doc-layout page: typeset, render, QA, score.""" + args, paras, offsets, keys = ctx["args"], ctx["paras"], ctx["offsets"], ctx["keys"] + i, j = page + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + _, font_key, variant = CONDITIONS[cond] + font = FONTS[font_key] + lines = ctx["lines"][cond][page] + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"exp19-doc-{font_key}-{variant}-{key}.png" + save_png(png, lambda: render_doc(lines, font, args.size, variant, CACHE)) + cols, rows, _ = capacity(font, args.size) + col_w = (cols - GUTTER) // 2 + prompt = load_prompt("exp19-qa-doc.md").format(col_w=col_w, rows=rows) + return qa_and_score(model, cond, prompt, png, questions, start, ctx) + + +def parse_answers(text: str, n: int) -> list[str]: + """parse_numbered, with a fallback for glm's intermittently unnumbered output. + + glm-4.6v often ignores the "numbered list" instruction and emits plain + answer lines in question order. When numbered parsing recovers fewer than + half the answers and the response is a clean <=n line list, map lines + positionally instead (preamble line ending in ':' is dropped). + """ + nums = squad.parse_numbered(text, n) + if sum(bool(a) for a in nums) >= max(1, n // 2): + return nums + lines = [ln.strip() for ln in text.splitlines() if ln.strip()] + if lines and lines[0].endswith(":"): + lines = lines[1:] + if 0 < len(lines) <= n: + return lines + [""] * (n - len(lines)) + return nums + + +def qa_and_score(model: str, cond: str, prompt: str, png: Path, questions: list[dict], start: int, ctx: dict) -> list[dict]: + args, keys = ctx["args"], ctx["keys"] + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [ + { + "role": "user", + "content": [ + {"text": prompt}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + qa = cached( + model, "exp19-qa", {"messages": messages, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, model, messages, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers = parse_answers(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true", help="render sample pages + capacity stats, no API") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / "exp19-bestglm" + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + for c in conditions: + if c not in CONDITIONS: + sys.exit(f"unknown condition: {c}") + + keys = {} + if not args.render_only: + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + capacity_stats = {} + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + # Doc page packing per font (max_lines differs with pitch). + page_lines: dict[str, dict] = {} + doc_stats = {} + for cond in conditions: + kind, font_key, _ = CONDITIONS[cond] + if kind != "doc": + continue + font = FONTS[font_key] + cols, rows, _cap = capacity(font, args.size) + col_w = (cols - GUTTER) // 2 + pages = pack_pages(paras, col_w, 2 * rows) + page_lines[cond] = {pg: layout_page(paras[pg[0] : pg[1]], col_w) for pg in pages} + chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + doc_stats[cond] = { + "pages": len(pages), + "mean_chars_page": round(sum(chars) / len(pages)), + "col_w": col_w, + "rows": rows, + } + capacity_stats[length] = { + "corpus_chars": len(flow), + "grid": { + fk: dict(zip(("cols", "rows", "chars"), capacity(FONTS[fk], args.size))) for fk in FONTS + }, + "doc": doc_stats, + } + ctx = { + "args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, + "length": length, "lines": page_lines, + } + for model in models: + for cond in conditions: + kind, font_key, _ = CONDITIONS[cond] + if kind == "grid": + budget = capacity(FONTS[font_key], args.size)[2] + for start in range(0, len(flow), budget): + tasks.append(("grid", model, cond, (start, min(start + budget, len(flow))), ctx)) + else: + for pg in page_lines[cond]: + tasks.append(("doc", model, cond, pg, ctx)) + + for length, st in capacity_stats.items(): + print(f"len {length}: corpus {st['corpus_chars']} chars") + for fk, g in st["grid"].items(): + print(f" grid {fk}: {g['cols']}x{g['rows']} = {g['chars']} chars/page") + for cond, d in st["doc"].items(): + print(f" {cond}: {d['pages']} pages, mean {d['mean_chars_page']} chars/page (2x{d['col_w']}w, {d['rows']} rows)") + + if args.render_only: + for length in lengths: + ctx = next(t[4] for t in tasks if t[4]["length"] == length) + for cond in conditions: + kind, font_key, variant = CONDITIONS[cond] + if kind == "grid": + budget = capacity(FONTS[font_key], args.size)[2] + chunk_text = ctx["flow"][:budget] + png = CACHE / f"exp19-{font_key}-{variant}-{sha8(chunk_text, str(args.size))}.png" + save_png(png, lambda: render(chunk_text, FONTS[font_key], CACHE, args.size, variant)) + else: + pg = next(iter(ctx["lines"][cond])) + i, j = pg + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in ctx["paras"][i:j]]), str(args.size)) + png = CACHE / f"exp19-doc-{font_key}-{variant}-{key}.png" + save_png(png, lambda: render_doc(ctx["lines"][cond][pg], FONTS[font_key], args.size, variant, CACHE)) + print(f" sample: {png}") + return + + print(f"grid: {len(tasks)} page/chunk tasks") + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [ + pool.submit(run_grid_chunk, m, c, u[0], u[1], ctx) if kind == "grid" else pool.submit(run_doc_page, m, c, u, ctx) + for kind, m, c, u, ctx in tasks + ] + for fut in futures: + records.extend(fut.result()) + done += 1 + print(f" {done}/{len(tasks)} tasks", flush=True) + + # Merge with prior invocations: replace the cells we just ran, keep the rest. + ran = {(m, ln, c) for ln in lengths for m in models for c in conditions} + rec_path = out_dir / "records.jsonl" + if rec_path.exists(): + old = [json.loads(ln) for ln in rec_path.read_text().splitlines() if ln.strip()] + records = [r for r in old if (r["model"], r["length"], r["cond"]) not in ran] + records + with rec_path.open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in sorted({r["model"] for r in records}): + for length in sorted({r["length"] for r in records}): + for cond in CONDITIONS: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text( + json.dumps({"args": vars(args), "capacity": capacity_stats, "cells": cells}, indent=1) + ) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for c in cells: + print( + f"{c['model']:<16} len {c['length']:<4} {c['condition']:<26} " + f"n={c['n']:<4} EM {c['em']:.3f} F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.4f}" + ) + print(f"\n-> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp20_8x8u.py b/packages/snapcompact/research/exp20_8x8u.py new file mode 100644 index 000000000..a23c063e5 --- /dev/null +++ b/packages/snapcompact/research/exp20_8x8u.py @@ -0,0 +1,362 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp20: swap the per-model round-2 winner font for unscii-8 (8x8u). + +For each model, re-run its best-known config with FontCfg("8x8u","unscii-8",8,8) +in place of the winning font, everything else unchanged (layout, variant, +1568px, seed 42, qpc 30). 8x8u @1568: grid 196x196 = 38,416 chars/page +(~2x doc-8on16's effective capacity); doc 2x96 cols x 196 rows. + + model config swapped from exp20 condition + gpt-5.5 doc-8on16-bw -> img-doc-8x8u-bw + google/gemini-3.5-flash doc-8on16-sent-dim -> img-doc-8x8u-sent-dim + moonshotai/kimi-k2.6 doc8on16-sent-dim -> img-doc-8x8u-sent-dim + z-ai/glm-4.6v doc-8on16-sent -> img-doc-8x8u-sent + claude-fable-5 grid 6x12-dim -> img-8x8u-dim + claude-opus-4-8 grid 8x13-bw -> img-8x8u-bw + +Usage: + uv run exp20_8x8u.py --model gpt-5.5 --render-only # sample PNG, no API + uv run exp20_8x8u.py --model gpt-5.5 # lengths 50,150,250 +""" + +import argparse +import os +import json +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import _DARK, _DIMMED, FontCfg, _stopword_mask, capacity, load_font, render # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp20" +OUT_DIR = RESULTS / f"{EXP}-8x8u" +FONT = FontCfg("8x8u", "unscii-8", 8, 8) +GUTTER = 3 +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) +_INK = (24, 24, 24) + +# model -> (cond, layout, variant, price_in, price_out, key_name) +CONFIGS = { + "gpt-5.5": ("img-doc-8x8u-bw", "doc", "bw", 2.0, 16.0, "openai"), + "google/gemini-3.5-flash": ("img-doc-8x8u-sent-dim", "doc", "sent-dim", 0.6, 4.0, "openrouter"), + "moonshotai/kimi-k2.6": ("img-doc-8x8u-sent-dim", "doc", "sent-dim", 0.68, 3.41, "openrouter"), + "z-ai/glm-4.6v": ("img-doc-8x8u-sent", "doc", "sent", 0.30, 0.90, "openrouter"), + "claude-fable-5": ("img-8x8u-dim", "grid", "dim", 10.0, 50.0, "anthropic"), + "claude-opus-4-8": ("img-8x8u-bw", "grid", "bw", 15.0, 75.0, "anthropic"), +} +KEY_ENV = {"openai": "OPENAI_API_KEY", "openrouter": "OPENROUTER_API_KEY", "anthropic": "ANTHROPIC_API_KEY"} + + +def slug(model: str) -> str: + return model.split("/")[-1] + + +def cached(model: str, payload: object, fn, fresh: bool) -> dict: + key = sha8(model, f"{EXP}-qa", json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if not fresh and path.exists(): + return json.loads(path.read_text()) + out = fn() + if out.get("stop") != "max_tokens": + tmp = path.with_suffix(".tmp.json") + tmp.write_text(json.dumps(out)) + tmp.replace(path) + else: + print(f"WARN truncated response, not cached ({model})", flush=True) + return out + + +# --- document layout (exp14's renderer, generalized for hex fonts + dim variants) --- + + +def wrap(text: str, width: int) -> list[str]: + lines, cur = [], "" + for word in text.split(): + while len(word) > width: + if cur: + lines.append(cur) + cur = "" + lines.append(word[:width]) + word = word[width:] + if not cur: + cur = word + elif len(cur) + 1 + len(word) <= width: + cur += " " + word + else: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def layout_page(paras: list[dict], col_w: int) -> list[dict]: + lines: list[dict] = [] + prev_title = None + for p in paras: + if p["title"] != prev_title: + if lines: + lines.append({"kind": "blank", "text": ""}) + for hl in wrap(p["title"].replace("_", " ").upper(), col_w): + lines.append({"kind": "heading", "text": hl}) + prev_title = p["title"] + elif lines: + lines.append({"kind": "blank", "text": ""}) + for bl in wrap(p["ctx"], col_w): + lines.append({"kind": "body", "text": bl}) + return lines + + +def pack_pages(paras: list[dict], col_w: int, max_lines: int) -> list[tuple[int, int]]: + pages = [] + i = 0 + while i < len(paras): + j = i + 1 + while j < len(paras) and len(layout_page(paras[i : j + 1], col_w)) <= max_lines: + j += 1 + pages.append((i, j)) + i = j + return pages + + +def _char_styles(lines: list[dict], variant: str) -> list[list[tuple[int, int, int]]]: + """Per-line per-char body glyph color for sent / dim composition.""" + joined = "\n".join(ln["text"] for ln in lines) + sent_idx = None + if "sent" in variant: + sent_idx, idx = [], 0 + for i, ch in enumerate(joined): + sent_idx.append(idx) + if ch in ".!?" and i + 1 < len(joined) and joined[i + 1] in " \n": + idx += 1 + dim = _stopword_mask(joined) if "dim" in variant else None + colors, pos = [], 0 + for ln in lines: + row = [] + for k in range(len(ln["text"])): + i = pos + k + if dim is not None and dim[i]: + row.append(_DIMMED) + elif sent_idx is not None: + row.append(_DARK[sent_idx[i] % 6]) + else: + row.append(_INK) + colors.append(row) + pos += len(ln["text"]) + 1 + return colors + + +def render_doc(lines: list[dict], size: int, variant: str, cache: Path) -> Image.Image: + glyphs, font_ascent = load_font(FONT, cache) + ascent = FONT.ascent if FONT.ascent is not None else font_ascent + cols, rows, _ = capacity(FONT, size) + col_w = (cols - GUTTER) // 2 + styles = _char_styles(lines, variant) + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for li, ln in enumerate(lines): + column, row = divmod(li, rows) + if column > 1: + break + x_origin = column * (col_w + GUTTER) * FONT.adv + y0 = row * FONT.pitch + for ci, ch in enumerate(ln["text"]): + glyph = glyphs.get(ord(ch)) + if glyph is None: + continue + fg = _BLACK if ln["kind"] == "heading" else styles[li][ci] + w, h, xoff, yoff = glyph["bbx"] + top = y0 + ascent - h - yoff + shift = 0x80 if w <= 8 else 0x8000 + strikes = (0, 1) if ln["kind"] == "heading" else (0,) + for dx in strikes: + for r, bits in enumerate(glyph["rows"]): + y = top + r + if not 0 <= y < size: + continue + for b in range(w): + if bits & (shift >> b): + x = x_origin + ci * FONT.adv + xoff + b + dx + if 0 <= x < size: + px[x, y] = fg + return img + + +# --- runner ------------------------------------------------------------------- + + +def atomic_save(img: Image.Image, png: Path) -> None: + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") # pid-unique: parallel models share sent-dim PNGs + img.save(tmp) + tmp.replace(png) + + +def parse_answers(text: str, n: int) -> list[str]: + """parse_numbered + exp19's positional fallback (glm drops numbering).""" + nums = squad.parse_numbered(text, n) + if sum(bool(a) for a in nums) >= max(1, n // 2): + return nums + lines = [ln.strip() for ln in text.splitlines() if ln.strip()] + if lines and lines[0].endswith(":"): + lines = lines[1:] + if 0 < len(lines) <= n: + return lines + [""] * (n - len(lines)) + return nums + + +def qa_unit(model: str, cond: str, prompt: str, png: Path, questions: list[dict], length: int, start: int, ctx: dict) -> list[dict]: + args, keys = ctx["args"], ctx["keys"] + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [{"role": "user", "content": [{"text": prompt}, {"image_path": png}, {"text": q_block}]}] + qa = cached( + model, {"messages": messages, "effort": None}, + lambda: dict(zip(("text", "usage", "stop"), llm_complete(keys, model, messages, max_tokens=args.max_tokens))), + args.fresh, + ) + answers = parse_answers(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append({ + "model": model, "length": length, "cond": cond, "chunk": start, + "pos_rel": q["pos_rel"], "q": q["q"], "answer": a, "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + }) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> list[dict]: + cells = {} + for r in records: + cells.setdefault((r["length"], r["cond"]), []).append(r) + out = [] + for (length, cond), recs in sorted(cells.items()): + n = len(recs) + f1s = [r["f1"] for r in recs] + mean = sum(f1s) / n + var = sum((x - mean) ** 2 for x in f1s) / (n - 1) if n > 1 else 0.0 + usage = [u for r in recs for u in r.get("usage", [])] + tin = sum(u["in"] for u in usage) + tout = sum(u["out"] for u in usage) + creads = sum(u.get("cache_r", 0) for u in usage) + rsn = sum(u.get("reasoning", 0) for u in usage) + cost = (tin + 0.1 * creads) * price_in / 1e6 + tout * price_out / 1e6 + out.append({ + "model": recs[0]["model"], "length": length, "condition": cond, "n": n, + "em": round(sum(r["em"] for r in recs) / n, 4), "f1": round(mean, 4), + "f1_se": round((var / n) ** 0.5, 4), "abstained": sum(r["abstained"] for r in recs), + "tok_in": tin, "tok_out": tout, "tok_cache_r": creads, "tok_reasoning": rsn, + "cost_usd": round(cost, 4), + }) + return out + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True, choices=sorted(CONFIGS)) + ap.add_argument("--lengths", default="50,150,250") + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + cond, layout, variant, price_in, price_out, key_name = CONFIGS[args.model] + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + cols, rows, grid_cap = capacity(FONT, args.size) + col_w = (cols - GUTTER) // 2 + max_lines = 2 * rows + print(f"{args.model}: {cond} ({layout}/{variant}); 8x8u grid {cols}x{rows}={grid_cap}, " + f"doc 2x{col_w}+g{GUTTER}, {max_lines} slots", flush=True) + + keys = {} if args.render_only else {key_name: load_env_key(KEY_ENV[key_name], args.env)} + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + cap_stats = {} + for length in (int(x) for x in args.lengths.split(",")): + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "keys": keys} + if layout == "doc": + pages = pack_pages(paras, col_w, max_lines) + page_chars = [offsets[j - 1] + len(paras[j - 1]["ctx"]) - offsets[i] for i, j in pages] + cap_stats[length] = {"pages": len(pages), "mean_chars_page": round(sum(page_chars) / len(pages))} + prompt = load_prompt("exp04-qa-image.md").format(col_w=col_w, rows=rows) + for i, j in pages: + start = offsets[i] + end = offsets[j - 1] + len(paras[j - 1]["ctx"]) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + continue + lines = layout_page(paras[i:j], col_w) + key = sha8(cond, json.dumps([(p["title"], p["ctx"]) for p in paras[i:j]]), str(args.size)) + png = CACHE / f"{EXP}-doc-{variant}-{key}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render_doc(lines, args.size, variant, CACHE), png) + tasks.append((args.model, cond, prompt, png, questions, length, start, ctx)) + else: + cap_stats[length] = {"pages": -(-len(flow) // grid_cap), "mean_chars_page": grid_cap} + prompt = load_prompt("qa-image.md").format(cols=cols, rows=rows) + for start in range(0, len(flow), grid_cap): + end = min(start + grid_cap, len(flow)) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + continue + png = CACHE / f"{EXP}-grid-{variant}-{sha8(flow[start:end], str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render(flow[start:end], FONT, CACHE, args.size, variant), png) + tasks.append((args.model, cond, prompt, png, questions, length, start, ctx)) + print(f" len {length}: {cap_stats[length]['pages']} pages, " + f"mean {cap_stats[length]['mean_chars_page']} chars/page, corpus {len(flow)}", flush=True) + + if args.render_only: + print(f"sample: {tasks[0][3]}" if tasks else "no tasks") + return + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(qa_unit, *t) for t in tasks] + for done, fut in enumerate(futures, 1): + records.extend(fut.result()) + print(f" {done}/{len(futures)}", flush=True) + + s = slug(args.model) + with (OUT_DIR / f"records-{s}.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + cells = aggregate(records, price_in, price_out) + hdr = "model,length,condition,n,em,f1,f1_se,abstained,tok_in,tok_out,tok_cache_r,tok_reasoning,cost_usd" + with (OUT_DIR / f"matrix-{s}.csv").open("w") as fh: + fh.write(hdr + "\n") + for c in cells: + fh.write(",".join(str(c[k]) for k in hdr.split(",")) + "\n") + (OUT_DIR / f"summary-{s}.json").write_text(json.dumps({"args": vars(args), "capacity": cap_stats, "cells": cells}, indent=1)) + for c in cells: + print(f"len {c['length']:<4} {c['condition']:<24} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f} " + f"out={c['tok_out']} rsn={c['tok_reasoning']}", flush=True) + print(f"-> {OUT_DIR}/records-{s}.jsonl", flush=True) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp21_braille.py b/packages/snapcompact/research/exp21_braille.py new file mode 100644 index 000000000..d553c7be0 --- /dev/null +++ b/packages/snapcompact/research/exp21_braille.py @@ -0,0 +1,276 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp21: render text as Grade-1 (uncontracted) 6-dot braille -> gemini-3.5-flash. + +Each character becomes a 2x3 dot matrix drawn directly (no font file): +lowercase letters, digits as number-sign + a-j (one sign per digit run), +a punctuation subset; everything else -> blank cell. Text is lowercased +(SQuAD scoring is case-insensitive; real braille capital signs would +waste cells). + +Conditions (gemini only, 1568px, bw): + img-braille-5x7 1px dots, cell 5x7 -> 313x224 = 70,112 cells/page + img-braille-7x10 2px dots, cell 7x10 -> 224x156 = 34,944 cells/page + +Usage: + uv run exp21_braille.py --render-only + uv run exp21_braille.py # lengths 50,150 +""" + +import argparse +import json +import os +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp21" +OUT_DIR = RESULTS / f"{EXP}-braille" +MODELS = { # model -> (price_in, price_out, key_name, key_env) + "google/gemini-3.5-flash": (0.6, 4.0, "openrouter", "OPENROUTER_API_KEY"), + "gpt-5.5": (2.0, 16.0, "openai", "OPENAI_API_KEY"), +} +MODEL = "google/gemini-3.5-flash" +PRICE_IN, PRICE_OUT = 0.6, 4.0 +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) + +# dots numbered 1-6: 1=top-left 2=mid-left 3=bottom-left 4=top-right 5=mid-right 6=bottom-right +# bitmask: bit0=dot1 .. bit5=dot6 (matches Unicode U+2800 offsets) +_L = { + "a": 0x01, "b": 0x03, "c": 0x09, "d": 0x19, "e": 0x11, "f": 0x0B, "g": 0x1B, + "h": 0x13, "i": 0x0A, "j": 0x1A, "k": 0x05, "l": 0x07, "m": 0x0D, "n": 0x1D, + "o": 0x15, "p": 0x0F, "q": 0x1F, "r": 0x17, "s": 0x0E, "t": 0x1E, "u": 0x25, + "v": 0x27, "w": 0x3A, "x": 0x2D, "y": 0x3D, "z": 0x35, +} +_PUNCT = { + ".": 0x32, ",": 0x02, "'": 0x04, "-": 0x24, ":": 0x12, ";": 0x06, + "?": 0x26, "!": 0x16, " ": 0x00, +} +_NUMSIGN = 0x3C # dots 3456 +_DIGIT = {d: _L["abcdefghij"[i]] for i, d in enumerate("1234567890")} + +# cell name -> (dot_px, adv, pitch); dot gap is 1px in both configs +CELLS = { + "5x7": (1, 5, 7), + "7x10": (2, 7, 10), +} + + +def braille_cells(text: str) -> tuple[list[int], list[int]]: + """(cell bitmasks, original char index per cell). Lowercases; digit runs share one number sign.""" + cells, origin = [], [] + in_num = False + for i, ch in enumerate(text): + c = ch.lower() + if c in _DIGIT: + if not in_num: + cells.append(_NUMSIGN) + origin.append(i) + in_num = True + cells.append(_DIGIT[c]) + origin.append(i) + continue + in_num = False + cells.append(_L.get(c, _PUNCT.get(c, 0x00))) + origin.append(i) + return cells, origin + + +def render_braille(cells: list[int], cell_name: str, size: int) -> Image.Image: + dpx, adv, pitch = CELLS[cell_name] + cols, rows = size // adv, size // pitch + step = dpx + 1 # dot pitch inside the cell + img = Image.new("RGB", (size, size), _WHITE) + px = img.load() + for idx, mask in enumerate(cells[: cols * rows]): + if not mask: + continue + row, col = divmod(idx, cols) + x0, y0 = col * adv, row * pitch + for bit in range(6): + if not mask & (1 << bit): + continue + dc, dr = divmod(bit, 3) # dots 1-3 left column, 4-6 right column + dx, dy = x0 + dc * step, y0 + dr * step + for yy in range(dy, dy + dpx): + for xx in range(dx, dx + dpx): + if xx < size and yy < size: + px[xx, yy] = _BLACK + return img + + +def atomic_save(img: Image.Image, png: Path) -> None: + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") + img.save(tmp) + tmp.replace(png) + + +def cached(payload: object, fn, fresh: bool) -> dict: + key = sha8(MODEL, f"{EXP}-qa", json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if not fresh and path.exists(): + return json.loads(path.read_text()) + out = fn() + if out.get("stop") != "max_tokens": + tmp = path.with_suffix(".tmp.json") + tmp.write_text(json.dumps(out)) + tmp.replace(path) + else: + print("WARN truncated response, not cached", flush=True) + return out + + +def qa_unit(cond: str, prompt: str, png: Path, questions: list[dict], length: int, start: int, ctx: dict) -> list[dict]: + args, keys = ctx["args"], ctx["keys"] + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [{"role": "user", "content": [{"text": prompt}, {"image_path": png}, {"text": q_block}]}] + payload = {"messages": messages} + if args.effort: + payload["effort"] = args.effort + qa = cached( + payload, + lambda: dict(zip(("text", "usage", "stop"), + llm_complete(keys, MODEL, messages, max_tokens=args.max_tokens, effort=args.effort))), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append({ + "model": MODEL, "length": length, "cond": cond, "chunk": start, + "pos_rel": q["pos_rel"], "q": q["q"], "answer": a, "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + }) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict]) -> list[dict]: + cells = {} + for r in records: + cells.setdefault((r["length"], r["cond"]), []).append(r) + out = [] + for (length, cond), recs in sorted(cells.items()): + n = len(recs) + f1s = [r["f1"] for r in recs] + mean = sum(f1s) / n + var = sum((x - mean) ** 2 for x in f1s) / (n - 1) if n > 1 else 0.0 + usage = [u for r in recs for u in r.get("usage", [])] + tin = sum(u["in"] for u in usage) + tout = sum(u["out"] for u in usage) + creads = sum(u.get("cache_r", 0) for u in usage) + rsn = sum(u.get("reasoning", 0) for u in usage) + cost = (tin + 0.1 * creads) * PRICE_IN / 1e6 + tout * PRICE_OUT / 1e6 + out.append({ + "model": MODEL, "length": length, "condition": cond, "n": n, + "em": round(sum(r["em"] for r in recs) / n, 4), "f1": round(mean, 4), + "f1_se": round((var / n) ** 0.5, 4), "abstained": sum(r["abstained"] for r in recs), + "tok_in": tin, "tok_out": tout, "tok_cache_r": creads, "tok_reasoning": rsn, + "cost_usd": round(cost, 4), + }) + return out + + +def main() -> None: + global MODEL, PRICE_IN, PRICE_OUT + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="google/gemini-3.5-flash", choices=sorted(MODELS)) + ap.add_argument("--cells", default="5x7,7x10") + ap.add_argument("--lengths", default="50,150") + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true") + ap.add_argument("--env", default="~/.env") + ap.add_argument("--effort", default=None) + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + MODEL = args.model + PRICE_IN, PRICE_OUT, key_name, key_env = MODELS[MODEL] + keys = {} if args.render_only else {key_name: load_env_key(key_env, args.env)} + all_paras = squad.load_paragraphs(CACHE) + prompt_tpl = load_prompt("exp21-qa-braille.md") + + tasks = [] + for cell_name in args.cells.split(","): + dpx, adv, pitch = CELLS[cell_name] + cols, rows = args.size // adv, args.size // pitch + cap = cols * rows + cond = f"img-braille-{cell_name}" + (f"+eff-{args.effort}" if args.effort else "") + for length in (int(x) for x in args.lengths.split(",")): + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + cells, origin = braille_cells(flow) + pages = [] + i = 0 + while i < len(cells): + j = min(i + cap, len(cells)) + pages.append((i, j)) + i = j + print(f"{cond} len {length}: {len(pages)} pages, {cap} cells/page " + f"({cols}x{rows}), {len(cells)} cells for {len(flow)} chars", flush=True) + ctx = {"args": args, "keys": keys} + for ci, cj in pages: + start = origin[ci] + end = origin[cj - 1] + 1 + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + continue + png = CACHE / f"{EXP}-{cell_name}-{sha8(flow[start:end], cell_name, str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render_braille(cells[ci:cj], cell_name, args.size), png) + prompt = prompt_tpl.format(cols=cols, rows=rows) + tasks.append((cond, prompt, png, questions, length, start, ctx)) + + if args.render_only: + for t in tasks[:2]: + print(f"sample: {t[2]}") + return + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(qa_unit, *t) for t in tasks] + for done, fut in enumerate(futures, 1): + records.extend(fut.result()) + print(f" {done}/{len(futures)}", flush=True) + + slug = MODEL.split("/")[-1] + (f"-eff{args.effort}" if args.effort else "") + with (OUT_DIR / f"records-{slug}.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + cells_out = aggregate(records) + hdr = "model,length,condition,n,em,f1,f1_se,abstained,tok_in,tok_out,tok_cache_r,tok_reasoning,cost_usd" + with (OUT_DIR / f"matrix-{slug}.csv").open("w") as fh: + fh.write(hdr + "\n") + for c in cells_out: + fh.write(",".join(str(c[k]) for k in hdr.split(",")) + "\n") + (OUT_DIR / f"summary-{slug}.json").write_text(json.dumps({"args": vars(args), "cells": cells_out}, indent=1)) + for c in cells_out: + print(f"len {c['length']:<4} {c['condition']:<20} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f} " + f"out={c['tok_out']} rsn={c['tok_reasoning']}", flush=True) + print(f"-> {OUT_DIR}/matrix-{slug}.csv", flush=True) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/exp22_ttf6pt.py b/packages/snapcompact/research/exp22_ttf6pt.py new file mode 100644 index 000000000..6d95db910 --- /dev/null +++ b/packages/snapcompact/research/exp22_ttf6pt.py @@ -0,0 +1,243 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""exp22: 6pt antialiased vector type vs 1-bit bitmap fonts. + +Every prior condition used hand-hinted 1-bit bitmap fonts. Here the corpus is +rasterized with a real monospace TTF (Monaco, fallback DejaVu Sans Mono) at +tiny em sizes WITH greyscale antialiasing — the hypothesis being that AA +preserves sub-pixel shape information a VLM can exploit below the bitmap-font +legibility floor. + +Conditions (bw, grid layout, 1568px): + img-ttf6-bw em 6px (6pt @ 72dpi) + img-ttf8-bw em 8px (6pt @ 96dpi) +Cell metrics are measured from the font (advance x line height). + +Usage: + uv run exp22_ttf6pt.py --render-only + uv run exp22_ttf6pt.py --model gpt-5.5 + uv run exp22_ttf6pt.py --model google/gemini-3.5-flash +""" + +import argparse +import json +import math +import os +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +from PIL import Image, ImageDraw, ImageFont + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, QA_CACHE, RESULTS, load_prompt, sha8 # noqa: E402 + +EXP = "exp22" +OUT_DIR = RESULTS / f"{EXP}-ttf6pt" +MODELS = { # model -> (price_in, price_out, key_name, key_env) + "google/gemini-3.5-flash": (0.6, 4.0, "openrouter", "OPENROUTER_API_KEY"), + "gpt-5.5": (2.0, 16.0, "openai", "OPENAI_API_KEY"), +} +MODEL = "google/gemini-3.5-flash" +PRICE_IN, PRICE_OUT = 0.6, 4.0 +_WHITE = (255, 255, 255) +_BLACK = (0, 0, 0) +TTF_CANDIDATES = [ + "/System/Library/Fonts/Monaco.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf", +] +EM_SIZES = (6, 8) + + +def mono_font(em: int) -> ImageFont.FreeTypeFont: + for path in TTF_CANDIDATES: + if Path(path).exists(): + return ImageFont.truetype(path, em) + raise FileNotFoundError("no monospace TTF found") + + +def metrics(em: int) -> tuple[float, int, int, int]: + """(advance, pitch, cols, rows) at 1568px for the em size.""" + f = mono_font(em) + adv = f.getlength("0") + ascent, descent = f.getmetrics() + pitch = ascent + descent # tight leading; AA keeps rows separable + cols = int(1568 // adv) + rows = 1568 // pitch + return adv, pitch, cols, rows + + +def render_ttf(text: str, em: int, size: int) -> Image.Image: + f = mono_font(em) + adv, pitch, cols, rows = metrics(em) + img = Image.new("RGB", (size, size), _WHITE) + draw = ImageDraw.Draw(img) + for r in range(rows): + line = text[r * cols : (r + 1) * cols] + if not line: + break + draw.text((0, r * pitch), line, font=f, fill=_BLACK) + return img + + +def atomic_save(img: Image.Image, png: Path) -> None: + tmp = png.with_suffix(f".{os.getpid()}.tmp.png") + img.save(tmp) + tmp.replace(png) + + +def cached(payload: object, fn, fresh: bool) -> dict: + key = sha8(MODEL, f"{EXP}-qa", json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if not fresh and path.exists(): + return json.loads(path.read_text()) + out = fn() + if out.get("stop") != "max_tokens": + tmp = path.with_suffix(".tmp.json") + tmp.write_text(json.dumps(out)) + tmp.replace(path) + else: + print("WARN truncated response, not cached", flush=True) + return out + + +def qa_unit(cond: str, prompt: str, png: Path, questions: list[dict], length: int, start: int, ctx: dict) -> list[dict]: + args, keys = ctx["args"], ctx["keys"] + q_block = "\n".join(f"{k + 1}. {q['q']}" for k, q in enumerate(questions)) + messages = [{"role": "user", "content": [{"text": prompt}, {"image_path": png}, {"text": q_block}]}] + payload = {"messages": messages} + if args.effort: + payload["effort"] = args.effort + qa = cached( + payload, + lambda: dict(zip(("text", "usage", "stop"), + llm_complete(keys, MODEL, messages, max_tokens=args.max_tokens, effort=args.effort))), + args.fresh, + ) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append({ + "model": MODEL, "length": length, "cond": cond, "chunk": start, + "pos_rel": q["pos_rel"], "q": q["q"], "answer": a, "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + }) + records[0]["usage"] = [{"phase": "qa", **qa["usage"]}] + return records + + +def aggregate(records: list[dict]) -> list[dict]: + cells = {} + for r in records: + cells.setdefault((r["length"], r["cond"]), []).append(r) + out = [] + for (length, cond), recs in sorted(cells.items()): + n = len(recs) + f1s = [r["f1"] for r in recs] + mean = sum(f1s) / n + var = sum((x - mean) ** 2 for x in f1s) / (n - 1) if n > 1 else 0.0 + usage = [u for r in recs for u in r.get("usage", [])] + tin = sum(u["in"] for u in usage) + tout = sum(u["out"] for u in usage) + creads = sum(u.get("cache_r", 0) for u in usage) + rsn = sum(u.get("reasoning", 0) for u in usage) + cost = (tin + 0.1 * creads) * PRICE_IN / 1e6 + tout * PRICE_OUT / 1e6 + out.append({ + "model": MODEL, "length": length, "condition": cond, "n": n, + "em": round(sum(r["em"] for r in recs) / n, 4), "f1": round(mean, 4), + "f1_se": round((var / n) ** 0.5, 4), "abstained": sum(r["abstained"] for r in recs), + "tok_in": tin, "tok_out": tout, "tok_cache_r": creads, "tok_reasoning": rsn, + "cost_usd": round(cost, 4), + }) + return out + + +def main() -> None: + global MODEL, PRICE_IN, PRICE_OUT + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="google/gemini-3.5-flash", choices=sorted(MODELS)) + ap.add_argument("--ems", default="6,8") + ap.add_argument("--lengths", default="50,150") + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=3) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--render-only", action="store_true") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + OUT_DIR.mkdir(parents=True, exist_ok=True) + + MODEL = args.model + PRICE_IN, PRICE_OUT, key_name, key_env = MODELS[MODEL] + keys = {} if args.render_only else {key_name: load_env_key(key_env, args.env)} + all_paras = squad.load_paragraphs(CACHE) + prompt_tpl = load_prompt("qa-image.md") + + tasks = [] + for em in (int(x) for x in args.ems.split(",")): + adv, pitch, cols, rows = metrics(em) + cap = cols * rows + cond = f"img-ttf{em}-bw" + (f"+eff-{args.effort}" if args.effort else "") + print(f"{cond}: adv {adv:.2f}px pitch {pitch}px -> {cols}x{rows} = {cap} chars/page", flush=True) + for length in (int(x) for x in args.lengths.split(",")): + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + n_pages = math.ceil(len(flow) / cap) + print(f" len {length}: {n_pages} pages, corpus {len(flow)}", flush=True) + ctx = {"args": args, "keys": keys} + for start in range(0, len(flow), cap): + end = min(start + cap, len(flow)) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + continue + png = CACHE / f"{EXP}-ttf{em}-{sha8(flow[start:end], str(em), str(args.size))}.png" + if not png.exists() or png.stat().st_size == 0: + atomic_save(render_ttf(flow[start:end], em, args.size), png) + prompt = prompt_tpl.format(cols=cols, rows=rows) + tasks.append((cond, prompt, png, questions, length, start, ctx)) + + if args.render_only: + for t in tasks[:2]: + print(f"sample: {t[2]}") + return + + records: list[dict] = [] + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(qa_unit, *t) for t in tasks] + for done, fut in enumerate(futures, 1): + records.extend(fut.result()) + print(f" {done}/{len(futures)}", flush=True) + + slug = MODEL.split("/")[-1] + (f"-eff{args.effort}" if args.effort else "") + with (OUT_DIR / f"records-{slug}.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + cells_out = aggregate(records) + hdr = "model,length,condition,n,em,f1,f1_se,abstained,tok_in,tok_out,tok_cache_r,tok_reasoning,cost_usd" + with (OUT_DIR / f"matrix-{slug}.csv").open("w") as fh: + fh.write(hdr + "\n") + for c in cells_out: + fh.write(",".join(str(c[k]) for k in hdr.split(",")) + "\n") + (OUT_DIR / f"summary-{slug}.json").write_text(json.dumps({"args": vars(args), "cells": cells_out}, indent=1)) + for c in cells_out: + print(f"len {c['length']:<4} {c['condition']:<20} n={c['n']:<4} EM {c['em']:.3f} " + f"F1 {c['f1']:.3f} ±{c['f1_se']:.3f} ${c['cost_usd']:.3f} " + f"out={c['tok_out']} rsn={c['tok_reasoning']}", flush=True) + print(f"-> {OUT_DIR}/matrix-{slug}.csv", flush=True) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/final.py b/packages/snapcompact/research/final.py new file mode 100644 index 000000000..4df51f1dc --- /dev/null +++ b/packages/snapcompact/research/final.py @@ -0,0 +1,331 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""snapcompact final: the blog-post dataset. One command, full grid, CSV out. + +Grid: lengths x models x techniques, SQuAD v1.1 dev QA recall (official EM/F1). + + lengths 50 / 150 / 250 passages (~30k / ~102k / ~170k chars) + models claude-fable-5, claude-opus-4-8 (Anthropic Messages API) + gpt-5.5 (OpenAI Responses API) + techniques text plain-text chunks (ceiling) + handoff agent handoff document, QA on the doc + compact Anthropic: agent compaction-summary prompt + OpenAI: remote /responses/compact window + img-{6x10,5x8}-{sent,bw} one 1568x1568 PNG per chunk + +Outputs in results/final/: records.jsonl (per question), matrix.csv (per cell), +summary.json. Responses are cached by payload hash: interrupted or re-scoped +runs only bill new cells. `--report` reprints from cache without API calls. + +Usage: uv run final.py [--models ...] [--lengths 50,150,250] [--conditions ...] +Keys: ANTHROPIC_API_KEY + OPENAI_API_KEY from ~/.env. +""" + +import argparse +import csv +import json +import re +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from providers import is_openai, llm_complete, load_env_key, openai_compact # noqa: E402 +from run import CACHE, FONTS, QA_CACHE, RESULTS, TEXT_CHUNK, agent_prompt, load_prompt, sha8 # noqa: E402 + +# (family display, $/M input, $/M output). Cached reads bill at 0.1x input, +# Anthropic cache writes at 1.25x. Edit prices here; `--report` recomputes. +MODELS = { + "claude-fable-5": (10.0, 50.0), + "claude-opus-4-8": (15.0, 75.0), + "gpt-5.5": (2.0, 16.0), + "google/gemini-3.5-flash": (0.6, 4.0), + "moonshotai/kimi-k2.6": (0.68, 3.41), + "z-ai/glm-4.6v": (0.30, 0.90), +} +LENGTHS = (50, 150, 250) +CONDITIONS = ("text", "handoff", "compact", "img-6x10-sent", "img-6x10-bw", "img-5x8-sent", "img-5x8-bw") +ACK = "Noted. I have read the passages and will keep them in mind." + + +def cached(model: str, tag: str, payload: object, fn, fresh: bool) -> dict: + """Disk-cache `fn() -> dict` keyed by (model, tag, payload). Truncated/empty outputs are not cached.""" + key = sha8(model, tag, json.dumps(payload, sort_keys=True, default=str)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens" and ("text" not in hit or hit["text"]): + return hit + out = fn() + if out.get("stop") == "max_tokens" or out.get("text") == "": + print(f" WARN truncated/empty, not cached: {model} {tag} {key}") + else: + path.write_text(json.dumps(out)) + return out + + +def parse_img_condition(name: str) -> tuple[str, str, int] | None: + if not name.startswith("img-"): + return None + _, font, variant = name.split("-", 2) + columns = 1 + m = re.match(r"(.+)-(\d+)col$", variant) + if m: + variant, columns = m.group(1), int(m.group(2)) + return font, variant, columns + + +def chunk_budget(cond: str, size: int) -> int: + img = parse_img_condition(cond) + return capacity(FONTS[img[0]], size, img[2])[2] if img else TEXT_CHUNK + + +def session_frame(chunk_text: str) -> list[dict]: + return [ + {"role": "user", "content": [{"text": load_prompt("session-frame.md").format(context=chunk_text)}]}, + {"role": "assistant", "content": [{"text": ACK}]}, + ] + + +def run_cell_chunk(model: str, cond: str, start: int, end: int, ctx: dict) -> list[dict]: + """One (model, condition, chunk) unit: build carrier, QA, score.""" + args, flow, paras, offsets, keys = ctx["args"], ctx["flow"], ctx["paras"], ctx["offsets"], ctx["keys"] + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(questions)) + usage_rows: list[tuple[str, dict]] = [] + img = parse_img_condition(cond) + extra_items: list[dict] | None = None + + if img: + font, variant, columns = img + tag = f"{font}-{variant}" if columns == 1 else f"{font}-{variant}-{columns}col" + # "dimv2" salts pure-dim renders: pre-fix PNGs (sticky-fg bug, glyphs after a + # row's first stopword all dimmed) and the QA cache entries keyed on their paths. + salt = ("dimv2",) if variant == "dim" else () + png = CACHE / f"img-{tag}-{sha8(chunk_text, str(args.size), *salt)}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{uuid.uuid4().hex[:8]}.tmp.png") + render(chunk_text, FONTS[font], CACHE, args.size, variant, columns=columns).save(tmp) + tmp.replace(png) + cols, rows, _ = capacity(FONTS[font], args.size, columns) + preamble = ( + load_prompt("qa-image-cols.md").format(cols=cols, rows=rows, columns=columns) + if columns > 1 + else load_prompt("qa-image.md").format(cols=cols, rows=rows) + ) + messages = [ + { + "role": "user", + "content": [ + {"text": preamble}, + {"image_path": png}, + {"text": q_block}, + ], + } + ] + elif cond == "compact" and is_openai(model): + comp = cached( + model, "remote-compact", {"chunk": chunk_text}, + lambda: dict(zip(("items", "usage"), openai_compact(keys["openai"], model, session_frame(chunk_text)))), + args.fresh, + ) + usage_rows.append(("compact", comp["usage"])) + extra_items = comp["items"] + messages = [ + {"role": "user", "content": [{"text": load_prompt("qa-remote-compact.md").format(questions=q_block)}]} + ] + elif cond in ("compact", "handoff"): + prompt_file = {"compact": "compaction-summary.md", "handoff": "handoff-document.md"}[cond] + gen = cached( + model, f"summary-{cond}", {"chunk": chunk_text}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete( + keys, model, + session_frame(chunk_text) + [{"role": "user", "content": [{"text": agent_prompt(prompt_file)}]}], + system=agent_prompt("summarization-system.md"), + max_tokens=args.max_tokens, + ), + ) + ), + args.fresh, + ) + usage_rows.append(("summarize", gen["usage"])) + messages = [ + { + "role": "user", + "content": [{"text": load_prompt("qa-text.md").format(context=gen["text"])}, {"text": q_block}], + } + ] + else: # text + messages = [ + { + "role": "user", + "content": [{"text": load_prompt("qa-text.md").format(context=chunk_text)}, {"text": q_block}], + } + ] + + qa = cached( + model, "qa", {"messages": messages, "extra": extra_items, "effort": args.effort}, + lambda: dict( + zip( + ("text", "usage", "stop"), + llm_complete( + keys, model, messages, + max_tokens=args.max_tokens, effort=args.effort, extra_input_items=extra_items, + ), + ) + ), + args.fresh, + ) + usage_rows.append(("qa", qa["usage"])) + answers = squad.parse_numbered(qa["text"], len(questions)) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "model": model, + "length": ctx["length"], + "cond": cond, + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + records[0]["usage"] = [{"phase": p, **u} for p, u in usage_rows] + return records + + +def aggregate(records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + us = [u for r in records if "usage" in r for u in r["usage"]] + tok = {k: sum(u.get(k, 0) for u in us) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + cost_in = (tok["in"] + 1.25 * tok["cache_w"] + 0.1 * tok["cache_r"]) / 1e6 * price_in + cost_out = tok["out"] / 1e6 * price_out + return { + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + **{f"tok_{k}": v for k, v in tok.items()}, + "cost_in_usd": round(cost_in, 4), + "cost_out_usd": round(cost_out, 4), + "cost_usd": round(cost_in + cost_out, 4), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--models", default=",".join(MODELS)) + ap.add_argument("--lengths", default=",".join(map(str, LENGTHS))) + ap.add_argument("--conditions", default=",".join(CONDITIONS)) + ap.add_argument("--qpc", type=int, default=30) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=6) + ap.add_argument("--max-tokens", type=int, default=16384) + ap.add_argument("--effort", default=None, help="reasoning effort; None = provider default") + ap.add_argument("--fresh", action="store_true") + ap.add_argument("--report", action="store_true", help="reprint from cache only") + ap.add_argument("--env", default="~/.env") + ap.add_argument("--out", default="final", help="results subdirectory (isolate concurrent runs)") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + out_dir = RESULTS / args.out + out_dir.mkdir(parents=True, exist_ok=True) + + models = [m.strip() for m in args.models.split(",") if m.strip()] + lengths = [int(x) for x in args.lengths.split(",") if x.strip()] + conditions = [c.strip() for c in args.conditions.split(",") if c.strip()] + for m in models: + if m not in MODELS: + raise SystemExit(f"unknown model {m}; add it to MODELS with prices") + + keys = {} + if not args.report: + keys["anthropic"] = load_env_key("ANTHROPIC_API_KEY", args.env) + keys["openai"] = load_env_key("OPENAI_API_KEY", args.env) + keys["openrouter"] = load_env_key("OPENROUTER_API_KEY", args.env) + + all_paras = squad.load_paragraphs(CACHE) + tasks = [] + for length in lengths: + paras = all_paras[:length] + flow, offsets = squad.build_flow(paras) + ctx = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "keys": keys, "length": length} + for model in models: + for cond in conditions: + budget = chunk_budget(cond, args.size) + for start in range(0, len(flow), budget): + tasks.append((model, cond, start, min(start + budget, len(flow)), ctx)) + print(f"grid: {len(models)} models x {len(lengths)} lengths x {len(conditions)} conditions = {len(tasks)} chunk tasks") + + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_cell_chunk, m, c, s, e, ctx) for m, c, s, e, ctx in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + if done % 25 == 0: + print(f" {done}/{len(tasks)} tasks", flush=True) + + with (out_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + cells = [] + for model in models: + for length in lengths: + for cond in conditions: + sub = [r for r in records if r["model"] == model and r["length"] == length and r["cond"] == cond] + if not sub: + continue + cells.append({"model": model, "length": length, "condition": cond, **aggregate(sub, *MODELS[model])}) + (out_dir / "summary.json").write_text(json.dumps({"args": vars(args), "cells": cells}, indent=1)) + with (out_dir / "matrix.csv").open("w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=list(cells[0].keys())) + writer.writeheader() + writer.writerows(cells) + + for length in lengths: + print(f"\n== {length} passages == (F1 / $carry-in / $decode-out)") + hdr = f"{'condition':<15}" + "".join(f"{m:>22}" for m in models) + print(hdr + "\n" + "-" * len(hdr)) + for cond in conditions: + row = f"{cond:<15}" + for model in models: + cell = next((c for c in cells if c["model"] == model and c["length"] == length and c["condition"] == cond), None) + row += ( + f"{cell['f1']:>10.3f} {cell['cost_in_usd']:>5.2f} {cell['cost_out_usd']:>5.2f}" + if cell + else f"{'-':>22}" + ) + print(row) + print(f"\ndataset -> {out_dir}/records.jsonl, matrix.csv, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/mono.py b/packages/snapcompact/research/mono.py new file mode 100644 index 000000000..db8071e63 --- /dev/null +++ b/packages/snapcompact/research/mono.py @@ -0,0 +1,154 @@ +"""Monolithic long-context probe: the WHOLE corpus in ONE request. + +final.py chunks every condition into ~10k-token QA calls, so it never tests +true in-request long-context retrieval. This runner stuffs an N-char SQuAD +flow (e.g. 800k chars ~ 200k text tokens) into a single request — either as +raw text or as a stack of dense-font images — with questions sampled evenly +across the whole span. Reports overall EM/F1 plus F1 by position quartile +(real lost-in-the-middle, which the chunked harness cannot see). + + uv run --with pillow python mono.py --model gpt-5.5 --chars 800000 \ + --conditions text,img-6x10-sent,img-6x8s-sent,img-8x8u-sent +""" + +import argparse +import json +import sys +import uuid +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from final import MODELS, cached, parse_img_condition # noqa: E402 +from providers import llm_complete, load_env_key # noqa: E402 +from run import CACHE, FONTS, RESULTS, load_prompt, sha8 # noqa: E402 + + +def build_content(cond: str, flow: str, size: int) -> tuple[list[dict], int]: + """Context blocks for the whole flow (questions appended per batch); returns (blocks, n_images).""" + img = parse_img_condition(cond) + if not img: + assert cond == "text", f"unsupported mono condition {cond!r}" + return [{"text": load_prompt("qa-text.md").format(context=flow), "cache": True}], 0 + font, variant, columns = img + cfg = FONTS[font] + cols, rows, cap = capacity(cfg, size, columns) + salt = ("dimv2",) if variant == "dim" else () + tag = f"{font}-{variant}" if columns == 1 else f"{font}-{variant}-{columns}col" + pngs = [] + for start in range(0, len(flow), cap): + chunk = flow[start : start + cap] + png = CACHE / f"img-{tag}-{sha8(chunk, str(size), *salt)}.png" + if not png.exists() or png.stat().st_size == 0: + tmp = png.with_suffix(f".{uuid.uuid4().hex[:8]}.tmp.png") + render(chunk, cfg, CACHE, size, variant, columns=columns).save(tmp) + tmp.replace(png) + pngs.append(png) + preamble = load_prompt("qa-image-multi.md").format(k=len(pngs), cols=cols, rows=rows) + if cfg.repeat > 1: + preamble += ( + f"\nNote: every text line is rendered {cfg.repeat} times consecutively - first on the plain " + "background, then repeated on a pale highlight band. The copies show identical characters; " + "cross-check between them when a glyph is hard to read, and do not treat copies as separate text." + ) + blocks = [{"text": preamble}, *({"image_path": p} for p in pngs), {"text": "End of images.", "cache": True}] + return blocks, len(pngs) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="gpt-5.5") + ap.add_argument("--chars", type=int, default=800_000) + ap.add_argument("--conditions", default="text,img-6x10-sent,img-6x8s-sent,img-8x8u-sent") + ap.add_argument("--questions", type=int, default=50, help="total questions sampled across the flow") + ap.add_argument("--qpb", type=int, default=5, help="questions per API call (context re-sent, prefix-cached)") + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--max-tokens", type=int, default=32768) + ap.add_argument("--effort", default=None) + ap.add_argument("--env", default="~/.env") + ap.add_argument("--out", default="mono") + ap.add_argument("--fresh", action="store_true") + args = ap.parse_args() + + keys = { + "anthropic": load_env_key("ANTHROPIC_API_KEY", args.env), + "openai": load_env_key("OPENAI_API_KEY", args.env), + "openrouter": load_env_key("OPENROUTER_API_KEY", args.env), + } + paras = squad.load_paragraphs(CACHE) + flow, offsets = squad.build_flow(paras, args.chars) + questions = squad.sample_chunk_questions(paras, offsets, 0, len(flow), args.questions, args.seed) + price_in, price_out = MODELS[args.model] + print( + f"flow: {len(flow):,} chars (~{len(flow) // 4 // 1000}k text tokens), " + f"{len(questions)} questions in batches of {args.qpb}" + ) + + out_dir = RESULTS / args.out + out_dir.mkdir(parents=True, exist_ok=True) + records, table = [], [] + for cond in [c.strip() for c in args.conditions.split(",") if c.strip()]: + ctx_blocks, n_imgs = build_content(cond, flow, args.size) + answers, usages, stops = [], [], [] + for b in range(0, len(questions), args.qpb): + batch = questions[b : b + args.qpb] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(batch)) + messages = [{"role": "user", "content": [*ctx_blocks, {"text": q_block}]}] + qa = cached( + args.model, "qa-mono", {"messages": messages, "effort": args.effort}, + lambda m=messages: dict( + zip( + ("text", "usage", "stop"), + llm_complete(keys, args.model, m, max_tokens=args.max_tokens, effort=args.effort), + ) + ), + args.fresh, + ) + answers.extend(squad.parse_numbered(qa["text"], len(batch))) + usages.append(qa["usage"]) + stops.append(qa["stop"]) + rows = [ + { + "model": args.model, "cond": cond, "pos_rel": q["pos_rel"], "q": q["q"], + "answer": a, "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + for q, a in zip(questions, answers) + ] + records.extend({**r, "usage": usages} if i == 0 else r for i, r in enumerate(rows)) + u = {k: sum(x[k] for x in usages) for k in ("in", "out", "cache_w", "cache_r", "reasoning")} + stop = next((s for s in stops if s == "max_tokens"), stops[-1] if stops else "") + cost = (u["in"] + 1.25 * u["cache_w"] + 0.1 * u["cache_r"]) / 1e6 * price_in + u["out"] / 1e6 * price_out + quart = [] + for lo, hi in ((0, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.01)): + qs = [r["f1"] for r in rows if lo <= r["pos_rel"] < hi] + quart.append(sum(qs) / len(qs) if qs else float("nan")) + table.append( + { + "cond": cond, "n": len(rows), "imgs": n_imgs, + "em": sum(r["em"] for r in rows) / len(rows), + "f1": sum(r["f1"] for r in rows) / len(rows), + "abst": sum(r["abstained"] for r in rows), + "tok_in": u["in"], "tok_cached": u["cache_r"], "tok_out": u["out"], "reas": u["reasoning"], + "cost": cost, "stop": stop, "q1": quart[0], "q2": quart[1], "q3": quart[2], "q4": quart[3], + } + ) + t = table[-1] + print( + f"{cond:<18} imgs={t['imgs']:>2} f1={t['f1']:.3f} em={t['em']:.3f} abst={t['abst']:>2} " + f"in={t['tok_in']:>7} cached={t['tok_cached']:>7} out={t['tok_out']:>6} reas={t['reas']:>6} " + f"${t['cost']:.2f} stop={t['stop']}" + ) + print(f"{'':<18} F1 by position quartile: " + " ".join(f"q{i + 1}={v:.3f}" for i, v in enumerate(quart))) + + (out_dir / "records.jsonl").write_text("\n".join(json.dumps(r) for r in records)) + (out_dir / "summary.json").write_text(json.dumps(table, indent=1)) + print(f"\ndataset -> {out_dir}/records.jsonl, summary.json") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/prompts/exp02-qa-image.md b/packages/snapcompact/research/prompts/exp02-qa-image.md new file mode 100644 index 000000000..edaaba915 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp02-qa-image.md @@ -0,0 +1,8 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +{extra} + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase from the text, written in normal full English spelling. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp04-qa-image.md b/packages/snapcompact/research/prompts/exp04-qa-image.md new file mode 100644 index 000000000..6bae75bff --- /dev/null +++ b/packages/snapcompact/research/prompts/exp04-qa-image.md @@ -0,0 +1,8 @@ +The attached image is a page of encyclopedia passages typeset like a newspaper: a monospace pixel font in TWO columns of {col_w} characters, {rows} rows each, separated by a vertical gutter of white space. + +Reading order: read the LEFT column top-to-bottom first, then continue at the top of the RIGHT column and read it top-to-bottom. Article titles appear as bold uppercase headings; paragraphs within an article are separated by blank lines, and text is word-wrapped (no words are split across lines or columns). + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp05-qa-image-ctl.md b/packages/snapcompact/research/prompts/exp05-qa-image-ctl.md new file mode 100644 index 000000000..ecf89b10c --- /dev/null +++ b/packages/snapcompact/research/prompts/exp05-qa-image-ctl.md @@ -0,0 +1,10 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +Questions follow after the image. Do NOT transcribe the image. For each question: +1. Locate the relevant region first (skim for the topic). +2. Read ONLY that narrow band of rows carefully. +3. Answer from what you read there. + +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp05-qa-image.md b/packages/snapcompact/research/prompts/exp05-qa-image.md new file mode 100644 index 000000000..50ab394ab --- /dev/null +++ b/packages/snapcompact/research/prompts/exp05-qa-image.md @@ -0,0 +1,13 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +The left margin of the image is a row-number ruler, not part of the text: every 5th row is labeled with its row index in blue digits (0, 5, 10, ..., {last_label}). Rows are numbered from 0 at the top. Use the ruler to address regions of the image like line numbers in a file. + +Questions follow after the image. Do NOT transcribe the image. For each question: +1. Locate the relevant region first (skim for the topic, use the blue ruler to pin down the row range). +2. Read ONLY that narrow band of rows carefully. +3. Answer from what you read there. + +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one line per question, in exactly this form: `N. | row≈` where is the ruler row where you found the answer. +- No other commentary. diff --git a/packages/snapcompact/research/prompts/exp06-prov-image.md b/packages/snapcompact/research/prompts/exp06-prov-image.md new file mode 100644 index 000000000..dac79c7eb --- /dev/null +++ b/packages/snapcompact/research/prompts/exp06-prov-image.md @@ -0,0 +1,6 @@ +The attached image contains a conversation transcript rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. The transcript interleaves messages from three roles: user, assistant, and tool. {encoding} + +For each numbered question below, do NOT answer the question itself. Instead identify which role's message contains the answer to it. +- Reply with exactly one word per line: user, assistant, or tool. +- You must choose one of the three roles for every question, even if uncertain — never reply UNREADABLE. +- Output a numbered list, one role per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp06-qa-image-tag.md b/packages/snapcompact/research/prompts/exp06-qa-image-tag.md new file mode 100644 index 000000000..12ebfe658 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp06-qa-image-tag.md @@ -0,0 +1,6 @@ +The attached image contains a conversation transcript rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. The transcript interleaves messages from three roles. Each message is preceded by a bracketed role tag: [user], [asst], or [tool]. The tags are part of the rendered text, not message content. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp06-qa-image.md b/packages/snapcompact/research/prompts/exp06-qa-image.md new file mode 100644 index 000000000..354aefb00 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp06-qa-image.md @@ -0,0 +1,6 @@ +The attached image contains a conversation transcript rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. The transcript interleaves messages from three roles. Glyph color encodes the author of each message: dark blue = user, dark green = assistant, dark red = tool. A message boundary is where the glyph color changes. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp07-answer-bands.md b/packages/snapcompact/research/prompts/exp07-answer-bands.md new file mode 100644 index 000000000..bc604ef41 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp07-answer-bands.md @@ -0,0 +1,4 @@ +Now answer the questions. For each question, read ONLY the row band you identified above (plus a row or two of margin). Do NOT transcribe the image or read the rest of it. +- Give short extractive answers: a word or phrase copied from the text in the image. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp07-locate.md b/packages/snapcompact/research/prompts/exp07-locate.md new file mode 100644 index 000000000..9a4846537 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp07-locate.md @@ -0,0 +1,5 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +Questions follow after the image. Do NOT answer them yet, and do NOT transcribe the image. For each question, skim the image coarsely (keywords, names, numbers) and guess which horizontal band of rows most likely contains the answer. +- Output a numbered list, one line per question, each line exactly of the form `rows A-B` (1-indexed, top row is 1; keep each band at most 25 rows tall). +- No answers, no commentary, nothing else. diff --git a/packages/snapcompact/research/prompts/exp07-qa-image.md b/packages/snapcompact/research/prompts/exp07-qa-image.md new file mode 100644 index 000000000..e20c37a9f --- /dev/null +++ b/packages/snapcompact/research/prompts/exp07-qa-image.md @@ -0,0 +1,7 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Do NOT transcribe the image or read it linearly. For each question, locate the region likely to contain the answer, read only that region, and answer. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp08-archive-eager.md b/packages/snapcompact/research/prompts/exp08-archive-eager.md new file mode 100644 index 000000000..d6731714b --- /dev/null +++ b/packages/snapcompact/research/prompts/exp08-archive-eager.md @@ -0,0 +1,10 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Rows are numbered 1 (top) to {rows} (bottom). Original paragraph breaks were collapsed to spaces. + +IMPORTANT: this font is rendered BELOW the size you can read reliably. Characters that look legible are often misread (digits, names, and dates especially). Be skeptical of your own reading. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give a direct short extractive answer ONLY when you are fully certain of every character in it. +- For any question where you are not fully certain — wrong-looking digits, a name you might be misreading, a region that is small or smudged — reply exactly `ZOOM rows A-B` for that question (e.g. `ZOOM rows 41-47`). That row band will be re-rendered for you at a large, comfortable font. Zooming is cheap and encouraged; guessing is penalized. +- Estimate the row range from the vertical position of the region (row 1 = top edge, row {rows} = bottom edge). If unsure of the exact rows, give a generous range — a margin around your estimate is included automatically. +- If you cannot even locate a relevant region, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp08-archive-phrase.md b/packages/snapcompact/research/prompts/exp08-archive-phrase.md new file mode 100644 index 000000000..6f5da8ea8 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp08-archive-phrase.md @@ -0,0 +1,9 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +IMPORTANT: this font is rendered BELOW the size you can read reliably. Characters that look legible are often misread (digits, names, and dates especially). Be skeptical of your own reading. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give a direct short extractive answer ONLY when you are fully certain of every character in it. +- For any question where you are not fully certain — wrong-looking digits, a name you might be misreading, a region that is small or smudged — reply exactly `ZOOM ""` for that question, where is 3-8 consecutive words you can partially make out in or right next to that region (your best-effort reading; a few wrong characters are fine). The matching region will be re-rendered for you at a large, comfortable font. Zooming is cheap and encouraged; guessing is penalized. +- If you cannot even locate a relevant region, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp08-archive.md b/packages/snapcompact/research/prompts/exp08-archive.md new file mode 100644 index 000000000..3042f313c --- /dev/null +++ b/packages/snapcompact/research/prompts/exp08-archive.md @@ -0,0 +1,7 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Rows are numbered 1 (top) to {rows} (bottom). Original paragraph breaks were collapsed to spaces. The font is deliberately tiny: most of it is readable with effort, but some regions may not be. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you locate the relevant region but it is too small or degraded to read reliably, do NOT guess: reply exactly `ZOOM rows A-B` for that question (e.g. `ZOOM rows 41-47`), giving the row range you need re-rendered at higher resolution. +- If you cannot even locate a relevant region, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp08-zoom.md b/packages/snapcompact/research/prompts/exp08-zoom.md new file mode 100644 index 000000000..e71e273f2 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp08-zoom.md @@ -0,0 +1,6 @@ +Below are high-resolution re-renderings of the archive row bands you requested. Each image is preceded by a label giving the archive row range it covers; the text inside is identical to those rows, re-flowed to the new line width. + +Answer your remaining questions, listed after the images, using the zoomed images plus anything you already read. Keep the same question numbers as before. +- Give short extractive answers: a word or phrase copied from the text. +- If you still cannot read the answer, reply exactly UNREADABLE. +- Output a numbered list (original numbering), one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp09-frame.md b/packages/snapcompact/research/prompts/exp09-frame.md new file mode 100644 index 000000000..e8c621c2d --- /dev/null +++ b/packages/snapcompact/research/prompts/exp09-frame.md @@ -0,0 +1,3 @@ +Reference material for this session arrives as bitmap image pages: monospace pixel font, {cols} characters per row, up to {rows} rows per page, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces; the text flows continuously across pages. More pages may be appended later in this session. I will ask questions about the material later. + +Here is page 1. diff --git a/packages/snapcompact/research/prompts/exp09-page.md b/packages/snapcompact/research/prompts/exp09-page.md new file mode 100644 index 000000000..2a3601496 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp09-page.md @@ -0,0 +1 @@ +Here is page {page}, continuing the reference material. diff --git a/packages/snapcompact/research/prompts/exp09-qa.md b/packages/snapcompact/research/prompts/exp09-qa.md new file mode 100644 index 000000000..baf6ff869 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp09-qa.md @@ -0,0 +1,6 @@ +Answer the following questions using ONLY text you can read in the reference pages above. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. + +{questions} diff --git a/packages/snapcompact/research/prompts/exp11-qa-hier.md b/packages/snapcompact/research/prompts/exp11-qa-hier.md new file mode 100644 index 000000000..85ed0bead --- /dev/null +++ b/packages/snapcompact/research/prompts/exp11-qa-hier.md @@ -0,0 +1,12 @@ +You hold a long stream of encyclopedia passages in a three-tier session memory. The stream is ordered oldest to newest and was split into three consecutive age bands, each stored at a different fidelity: + +- TIER 3 (oldest third): available only as a narrative summary that was written earlier{appendix_note}. +- TIER 2 (middle third): rendered as {n_pages} bitmap page image(s) — monospace pixel font, {cols} characters per row, {rows} rows per page, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. +- TIER 1 (newest third): plain verbatim text. + +The tiers appear below in age order (oldest first), each clearly labeled. Questions follow at the end; any question may target material from any tier. + +Answer the questions using ONLY the material provided here. +- Give short extractive answers: a word or phrase copied from the text or read out of an image. For facts that only appear in the TIER 3 summary, copy the best short phrase from the summary. +- If the needed fact is not present in any tier and not legible in any image, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp13-extract.md b/packages/snapcompact/research/prompts/exp13-extract.md new file mode 100644 index 000000000..d5df0a542 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp13-extract.md @@ -0,0 +1,10 @@ +You are compacting the reference passages above for a later question-answering session. Produce an EXTRACTIVE digest: copy out the most load-bearing sentences and spans VERBATIM. + +Rules: +- Copy text exactly as it appears in the passages: no paraphrasing, no summarizing, no rewording, no truncation of sentences. Keep numbers, dates, names, units, capitalization, and spelling character-for-character. +- Preserve the original order of the source; never reorder spans. +- Prefer complete sentences dense with facts: named entities, quantities, dates, definitions, titles, causal claims. Drop filler, transitions, and rhetorical framing. +- Cover the FULL span of the material evenly — do not spend the budget on the earliest passages and starve the later ones. +- Separate non-adjacent extracts with a line containing only: ... +- Hard output budget: at most {budget} characters in total. Use most of it; never exceed it. +- Output ONLY the extracted text (and the ... separators). No headers, no commentary, no explanations. diff --git a/packages/snapcompact/research/prompts/exp19-qa-doc.md b/packages/snapcompact/research/prompts/exp19-qa-doc.md new file mode 100644 index 000000000..1a08492da --- /dev/null +++ b/packages/snapcompact/research/prompts/exp19-qa-doc.md @@ -0,0 +1,8 @@ +The attached image is a page of encyclopedia passages typeset like a newspaper: a monospace pixel font in TWO columns of {col_w} characters, {rows} rows each, separated by a blank vertical gutter. + +Reading order: read the LEFT column top-to-bottom first, then continue at the top of the RIGHT column and read it top-to-bottom. Article titles appear as bold uppercase headings; paragraphs within an article are separated by blank lines, and text is word-wrapped (no words are split across lines or columns). + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/exp21-qa-braille.md b/packages/snapcompact/research/prompts/exp21-qa-braille.md new file mode 100644 index 000000000..b68515ca2 --- /dev/null +++ b/packages/snapcompact/research/prompts/exp21-qa-braille.md @@ -0,0 +1,11 @@ +The attached image contains English text encoded in Grade-1 (uncontracted, letter-by-letter) 6-dot braille, drawn as a dense grid of dot-matrix cells: {cols} cells per row, {rows} rows, read left-to-right, top-to-bottom, wrapping at row ends mid-word (no hyphenation marks). + +Encoding details: +- Standard braille letter patterns a-z. All text is lowercase; there are NO capital signs. +- Digits: the number sign (dots 3456) starts a digit run, then digits use a-j (1=a ... 9=i, 0=j). One number sign covers the whole run. +- Punctuation: period (dots 256), comma (dot 2), apostrophe (dot 3), hyphen (dots 36), colon (dots 25), semicolon (dots 23), question mark (dots 236), exclamation (dots 235). A blank cell is a space. All other symbols were dropped as spaces. + +Questions follow after the image. Decode only the regions you need and answer using ONLY text encoded in the image. +- Give short extractive answers in ordinary English: a word or phrase from the text. +- If you cannot decode the relevant region, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/qa-image-cols.md b/packages/snapcompact/research/prompts/qa-image-cols.md new file mode 100644 index 000000000..2cb675006 --- /dev/null +++ b/packages/snapcompact/research/prompts/qa-image-cols.md @@ -0,0 +1,6 @@ +The attached image contains encyclopedia passages rendered as a bitmap in a newspaper-style layout: monospace pixel font, {columns} columns of text separated by a vertical gutter line. Each column holds {rows} lines of {cols} characters. Reading order: start with the leftmost column, read its lines left-to-right from top to bottom, then continue at the top of the next column. Original paragraph breaks were collapsed to spaces. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/qa-image-multi.md b/packages/snapcompact/research/prompts/qa-image-multi.md new file mode 100644 index 000000000..fe3cc7380 --- /dev/null +++ b/packages/snapcompact/research/prompts/qa-image-multi.md @@ -0,0 +1,6 @@ +The attached {k} images contain encyclopedia passages rendered as dense bitmaps: monospace pixel font, {cols} characters per row, {rows} rows per image. The text flows continuously across the images: read each image left-to-right, top-to-bottom, then continue with the next image in order (image 1 first, image {k} last). Original paragraph breaks were collapsed to spaces. + +Questions follow after the images. Answer them using ONLY text you can read in the images. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/qa-image.md b/packages/snapcompact/research/prompts/qa-image.md new file mode 100644 index 000000000..085061a4b --- /dev/null +++ b/packages/snapcompact/research/prompts/qa-image.md @@ -0,0 +1,6 @@ +The attached image contains encyclopedia passages rendered as a bitmap: monospace pixel font, {cols} characters per row, {rows} rows, read left-to-right then top-to-bottom. Original paragraph breaks were collapsed to spaces. + +Questions follow after the image. Answer them using ONLY text you can read in the image. +- Give short extractive answers: a word or phrase copied from the text. +- If you cannot read the relevant region well enough to answer, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/qa-remote-compact.md b/packages/snapcompact/research/prompts/qa-remote-compact.md new file mode 100644 index 000000000..504c3c9bf --- /dev/null +++ b/packages/snapcompact/research/prompts/qa-remote-compact.md @@ -0,0 +1,8 @@ +Earlier in this conversation you were given a set of reference passages to keep in mind. The conversation context may have been compacted since. + +Answer the questions below using ONLY what you retained from those passages. +- Give short extractive answers: a word or phrase from the passages. +- If the information did not survive in your context, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. + +{questions} diff --git a/packages/snapcompact/research/prompts/qa-text.md b/packages/snapcompact/research/prompts/qa-text.md new file mode 100644 index 000000000..d4606aa1f --- /dev/null +++ b/packages/snapcompact/research/prompts/qa-text.md @@ -0,0 +1,10 @@ +Below is reference material. Questions follow in the next block. + + +{context} + + +Answer the questions using ONLY the reference material above. +- Give short extractive answers: a word or phrase copied from the text. +- If the reference does not contain the information, reply exactly UNREADABLE for that question. +- Output a numbered list, one answer per line, no commentary. diff --git a/packages/snapcompact/research/prompts/session-frame.md b/packages/snapcompact/research/prompts/session-frame.md new file mode 100644 index 000000000..0941b56bd --- /dev/null +++ b/packages/snapcompact/research/prompts/session-frame.md @@ -0,0 +1,5 @@ +Here is a set of reference passages I need you to keep in mind for this session. I will ask questions about them later. + + +{context} + diff --git a/packages/snapcompact/research/providers.py b/packages/snapcompact/research/providers.py new file mode 100644 index 000000000..270cefd5d --- /dev/null +++ b/packages/snapcompact/research/providers.py @@ -0,0 +1,261 @@ +"""Provider-neutral LLM client: Anthropic Messages + OpenAI Responses. + +Neutral message shape: [{"role": str, "content": [block, ...]}] where block is + {"text": str} - text block + {"image_path": Path, "cache": bool} - PNG by path; cache marks the prompt-cache + breakpoint (Anthropic only; OpenAI caches + automatically) + +Normalized usage: {"in", "out", "cache_w", "cache_r", "reasoning"}. +""" + +import base64 +import json +import time +import urllib.error +import urllib.request +from pathlib import Path + +ANTHROPIC_URL = "https://api.anthropic.com/v1/messages" +OPENAI_URL = "https://api.openai.com/v1/responses" +ANTHROPIC_VERSION = "2023-06-01" + + +def load_env_key(var: str, env_path: str = "~/.env") -> str: + """Last assignment wins (mirrors shell sourcing).""" + key = None + for line in Path(env_path).expanduser().read_text().splitlines(): + line = line.strip() + if line.startswith("export "): + line = line[len("export ") :] + if line.startswith(f"{var}="): + key = line.split("=", 1)[1].strip().strip("'\"") + if not key: + raise SystemExit(f"no {var} in {env_path}") + return key + + +def _post(url: str, body: dict, headers: dict, retries: int = 4) -> dict: + payload = json.dumps(body).encode() + req = urllib.request.Request(url, data=payload, headers={"content-type": "application/json", **headers}) + for attempt in range(retries + 1): + try: + with urllib.request.urlopen(req, timeout=600) as resp: + raw = resp.read() + return json.loads(raw) + except urllib.error.HTTPError as err: + detail = err.read().decode(errors="replace")[:500] + if err.code in (408, 429, 500, 502, 503, 529) and attempt < retries: + wait = 2.0 * 2**attempt + print(f" HTTP {err.code}, retrying in {wait:.0f}s: {detail[:120]}") + time.sleep(wait) + continue + raise SystemExit(f"API error {err.code} ({url}): {detail}") from err + except (json.JSONDecodeError, TimeoutError, urllib.error.URLError) as err: + if attempt < retries: + wait = 2.0 * 2**attempt + print(f" bad response ({type(err).__name__}), retrying in {wait:.0f}s") + time.sleep(wait) + continue + raise + raise AssertionError("unreachable") + + +def _png_b64(path: Path) -> str: + return base64.b64encode(path.read_bytes()).decode() + + +# --- Anthropic --- + + +def _anthropic_blocks(blocks: list[dict]) -> list[dict]: + out = [] + for b in blocks: + if "text" in b: + item: dict = {"type": "text", "text": b["text"]} + else: + item = { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": _png_b64(b["image_path"])}, + } + if b.get("cache"): + item["cache_control"] = {"type": "ephemeral"} + out.append(item) + return out + + +def _anthropic_complete( + api_key: str, model: str, messages: list[dict], system: str | None, max_tokens: int, effort: str | None +) -> tuple[str, dict, str]: + body: dict = { + "model": model, + "max_tokens": max_tokens, + "messages": [{"role": m["role"], "content": _anthropic_blocks(m["content"])} for m in messages], + } + if system: + body["system"] = system + headers = {"x-api-key": api_key, "anthropic-version": ANTHROPIC_VERSION} + if effort: + body["output_config"] = {"effort": effort} + headers["anthropic-beta"] = "effort-2025-11-24" + out = _post(ANTHROPIC_URL, body, headers) + text = "".join(b.get("text", "") for b in out["content"] if b.get("type") == "text") + u = out.get("usage", {}) + usage = { + "in": u.get("input_tokens", 0), + "out": u.get("output_tokens", 0), + "cache_w": u.get("cache_creation_input_tokens", 0), + "cache_r": u.get("cache_read_input_tokens", 0), + "reasoning": (u.get("output_tokens_details") or {}).get("thinking_tokens", 0), + } + return text, usage, out.get("stop_reason", "") + + +# --- OpenAI (Responses API) --- + + +def _openai_content(blocks: list[dict], role: str) -> list[dict]: + text_type = "output_text" if role == "assistant" else "input_text" + out = [] + for b in blocks: + if "text" in b: + out.append({"type": text_type, "text": b["text"]}) + else: + out.append( + { + "type": "input_image", + "image_url": f"data:image/png;base64,{_png_b64(b['image_path'])}", + "detail": "original", + } + ) + return out + + +def _openai_usage(out: dict) -> dict: + u = out.get("usage", {}) + cached = (u.get("input_tokens_details") or {}).get("cached_tokens", 0) + return { + "in": u.get("input_tokens", 0) - cached, + "out": u.get("output_tokens", 0), + "cache_w": 0, + "cache_r": cached, + "reasoning": (u.get("output_tokens_details") or {}).get("reasoning_tokens", 0), + } + + +def _openai_output_text(out: dict) -> str: + parts = [] + for item in out.get("output", []): + if item.get("type") == "message": + for c in item.get("content", []): + if c.get("type") == "output_text": + parts.append(c.get("text", "")) + return "".join(parts) + + +def _openai_complete( + api_key: str, + model: str, + messages: list[dict], + system: str | None, + max_tokens: int, + effort: str | None, + extra_input_items: list[dict] | None = None, +) -> tuple[str, dict, str]: + input_items: list[dict] = list(extra_input_items or []) + input_items += [{"role": m["role"], "content": _openai_content(m["content"], m["role"])} for m in messages] + body: dict = {"model": model, "input": input_items, "max_output_tokens": max_tokens, "store": False} + if system: + body["instructions"] = system + if effort: + body["reasoning"] = {"effort": "high" if effort in ("xhigh", "max") else effort} + out = _post(OPENAI_URL, body, {"authorization": f"Bearer {api_key}"}) + status = out.get("status", "") + stop = "max_tokens" if (out.get("incomplete_details") or {}).get("reason") == "max_output_tokens" else status + return _openai_output_text(out), _openai_usage(out), stop + + +def openai_compact(api_key: str, model: str, messages: list[dict]) -> tuple[list[dict], dict]: + """POST /responses/compact: returns (compacted output items, usage).""" + body = { + "model": model, + "input": [{"role": m["role"], "content": _openai_content(m["content"], m["role"])} for m in messages], + } + out = _post(f"{OPENAI_URL}/compact", body, {"authorization": f"Bearer {api_key}"}) + return out.get("output", []), _openai_usage(out) + + +# --- OpenRouter (chat completions) --- + +OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions" + + +def _openrouter_complete( + api_key: str, model: str, messages: list[dict], system: str | None, max_tokens: int, effort: str | None +) -> tuple[str, dict, str]: + def content(blocks: list[dict]) -> list[dict]: + out = [] + for b in blocks: + if "text" in b: + out.append({"type": "text", "text": b["text"]}) + else: + out.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_png_b64(b['image_path'])}"}}) + return out + + chat_messages = [{"role": m["role"], "content": content(m["content"])} for m in messages] + if system: + chat_messages.insert(0, {"role": "system", "content": system}) + body: dict = {"model": model, "messages": chat_messages, "max_tokens": max_tokens} + if effort == "none": + body["reasoning"] = {"enabled": False} # OpenRouter's disable switch; effort "none" is not a valid level + elif effort: + body["reasoning"] = {"effort": "high" if effort in ("xhigh", "max") else effort} + out = _post(OPENROUTER_URL, body, {"authorization": f"Bearer {api_key}"}) + choice = (out.get("choices") or [{}])[0] + text = (choice.get("message") or {}).get("content") or "" + if isinstance(text, list): # some providers return content parts + text = "".join(p.get("text", "") for p in text if isinstance(p, dict)) + u = out.get("usage", {}) + usage = { + "in": u.get("prompt_tokens", 0) - (u.get("prompt_tokens_details") or {}).get("cached_tokens", 0), + "out": u.get("completion_tokens", 0), + "cache_w": 0, + "cache_r": (u.get("prompt_tokens_details") or {}).get("cached_tokens", 0), + "reasoning": (u.get("completion_tokens_details") or {}).get("reasoning_tokens", 0), + } + stop = "max_tokens" if choice.get("finish_reason") == "length" else (choice.get("finish_reason") or "") + return text, usage, stop + + +# --- dispatch --- + + +def is_openai(model: str) -> bool: + return model.startswith(("gpt-", "o3", "o4", "codex")) + + +def is_openrouter(model: str) -> bool: + return "/" in model + + +def llm_complete( + api_keys: dict[str, str], + model: str, + messages: list[dict], + system: str | None = None, + max_tokens: int = 16384, + effort: str | None = None, + extra_input_items: list[dict] | None = None, +) -> tuple[str, dict, str]: + """Returns (text, normalized usage, stop). stop == "max_tokens" means truncated.""" + if is_openrouter(model): + if extra_input_items: + raise ValueError("extra_input_items is OpenAI-only (compacted window replay)") + return _openrouter_complete(api_keys["openrouter"], model, messages, system, max_tokens, effort) + if is_openai(model): + return _openai_complete( + api_keys["openai"], model, messages, system, max_tokens, effort, extra_input_items + ) + if extra_input_items: + raise ValueError("extra_input_items is OpenAI-only (compacted window replay)") + return _anthropic_complete(api_keys["anthropic"], model, messages, system, max_tokens, effort) diff --git a/packages/snapcompact/research/run.py b/packages/snapcompact/research/run.py new file mode 100644 index 000000000..efc570d75 --- /dev/null +++ b/packages/snapcompact/research/run.py @@ -0,0 +1,353 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""snapcompact: QA recall per context-compression strategy, over the full SQuAD dev set. + +The corpus (all SQuAD v1.1 dev passages, space-joined; ~1.5M chars) is processed +per condition in chunks sized to that condition's carrying capacity: + + text plain text, fixed 40,716-char chunks (= img-6x10 capacity) + compact agent compaction summary of each text chunk, QA on the summary + handoff agent handoff document of each text chunk, QA on the summary + img-- one 1568x1568 image per chunk; font in {8x13,6x10,5x8,5x7, + 4x6tt,4x5tt}, render variant v in {color,zebra,bw} + +Per chunk, up to --qpc questions are sampled (seeded, evenly spread across the +chunk so answers land at every image row band; pos_rel is recorded for position +analysis). Scoring is official SQuAD EM/F1. Responses are cached by payload +hash, so interrupted runs resume for free. + +Usage examples: + uv run run.py # default condition set, full corpus + uv run run.py --limit-chars 200000 # quick pass on a corpus prefix + uv run run.py --conditions img-6x10-bw # one condition + uv run run.py --report # re-print tables from cache, no API + +Key: ANTHROPIC_API_KEY from ~/.env (last assignment wins). +""" + +import argparse +import hashlib +import json +import re +import sys +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +HERE = Path(__file__).resolve().parent + +def find_agent_prompts() -> Path: + for parent in HERE.parents: + for candidate in ( + parent / "packages" / "agent" / "src" / "compaction" / "prompts", + parent / "agent" / "src" / "compaction" / "prompts", + ): + if candidate.exists(): + return candidate + raise FileNotFoundError("Could not find agent compaction prompts") + + +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from anthropic_api import complete, image_block, load_api_key # noqa: E402 +from bdf import VARIANTS, FontCfg, capacity, render # noqa: E402 + +AGENT_PROMPTS = find_agent_prompts() +CACHE = HERE / ".cache" +QA_CACHE = CACHE / "qa" +RESULTS = HERE / "results" + +FONTS = { + "8x13": FontCfg("8x13", "8x13", 8, 13), + "7x13": FontCfg("7x13", "7x13", 7, 13), + "6x12": FontCfg("6x12", "6x12", 6, 12), + "6x10": FontCfg("6x10", "6x10", 6, 10), + "6x9": FontCfg("6x9", "6x9", 6, 9), + # Anisotropic squashes: same glyphs, tighter grid. Crop flavor (6x6/6x8/3x10): + # next row's band paints over the vertical overhang, halved advance fuses ink. + # Stretch flavor (*s): rasterize native 6x10, Lanczos-resize to the target cell. + "6x8s": FontCfg("6x8s", "6x10", 6, 8, native=(6, 10)), + "6x6s": FontCfg("6x6s", "6x10", 6, 6, native=(6, 10)), + "3x10s": FontCfg("3x10s", "6x10", 3, 10, native=(6, 10)), + # Human-designed square cell: unscii-8 native, and stretched to a 6x6 cell. + "8x8u": FontCfg("8x8u", "unscii-8", 8, 8), + "6x6u": FontCfg("6x6u", "unscii-8", 6, 6, native=(8, 8)), + "6x8": FontCfg("6x8", "6x9", 6, 8, ascent=7), + "6x6": FontCfg("6x6", "6x9", 6, 6, ascent=6), + "3x10": FontCfg("3x10", "6x10", 3, 10), + # Redundancy coding: every line rendered twice, duplicate on a pale highlight. + "8x8r": FontCfg("8x8r", "unscii-8", 8, 8, repeat=2), + # Downsample-survivors: unscii-8 upscaled so gateway resizes leave legible glyphs. + "12x12u": FontCfg("12x12u", "unscii-8", 12, 12, native=(8, 8)), + "16x16u": FontCfg("16x16u", "unscii-8", 16, 16, native=(8, 8)), + "5x8": FontCfg("5x8", "5x8", 5, 8), + "5x7": FontCfg("5x7", "5x7", 5, 7), + "4x6tt": FontCfg("4x6tt", "tom-thumb", 4, 6, ascent=5), + "4x5tt": FontCfg("4x5tt", "tom-thumb", 4, 5, ascent=5), +} +TEXT_CHUNK = 40716 # = img-6x10 capacity; keeps text/summary chunks comparable +DEFAULT_CONDITIONS = ( + "text,compact,handoff," + "img-8x13-color,img-6x10-color,img-5x8-color,img-5x7-color," + "img-6x10-zebra,img-5x8-zebra,img-6x10-bw,img-5x8-bw" +) + + +def sha8(*parts: str) -> str: + return hashlib.sha1("\x00".join(parts).encode()).hexdigest()[:8] + + +def load_prompt(name: str) -> str: + return (HERE / "prompts" / name).read_text() + + +def agent_prompt(name: str) -> str: + text = (AGENT_PROMPTS / name).read_text() + # Drop unused Handlebars conditionals (no custom focus in this eval). + return re.sub(r"\{\{#if .*?\{\{/if\}\}\n?", "", text, flags=re.DOTALL) + + +def cached_complete(api_key: str, model: str, messages: list[dict], fresh: bool, **kw) -> tuple[str, dict]: + """complete() with response caching keyed on the full request payload. + + Truncated responses (stop_reason == max_tokens) are never cached and never + served from cache, so re-runs with a larger budget repair them. + """ + key = sha8(model, kw.get("effort") or "", json.dumps(messages, sort_keys=True)) + path = QA_CACHE / f"{key}.json" + if path.exists() and not fresh: + hit = json.loads(path.read_text()) + if hit.get("stop") != "max_tokens" and hit["text"]: + return hit["text"], hit["usage"] + text, usage, stop = complete(api_key, model, messages, **kw) + if stop == "max_tokens": + print(f" WARN truncated response (stop=max_tokens), not cached: {key}") + else: + path.write_text(json.dumps({"text": text, "usage": usage, "stop": stop})) + return text, usage + + +def parse_condition(name: str) -> dict: + if name in ("text", "compact", "handoff"): + return {"name": name, "kind": name} + m = re.fullmatch(r"img-([a-z0-9]+)-([a-z-]+)", name) + if not m or m.group(1) not in FONTS or m.group(2) not in VARIANTS: + raise SystemExit(f"bad condition {name!r}; expected text|compact|handoff|img--") + return {"name": name, "kind": "image", "font": FONTS[m.group(1)], "variant": m.group(2)} + + +def run_chunk(cond: dict, start: int, end: int, ctx_args: dict) -> list[dict]: + """Execute one (condition, chunk) task; returns per-question records.""" + args, flow, paras, offsets, api_key = ( + ctx_args["args"], + ctx_args["flow"], + ctx_args["paras"], + ctx_args["offsets"], + ctx_args["api_key"], + ) + questions = squad.sample_chunk_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + return [] + chunk_text = flow[start:end] + usage_rows: list[tuple[str, dict]] = [] + + png = cols = rows = None + context = chunk_text + if cond["kind"] == "image": + salt = ("dimv2",) if cond["variant"] == "dim" else () # cache-bust pre-fix sticky-fg dim renders + png = CACHE / f"img-{cond['font'].name}-{cond['variant']}-{sha8(chunk_text, str(args.size), *salt)}.png" + if not png.exists(): + render(chunk_text, cond["font"], CACHE, args.size, cond["variant"]).save(png) + cols, rows, _ = capacity(cond["font"], args.size) + elif cond["kind"] in ("compact", "handoff"): + prompt_file = {"compact": "compaction-summary.md", "handoff": "handoff-document.md"}[cond["kind"]] + gen_messages = [ + {"role": "user", "content": load_prompt("session-frame.md").format(context=chunk_text)}, + {"role": "assistant", "content": "Noted. I have read the passages and will keep them in mind."}, + {"role": "user", "content": agent_prompt(prompt_file)}, + ] + context, gen_usage = cached_complete( + api_key, args.model, gen_messages, args.fresh, + system=agent_prompt("summarization-system.md"), max_tokens=4096, + ) + usage_rows.append(("summarize", gen_usage)) + + use_cache = args.cache == "on" or (args.cache == "auto" and args.qpb > 0) + batch_size = args.qpb or len(questions) + answers: list[str] = [] + for b in range(0, len(questions), batch_size): + batch = questions[b : b + batch_size] + q_block = "\n".join(f"{i + 1}. {q['q']}" for i, q in enumerate(batch)) + if cond["kind"] == "image": + carrier = image_block(png) + preamble = {"type": "text", "text": load_prompt("qa-image.md").format(cols=cols, rows=rows)} + else: + carrier = {"type": "text", "text": load_prompt("qa-text.md").format(context=context)} + preamble = None + if use_cache: + carrier["cache_control"] = {"type": "ephemeral"} + content = ([preamble] if preamble else []) + [carrier, {"type": "text", "text": q_block}] + messages = [{"role": "user", "content": content}] + text, usage = cached_complete( + api_key, args.model, messages, args.fresh, max_tokens=args.max_tokens, effort=args.effort + ) + usage_rows.append(("qa", usage)) + answers.extend(squad.parse_numbered(text, len(batch))) + records = [] + for q, a in zip(questions, answers): + records.append( + { + "cond": cond["name"], + "chunk": start, + "pos_rel": q["pos_rel"], + "q": q["q"], + "answer": a, + "golds": q["golds"], + "em": squad.exact_match(a, q["golds"]), + "f1": squad.f1(a, q["golds"]), + "abstained": "unreadable" in a.lower(), + } + ) + # Attach token usage to the first record of the chunk (avoids double counting). + records[0]["usage"] = [ + { + "phase": phase, + "in": u.get("input_tokens", 0), + "out": u.get("output_tokens", 0), + "cache_w": u.get("cache_creation_input_tokens", 0), + "cache_r": u.get("cache_read_input_tokens", 0), + } + for phase, u in usage_rows + ] + return records + + +def aggregate(name: str, records: list[dict], price_in: float, price_out: float) -> dict: + n = len(records) + f1s = [r["f1"] for r in records] + mean_f1 = sum(f1s) / n + se = (sum((x - mean_f1) ** 2 for x in f1s) / (n * (n - 1))) ** 0.5 if n > 1 else 0.0 + usages = [u for r in records if "usage" in r for u in r["usage"]] + tok_in = sum(u["in"] for u in usages) + tok_out = sum(u["out"] for u in usages) + cache_w = sum(u.get("cache_w", 0) for u in usages) + cache_r = sum(u.get("cache_r", 0) for u in usages) + quart = [] + for lo, hi in ((0.0, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.01)): + qs = [r["f1"] for r in records if lo <= r["pos_rel"] < hi] + quart.append(sum(qs) / len(qs) if qs else None) + return { + "name": name, + "n": n, + "em": sum(r["em"] for r in records) / n, + "f1": mean_f1, + "f1_se": se, + "abstained": sum(r["abstained"] for r in records), + "tokens_in": tok_in, + "tokens_out": tok_out, + "cache_w": cache_w, + "cache_r": cache_r, + # Anthropic pricing: cache write 1.25x input, cache read 0.1x input (5m TTL). + "cost_usd": (tok_in + 1.25 * cache_w + 0.1 * cache_r) / 1e6 * price_in + tok_out / 1e6 * price_out, + "f1_by_quartile": quart, + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="claude-fable-5") + ap.add_argument("--conditions", default=DEFAULT_CONDITIONS) + ap.add_argument("--qpc", type=int, default=30, help="questions sampled per chunk") + ap.add_argument("--qpb", type=int, default=0, help="questions per API call (batches the chunk); 0 = all at once") + ap.add_argument("--cache", choices=["auto", "on", "off"], default="auto", + help="prompt-cache the carrier block; auto = on when --qpb is set") + ap.add_argument("--max-tokens", type=int, default=8192, help="output budget per QA call (incl. thinking)") + ap.add_argument("--effort", choices=["low", "medium", "high", "xhigh", "max"], default=None, + help="adaptive-thinking effort for QA calls; default = provider default") + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--workers", type=int, default=4) + ap.add_argument("--limit-chars", type=int, default=0, help="cap corpus size; 0 = full dev set") + ap.add_argument("--limit-paras", type=int, default=0, help="cap corpus to first N passages; 0 = all") + ap.add_argument("--fresh", action="store_true", help="ignore cached responses") + ap.add_argument("--report", action="store_true", help="aggregate cached records only; no API calls") + ap.add_argument("--price-in", type=float, default=10.0, help="$ per 1M input tokens") + ap.add_argument("--price-out", type=float, default=50.0, help="$ per 1M output tokens") + ap.add_argument("--env", default="~/.env") + args = ap.parse_args() + + CACHE.mkdir(exist_ok=True) + QA_CACHE.mkdir(exist_ok=True) + scope = f"p{args.limit_paras}" if args.limit_paras else (args.limit_chars or "full") + tag = "".join( + [ + f"-qpb{args.qpb}" if args.qpb else "", + f"-s{args.size}" if args.size != 1568 else "", + f"-e{args.effort}" if args.effort else "", + ] + ) + run_dir = RESULTS / f"{args.model}-seed{args.seed}-qpc{args.qpc}-{scope}{tag}" + run_dir.mkdir(parents=True, exist_ok=True) + + paras = squad.load_paragraphs(CACHE) + if args.limit_paras: + paras = paras[: args.limit_paras] + flow, offsets = squad.build_flow(paras, args.limit_chars or None) + conditions = [parse_condition(c.strip()) for c in args.conditions.split(",") if c.strip()] + + tasks: list[tuple[dict, int, int]] = [] + for cond in conditions: + budget = capacity(cond["font"], args.size)[2] if cond["kind"] == "image" else TEXT_CHUNK + for start in range(0, len(flow), budget): + tasks.append((cond, start, min(start + budget, len(flow)))) + calls = len(tasks) + sum(1 for c, *_ in tasks if c["kind"] in ("compact", "handoff")) + print( + f"corpus={len(flow):,} chars ({len(offsets):,} passages), {len(conditions)} conditions, " + f"{len(tasks)} chunks, <= {calls} API calls, qpc={args.qpc}, model={args.model}" + ) + + api_key = "" if args.report else load_api_key(args.env) + ctx_args = {"args": args, "flow": flow, "paras": paras, "offsets": offsets, "api_key": api_key} + records: list[dict] = [] + done = 0 + with ThreadPoolExecutor(args.workers) as pool: + futures = [pool.submit(run_chunk, cond, start, end, ctx_args) for cond, start, end in tasks] + for fut in futures: + records.extend(fut.result()) + done += 1 + if done % 20 == 0: + print(f" {done}/{len(tasks)} chunks", flush=True) + + with (run_dir / "records.jsonl").open("w") as fh: + for r in records: + fh.write(json.dumps(r) + "\n") + + rows = [ + aggregate(cond["name"], [r for r in records if r["cond"] == cond["name"]], args.price_in, args.price_out) + for cond in conditions + if any(r["cond"] == cond["name"] for r in records) + ] + rows.sort(key=lambda r: -r["f1"]) + (run_dir / "summary.json").write_text(json.dumps({"args": vars(args), "rows": rows}, indent=1)) + + hdr = ( + f"{'condition':<15}{'n':>6}{'EM':>7}{'F1':>7}{'±se':>6}{'abst':>6}" + f"{'in tok':>10}{'cache w':>9}{'cache r':>9}{'out tok':>9}{'$':>7}" + ) + print("\n" + hdr + "\n" + "-" * len(hdr)) + for r in rows: + print( + f"{r['name']:<15}{r['n']:>6}{r['em']:>7.3f}{r['f1']:>7.3f}{r['f1_se']:>6.3f}{r['abstained']:>6}" + f"{r['tokens_in']:>10,}{r['cache_w']:>9,}{r['cache_r']:>9,}{r['tokens_out']:>9,}{r['cost_usd']:>7.2f}" + ) + print(f"\n{'condition':<15} F1 by position quartile (Q1..Q4)") + for r in rows: + cells = " ".join(" - " if q is None else f"{q:.3f}" for q in r["f1_by_quartile"]) + print(f"{r['name']:<15} {cells}") + print(f"\nresults -> {run_dir}/") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_3d_activation_html.py b/packages/snapcompact/research/snapcompact_3d_activation_html.py new file mode 100644 index 000000000..f29726405 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_3d_activation_html.py @@ -0,0 +1,192 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["numpy", "plotly"] +# /// +"""Build an embeddable interactive 3D activation terrain for snapcompact.""" + +from __future__ import annotations + +import argparse +import base64 +import json +from pathlib import Path + +import numpy as np +import plotly.graph_objects as go +from plotly.subplots import make_subplots + +HERE = Path(__file__).resolve().parent + + +def downsample(arr: np.ndarray, cols: int) -> np.ndarray: + if arr.shape[1] <= cols: + return arr + edges = np.linspace(0, arr.shape[1], cols + 1).round().astype(int) + out = np.zeros((arr.shape[0], cols), dtype=np.float32) + for i in range(cols): + lo = edges[i] + hi = max(lo + 1, edges[i + 1]) + out[:, i] = arr[:, lo:hi].mean(axis=1) + return out + + +def norm(arr: np.ndarray, q: float = 0.985) -> np.ndarray: + scale = float(np.quantile(arr, q)) + if scale <= 0: + scale = 1.0 + return np.clip(arr / scale, 0, 1) + + +def image_data_uri(path: Path) -> str: + return "data:image/png;base64," + base64.b64encode(path.read_bytes()).decode() + + +def add_surface(fig: go.Figure, z: np.ndarray, row: int, col: int, name: str, colorscale: str, showscale: bool = False) -> None: + y = np.arange(z.shape[0]) + x = np.arange(z.shape[1]) + fig.add_trace( + go.Surface( + x=x, + y=y, + z=z, + name=name, + colorscale=colorscale, + cmin=0, + cmax=1, + showscale=showscale, + lighting={"ambient": 0.58, "diffuse": 0.72, "specular": 0.28, "roughness": 0.52}, + contours={ + "z": {"show": True, "usecolormap": True, "highlightcolor": "#fff0a8", "project_z": True}, + }, + hovertemplate="layer %{y}
image bin %{x}
Δ %{z:.3f}" + name + "", + ), + row=row, + col=col, + ) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "tensor-heatmap-paddleocr-q7")) + ap.add_argument("--out", default=str(HERE / "results" / "snapcompact-activation-terrain.html")) + ap.add_argument("--bins", type=int, default=150) + args = ap.parse_args() + + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "heatmaps.npz") + answer = norm(downsample(data["answer_binned"], args.bins)) + random = norm(downsample(data["random_binned"], args.bins)) + ratio = norm(downsample(data["ratio_binned"], args.bins), 0.97) + + fig = make_subplots( + rows=2, + cols=2, + specs=[[{"type": "surface"}, {"type": "surface"}], [{"type": "surface", "colspan": 2}, None]], + horizontal_spacing=0.02, + vertical_spacing=0.03, + subplot_titles=("Gold answer erased", "Random equal-size erase", "Answer / random residual scar"), + ) + add_surface(fig, answer, 1, 1, "gold answer mask", "Magma") + add_surface(fig, random, 1, 2, "random mask", "Viridis") + add_surface(fig, ratio, 2, 1, "answer/random ratio", "Inferno", True) + + camera = {"eye": {"x": 1.65, "y": -1.75, "z": 0.82}, "center": {"x": 0, "y": 0, "z": -0.08}} + scene_common = { + "bgcolor": "rgba(0,0,0,0)", + "camera": camera, + "xaxis": {"title": "image-token bins", "gridcolor": "rgba(140,170,180,0.18)", "color": "#94a3aa", "zeroline": False}, + "yaxis": {"title": "decoder layer", "gridcolor": "rgba(140,170,180,0.18)", "color": "#94a3aa", "autorange": "reversed", "dtick": 4}, + "zaxis": {"title": "Δ hidden", "gridcolor": "rgba(140,170,180,0.18)", "color": "#94a3aa", "range": [0, 1]}, + "aspectratio": {"x": 2.6, "y": 0.78, "z": 0.52}, + } + fig.update_layout( + template="plotly_dark", + paper_bgcolor="rgba(0,0,0,0)", + plot_bgcolor="rgba(0,0,0,0)", + height=920, + margin={"l": 0, "r": 0, "t": 58, "b": 0}, + font={"family": "Arial, sans-serif", "color": "#efeede"}, + scene=scene_common, + scene2=scene_common, + scene3={**scene_common, "aspectratio": {"x": 3.2, "y": 0.78, "z": 0.58}}, + coloraxis_showscale=False, + ) + fig.update_annotations(font={"size": 18, "color": "#efeede"}) + + q = summary["question"] + original_uri = image_data_uri(result_dir / "images" / "original.png") + masked_uri = image_data_uri(result_dir / "images" / "answer-mask.png") + graph_html = fig.to_html(full_html=False, include_plotlyjs="cdn", config={"displayModeBar": False, "responsive": True}) + html = f""" + + + + +Snapcompact activation terrain + + + +
+
SNAPCOMPACT WHITEBOX
+

Activation terrain from a missing answer

+
Actual decoder hidden states: layer × image-token bin × ||original − masked||. Drag the terrain to inspect the residual-stream scar.
+
+ +
+
interactive 3D residual terraingold-mask spikes rise where the model reacts to losing the answer glyphs
+
{graph_html}
+
+
+
+ + +""" + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(html) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_3d_activation_viz.py b/packages/snapcompact/research/snapcompact_3d_activation_viz.py new file mode 100644 index 000000000..d002182d2 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_3d_activation_viz.py @@ -0,0 +1,230 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["matplotlib", "numpy", "pillow"] +# /// +"""Render a 3D blog visualization of snapcompact hidden-state deltas.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.colors as mcolors +import matplotlib.pyplot as plt +import numpy as np +from matplotlib import cm +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent + +BG = (5, 7, 10) +PANEL = (13, 18, 23) +INK = (241, 239, 224) +MUTED = (139, 151, 156) +CYAN = (80, 220, 255) +RED = (255, 83, 62) +GREEN = (145, 255, 112) +AMBER = (255, 194, 65) + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if path and Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def norm_quantile(arr: np.ndarray, q: float = 0.985) -> np.ndarray: + scale = float(np.quantile(arr, q)) + if scale <= 0: + scale = 1.0 + return np.clip(arr / scale, 0, 1) + + +def downsample(arr: np.ndarray, cols: int) -> np.ndarray: + if arr.shape[1] <= cols: + return arr + edges = np.linspace(0, arr.shape[1], cols + 1).round().astype(int) + out = np.zeros((arr.shape[0], cols), dtype=np.float32) + for i in range(cols): + lo = edges[i] + hi = max(lo + 1, edges[i + 1]) + out[:, i] = arr[:, lo:hi].mean(axis=1) + return out + + +def style_3d(ax, title: str, subtitle: str, color: str) -> None: + ax.set_facecolor((0.02, 0.025, 0.035, 1)) + ax.xaxis.pane.set_facecolor((0.02, 0.025, 0.035, 0.0)) + ax.yaxis.pane.set_facecolor((0.02, 0.025, 0.035, 0.0)) + ax.zaxis.pane.set_facecolor((0.02, 0.025, 0.035, 0.0)) + for axis in (ax.xaxis, ax.yaxis, ax.zaxis): + axis._axinfo["grid"]["color"] = (0.28, 0.34, 0.38, 0.20) + axis._axinfo["tick"]["color"] = (0.80, 0.84, 0.84, 0.55) + ax.tick_params(colors="#8b979c", labelsize=8, pad=0) + ax.set_xlabel("image-token bins", color="#8b979c", labelpad=6) + ax.set_ylabel("decoder layer", color="#8b979c", labelpad=6) + ax.set_zlabel("Δ hidden", color="#8b979c", labelpad=5) + ax.set_title(title, color=color, fontsize=16, fontweight="bold", loc="left", pad=10) + ax.text2D(0.0, 0.94, subtitle, transform=ax.transAxes, color="#8b979c", fontsize=9) + ax.view_init(elev=31, azim=-58) + ax.set_box_aspect((2.7, 0.85, 0.55)) + + +def draw_surface(ax, arr: np.ndarray, cmap_name: str, title: str, subtitle: str, color: str, zmax: float = 1.0) -> None: + y = np.arange(arr.shape[0]) + x = np.arange(arr.shape[1]) + X, Y = np.meshgrid(x, y) + Z = arr * zmax + cmap = cm.get_cmap(cmap_name) + ax.plot_surface( + X, + Y, + Z, + rstride=1, + cstride=1, + facecolors=cmap(arr), + linewidth=0, + antialiased=True, + shade=False, + alpha=0.98, + ) + # A dark floor with projected contour lines makes the shape read as 3D. + ax.contour(X, Y, Z, zdir="z", offset=-0.05, levels=9, cmap=cmap, linewidths=0.8, alpha=0.72) + ax.set_zlim(-0.05, zmax) + ax.set_ylim(arr.shape[0] - 1, 0) + style_3d(ax, title, subtitle, color) + + +def crop_with_box(img: Image.Image, start: int, end: int, cols: int, adv: int, pitch: int, pad_cells: int = 34) -> Image.Image: + row0 = max(0, start // cols - 5) + row1 = min(img.height // pitch, end // cols + 6) + col0 = max(0, start % cols - pad_cells) + col1 = min(cols, end % cols + pad_cells) + if col1 <= col0: + col1 = min(cols, col0 + 72) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=RED, width=3) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.NEAREST) + canvas.paste(resized, (x0 + (x1 - x0 - resized.width) // 2, y0 + (y1 - y0 - resized.height) // 2)) + + +def render_matplotlib_panel(answer: np.ndarray, random: np.ndarray, ratio: np.ndarray) -> Image.Image: + fig = plt.figure(figsize=(16.6, 9.0), dpi=170) + fig.patch.set_facecolor("#05070a") + gs = fig.add_gridspec(2, 2, left=0.02, right=0.99, top=0.96, bottom=0.04, wspace=0.03, hspace=0.08) + ax1 = fig.add_subplot(gs[0, 0], projection="3d") + ax2 = fig.add_subplot(gs[0, 1], projection="3d") + ax3 = fig.add_subplot(gs[1, :], projection="3d") + draw_surface(ax1, answer, "magma", "Gold answer mask", "true answer cells erased", "#ff533e") + draw_surface(ax2, random, "viridis", "Random control mask", "same-sized blank elsewhere", "#91ff70") + draw_surface(ax3, ratio, "inferno", "Answer / random ratio", "where the missing answer leaves a larger residual-stream scar", "#ffc241", zmax=1.08) + tmp = HERE / "results" / ".snapcompact-3d-panel.png" + fig.savefig(tmp, facecolor=fig.get_facecolor(), transparent=False) + plt.close(fig) + img = Image.open(tmp).convert("RGB") + tmp.unlink(missing_ok=True) + return img + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "tensor-heatmap-paddleocr-q7")) + ap.add_argument("--out", default=str(HERE / "results" / "snapcompact-3d-activation-terrain.png")) + ap.add_argument("--bins", type=int, default=128) + args = ap.parse_args() + + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "heatmaps.npz") + answer = norm_quantile(downsample(data["answer_binned"], args.bins)) + random = norm_quantile(downsample(data["random_binned"], args.bins)) + ratio = norm_quantile(downsample(data["ratio_binned"], args.bins), 0.97) + + panel = render_matplotlib_panel(answer, random, ratio) + w, h = 2200, 1320 + canvas = Image.new("RGB", (w, h), BG) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(8, 10 + y % 10, 14 + y % 12)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-260, -160, 950, 600), fill=(255, 83, 62, 34)) + gd.ellipse((1100, 220, 2500, 1500), fill=(80, 220, 255, 28)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(82))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 42), "SNAPCOMPACT WHITEBOX", fill=AMBER, font=font(24, True)) + draw.text((64, 82), "Activation terrain from a missing answer", fill=INK, font=font(68, True)) + draw.text((66, 168), "Actual decoder hidden states: layer × image-token bin × ||original − masked||. A blog-friendly 3D tensor slice, not a schematic.", fill=MUTED, font=font(26)) + + # Left evidence strip. + draw.rounded_rectangle((64, 238, 600, 1234), radius=30, fill=PANEL, outline=(31, 42, 50), width=1) + q = summary["question"] + draw.text((96, 270), "visual intervention", fill=INK, font=font(32, True)) + draw.text((96, 310), "answer cells are blanked", fill=MUTED, font=font(18)) + base = Image.open(result_dir / "images" / "original.png").convert("RGB") + masked = Image.open(result_dir / "images" / "answer-mask.png").convert("RGB") + cols = summary["geometry"]["cols"] + crop = crop_with_box(base, q["answer_start"], q["answer_end"], cols, 8, 13) + masked_crop = crop_with_box(masked, q["answer_start"], q["answer_end"], cols, 8, 13) + draw.text((96, 366), "ORIGINAL", fill=CYAN, font=font(17, True)) + draw.rounded_rectangle((96, 394, 568, 560), radius=14, fill=(244, 242, 230), outline=CYAN, width=3) + paste_fit(canvas, crop, (112, 410, 552, 544)) + draw.text((96, 618), "ANSWER ERASED", fill=RED, font=font(17, True)) + draw.rounded_rectangle((96, 646, 568, 812), radius=14, fill=(244, 242, 230), outline=RED, width=3) + paste_fit(canvas, masked_crop, (112, 662, 552, 796)) + question = q["q"] + if len(question) > 54: + question = question[:51] + "…" + draw.text((96, 890), "question", fill=MUTED, font=font(16, True)) + draw.text((96, 920), question, fill=INK, font=font(22)) + draw.text((96, 990), "gold answer", fill=MUTED, font=font(16, True)) + draw.text((96, 1024), str(q["answer_text"]), fill=AMBER, font=font(42, True)) + draw.text((96, 1110), f"{summary['layers']} layers", fill=MUTED, font=font(20)) + draw.text((96, 1142), f"{summary['image_tokens']} image tokens", fill=MUTED, font=font(20)) + draw.text((96, 1174), f"answer/random Δ = {summary['answer_over_random_delta']:.2f}×", fill=INK, font=font(22, True)) + + # Main 3D panel. + draw.rounded_rectangle((632, 238, 2134, 1234), radius=30, fill=PANEL, outline=(31, 42, 50), width=1) + panel = panel.resize((1450, 786), Image.Resampling.LANCZOS) + canvas.paste(panel, (660, 330)) + draw.text((672, 268), "3D residual-stream delta terrain", fill=INK, font=font(36, True)) + draw.text((672, 311), "Gold-mask spikes rise where the model’s image-token activations react to losing the answer glyphs.", fill=MUTED, font=font(20)) + + # Color scale. + cmap = cm.get_cmap("magma") + for i in range(260): + rgb = tuple(int(v * 255) for v in cmap(i / 259)[:3]) + draw.rectangle((1810 + i, 274, 1811 + i, 292), fill=rgb) + draw.text((1810, 246), "low Δ", fill=MUTED, font=font(14)) + draw.text((2018, 246), "high Δ", fill=MUTED, font=font(14)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_activation_probe.py b/packages/snapcompact/research/snapcompact_activation_probe.py new file mode 100644 index 000000000..b65036dfe --- /dev/null +++ b/packages/snapcompact/research/snapcompact_activation_probe.py @@ -0,0 +1,225 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers", "sentencepiece", "protobuf", "einops"] +# /// +"""White-box snapcompact activation pilot for a local Hugging Face VLM. + +Runs a tiny dense-bitmap/text paired corpus through a local VLM and compares +hidden states across carriers. The default targets the PaddleOCR-VL snapshot +available on spark.internal because the served Qwen2.5-VL NVFP4 checkpoint is a +vLLM/modelopt artifact that Transformers cannot load directly. +""" + +from __future__ import annotations + +import argparse +import json +import math +import random +import sys +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import mask_cells, random_span, sample_answer_questions # noqa: E402 + +DEFAULT_MODEL_DIR = ( + "/home/can/.cache/huggingface/hub/models--PaddlePaddle--PaddleOCR-VL/" + "snapshots/2b77538ef936207f60c16b45082841068987d08c" +) + + +def centered_gram(x: np.ndarray) -> np.ndarray: + gram = x @ x.T + row_mean = gram.mean(axis=1, keepdims=True) + col_mean = gram.mean(axis=0, keepdims=True) + return gram - row_mean - col_mean + gram.mean() + + +def linear_cka(x: np.ndarray, y: np.ndarray) -> float: + if x.shape[0] < 2 or y.shape[0] < 2: + return float("nan") + x = x - x.mean(axis=0, keepdims=True) + y = y - y.mean(axis=0, keepdims=True) + k = centered_gram(x) + l = centered_gram(y) + denom = math.sqrt(float((k * k).sum()) * float((l * l).sum())) + if denom == 0: + return float("nan") + return float((k * l).sum() / denom) + + +def paired_cosine(x: np.ndarray, y: np.ndarray) -> float: + dot = (x * y).sum(axis=1) + denom = np.linalg.norm(x, axis=1) * np.linalg.norm(y, axis=1) + valid = denom > 0 + if not valid.any(): + return float("nan") + return float((dot[valid] / denom[valid]).mean()) + + +def make_prompt(q: str, cols: int, rows: int) -> str: + return ( + load_prompt("qa-image.md").format(cols=cols, rows=rows) + + f"\n\nQuestion: {q}\nAnswer with only the shortest extractive answer." + ) + + +def to_device(batch: dict[str, Any], device: Any) -> dict[str, Any]: + return {k: (v.to(device) if hasattr(v, "to") else v) for k, v in batch.items()} + + +def hidden_features(model: Any, processor: Any, *, image: Image.Image | None, text: str, device: Any) -> list[np.ndarray]: + import torch + + if image is None: + messages = [{"role": "user", "content": [{"type": "text", "text": text}]}] + templated = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + batch = processor(text=templated, return_tensors="pt") + else: + messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": text}]}] + templated = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + batch = processor(images=image, text=templated, return_tensors="pt") + batch = to_device(batch, device) + with torch.no_grad(): + out = model(**batch, output_hidden_states=True, output_attentions=False, use_cache=False) + feats: list[np.ndarray] = [] + for h in out.hidden_states: + # Mean-pool the prompt sequence. This avoids brittle alignment between + # image-token and text-token positions while preserving layer geometry. + pooled = h.float().mean(dim=1).detach().cpu().numpy()[0] + feats.append(pooled.astype(np.float32, copy=False)) + return feats + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default=DEFAULT_MODEL_DIR) + ap.add_argument("--font", default="5x8", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=512) + ap.add_argument("--limit-paras", type=int, default=20) + ap.add_argument("--qpc", type=int, default=8) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--out", default="snapcompact-activation") + args = ap.parse_args() + + import torch + from transformers import AutoModel, AutoProcessor + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "activation-images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + start, end = 0, min(len(flow), budget) + chunk = flow[start:end] + questions = sample_answer_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + raise SystemExit("no sampled questions fit in the activation chunk") + + base_img = render(chunk, cfg, CACHE, args.size, args.variant) + base_path = img_dir / "base.png" + base_img.save(base_path) + fill = (255, 255, 255) if args.variant not in ("dark", "dark-sent") else (0, 0, 0) + + print(f"loading {args.model_dir}", flush=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 + model = AutoModel.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=dtype).to(device).eval() + + feature_sets: dict[str, list[list[np.ndarray]]] = {"text": [], "image": [], "answer_mask": [], "random_mask": []} + records: list[dict[str, Any]] = [] + for qi, q in enumerate(questions): + span_len = max(1, q["answer_end"] - q["answer_start"]) + rng = random.Random(args.seed * 31 + qi) + rand_start, rand_end = random_span(rng, len(chunk), span_len, q["answer_start"], q["answer_end"]) + answer_img = mask_cells(base_img, q["answer_start"], q["answer_end"], cols, cfg.adv, cfg.pitch, fill) + random_img = mask_cells(base_img, rand_start, rand_end, cols, cfg.adv, cfg.pitch, fill) + answer_path = img_dir / f"q{qi}-answer-mask.png" + random_path = img_dir / f"q{qi}-random-mask.png" + answer_img.save(answer_path) + random_img.save(random_path) + + img_prompt = make_prompt(q["q"], cols, rows) + text_prompt = ( + "Below is reference material. Answer the question using only it.\n\n" + f"{chunk}\n\nQuestion: {q['q']}\n" + "Answer with only the shortest extractive answer." + ) + feature_sets["text"].append(hidden_features(model, processor, image=None, text=text_prompt, device=device)) + feature_sets["image"].append(hidden_features(model, processor, image=base_img, text=img_prompt, device=device)) + feature_sets["answer_mask"].append(hidden_features(model, processor, image=answer_img, text=img_prompt, device=device)) + feature_sets["random_mask"].append(hidden_features(model, processor, image=random_img, text=img_prompt, device=device)) + records.append( + { + "question_index": qi, + "q": q["q"], + "golds": q["golds"], + "answer_text": q["answer_text"], + "answer_start": q["answer_start"], + "answer_end": q["answer_end"], + "random_start": rand_start, + "random_end": rand_end, + } + ) + print(f"captured {qi + 1}/{len(questions)}", flush=True) + + layer_count = len(feature_sets["image"][0]) + layers = [] + for layer in range(layer_count): + arrays = { + name: np.stack([sample[layer] for sample in samples], axis=0) + for name, samples in feature_sets.items() + } + img = arrays["image"] + ans = arrays["answer_mask"] + rnd = arrays["random_mask"] + answer_delta = np.linalg.norm(img - ans, axis=1) + random_delta = np.linalg.norm(img - rnd, axis=1) + layers.append( + { + "layer": layer, + "cka_text_image": linear_cka(arrays["text"], img), + "cka_image_answer_mask": linear_cka(img, ans), + "cka_image_random_mask": linear_cka(img, rnd), + "cos_text_image": paired_cosine(arrays["text"], img), + "cos_image_answer_mask": paired_cosine(img, ans), + "cos_image_random_mask": paired_cosine(img, rnd), + "answer_delta_norm": float(answer_delta.mean()), + "random_delta_norm": float(random_delta.mean()), + "answer_over_random_delta": float(answer_delta.mean() / random_delta.mean()) if random_delta.mean() else float("inf"), + } + ) + + summary = { + "args": vars(args), + "model_dir": args.model_dir, + "device": str(device), + "geometry": {"cols": cols, "rows": rows, "capacity": budget}, + "n": len(records), + "layers": layers, + } + with (out_dir / "records.jsonl").open("w") as fh: + for record in records: + fh.write(json.dumps(record) + "\n") + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + print(json.dumps(summary, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_blackbox_occlusion.py b/packages/snapcompact/research/snapcompact_blackbox_occlusion.py new file mode 100644 index 000000000..15faac599 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_blackbox_occlusion.py @@ -0,0 +1,261 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""Black-box snapcompact occlusion probe against an OpenAI-compatible VLM endpoint. + +For sampled SQuAD questions, render the carrier as a dense bitmap, then compare +QA on the original image, an image with the gold answer cells masked, and an +image with an equal-sized random mask. A real visual-retrieval mechanism should +show a larger F1 drop for answer masks than random masks. +""" + +from __future__ import annotations + +import argparse +import base64 +import hashlib +import json +import random +import sys +import urllib.error +import urllib.request +from pathlib import Path +from typing import Any + +from PIL import Image, ImageDraw + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt, sha8 # noqa: E402 + + +def sample_answer_questions(paras: list[dict], offsets: list[int], start: int, end: int, n: int, seed: int) -> list[dict]: + """Sample questions like squad.sample_chunk_questions, preserving answer offsets.""" + rng = random.Random(seed * 1_000_003 + start) + eligible = [ + i + for i in range(len(offsets)) + if offsets[i] >= start and offsets[i] + len(paras[i]["ctx"]) <= end and paras[i].get("qas") + ] + if not eligible: + return [] + n = min(n, len(eligible)) + step = len(eligible) / n + picked: list[dict] = [] + for k in range(n): + pi = eligible[int(k * step)] + qa = rng.choice(paras[pi]["qas"]) + answers = qa.get("answers") or [] + if not answers: + continue + answer = answers[0] + picked.append( + { + "q": " ".join(qa["question"].split()), + "golds": sorted({a["text"] for a in answers}), + "answer_text": answer["text"], + "answer_start": offsets[pi] - start + int(answer["answer_start"]), + "answer_end": offsets[pi] - start + int(answer["answer_start"]) + len(answer["text"]), + "pos_rel": (offsets[pi] - start) / (end - start), + } + ) + return picked + + +def mask_cells(img: Image.Image, start: int, end: int, cols: int, adv: int, pitch: int, fill: tuple[int, int, int]) -> Image.Image: + out = img.copy() + draw = ImageDraw.Draw(out) + start = max(0, start) + end = max(start + 1, end) + first_row = start // cols + last_row = (end - 1) // cols + for row in range(first_row, last_row + 1): + c0 = start % cols if row == first_row else 0 + c1 = (end - 1) % cols if row == last_row else cols - 1 + x0 = max(0, c0 * adv - adv) + y0 = max(0, row * pitch - 1) + x1 = min(out.width, (c1 + 2) * adv) + y1 = min(out.height, (row + 1) * pitch + 1) + draw.rectangle((x0, y0, x1, y1), fill=fill) + return out + + +def random_span(rng: random.Random, text_len: int, span_len: int, avoid_start: int, avoid_end: int) -> tuple[int, int]: + if text_len <= span_len: + return 0, text_len + for _ in range(100): + start = rng.randrange(0, text_len - span_len) + end = start + span_len + if end < avoid_start - span_len or start > avoid_end + span_len: + return start, end + start = 0 if avoid_start > text_len // 2 else max(0, text_len - span_len) + return start, min(text_len, start + span_len) + + +def post_chat(endpoint: str, model: str, image_path: Path, prompt: str, max_tokens: int, cache_dir: Path, fresh: bool) -> tuple[str, dict]: + payload_key = sha8(model, prompt, hashlib.sha1(image_path.read_bytes()).hexdigest()) + cache_path = cache_dir / f"{payload_key}.json" + if cache_path.exists() and not fresh: + cached = json.loads(cache_path.read_text()) + return cached["text"], cached.get("usage", {}) + + image_b64 = base64.b64encode(image_path.read_bytes()).decode() + body = { + "model": model, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_b64}"}}, + ], + } + ], + "temperature": 0, + "max_tokens": max_tokens, + } + req = urllib.request.Request( + endpoint, + data=json.dumps(body).encode(), + headers={"content-type": "application/json"}, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=180) as resp: + out = json.loads(resp.read().decode()) + except urllib.error.HTTPError as err: + raise RuntimeError(err.read().decode()) from err + choice = (out.get("choices") or [{}])[0] + text = ((choice.get("message") or {}).get("content") or "").strip() + usage = out.get("usage") or {} + cache_path.write_text(json.dumps({"text": text, "usage": usage}, indent=1)) + return text, usage + + +def aggregate(records: list[dict]) -> dict: + by_variant: dict[str, list[dict]] = {} + for record in records: + by_variant.setdefault(record["variant"], []).append(record) + out: dict[str, Any] = {"n": len(records) // 3, "variants": {}} + for name, rows in sorted(by_variant.items()): + out["variants"][name] = { + "n": len(rows), + "em": sum(r["em"] for r in rows) / max(1, len(rows)), + "f1": sum(r["f1"] for r in rows) / max(1, len(rows)), + "abstained": sum(1 for r in rows if "unreadable" in r["answer"].lower()), + "prompt_tokens": sum((r.get("usage") or {}).get("prompt_tokens", 0) for r in rows), + "completion_tokens": sum((r.get("usage") or {}).get("completion_tokens", 0) for r in rows), + } + base = out["variants"].get("original", {}).get("f1", 0.0) + out["drops"] = { + name: base - row["f1"] for name, row in out["variants"].items() if name != "original" + } + return out + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--endpoint", default="http://spark.internal:8000/v1/chat/completions") + ap.add_argument("--model", default="Qwen2.5-VL-7B-Instruct-NVFP4") + ap.add_argument("--font", default="5x8", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--limit-paras", type=int, default=50) + ap.add_argument("--qpc", type=int, default=12) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--max-tokens", type=int, default=48) + ap.add_argument("--out", default="snapcompact-occlusion") + ap.add_argument("--fresh", action="store_true") + args = ap.parse_args() + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + cache_dir = out_dir / "api-cache" + img_dir.mkdir(parents=True, exist_ok=True) + cache_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + prompt_base = load_prompt("qa-image.md").format(cols=cols, rows=rows) + mask_fill = (255, 255, 255) if args.variant not in ("dark", "dark-sent") else (0, 0, 0) + + tasks = [] + for start in range(0, len(flow), budget): + end = min(start + budget, len(flow)) + chunk = flow[start:end] + questions = sample_answer_questions(paras, offsets, start, end, args.qpc, args.seed) + if questions: + tasks.append((start, end, chunk, questions)) + + records: list[dict] = [] + for chunk_index, (start, end, chunk, questions) in enumerate(tasks): + base_img = render(chunk, cfg, CACHE, args.size, args.variant) + base_path = img_dir / f"chunk-{start}-{args.font}-{args.variant}.png" + if not base_path.exists(): + base_img.save(base_path) + for qi, q in enumerate(questions): + span_len = max(1, q["answer_end"] - q["answer_start"]) + rng = random.Random(args.seed * 17 + start + qi) + rand_start, rand_end = random_span(rng, len(chunk), span_len, q["answer_start"], q["answer_end"]) + answer_path = img_dir / f"chunk-{start}-q{qi}-answer-mask.png" + random_path = img_dir / f"chunk-{start}-q{qi}-random-mask.png" + if not answer_path.exists(): + mask_cells(base_img, q["answer_start"], q["answer_end"], cols, cfg.adv, cfg.pitch, mask_fill).save(answer_path) + if not random_path.exists(): + mask_cells(base_img, rand_start, rand_end, cols, cfg.adv, cfg.pitch, mask_fill).save(random_path) + + prompt = ( + f"{prompt_base}\n\nQuestion: {q['q']}\n" + "Answer with only the shortest extractive answer copied from the image. " + "If the answer is unreadable, reply exactly UNREADABLE." + ) + for variant_name, path in ( + ("original", base_path), + ("answer_mask", answer_path), + ("random_mask", random_path), + ): + answer, usage = post_chat(args.endpoint, args.model, path, prompt, args.max_tokens, cache_dir, args.fresh) + records.append( + { + "chunk": start, + "chunk_index": chunk_index, + "question_index": qi, + "variant": variant_name, + "q": q["q"], + "answer": answer, + "golds": q["golds"], + "answer_text": q["answer_text"], + "answer_start": q["answer_start"], + "answer_end": q["answer_end"], + "random_start": rand_start, + "random_end": rand_end, + "pos_rel": q["pos_rel"], + "em": squad.exact_match(answer, q["golds"]), + "f1": squad.f1(answer, q["golds"]), + "usage": usage, + } + ) + print(f"{len(records):04d} {variant_name:<11} f1={records[-1]['f1']:.3f} answer={answer[:80]!r}", flush=True) + + summary = { + "args": vars(args), + "geometry": {"cols": cols, "rows": rows, "capacity": budget}, + **aggregate(records), + } + with (out_dir / "records.jsonl").open("w") as fh: + for record in records: + fh.write(json.dumps(record) + "\n") + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + print(json.dumps(summary, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_blog_viz.py b/packages/snapcompact/research/snapcompact_blog_viz.py new file mode 100644 index 000000000..c431c2205 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_blog_viz.py @@ -0,0 +1,231 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""Render a blog-ready snapcompact white-box visualization from pilot outputs.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent + +PALETTE = { + "bg": (8, 11, 14), + "panel": (18, 24, 30), + "panel2": (13, 18, 24), + "grid": (47, 61, 72), + "text": (234, 238, 229), + "muted": (139, 151, 156), + "accent": (255, 104, 72), + "accent2": (67, 210, 255), + "green": (158, 255, 121), + "amber": (255, 197, 74), + "red": (255, 70, 70), +} + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Monaco.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if path and Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def rounded(draw: ImageDraw.ImageDraw, xy: tuple[int, int, int, int], fill: tuple[int, int, int], outline=None, radius=24, width=1) -> None: + draw.rounded_rectangle(xy, radius=radius, fill=fill, outline=outline, width=width) + + +def draw_label(draw: ImageDraw.ImageDraw, xy: tuple[int, int], text: str, color: tuple[int, int, int], size: int = 24, bold: bool = False) -> None: + draw.text(xy, text, fill=color, font=font(size, bold=bold)) + + +def chart( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], + series: list[tuple[str, list[float], tuple[int, int, int]]], + y_min: float, + y_max: float, + title: str, + subtitle: str, +) -> None: + x0, y0, x1, y1 = box + rounded(draw, box, PALETTE["panel"], outline=(35, 47, 56), radius=22) + draw_label(draw, (x0 + 28, y0 + 22), title, PALETTE["text"], 28, True) + draw_label(draw, (x0 + 28, y0 + 57), subtitle, PALETTE["muted"], 17) + gx0, gy0, gx1, gy1 = x0 + 58, y0 + 98, x1 - 30, y1 - 58 + for i in range(5): + y = gy0 + round((gy1 - gy0) * i / 4) + draw.line((gx0, y, gx1, y), fill=PALETTE["grid"], width=1) + value = y_max - (y_max - y_min) * i / 4 + draw.text((x0 + 18, y - 9), f"{value:.2f}", fill=PALETTE["muted"], font=font(13)) + n = len(series[0][1]) + for label, values, color in series: + pts = [] + for i, value in enumerate(values): + x = gx0 + (gx1 - gx0) * i / max(1, n - 1) + y = gy1 - (gy1 - gy0) * (value - y_min) / (y_max - y_min) + pts.append((round(x), round(y))) + draw.line(pts, fill=color, width=4, joint="curve") + for p in pts[:: max(1, n // 6)]: + draw.ellipse((p[0] - 4, p[1] - 4, p[0] + 4, p[1] + 4), fill=color) + lx = gx0 + ly = y1 - 36 + for label, _values, color in series: + draw.rounded_rectangle((lx, ly, lx + 20, ly + 10), radius=5, fill=color) + draw.text((lx + 28, ly - 5), label, fill=PALETTE["muted"], font=font(15)) + lx += 210 + + +def crop_with_box(img: Image.Image, start: int, end: int, cols: int, adv: int, pitch: int, pad_cells: int = 24) -> Image.Image: + row0 = max(0, start // cols - 4) + row1 = min(img.height // pitch, end // cols + 5) + col0 = max(0, start % cols - pad_cells) + col1 = min(cols, end % cols + pad_cells) + if row1 <= row0: + row1 = min(img.height // pitch, row0 + 8) + if col1 <= col0: + col1 = min(cols, col0 + 48) + x0, y0, x1, y1 = col0 * adv, row0 * pitch, col1 * adv, row1 * pitch + crop = img.crop((x0, y0, x1, y1)).convert("RGB") + draw = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + draw.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=PALETTE["red"], width=3) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.NEAREST) + px = x0 + (x1 - x0 - resized.width) // 2 + py = y0 + (y1 - y0 - resized.height) // 2 + canvas.paste(resized, (px, py)) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--activation", default=str(HERE / "results" / "activation-paddleocr-8x13-n16")) + ap.add_argument("--occlusion", default=str(HERE / "results" / "snapcompact-occlusion-qwen-8x13")) + ap.add_argument("--out", default=str(HERE / "results" / "snapcompact-blog-whitebox.png")) + args = ap.parse_args() + + act_dir = Path(args.activation) + occ_dir = Path(args.occlusion) + summary = json.loads((act_dir / "summary.json").read_text()) + occ = json.loads((occ_dir / "summary.json").read_text()) + records = [json.loads(line) for line in (act_dir / "records.jsonl").read_text().splitlines() if line] + layers = summary["layers"] + + w, h = 1800, 1040 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + # Background texture. + for y in range(0, h, 18): + color = (10 + y % 17, 13 + y % 13, 17 + y % 11) + draw.line((0, y, w, y), fill=color) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-160, -220, 760, 520), fill=(255, 104, 72, 34)) + gd.ellipse((1100, 130, 2100, 1160), fill=(67, 210, 255, 28)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(70))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw_label(draw, (56, 38), "SNAPCOMPACT UNDER THE MICROSCOPE", PALETTE["amber"], 21, True) + draw_label(draw, (56, 78), "Dense text-images leave a white-box trace", PALETTE["text"], 54, True) + draw_label(draw, (58, 145), "Text and image prompts converge late; blanking the gold answer region perturbs hidden states far more than an equal random blank.", PALETTE["muted"], 24) + + # Big stat cards. + stats = [ + ("Qwen black-box F1", f"{occ['variants']['original']['f1']:.2f}", "original image"), + ("gold-mask drop", f"−{occ['drops']['answer_mask']:.2f}", "answer region blanked"), + ("random-mask drop", f"−{occ['drops']['random_mask']:.2f}", "same-size random blank"), + ] + sx = 56 + card_w = 258 + for title, value, caption in stats: + rounded(draw, (sx, 205, sx + card_w, 330), PALETTE["panel2"], outline=(36, 48, 56), radius=22) + draw_label(draw, (sx + 20, 226), title, PALETTE["muted"], 16) + draw_label(draw, (sx + 20, 252), value, PALETTE["text"], 42, True) + draw_label(draw, (sx + 20, 300), caption, PALETTE["muted"], 14) + sx += card_w + 22 + + chart( + draw, + (56, 368, 872, 668), + [ + ("text ↔ image CKA", [x["cka_text_image"] for x in layers], PALETTE["accent2"]), + ("answer-mask CKA", [x["cka_image_answer_mask"] for x in layers], PALETTE["accent"]), + ("random-mask CKA", [x["cka_image_random_mask"] for x in layers], PALETTE["green"]), + ], + 0.2, + 1.0, + "Layer geometry", + "PaddleOCR-VL hidden-state similarity across 19 decoder layers", + ) + chart( + draw, + (56, 698, 872, 990), + [ + ("answer / random perturbation", [x["answer_over_random_delta"] for x in layers], PALETTE["amber"]), + ], + 1.0, + 1.6, + "Causal-ish scar", + "Mean hidden-state shift: gold answer mask divided by random mask", + ) + + # Visual crop panel. + panel = (920, 205, 1744, 990) + rounded(draw, panel, PALETTE["panel"], outline=(35, 47, 56), radius=26) + draw_label(draw, (950, 232), "What the mask test looks like", PALETTE["text"], 34, True) + draw_label(draw, (950, 274), "Same question, same bitmap. Only the gold answer cells are erased.", PALETTE["muted"], 19) + + base = Image.open(act_dir / "activation-images" / "base.png").convert("RGB") + # Use a later question if possible because it gives a better-looking crop. + rec = records[min(7, len(records) - 1)] + ans = Image.open(act_dir / "activation-images" / f"q{rec['question_index']}-answer-mask.png").convert("RGB") + rnd = Image.open(act_dir / "activation-images" / f"q{rec['question_index']}-random-mask.png").convert("RGB") + cols = summary["geometry"]["cols"] + adv = 8 + pitch = 13 + crops = [ + ("original", crop_with_box(base, rec["answer_start"], rec["answer_end"], cols, adv, pitch), PALETTE["accent2"]), + ("answer masked", crop_with_box(ans, rec["answer_start"], rec["answer_end"], cols, adv, pitch), PALETTE["accent"]), + ("random masked", crop_with_box(rnd, rec["random_start"], rec["random_end"], cols, adv, pitch), PALETTE["green"]), + ] + y = 334 + for label, img, color in crops: + draw_label(draw, (950, y - 31), label.upper(), color, 17, True) + rounded(draw, (950, y, 1714, y + 150), (244, 242, 230), outline=color, radius=14, width=3) + paste_fit(canvas, img, (966, y + 16, 1698, y + 134)) + y += 198 + + q = rec["q"] + if len(q) > 92: + q = q[:89] + "…" + draw_label(draw, (950, 916), "sample question", PALETTE["muted"], 17, True) + draw_label(draw, (950, 940), q, PALETTE["text"], 20) + draw_label(draw, (950, 966), f"gold answer: {rec['answer_text']}", PALETTE["amber"], 18, True) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_carrier_convergence.py b/packages/snapcompact/research/snapcompact_carrier_convergence.py new file mode 100644 index 000000000..cfdb21af2 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_carrier_convergence.py @@ -0,0 +1,212 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers"] +# /// +"""Do text and image carriers converge to the same internal state? + +For N questions over the same chunk, capture the last-prompt-token hidden state +(the model's "about to answer" summary) per decoder layer, once with the chunk +as raw text and once as a snapcompact bitmap. Carrier-specific means are +subtracted per layer so prompt boilerplate and modality signatures cancel out. + +Evidence of convergence: +1. matched pairs (same question, different carrier) >> mismatched pairs +2. the question-by-question similarity geometry (RSA) is shared across carriers +3. both carriers generate the same answers +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import sample_answer_questions # noqa: E402 + + +def make_text_prompt(chunk: str, question: str) -> str: + return ( + "Below is reference material. Answer the question using only it.\n\n" + f"{chunk}\n\nQuestion: {question}\n" + "Answer with only the shortest extractive answer." + ) + + +def make_image_prompt(cols: int, rows: int, question: str) -> str: + return load_prompt("qa-image.md").format(cols=cols, rows=rows) + f"\n\nQuestion: {question}\nAnswer with only the shortest extractive answer." + + +def capture_last_token(model: Any, processor: Any, device: Any, text: str, image: Image.Image | None) -> tuple[np.ndarray, str]: + """Return per-layer hidden state at the final prompt position plus a short generation.""" + import torch + + content: list[dict[str, Any]] = [] + if image is not None: + content.append({"type": "image", "image": image}) + content.append({"type": "text", "text": text}) + templated = processor.apply_chat_template([{"role": "user", "content": content}], tokenize=False, add_generation_prompt=True) + if image is not None: + batch = processor(images=image, text=templated, return_tensors="pt") + else: + batch = processor(text=templated, return_tensors="pt") + batch = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in batch.items()} + with torch.no_grad(): + out = model(**batch, output_hidden_states=True, use_cache=False) + generated = model.generate(**batch, max_new_tokens=16, do_sample=False) + states = np.stack([h[0, -1, :].float().detach().cpu().numpy() for h in out.hidden_states], axis=0) + answer = processor.batch_decode(generated[:, batch["input_ids"].shape[1] :], skip_special_tokens=True)[0].strip() + return states.astype(np.float32, copy=False), answer + + +def cosine_rows(a: np.ndarray, b: np.ndarray) -> np.ndarray: + a_n = a / np.maximum(np.linalg.norm(a, axis=-1, keepdims=True), 1e-6) + b_n = b / np.maximum(np.linalg.norm(b, axis=-1, keepdims=True), 1e-6) + return a_n @ b_n.T + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default="Qwen/Qwen2.5-VL-7B-Instruct") + ap.add_argument("--font", default="8x13", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--limit-paras", type=int, default=80) + ap.add_argument("--questions", type=int, default=12) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--out", default="qwen-carrier-convergence") + args = ap.parse_args() + + import torch + from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + chunk = flow[: min(len(flow), budget)] + questions = sample_answer_questions(paras, offsets, 0, len(chunk), args.questions * 2, args.seed) + # Deduplicate gold answers so the RSA geometry has distinct content per row. + seen: set[str] = set() + picked: list[dict[str, Any]] = [] + for q in questions: + key = q["answer_text"].lower() + if key not in seen: + seen.add(key) + picked.append(q) + if len(picked) >= args.questions: + break + if len(picked) < 4: + raise SystemExit("not enough distinct questions in chunk") + img = render(chunk, cfg, CACHE, args.size, args.variant) + img.save(img_dir / "image-carrier.png") + + print(f"loading {args.model_dir}", flush=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + model = Qwen2_5_VLForConditionalGeneration.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto").eval() + device = next(model.parameters()).device + + text_states: list[np.ndarray] = [] + image_states: list[np.ndarray] = [] + records: list[dict[str, Any]] = [] + for qi, q in enumerate(picked): + t_states, t_answer = capture_last_token(model, processor, device, make_text_prompt(chunk, q["q"]), None) + i_states, i_answer = capture_last_token(model, processor, device, make_image_prompt(cols, rows, q["q"]), img) + text_states.append(t_states) + image_states.append(i_states) + records.append( + { + "question_index": qi, + "q": q["q"], + "gold": q["answer_text"], + "golds": q["golds"], + "text_answer": t_answer, + "image_answer": i_answer, + "text_em": squad.exact_match(t_answer, q["golds"]), + "image_em": squad.exact_match(i_answer, q["golds"]), + "agree": squad.f1(t_answer, [i_answer]) >= 0.99, + } + ) + print(f"{qi + 1}/{len(picked)} text={t_answer!r} image={i_answer!r} gold={q['answer_text']!r}", flush=True) + + text_arr = np.stack(text_states, axis=0) # [Q, L, D] + image_arr = np.stack(image_states, axis=0) + n_q, n_layers, _dim = text_arr.shape + + layers: list[dict[str, Any]] = [] + text_sim_by_layer = np.zeros((n_layers, n_q, n_q), dtype=np.float32) + image_sim_by_layer = np.zeros((n_layers, n_q, n_q), dtype=np.float32) + cross_sim_by_layer = np.zeros((n_layers, n_q, n_q), dtype=np.float32) + off_diag = ~np.eye(n_q, dtype=bool) + for layer in range(n_layers): + text_l = text_arr[:, layer, :] + image_l = image_arr[:, layer, :] + # Carrier-centering removes modality/prompt signature; what remains is + # per-question content variation within each carrier. + text_c = text_l - text_l.mean(axis=0, keepdims=True) + image_c = image_l - image_l.mean(axis=0, keepdims=True) + cross = cosine_rows(text_c, image_c) + text_sim = cosine_rows(text_c, text_c) + image_sim = cosine_rows(image_c, image_c) + matched = float(np.diag(cross).mean()) + mismatched = float(cross[off_diag].mean()) + rsa = float(np.corrcoef(text_sim[off_diag], image_sim[off_diag])[0, 1]) + layers.append( + { + "layer": layer, + "matched_cosine": matched, + "mismatched_cosine": mismatched, + "separation": matched - mismatched, + "rsa_pearson": rsa, + "match_rank_accuracy": float((np.argmax(cross, axis=1) == np.arange(n_q)).mean()), + } + ) + text_sim_by_layer[layer] = text_sim + image_sim_by_layer[layer] = image_sim + cross_sim_by_layer[layer] = cross + + best_layer = int(np.argmax([l["separation"] for l in layers])) + summary = { + "args": vars(args), + "geometry": {"cols": cols, "rows": rows, "capacity": budget}, + "n_questions": n_q, + "layers": len(layers), + "per_layer": layers, + "best_layer": best_layer, + "best": layers[best_layer], + "final": layers[-1], + "answer_agreement": float(np.mean([r["agree"] for r in records])), + "text_em": float(np.mean([r["text_em"] for r in records])), + "image_em": float(np.mean([r["image_em"] for r in records])), + "records": records, + } + np.savez_compressed( + out_dir / "carrier_convergence.npz", + text_states=text_arr, + image_states=image_arr, + text_sim=text_sim_by_layer, + image_sim=image_sim_by_layer, + cross_sim=cross_sim_by_layer, + ) + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + print(json.dumps({k: v for k, v in summary.items() if k not in ("per_layer", "records")}, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_convergence_3d.py b/packages/snapcompact/research/snapcompact_convergence_3d.py new file mode 100644 index 000000000..ec4f8ed28 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_convergence_3d.py @@ -0,0 +1,187 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["matplotlib", "numpy", "pillow"] +# /// +"""3D convergence strands: text and image trajectories fusing through depth.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +BG = (5, 7, 10) +PANEL = (12, 17, 23) +INK = (241, 239, 224) +MUTED = (143, 154, 160) +AMBER = (255, 196, 68) +CYAN = (75, 220, 255) +ORANGE = (255, 112, 72) + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def question_hue(i: int, n: int) -> tuple[float, float, float]: + h = i / n + r = 0.5 + 0.5 * math.cos(2 * math.pi * (h + 0.00)) + g = 0.5 + 0.5 * math.cos(2 * math.pi * (h + 0.33)) + b = 0.5 + 0.5 * math.cos(2 * math.pi * (h + 0.67)) + return (0.28 + 0.72 * r, 0.28 + 0.72 * g, 0.28 + 0.72 * b) + + +def center(arr: np.ndarray) -> np.ndarray: + return arr - arr.mean(axis=0, keepdims=True) + + +def smooth_path(path: np.ndarray, passes: int = 2) -> np.ndarray: + out = path.copy() + for _ in range(passes): + mid = (out[:-2] + out[1:-1] * 2 + out[2:]) / 4 + out[1:-1] = mid + return out + + +def render_strands(text_arr: np.ndarray, image_arr: np.ndarray, best_layer: int) -> Image.Image: + n_q, n_layers, _ = text_arr.shape + ref = np.concatenate([center(text_arr[:, best_layer, :]), center(image_arr[:, best_layer, :])], axis=0) + _, _, vt = np.linalg.svd(ref, full_matrices=False) + basis = vt[:2].T + + # Per-layer projections, per-layer scale normalization so depth shows shape, + # not raw norm growth across layers. + t_proj = np.zeros((n_q, n_layers, 2), dtype=np.float64) + i_proj = np.zeros((n_q, n_layers, 2), dtype=np.float64) + for layer in range(n_layers): + t = center(text_arr[:, layer, :]) @ basis + i = center(image_arr[:, layer, :]) @ basis + scale = max(1e-6, float(np.abs(np.concatenate([t, i], axis=0)).max())) + t_proj[:, layer] = t / scale + i_proj[:, layer] = i / scale + + fig = plt.figure(figsize=(15.2, 9.4), dpi=170) + fig.patch.set_facecolor("#05070a") + ax = fig.add_subplot(111, projection="3d") + ax.set_facecolor((0.02, 0.025, 0.035, 1)) + for axis in (ax.xaxis, ax.yaxis, ax.zaxis): + axis.pane.set_facecolor((0.02, 0.025, 0.035, 0.0)) + axis._axinfo["grid"]["color"] = (0.32, 0.42, 0.48, 0.16) + ax.tick_params(colors="#8f9aa0", labelsize=8) + + layers_axis = np.arange(n_layers) + for qi in range(n_q): + color = question_hue(qi, n_q) + tp = smooth_path(np.column_stack([layers_axis, t_proj[qi, :, 0], t_proj[qi, :, 1]])) + ip = smooth_path(np.column_stack([layers_axis, i_proj[qi, :, 0], i_proj[qi, :, 1]])) + ax.plot(tp[:, 0], tp[:, 1], tp[:, 2], color=color, linewidth=2.6, alpha=0.95) + ax.plot(ip[:, 0], ip[:, 1], ip[:, 2], color=color, linewidth=2.6, alpha=0.55, linestyle=(0, (4, 2))) + # tie-lines every few layers showing the closing gap + for layer in range(1, n_layers, 4): + ax.plot( + [layer, layer], + [t_proj[qi, layer, 0], i_proj[qi, layer, 0]], + [t_proj[qi, layer, 1], i_proj[qi, layer, 1]], + color=color, + linewidth=0.9, + alpha=0.38, + ) + ax.scatter([0], [t_proj[qi, 0, 0]], [t_proj[qi, 0, 1]], color=color, s=26, marker="o", depthshade=False) + ax.scatter([0], [i_proj[qi, 0, 0]], [i_proj[qi, 0, 1]], color=color, s=30, marker="D", depthshade=False) + ax.scatter([best_layer], [t_proj[qi, best_layer, 0]], [t_proj[qi, best_layer, 1]], color=color, s=46, marker="o", edgecolors="white", linewidths=0.6, depthshade=False) + + # Peak-layer plane. + yy, zz = np.meshgrid(np.linspace(-1.05, 1.05, 2), np.linspace(-1.05, 1.05, 2)) + ax.plot_surface(np.full_like(yy, best_layer), yy, zz, color=(1.0, 0.77, 0.27, 0.10), shade=False) + + ax.set_xlim(0, n_layers - 1) + ax.set_ylim(-1.1, 1.1) + ax.set_zlim(-1.1, 1.1) + ax.set_xlabel("decoder layer →", color="#8f9aa0", labelpad=12) + ax.set_ylabel("content PC1", color="#8f9aa0", labelpad=10) + ax.set_zlabel("content PC2", color="#8f9aa0", labelpad=8) + ax.view_init(elev=18, azim=-66) + ax.set_box_aspect((2.9, 1.0, 0.9)) + tmp = HERE / "results" / ".convergence-3d-panel.png" + fig.subplots_adjust(left=0, right=1, top=1, bottom=0) + fig.savefig(tmp, facecolor=fig.get_facecolor(), transparent=False, bbox_inches="tight", pad_inches=0.05) + plt.close(fig) + img = Image.open(tmp).convert("RGB") + tmp.unlink(missing_ok=True) + return img + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-carrier-convergence-n12")) + ap.add_argument("--out", default=str(HERE / "results" / "qwen-carrier-convergence-n12" / "convergence-strands-3d.png")) + args = ap.parse_args() + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "carrier_convergence.npz") + best_layer = summary["best_layer"] + best = summary["best"] + panel = render_strands(data["text_states"], data["image_states"], best_layer) + + w, h = 2200, 1300 + canvas = Image.new("RGB", (w, h), BG) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-260, -220, 900, 700), fill=(75, 220, 255, 27)) + gd.ellipse((1240, 160, 2460, 1360), fill=(255, 112, 72, 25)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(84))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 42), "QWEN CARRIER CONVERGENCE — 3D STRANDS", fill=AMBER, font=ui_font(24, True)) + draw.text((64, 84), "Twelve thoughts, two doors, one room", fill=INK, font=ui_font(64, True)) + draw.text( + (66, 164), + "Each color is one question travelling through the decoder. Solid strand entered as text; dashed strand entered as pixels. Strand pairs braid together by depth.", + fill=MUTED, + font=ui_font(23), + ) + + draw.rounded_rectangle((64, 234, 2136, 1146), radius=30, fill=PANEL, outline=(35, 49, 59), width=1) + panel = panel.resize((1980, 832), Image.Resampling.LANCZOS) + canvas.paste(panel, (104, 286)) + draw.text((96, 252), f"PCA frame fixed at peak layer {best_layer}; per-layer scale normalized", fill=MUTED, font=ui_font(17)) + + stats = [ + ("matched cosine", f"{best['matched_cosine']:.2f}"), + ("mismatched", f"{best['mismatched_cosine']:.2f}"), + ("RSA geometry", f"{best['rsa_pearson']:.2f}"), + ("pair retrieval", f"{best['match_rank_accuracy'] * 100:.0f}%"), + ] + sx = 64 + for title, value in stats: + draw.rounded_rectangle((sx, 1170, sx + 320, 1262), radius=18, fill=PANEL, outline=(35, 49, 59), width=1) + draw.text((sx + 22, 1184), title, fill=MUTED, font=ui_font(16)) + draw.text((sx + 22, 1208), value, fill=INK, font=ui_font(34, True)) + sx += 344 + draw.text((sx + 20, 1196), "solid = text carrier dashed = image carrier thin rungs = pair gap", fill=MUTED, font=ui_font(18)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_convergence_extras.py b/packages/snapcompact/research/snapcompact_convergence_extras.py new file mode 100644 index 000000000..7ab101d6b --- /dev/null +++ b/packages/snapcompact/research/snapcompact_convergence_extras.py @@ -0,0 +1,207 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Extra convergence graphics: PCA funnel snapshots and an animated diagonal GIF.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "amber": (255, 196, 68), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "grid": (38, 49, 58), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def question_hue(i: int, n: int) -> tuple[int, int, int]: + """Distinct, bright hue per question.""" + h = i / n + r = 0.5 + 0.5 * math.cos(2 * math.pi * (h + 0.00)) + g = 0.5 + 0.5 * math.cos(2 * math.pi * (h + 0.33)) + b = 0.5 + 0.5 * math.cos(2 * math.pi * (h + 0.67)) + return (round(70 + 185 * r), round(70 + 185 * g), round(70 + 185 * b)) + + +def center(arr: np.ndarray) -> np.ndarray: + return arr - arr.mean(axis=0, keepdims=True) + + +def diverging_color(t: float) -> tuple[int, int, int]: + t = max(-1.0, min(1.0, t)) + if t < 0: + u = -t + return (round(8 + 12 * u), round(20 + 90 * u), round(34 + 190 * u)) + return (round(8 + 247 * t), round(20 + 130 * t), round(34 + 20 * t)) + + +def background(w: int, h: int) -> Image.Image: + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -200, 880, 680), fill=(75, 220, 255, 25)) + gd.ellipse((w - 1000, h - 760, w + 240, h + 220), fill=(255, 112, 72, 25)) + return Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(84))).convert("RGB") + + +def render_funnel(out_path: Path, text_arr: np.ndarray, image_arr: np.ndarray, layers_meta: list[dict[str, Any]], best_layer: int, records: list[dict[str, Any]]) -> None: + n_q, n_layers, _ = text_arr.shape + snapshots = [1, max(2, best_layer // 2), best_layer] + # Shared PCA frame from the peak layer keeps the panels comparable. + ref = np.concatenate([center(text_arr[:, best_layer, :]), center(image_arr[:, best_layer, :])], axis=0) + _, _, vt = np.linalg.svd(ref, full_matrices=False) + basis = vt[:2].T # [D, 2] + + w, h = 2200, 1240 + canvas = background(w, h) + draw = ImageDraw.Draw(canvas) + draw.text((64, 42), "QWEN CARRIER CONVERGENCE — TRAJECTORY VIEW", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((64, 84), "Watch the two carriers fuse", fill=PALETTE["ink"], font=ui_font(64, True)) + draw.text((66, 164), "Each color is one question; ● came in as text, ◆ came in as pixels. Same 2D projection at every depth. The tie-lines shrink as carriers converge.", fill=PALETTE["muted"], font=ui_font(23)) + + panel_w = 660 + titles = ["early (layer {})", "middle (layer {})", "peak (layer {})"] + for pi, (layer, title) in enumerate(zip(snapshots, titles)): + x0 = 64 + pi * (panel_w + 44) + box = (x0, 232, x0 + panel_w, 952) + draw.rounded_rectangle(box, radius=24, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((x0 + 26, 252), title.format(layer), fill=PALETTE["ink"], font=ui_font(27, True)) + t_proj = center(text_arr[:, layer, :]) @ basis + i_proj = center(image_arr[:, layer, :]) @ basis + both = np.concatenate([t_proj, i_proj], axis=0) + lim = float(np.abs(both).max()) * 1.15 or 1.0 + gx0, gy0, gx1, gy1 = x0 + 36, 306, x0 + panel_w - 36, 912 + def to_px(p: np.ndarray) -> tuple[int, int]: + return ( + round(gx0 + (p[0] + lim) / (2 * lim) * (gx1 - gx0)), + round(gy0 + (1 - (p[1] + lim) / (2 * lim)) * (gy1 - gy0)), + ) + draw.line((gx0, (gy0 + gy1) // 2, gx1, (gy0 + gy1) // 2), fill=PALETTE["grid"], width=1) + draw.line(((gx0 + gx1) // 2, gy0, (gx0 + gx1) // 2, gy1), fill=PALETTE["grid"], width=1) + pair_dist = 0.0 + for qi in range(n_q): + color = question_hue(qi, n_q) + tp = to_px(t_proj[qi]) + ip = to_px(i_proj[qi]) + draw.line((tp, ip), fill=(*color, 0)[:3], width=3) + r = 11 + draw.ellipse((tp[0] - r, tp[1] - r, tp[0] + r, tp[1] + r), fill=color, outline=(8, 10, 12), width=2) + d = ImageDraw.Draw(canvas) + d.polygon([(ip[0], ip[1] - r - 2), (ip[0] + r + 2, ip[1]), (ip[0], ip[1] + r + 2), (ip[0] - r - 2, ip[1])], fill=color, outline=(8, 10, 12)) + pair_dist += float(np.linalg.norm(t_proj[qi] - i_proj[qi])) + pair_dist /= n_q + norm_dist = pair_dist / (2 * lim) + meta = layers_meta[layer] + draw.text((x0 + 26, 916), f"mean pair gap: {norm_dist * 100:.0f}% of frame · matched cos {meta['matched_cosine']:.2f}", fill=PALETTE["muted"], font=ui_font(17)) + + # Pair-distance by layer strip. + strip = (64, 996, 2136, 1190) + draw.rounded_rectangle(strip, radius=24, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 1014), "matched-pair separation by layer (lower = carriers agree)", fill=PALETTE["ink"], font=ui_font(22, True)) + gx0, gy0, gx1, gy1 = 110, 1062, 2100, 1162 + gaps = [] + for layer in range(n_layers): + t_proj = center(text_arr[:, layer, :]) + i_proj = center(image_arr[:, layer, :]) + t_n = t_proj / np.maximum(np.linalg.norm(t_proj, axis=1, keepdims=True), 1e-6) + i_n = i_proj / np.maximum(np.linalg.norm(i_proj, axis=1, keepdims=True), 1e-6) + gaps.append(1.0 - float((t_n * i_n).sum(axis=1).mean())) + hi = max(gaps) + bw = (gx1 - gx0) / n_layers + for layer, gap in enumerate(gaps): + xa = gx0 + layer * bw + 3 + xb = gx0 + (layer + 1) * bw - 3 + bh = (gy1 - gy0) * gap / hi + color = PALETTE["orange"] if layer == best_layer else (62, 86, 102) + draw.rounded_rectangle((round(xa), round(gy1 - bh), round(xb), gy1), radius=5, fill=color) + draw.text((gx0, gy1 + 6), "layer 0", fill=PALETTE["muted"], font=ui_font(13)) + draw.text((gx1 - 70, gy1 + 6), f"layer {n_layers - 1}", fill=PALETTE["muted"], font=ui_font(13)) + + out_path.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out_path) + + +def render_gif(out_path: Path, cross_sim: np.ndarray, layers_meta: list[dict[str, Any]]) -> None: + n_layers, n_q, _ = cross_sim.shape + cell = 46 + pad = 36 + header = 132 + w = n_q * cell + pad * 2 + h = n_q * cell + header + pad + 64 + frames: list[Image.Image] = [] + for layer in range(n_layers): + frame = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(frame) + for y in range(0, h, 14): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + draw.text((pad, 22), "cross-carrier matching", fill=PALETTE["ink"], font=ui_font(30, True)) + draw.text((pad, 62), "text question i × image question j", fill=PALETTE["muted"], font=ui_font(17)) + meta = layers_meta[layer] + draw.text((pad, 92), f"layer {layer:02d} matched {meta['matched_cosine']:+.2f} others {meta['mismatched_cosine']:+.2f}", fill=PALETTE["amber"], font=ui_font(19, True)) + for r in range(n_q): + for c in range(n_q): + xa = pad + c * cell + ya = header + r * cell + draw.rounded_rectangle((xa, ya, xa + cell - 4, ya + cell - 4), radius=7, fill=diverging_color(float(cross_sim[layer, r, c]))) + # progress bar + bar_y = header + n_q * cell + 18 + draw.rounded_rectangle((pad, bar_y, w - pad, bar_y + 10), radius=5, fill=(30, 40, 48)) + draw.rounded_rectangle((pad, bar_y, pad + (w - 2 * pad) * (layer + 1) // n_layers, bar_y + 10), radius=5, fill=PALETTE["cyan"]) + frames.append(frame) + durations = [240] * n_layers + durations[-1] = 2200 + out_path.parent.mkdir(parents=True, exist_ok=True) + frames[0].save(out_path, save_all=True, append_images=frames[1:], duration=durations, loop=0) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-carrier-convergence-n12")) + args = ap.parse_args() + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "carrier_convergence.npz") + text_arr = data["text_states"] + image_arr = data["image_states"] + cross_sim = data["cross_sim"] + layers_meta = summary["per_layer"] + best_layer = summary["best_layer"] + + funnel_path = result_dir / "convergence-funnel.png" + gif_path = result_dir / "diagonal-emerges.gif" + render_funnel(funnel_path, text_arr, image_arr, layers_meta, best_layer, summary["records"]) + render_gif(gif_path, cross_sim, layers_meta) + print(funnel_path) + print(gif_path) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_convergence_viz.py b/packages/snapcompact/research/snapcompact_convergence_viz.py new file mode 100644 index 000000000..b28b33088 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_convergence_viz.py @@ -0,0 +1,219 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render the carrier-convergence geometry figure: same content, same nothings.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "red": (255, 76, 62), + "grid": (38, 49, 58), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def diverging_color(t: float) -> tuple[int, int, int]: + """-1..1 → blue-black-orange diverging scale.""" + t = max(-1.0, min(1.0, t)) + if t < 0: + u = -t + return (round(8 + 12 * u), round(20 + 90 * u), round(34 + 190 * u)) + u = t + return (round(8 + 247 * u), round(20 + 130 * u), round(34 + 20 * u)) + + +def draw_matrix(draw: ImageDraw.ImageDraw, mat: np.ndarray, box: tuple[int, int, int, int], title: str, subtitle: str, color: tuple[int, int, int], highlight_diag: bool = False) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=20, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((x0 + 20, y0 + 16), title, fill=color, font=ui_font(24, True)) + draw.text((x0 + 20, y0 + 48), subtitle, fill=PALETTE["muted"], font=ui_font(15)) + n = mat.shape[0] + gx0, gy0 = x0 + 34, y0 + 84 + side = min(x1 - 34 - gx0, y1 - 30 - gy0) + cell = side / n + for r in range(n): + for c in range(n): + xa = round(gx0 + c * cell) + xb = round(gx0 + (c + 1) * cell) - 2 + ya = round(gy0 + r * cell) + yb = round(gy0 + (r + 1) * cell) - 2 + draw.rounded_rectangle((xa, ya, xb, yb), radius=4, fill=diverging_color(float(mat[r, c]))) + if highlight_diag: + for r in range(n): + xa = round(gx0 + r * cell) + ya = round(gy0 + r * cell) + draw.rounded_rectangle((xa - 1, ya - 1, round(xa + cell) - 1, round(ya + cell) - 1), radius=5, outline=PALETTE["amber"], width=2) + draw.text((gx0, round(gy0 + side) + 6), "questions →", fill=PALETTE["muted"], font=ui_font(13)) + + +def draw_curves(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], layers: list[dict[str, Any]]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=20, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((x0 + 22, y0 + 16), "convergence by depth", fill=PALETTE["ink"], font=ui_font(24, True)) + draw.text((x0 + 22, y0 + 48), "carrier-centered cosine: same question across carriers vs different questions", fill=PALETTE["muted"], font=ui_font(15)) + gx0, gy0, gx1, gy1 = x0 + 52, y0 + 92, x1 - 26, y1 - 56 + lo, hi = -0.15, 1.0 + for i in range(5): + y = gy0 + (gy1 - gy0) * i / 4 + draw.line((gx0, y, gx1, y), fill=PALETTE["grid"], width=1) + value = hi - (hi - lo) * i / 4 + draw.text((x0 + 12, y - 8), f"{value:.1f}", fill=PALETTE["muted"], font=ui_font(12)) + series = [ + ("matched_cosine", PALETTE["amber"], 6), + ("mismatched_cosine", PALETTE["muted"], 4), + ("rsa_pearson", PALETTE["cyan"], 4), + ] + n = len(layers) + for key, color, width in series: + pts = [] + for i, row in enumerate(layers): + value = float(row[key]) + if value != value: # NaN guard (layer-0 RSA is undefined) + continue + value = max(lo, min(hi, value)) + x = gx0 + (gx1 - gx0) * i / max(1, n - 1) + y = gy1 - (gy1 - gy0) * (value - lo) / (hi - lo) + pts.append((round(x), round(y))) + if len(pts) < 2: + continue + draw.line(pts, fill=color, width=width, joint="curve") + draw.text((gx0, gy1 + 14), "layer 0", fill=PALETTE["muted"], font=ui_font(13)) + draw.text((gx1 - 64, gy1 + 14), f"layer {n - 1}", fill=PALETTE["muted"], font=ui_font(13)) + legend = [("same question, text↔image", PALETTE["amber"]), ("different questions", PALETTE["muted"]), ("RSA geometry corr", PALETTE["cyan"])] + lx = gx0 + for label, color in legend: + draw.rounded_rectangle((lx, y0 + 70, lx + 18, y0 + 78), radius=4, fill=color) + draw.text((lx + 24, y0 + 62), label, fill=PALETTE["muted"], font=ui_font(13)) + lx += 232 + + +def draw_answers(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], records: list[dict[str, Any]]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=20, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((x0 + 22, y0 + 16), "behavioral check: both carriers answer alike", fill=PALETTE["ink"], font=ui_font(24, True)) + draw.text((x0 + 240, y0 + 56), "text carrier", fill=PALETTE["cyan"], font=ui_font(15, True)) + draw.text((x0 + 470, y0 + 56), "image carrier", fill=PALETTE["orange"], font=ui_font(15, True)) + y = y0 + 84 + row_h = (y1 - y0 - 96) // len(records) + fnt = mono_font(15) + for r in records: + gold = r["gold"][:22] + draw.text((x0 + 22, y), gold, fill=PALETTE["muted"], font=fnt) + draw.text((x0 + 240, y), r["text_answer"][:22], fill=PALETTE["ink"], font=fnt) + draw.text((x0 + 470, y), r["image_answer"][:22], fill=PALETTE["ink"], font=fnt) + mark = "=" if r["agree"] else "≠" + draw.text((x1 - 44, y), mark, fill=PALETTE["green"] if r["agree"] else PALETTE["red"], font=ui_font(17, True)) + y += row_h + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-carrier-convergence-n12")) + ap.add_argument("--out", default=str(HERE / "results" / "qwen-carrier-convergence-n12" / "carrier-convergence.png")) + args = ap.parse_args() + + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "carrier_convergence.npz") + layers = summary["per_layer"] + best_layer = summary["best_layer"] + text_sim = data["text_sim"][best_layer] + image_sim = data["image_sim"][best_layer] + cross_sim = data["cross_sim"][best_layer] + records = summary["records"] + + w, h = 2200, 1320 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -200, 900, 700), fill=(75, 220, 255, 26)) + gd.ellipse((1300, 180, 2480, 1380), fill=(255, 112, 72, 24)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(86))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + best = summary["best"] + draw.text((64, 42), "QWEN CARRIER CONVERGENCE", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((64, 84), "Two carriers, one thought", fill=PALETTE["ink"], font=ui_font(66, True)) + draw.text( + (66, 166), + "Hidden state at the answer position, carrier means removed. Same question through text or bitmap lands in the same place; different questions do not.", + fill=PALETTE["muted"], + font=ui_font(23), + ) + + stats = [ + ("matched pairs", f"{best['matched_cosine']:.2f}", "same Q, text ↔ image"), + ("mismatched pairs", f"{best['mismatched_cosine']:.2f}", "different questions"), + ("RSA geometry corr", f"{best['rsa_pearson']:.2f}", f"layer {best['layer']}"), + ("pair retrieval", f"{best['match_rank_accuracy'] * 100:.0f}%", "nearest cross-carrier match"), + ("answer agreement", f"{summary['answer_agreement'] * 100:.0f}%", "text vs image generations"), + ] + sx = 64 + for title, value, caption in stats: + draw.rounded_rectangle((sx, 222, sx + 396, 332), radius=20, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((sx + 22, 240), title, fill=PALETTE["muted"], font=ui_font(16)) + draw.text((sx + 22, 264), value, fill=PALETTE["ink"], font=ui_font(40, True)) + draw.text((sx + 226, 290), caption, fill=PALETTE["muted"], font=ui_font(13)) + sx += 420 + + n = text_sim.shape[0] + draw_matrix(draw, text_sim, (64, 376, 600, 952), "text-carrier geometry", f"{n}×{n} question similarity, layer {best_layer}", PALETTE["cyan"]) + draw_matrix(draw, image_sim, (628, 376, 1164, 952), "image-carrier geometry", "same questions through the bitmap — same shape", PALETTE["orange"]) + draw_matrix(draw, cross_sim, (1192, 376, 1728, 952), "cross-carrier matching", "text question i × image question j — bright diagonal", PALETTE["green"], highlight_diag=True) + + draw_curves(draw, (64, 996, 1164, 1264), layers) + draw_answers(draw, (1192, 996, 2136, 1264), records) + + # Color scale. + for i in range(240): + t = 1 - i / 239 * 2 + draw.rectangle((1816, 420 + i * 2, 1836, 422 + i * 2), fill=diverging_color(t)) + draw.text((1848, 412), "+1 similar", fill=PALETTE["muted"], font=ui_font(14)) + draw.text((1848, 884), "−1 opposite", fill=PALETTE["muted"], font=ui_font(14)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_lockon_anatomy_viz.py b/packages/snapcompact/research/snapcompact_lockon_anatomy_viz.py new file mode 100644 index 000000000..357c63173 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_lockon_anatomy_viz.py @@ -0,0 +1,288 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""The lock-on instrument: how we decide WHERE the answer materializes. + +Renders the measurement methodology as a depth-gauge diagram: the answer patch +descends the decoder shaft; at every layer a logit-lens probe (final norm + LM +head) reads the vocabulary distribution; lock-on is the first layer whose top-1 +token is a BPE piece of the answer. All readouts are real sweep data. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +P = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "faint": (90, 101, 108), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "purple": (188, 112, 255), + "grid": (38, 49, 58), +} +COND_COLORS = { + "base-8x13": (143, 154, 160), + "repeat2-color": (255, 196, 68), + "align-7x14": (148, 255, 117), + "align-14x28": (75, 220, 255), + "align-28x28": (255, 112, 72), + "repeat2-align-14x28": (188, 112, 255), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def label_font(label: str, size: int) -> ImageFont.ImageFont: + """Monaco for ASCII; Arial Unicode for anything it cannot shape (CJK).""" + if all(ord(ch) < 0x2000 for ch in label): + return mono_font(size) + unicode_path = "/System/Library/Fonts/Supplemental/Arial Unicode.ttf" + if Path(unicode_path).exists(): + return ImageFont.truetype(unicode_path, size) + return mono_font(size) + + +def crosshair(draw: ImageDraw.ImageDraw, cx: int, cy: int, r: int, color: tuple[int, int, int], width: int = 4) -> None: + draw.ellipse((cx - r, cy - r, cx + r, cy + r), outline=color, width=width) + draw.ellipse((cx - r // 2, cy - r // 2, cx + r // 2, cy + r // 2), outline=color, width=2) + for dx, dy in ((-1, 0), (1, 0), (0, -1), (0, 1)): + draw.line((cx + dx * (r - 6), cy + dy * (r - 6), cx + dx * (r + 14), cy + dy * (r + 14)), fill=color, width=width) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-materialize-sweep-q3")) + ap.add_argument("--condition", default="base-8x13") + ap.add_argument("--out", default=str(HERE / "results" / "qwen-materialize-sweep-q3" / "lockon-anatomy.png")) + args = ap.parse_args() + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + conditions = {c["name"]: c for c in summary["conditions"]} + cond = conditions[args.condition] + q = summary["question"] + answer_strs = summary["answer_token_strs"] + answer_set = {s.strip() for s in answer_strs} + layers = cond["layers"] + n_layers = len(layers) + lock_on = cond["lock_on_layer"] + + w, h = 2200, 1420 + canvas = Image.new("RGB", (w, h), P["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((520, 620, 1280, 1180), fill=(255, 196, 68, 36)) + gd.ellipse((-260, -240, 760, 560), fill=(75, 220, 255, 26)) + gd.ellipse((1500, -100, 2480, 700), fill=(255, 112, 72, 20)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(90))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 40), "THE LOCK-ON INSTRUMENT", fill=P["amber"], font=ui_font(24, True)) + draw.text((64, 80), "How we decide where the answer materializes", fill=P["ink"], font=ui_font(58, True)) + draw.text( + (66, 154), + "At every layer, a logit-lens probe taps the answer patch's residual stream: final RMSNorm → LM head → softmax over 152k vocabulary entries.", + fill=P["muted"], + font=ui_font(22), + ) + draw.text( + (66, 186), + "LOCK-ON = the first layer whose #1 vocabulary entry is a BPE piece of the answer. Past this depth the fact is settled — remaining layers are free for reasoning.", + fill=P["amber"], + font=ui_font(22, True), + ) + + # ---- Probe pipeline card (top left). + pipe = (64, 248, 700, 420) + draw.rounded_rectangle(pipe, radius=22, fill=P["panel"], outline=(35, 49, 59), width=1) + draw.text((92, 268), "the probe, applied at every layer ℓ", fill=P["ink"], font=ui_font(22, True)) + stages = ["h(patch)", "RMSNorm", "LM head", "softmax", "top-1?"] + sx = 92 + for si, stage in enumerate(stages): + color = P["amber"] if si == len(stages) - 1 else P["cyan"] + tw = int(draw.textlength(stage, font=mono_font(16))) + 24 + draw.rounded_rectangle((sx, 318, sx + tw, 356), radius=10, fill=P["panel2"], outline=color, width=2) + draw.text((sx + 12, 327), stage, fill=color, font=mono_font(16)) + if si < len(stages) - 1: + draw.text((sx + tw + 4, 327), "→", fill=P["faint"], font=ui_font(18)) + sx += tw + 28 + draw.text((92, 376), f"vocabulary = 152k entries · answer BPEs = {answer_strs}", fill=P["muted"], font=mono_font(14)) + + # ---- The patch under test (left). + patch_card = (64, 460, 380, 760) + draw.rounded_rectangle(patch_card, radius=22, fill=P["panel"], outline=(35, 49, 59), width=1) + draw.text((92, 480), "specimen", fill=P["orange"], font=ui_font(21, True)) + carrier = Image.open(result_dir / "images" / f"{args.condition}.png").convert("RGB") + rw = 1568 + grid = 56 + px = 28 + lock_entry = layers[lock_on] + tok_idx = lock_entry["best_token_index"] + r0, c0 = tok_idx // grid, tok_idx % grid + cell = carrier.resize((rw, rw), Image.Resampling.LANCZOS).crop((c0 * px, r0 * px, (c0 + 1) * px, (r0 + 1) * px)) + big = cell.resize((196, 196), Image.Resampling.NEAREST) + draw.rounded_rectangle((118, 516, 326, 724), radius=12, fill=(244, 242, 230), outline=P["orange"], width=4) + canvas.paste(big, (124, 522)) + draw.text((118, 730), f"visual token #{tok_idx} · 28×28 px", fill=P["muted"], font=mono_font(13)) + + # ---- Depth shaft. + shaft_x = 470 + shaft_top, shaft_bot = 470, 1340 + draw.rounded_rectangle((shaft_x - 7, shaft_top, shaft_x + 7, shaft_bot), radius=7, fill=(20, 28, 35), outline=(40, 54, 64), width=1) + + def layer_y(layer: int) -> int: + return round(shaft_top + (shaft_bot - shaft_top) * layer / (n_layers - 1)) + + # p(answer) trajectory along the shaft. + traj = [(shaft_x + 14 + 230 * min(1.0, e["best_answer_p"]), layer_y(e["layer"])) for e in layers] + for i in range(len(traj) - 1): + draw.line((traj[i], traj[i + 1]), fill=(120, 96, 40), width=3) + draw.text((shaft_x + 30, shaft_bot + 10), "p(answer BPE) →", fill=(150, 124, 60), font=ui_font(14)) + + for layer in range(n_layers): + y = layer_y(layer) + major = layer % 4 == 0 or layer == n_layers - 1 + draw.line((shaft_x - (16 if major else 10), y, shaft_x + (16 if major else 10), y), fill=P["faint"] if major else (52, 64, 73), width=2) + if major: + draw.text((shaft_x - 58, y - 9), f"L{layer:02d}", fill=P["muted"], font=mono_font(13)) + # Patch entering the shaft. + draw.line((326, 620, shaft_x - 18, shaft_top + 6), fill=P["orange"], width=3) + draw.polygon([(shaft_x - 14, shaft_top + 2), (shaft_x - 30, shaft_top - 4), (shaft_x - 26, shaft_top + 16)], fill=P["orange"]) + + # ---- Readout cards at sampled depths (real top-5). + samples = [0, 10, 18, lock_on, n_layers - 1] + card_x = 790 + card_w = 620 + card_h = 128 + gap = 14 + desired = [layer_y(layer) - card_h // 2 for layer in samples] + card_ys = [0] * len(samples) + # Bottom-up pass: clamp the last card into the canvas, then keep every + # earlier card fully above its successor; final top clamp at 440. + card_ys[-1] = min(desired[-1], h - card_h - 70) + for i in range(len(samples) - 2, -1, -1): + card_ys[i] = min(desired[i], card_ys[i + 1] - card_h - gap) + shift = max(0, 440 - card_ys[0]) + card_ys = [cy + shift for cy in card_ys] + for layer, cy in zip(samples, card_ys): + entry = layers[layer] + is_lock = layer == lock_on + accent = P["amber"] if is_lock else P["cyan"] if entry["best_answer_p"] > 0.01 else P["faint"] + # Connector. + ly = layer_y(layer) + draw.line((shaft_x + 16, ly, card_x - 18, cy + card_h // 2), fill=accent, width=3 if is_lock else 2) + draw.ellipse((shaft_x + 12, ly - 5, shaft_x + 22, ly + 5), fill=accent) + draw.rounded_rectangle((card_x, cy, card_x + card_w, cy + card_h), radius=16, fill=P["panel2"], outline=accent, width=3 if is_lock else 1) + title = f"L{layer:02d} readout" + (" LOCK-ON" if is_lock else "") + draw.text((card_x + 20, cy + 10), title, fill=accent, font=ui_font(19, True)) + if is_lock: + tx = card_x + 20 + draw.textlength(f"L{layer:02d} readout ", font=ui_font(19, True)) + draw.ellipse((tx - 8, cy + 14, tx + 4, cy + 26), outline=accent, width=3) + bx = card_x + 20 + by = cy + 44 + for k, t in enumerate(entry["best_token_top"]): + label = t["str"].strip() or "␣" + if len(label) > 9: + label = label[:8] + "…" + hit = t["str"].strip() in answer_set + pill_w = 108 + fill = (66, 92, 36) if hit else (16, 22, 28) + outline = P["green"] if hit else (38, 52, 61) + draw.rounded_rectangle((bx, by, bx + pill_w, by + 30), radius=8, fill=fill, outline=outline, width=2) + draw.text((bx + 8, by + 6), label, fill=(220, 255, 190) if hit else P["ink"], font=label_font(label, 13)) + bar = round(min(1.0, t["p"] / 0.4) * pill_w) + draw.rounded_rectangle((bx, by + 36, bx + max(3, bar), by + 42), radius=3, fill=P["amber"] if hit else (60, 76, 88)) + draw.text((bx, by + 46, ), f"{t['p']:.3f}", fill=P["muted"], font=mono_font(10)) + bx += pill_w + 12 + if is_lock: + crosshair(draw, shaft_x, ly, 26, P["amber"], 4) + draw.text((shaft_x + 44, ly + 26), f"first top-1 hit: “{entry['best_token_top'][0]['str'].strip()}” p={entry['best_token_top'][0]['p']:.2f}", fill=P["amber"], font=ui_font(16, True)) + + # ---- Why it matters (right column). + why = (1460, 248, 2136, 716) + draw.rounded_rectangle(why, radius=22, fill=P["panel"], outline=(35, 49, 59), width=1) + draw.text((1492, 270), "why lock-on is the metric", fill=P["ink"], font=ui_font(26, True)) + lines = [ + ("It separates decoding from reasoning.", P["ink"]), + ("Layers before lock-on are spent turning", P["muted"]), + ("pixels into words; layers after are free to", P["muted"]), + ("reason about them. Earlier, harder lock-on", P["muted"]), + ("= more of the network left for thinking.", P["muted"]), + ] + ty = 314 + for text, color in lines: + draw.text((1492, ty), text, fill=color, font=ui_font(19)) + ty += 30 + draw.line((1492, ty + 8, 2104, ty + 8), fill=P["grid"], width=1) + ty += 26 + draw.text((1492, ty), "reasoning budget after lock-on", fill=P["muted"], font=ui_font(16, True)) + ty += 30 + for name, color in COND_COLORS.items(): + c = conditions.get(name) + if not c or c["lock_on_layer"] is None: + continue + budget = n_layers - 1 - c["lock_on_layer"] + bw_px = round(budget / (n_layers - 1) * 430) + draw.text((1492, ty), name, fill=color, font=mono_font(13)) + draw.rounded_rectangle((1492, ty + 20, 1492 + bw_px, ty + 32), radius=6, fill=color) + draw.text((1492 + bw_px + 10, ty + 17), f"{budget} layers · p {c['max_answer_p']:.2f}", fill=P["muted"], font=mono_font(12)) + ty += 44 + + # ---- Rule plate (bottom right). + plate = (1460, 740, 2136, 1000) + draw.rounded_rectangle(plate, radius=22, fill=P["panel"], outline=(255, 196, 68), width=2) + draw.text((1492, 762), "the rule", fill=P["amber"], font=ui_font(24, True)) + rule_lines = [ + "lock_on(patch) = min L such that", + " argmax softmax(W * norm(h_L))", + f" in {{{answer_strs[0]!r}, {answer_strs[1]!r}}}", + "", + f"here: L = {lock_on}, decoded “{layers[lock_on]['best_token_top'][0]['str'].strip()}”", + ] + ry = 806 + for line in rule_lines: + draw.text((1492, ry), line, fill=P["ink"] if line else P["muted"], font=mono_font(17)) + ry += 32 + draw.text((1492, 1014), f"question: {q['q'][:60]}…", fill=P["muted"], font=ui_font(15)) + draw.text((1492, 1040), f"answer: “{q['answer_text']}” · condition: {args.condition} · generation: “{cond['generation']}”", fill=P["muted"], font=ui_font(15)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_logit_lens_dump.py b/packages/snapcompact/research/snapcompact_logit_lens_dump.py new file mode 100644 index 000000000..095c5de17 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_logit_lens_dump.py @@ -0,0 +1,143 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers"] +# /// +"""Logit-lens dump: what vocabulary word does each visual token become, per layer? + +For the visual tokens covering the answer word in a snapcompact bitmap, decode +every layer's hidden state through the final norm + lm_head and record the +top-k vocabulary tokens. If the bitmap is truly read into text space, the +patches' hidden states should decode to the answer's BPE tokens mid-stack. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import sample_answer_questions # noqa: E402 +from snapcompact_text_image_compare import image_answer_token_indices # noqa: E402 + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default="Qwen/Qwen2.5-VL-7B-Instruct") + ap.add_argument("--font", default="8x13", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--limit-paras", type=int, default=80) + ap.add_argument("--question-index", type=int, default=3) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--topk", type=int, default=5) + ap.add_argument("--control-tokens", type=int, default=2) + ap.add_argument("--out", default="qwen-logit-lens") + args = ap.parse_args() + + import torch + from transformers import AutoProcessor, AutoTokenizer, Qwen2_5_VLForConditionalGeneration + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + chunk = flow[: min(len(flow), budget)] + questions = sample_answer_questions(paras, offsets, 0, len(chunk), 24, args.seed) + q = questions[min(args.question_index, len(questions) - 1)] + img = render(chunk, cfg, CACHE, args.size, args.variant) + img.save(img_dir / "image-carrier.png") + + print(f"loading {args.model_dir}", flush=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + tokenizer = AutoTokenizer.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True) + model = Qwen2_5_VLForConditionalGeneration.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto").eval() + device = next(model.parameters()).device + + prompt = load_prompt("qa-image.md").format(cols=cols, rows=rows) + f"\n\nQuestion: {q['q']}\nAnswer with only the shortest extractive answer." + messages = [{"role": "user", "content": [{"type": "image", "image": img}, {"type": "text", "text": prompt}]}] + templated = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + batch = processor(images=img, text=templated, return_tensors="pt") + image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token) + ids = batch["input_ids"][0].tolist() + image_positions = [i for i, token_id in enumerate(ids) if token_id == image_token_id] + n_tokens = len(image_positions) + grid = int(round(n_tokens**0.5)) + answer_indices = image_answer_token_indices(q["answer_start"], q["answer_end"], cols, cfg.adv, cfg.pitch, img.width, img.height, n_tokens) + + # Controls: blank-region tokens far from any text row boundary effects. + control_indices = [] + if answer_indices: + row_far = (answer_indices[0] // grid + grid // 2) % grid + for k in range(args.control_tokens): + control_indices.append(row_far * grid + (answer_indices[0] % grid + k)) + track = [("answer", idx) for idx in answer_indices] + [("control", idx) for idx in control_indices] + track_positions = [image_positions[idx] for _kind, idx in track] + + batch = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in batch.items()} + with torch.no_grad(): + out = model(**batch, output_hidden_states=True, use_cache=False) + + norm = model.model.language_model.norm + lm_head = model.lm_head + answer_token_ids = tokenizer(q["answer_text"], add_special_tokens=False)["input_ids"] + answer_token_strs = [tokenizer.decode([t]) for t in answer_token_ids] + + lens: list[dict[str, Any]] = [] + with torch.no_grad(): + for layer, hidden in enumerate(out.hidden_states): + states = hidden[0, track_positions, :] + logits = lm_head(norm(states)).float() + probs = torch.softmax(logits, dim=-1) + topv, topi = probs.topk(args.topk, dim=-1) + for ti, (kind, idx) in enumerate(track): + entry = { + "layer": layer, + "kind": kind, + "token_index": int(idx), + "grid_rc": [int(idx // grid), int(idx % grid)], + "top": [ + {"str": tokenizer.decode([int(topi[ti, k])]), "id": int(topi[ti, k]), "p": round(float(topv[ti, k]), 5)} + for k in range(args.topk) + ], + "answer_token_p": [round(float(probs[ti, t]), 6) for t in answer_token_ids], + } + lens.append(entry) + print(f"layer {layer} done", flush=True) + + dump = { + "args": vars(args), + "question": {"q": q["q"], "answer_text": q["answer_text"], "answer_start": q["answer_start"], "answer_end": q["answer_end"]}, + "geometry": {"cols": cols, "rows": rows, "image_w": img.width, "image_h": img.height}, + "image_tokens": n_tokens, + "image_grid": grid, + "token_pixel_size": 28, + "answer_token_ids": answer_token_ids, + "answer_token_strs": answer_token_strs, + "answer_indices": [int(i) for i in answer_indices], + "control_indices": [int(i) for i in control_indices], + "layers": len(out.hidden_states), + "lens": lens, + } + (out_dir / "logit_lens.json").write_text(json.dumps(dump, indent=1)) + # Quick console summary: best layer per answer token. + for kind, idx in track: + best = max((e for e in lens if e["token_index"] == idx), key=lambda e: max(e["answer_token_p"])) + print(kind, idx, "best layer", best["layer"], "p", max(best["answer_token_p"]), "top1", best["top"][0]["str"]) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_logit_lens_viz.py b/packages/snapcompact/research/snapcompact_logit_lens_viz.py new file mode 100644 index 000000000..9c9b4eb15 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_logit_lens_viz.py @@ -0,0 +1,169 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render the logit-lens grid: pixel patches morphing into BPE tokens by layer.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "grid": (38, 49, 58), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def heat_fill(p: float, hit: bool) -> tuple[int, int, int]: + """Cell background: brightness by confidence, hue by answer-match.""" + t = min(1.0, max(0.0, math.log10(max(p, 1e-6)) / 3 + 1)) # p=1 -> 1, p=1e-3 -> 0 + if hit: + return (round(30 + 130 * t), round(48 + 130 * t), round(18 + 40 * t)) + return (round(14 + 26 * t), round(19 + 30 * t), round(26 + 36 * t)) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-logit-lens-q3")) + ap.add_argument("--out", default=str(HERE / "results" / "qwen-logit-lens-q3" / "logit-lens-grid.png")) + ap.add_argument("--layer-step", type=int, default=1) + args = ap.parse_args() + result_dir = Path(args.result_dir) + dump = json.loads((result_dir / "logit_lens.json").read_text()) + carrier = Image.open(result_dir / "images" / "image-carrier.png").convert("RGB") + + q = dump["question"] + answer = q["answer_text"] + answer_ids = set(dump["answer_token_ids"]) + grid = dump["image_grid"] + px = dump["token_pixel_size"] + rw = grid * px + resized = carrier.resize((rw, rw), Image.Resampling.LANCZOS) + + track_indices = dump["answer_indices"] + dump["control_indices"] + by_token: dict[int, list[dict[str, Any]]] = {} + for e in dump["lens"]: + by_token.setdefault(e["token_index"], []).append(e) + for entries in by_token.values(): + entries.sort(key=lambda e: e["layer"]) + n_layers = dump["layers"] + layer_rows = list(range(0, n_layers, args.layer_step)) + + cell_w, cell_h = 150, 34 + header_h = 210 + left_w = 120 + n_cols = len(track_indices) + grid_w = left_w + n_cols * cell_w + margin = 64 + title_h = 200 + w = max(1900, grid_w + margin * 2) + h = title_h + header_h + len(layer_rows) * cell_h + 160 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -200, 900, 700), fill=(75, 220, 255, 26)) + gd.ellipse((w - 1000, h - 800, w + 240, h + 200), fill=(255, 112, 72, 24)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(86))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((margin, 42), "QWEN LOGIT LENS — PIXELS BECOMING WORDS", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((margin, 84), "Watch each patch decode into vocabulary", fill=PALETTE["ink"], font=ui_font(58, True)) + draw.text( + (margin + 2, 156), + f"Each column is one 28×28px visual token; each row is a decoder layer projected through the LM head. Green cells decode to a BPE piece of “{answer}”.", + fill=PALETTE["muted"], + font=ui_font(22), + ) + + gx0 = margin + left_w + gy0 = title_h + header_h + # Column headers: actual patch pixels. + patch_size = 108 + for ci, idx in enumerate(track_indices): + r, c = idx // grid, idx % grid + cell = resized.crop((c * px, r * px, (c + 1) * px, (r + 1) * px)).resize((patch_size, patch_size), Image.Resampling.NEAREST) + cx = gx0 + ci * cell_w + (cell_w - patch_size) // 2 + is_control = idx in dump["control_indices"] + color = PALETTE["muted"] if is_control else PALETTE["orange"] + draw.rounded_rectangle((cx - 4, title_h + 26, cx + patch_size + 4, title_h + 34 + patch_size), radius=8, fill=(244, 242, 230), outline=color, width=3) + canvas.paste(cell, (cx, title_h + 30)) + label = "control" if is_control else f"tok[{idx}]" + tw = draw.textlength(label, font=mono_font(13)) + draw.text((cx + (patch_size - tw) / 2, title_h + 42 + patch_size), label, fill=color, font=mono_font(13)) + draw.text((margin, title_h + 30 + patch_size // 2 - 10), "input\npixels", fill=PALETTE["muted"], font=ui_font(15, True)) + + # Grid rows. + fnt = mono_font(14) + for ri, layer in enumerate(layer_rows): + y = gy0 + ri * cell_h + draw.text((margin + 24, y + 8), f"L{layer:02d}", fill=PALETTE["muted"], font=mono_font(13)) + for ci, idx in enumerate(track_indices): + e = by_token[idx][layer] + top = e["top"][0] + hit = top["id"] in answer_ids + p_ans = max(e["answer_token_p"]) + x = gx0 + ci * cell_w + fill = heat_fill(top["p"] if not hit else max(top["p"], p_ans), hit) + draw.rounded_rectangle((x + 2, y + 2, x + cell_w - 6, y + cell_h - 4), radius=6, fill=fill, outline=(32, 44, 53), width=1) + label = top["str"].replace("\n", "⏎").strip() or "·" + if len(label) > 12: + label = label[:11] + "…" + color = (220, 255, 190) if hit else PALETTE["ink"] if top["p"] > 0.05 else PALETTE["muted"] + draw.text((x + 10, y + 8), label, fill=color, font=fnt) + if hit: + draw.text((x + cell_w - 52, y + 9), f"{p_ans:.2f}", fill=PALETTE["green"], font=mono_font(11)) + + # Footer. + fy = gy0 + len(layer_rows) * cell_h + 22 + draw.rounded_rectangle((margin, fy, w - margin, fy + 88), radius=18, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((margin + 28, fy + 16), f"question: {q['q'][:88]}", fill=PALETTE["ink"], font=ui_font(19)) + draw.text( + (margin + 28, fy + 50), + f"gold answer “{answer}” = BPE {dump['answer_token_strs']} · logit lens = hidden state → final norm → LM head · {dump['image_tokens']:,} visual tokens total, showing the {len(dump['answer_indices'])} covering the answer + {len(dump['control_indices'])} blank-region controls", + fill=PALETTE["muted"], + font=ui_font(16), + ) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_materialize_sweep.py b/packages/snapcompact/research/snapcompact_materialize_sweep.py new file mode 100644 index 000000000..68ab4cfc5 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_materialize_sweep.py @@ -0,0 +1,226 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers"] +# /// +"""Can rendering choices make visual tokens decode to vocabulary EARLIER? + +Sweeps rendering conditions (baseline, line-repeat-in-color, patch-aligned +glyph grids) over the same content/question and measures, per condition, the +layer at which the answer word materializes in logit-lens vocabulary space. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import FontCfg, capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import sample_answer_questions # noqa: E402 + +TOKEN_PX = 28 # Qwen2.5-VL merged visual token size at native resolution + + +@dataclass(frozen=True) +class Condition: + name: str + cfg: FontCfg + variant: str + repeat: int # each text line rendered this many times consecutively + note: str + + +CONDITIONS = [ + Condition("base-8x13", FONTS["8x13"], "bw", 1, "baseline: glyphs straddle token cells on both axes"), + Condition("repeat2-color", FONTS["8x13"], "color", 2, "every line twice, consecutive rows in different hues"), + Condition("align-7x14", FontCfg("7x14a", "7x13", 7, 14), "bw", 1, "4 chars x 2 rows per token, no straddling"), + Condition("align-14x28", FontCfg("14x28a", "7x13", 14, 28, native=(7, 14)), "bw", 1, "2 chars x 1 row per token"), + Condition("align-28x28", FontCfg("28x28a", "8x13", 28, 28, native=(8, 13)), "bw", 1, "1 char per token"), + Condition("repeat2-align-14x28", FontCfg("14x28a", "7x13", 14, 28, native=(7, 14)), "color", 2, "aligned + repeated lines in hues"), +] + + +def build_layout(chunk: str, cols: int, rows: int, repeat: int) -> tuple[str, int]: + """Row-major render string with each line repeated `repeat` times. + + Returns (render_text, usable_chars) where usable_chars is how much of + `chunk` actually fits. + """ + if repeat == 1: + usable = min(len(chunk), cols * rows) + return chunk[:usable], usable + lines = rows // repeat + usable = min(len(chunk), cols * lines) + out: list[str] = [] + for li in range(lines): + line = chunk[li * cols : (li + 1) * cols].ljust(cols) + out.append(line * repeat) + return "".join(out), usable + + +def answer_token_indices(start: int, end: int, cols: int, adv: int, pitch: int, repeat: int, image_size: int, grid: int) -> list[int]: + """Visual-token indices covering chars [start, end) under the layout.""" + indices: set[int] = set() + for i in range(start, end): + row = i // cols + col = i % cols + for copy in range(repeat): + render_row = row * repeat + copy + x0 = col * adv + x1 = min(image_size - 1, (col + 1) * adv - 1) + y0 = render_row * pitch + y1 = min(image_size - 1, (render_row + 1) * pitch - 1) + if y0 >= image_size: + continue + for x in (x0, x1): + for y in (y0, y1): + indices.add((y // TOKEN_PX) * grid + (x // TOKEN_PX)) + return sorted(indices) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default="Qwen/Qwen2.5-VL-7B-Instruct") + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--limit-paras", type=int, default=80) + ap.add_argument("--question-index", type=int, default=3) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--topk", type=int, default=5) + ap.add_argument("--out", default="qwen-materialize-sweep") + args = ap.parse_args() + + import torch + from transformers import AutoProcessor, AutoTokenizer, Qwen2_5_VLForConditionalGeneration + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + # The question is selected on the baseline layout; every condition renders a + # prefix of the same flow, so chunk-relative answer offsets are unchanged. + base_cfg = FONTS["8x13"] + base_cols, base_rows, base_budget = capacity(base_cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + base_chunk = flow[: min(len(flow), base_budget)] + questions = sample_answer_questions(paras, offsets, 0, len(base_chunk), 24, args.seed) + q = questions[min(args.question_index, len(questions) - 1)] + print(f"question: {q['q']!r} answer: {q['answer_text']!r} @ {q['answer_start']}", flush=True) + + print(f"loading {args.model_dir}", flush=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + tokenizer = AutoTokenizer.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True) + model = Qwen2_5_VLForConditionalGeneration.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto").eval() + device = next(model.parameters()).device + image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token) + answer_token_ids = tokenizer(q["answer_text"], add_special_tokens=False)["input_ids"] + answer_id_set = set(answer_token_ids) + answer_token_strs = [tokenizer.decode([t]) for t in answer_token_ids] + norm = model.model.language_model.norm + lm_head = model.lm_head + + conditions_out: list[dict[str, Any]] = [] + for cond in CONDITIONS: + cols, rows, _cap = capacity(cond.cfg, args.size) + render_text, usable = build_layout(flow[: cols * rows], cols, rows, cond.repeat) + if q["answer_end"] > usable: + print(f"SKIP {cond.name}: answer beyond capacity ({usable})", flush=True) + continue + img = render(render_text, cond.cfg, CACHE, args.size, cond.variant) + img.save(img_dir / f"{cond.name}.png") + + prompt = load_prompt("qa-image.md").format(cols=cols, rows=rows) + f"\n\nQuestion: {q['q']}\nAnswer with only the shortest extractive answer." + messages = [{"role": "user", "content": [{"type": "image", "image": img}, {"type": "text", "text": prompt}]}] + templated = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + batch = processor(images=img, text=templated, return_tensors="pt") + ids = batch["input_ids"][0].tolist() + image_positions = [i for i, token_id in enumerate(ids) if token_id == image_token_id] + grid = int(round(len(image_positions) ** 0.5)) + track = answer_token_indices(q["answer_start"], q["answer_end"], cols, cond.cfg.adv, cond.cfg.pitch, cond.repeat, args.size, grid) + track_positions = [image_positions[idx] for idx in track] + batch = {k: (v.to(device) if hasattr(v, "to") else v) for k, v in batch.items()} + + with torch.no_grad(): + fwd = model(**batch, output_hidden_states=True, use_cache=False) + generated = model.generate(**batch, max_new_tokens=16, do_sample=False) + answer_gen = processor.batch_decode(generated[:, batch["input_ids"].shape[1] :], skip_special_tokens=True)[0].strip() + + layers_data: list[dict[str, Any]] = [] + lock_on_layer: int | None = None + soft_layer: int | None = None + with torch.no_grad(): + for layer, hidden in enumerate(fwd.hidden_states): + states = hidden[0, track_positions, :] + logits = lm_head(norm(states)).float() + probs = torch.softmax(logits, dim=-1) + top1 = probs.argmax(dim=-1) + answer_p = probs[:, answer_token_ids] # [n_track, n_answer_tokens] + best_p = float(answer_p.max()) + top1_hit = any(int(t) in answer_id_set for t in top1) + best_idx = int(answer_p.max(dim=1).values.argmax()) + tv, ti = probs[best_idx].topk(args.topk) + layers_data.append( + { + "layer": layer, + "best_answer_p": round(best_p, 6), + "top1_hit": bool(top1_hit), + "best_token_index": track[best_idx], + "best_token_top": [ + {"str": tokenizer.decode([int(ti[k])]), "p": round(float(tv[k]), 5)} for k in range(args.topk) + ], + } + ) + if top1_hit and lock_on_layer is None: + lock_on_layer = layer + if best_p > 0.1 and soft_layer is None: + soft_layer = layer + del fwd + torch.cuda.empty_cache() + + result = { + "name": cond.name, + "note": cond.note, + "variant": cond.variant, + "repeat": cond.repeat, + "adv": cond.cfg.adv, + "pitch": cond.cfg.pitch, + "cols": cols, + "rows": rows, + "usable_chars": usable, + "chars_per_token": round(cols * rows / (cond.repeat * grid * grid), 2), + "tracked_tokens": track, + "generation": answer_gen, + "generation_correct": q["answer_text"].lower() in answer_gen.lower(), + "lock_on_layer": lock_on_layer, + "soft_layer_p10": soft_layer, + "max_answer_p": max(l["best_answer_p"] for l in layers_data), + "layers": layers_data, + } + conditions_out.append(result) + print( + f"{cond.name}: lock_on={lock_on_layer} soft={soft_layer} max_p={result['max_answer_p']:.3f} gen={answer_gen!r}", + flush=True, + ) + + summary = { + "args": vars(args), + "question": {"q": q["q"], "answer_text": q["answer_text"], "answer_start": q["answer_start"], "answer_end": q["answer_end"]}, + "answer_token_ids": answer_token_ids, + "answer_token_strs": answer_token_strs, + "conditions": conditions_out, + } + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_materialize_viz.py b/packages/snapcompact/research/snapcompact_materialize_viz.py new file mode 100644 index 000000000..0e5414f82 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_materialize_viz.py @@ -0,0 +1,171 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render the materialization sweep: rendering choices vs logit-lens confidence.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "grid": (38, 49, 58), +} +SERIES = [ + ("base-8x13", (143, 154, 160)), + ("repeat2-color", (255, 196, 68)), + ("align-7x14", (148, 255, 117)), + ("align-14x28", (75, 220, 255)), + ("align-28x28", (255, 112, 72)), + ("repeat2-align-14x28", (188, 112, 255)), +] + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def crop_answer_region(img_path: Path, cond: dict[str, Any], answer_start: int, answer_end: int, image_size: int = 1568) -> Image.Image: + img = Image.open(img_path).convert("RGB") + cols = cond["cols"] + adv = cond["adv"] + pitch = cond["pitch"] + repeat = cond["repeat"] + row = answer_start // cols + c0 = answer_start % cols + c1 = min(cols - 1, (answer_end - 1) % cols) + y0 = max(0, row * repeat * pitch - pitch) + y1 = min(image_size, (row * repeat + repeat) * pitch + pitch) + x0 = max(0, c0 * adv - 10 * adv) + x1 = min(image_size, (c1 + 1) * adv + 10 * adv) + crop = img.crop((x0, y0, x1, y1)) + d = ImageDraw.Draw(crop) + d.rectangle((c0 * adv - x0 - 2, row * repeat * pitch - y0 - 1, (c1 + 1) * adv - x0 + 2, (row * repeat + repeat) * pitch - y0 + 1), outline=(255, 112, 72), width=3) + return crop + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-materialize-sweep-q3")) + ap.add_argument("--out", default=str(HERE / "results" / "qwen-materialize-sweep-q3" / "materialize-sweep.png")) + args = ap.parse_args() + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + conditions = {c["name"]: c for c in summary["conditions"]} + q = summary["question"] + + w, h = 2200, 1380 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -200, 900, 700), fill=(75, 220, 255, 25)) + gd.ellipse((1240, 540, 2460, 1480), fill=(255, 112, 72, 24)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(86))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 42), "QWEN MATERIALIZATION SWEEP — CAN RENDERING MOVE THE LAYER?", fill=(255, 196, 68), font=ui_font(24, True)) + draw.text((64, 84), "The depth is the model's; the clarity is yours", fill=PALETTE["ink"], font=ui_font(58, True)) + draw.text( + (66, 158), + "Six renderings of the same passage. Logit-lens p(answer BPE) at the answer patch, by layer.\n" + "Alignment and repetition barely move WHERE it materializes — they transform HOW HARD.", + fill=PALETTE["muted"], + font=ui_font(22), + ) + + # Main curve panel. + panel = (64, 226, 1380, 900) + draw.rounded_rectangle(panel, radius=26, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 248), "p(answer BPE) at the best answer patch, per layer", fill=PALETTE["ink"], font=ui_font(26, True)) + gx0, gy0, gx1, gy1 = 150, 320, 1330, 800 + n_layers = len(conditions["base-8x13"]["layers"]) + for i in range(6): + yy = gy0 + (gy1 - gy0) * i / 5 + draw.line((gx0, yy, gx1, yy), fill=PALETTE["grid"], width=1) + draw.text((96, yy - 9), f"{1.0 - i / 5:.1f}", fill=PALETTE["muted"], font=ui_font(14)) + for name, color in SERIES: + cond = conditions.get(name) + if not cond: + continue + pts = [] + for e in cond["layers"]: + x = gx0 + (gx1 - gx0) * e["layer"] / (n_layers - 1) + y = gy1 - (gy1 - gy0) * min(1.0, e["best_answer_p"]) + pts.append((round(x), round(y))) + draw.line(pts, fill=color, width=5 if name != "base-8x13" else 4, joint="curve") + if cond["lock_on_layer"] is not None: + lx = gx0 + (gx1 - gx0) * cond["lock_on_layer"] / (n_layers - 1) + draw.ellipse((lx - 7, gy1 - (gy1 - gy0) * min(1.0, cond["layers"][cond["lock_on_layer"]]["best_answer_p"]) - 7, lx + 7, gy1 - (gy1 - gy0) * min(1.0, cond["layers"][cond["lock_on_layer"]]["best_answer_p"]) + 7), outline=color, width=3) + draw.text((gx0, gy1 + 16), "layer 0", fill=PALETTE["muted"], font=ui_font(15)) + draw.text((gx1 - 76, gy1 + 16), f"layer {n_layers - 1}", fill=PALETTE["muted"], font=ui_font(15)) + draw.text((gx0 + 320, gy1 + 16), "rings mark lock-on (top-1 becomes an answer BPE)", fill=PALETTE["muted"], font=ui_font(15)) + legend_box = (1420, 226, 2136, 900) + draw.rounded_rectangle(legend_box, radius=26, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((1452, 248), "conditions", fill=PALETTE["ink"], font=ui_font(26, True)) + ly = 304 + for name, color in SERIES: + cond = conditions.get(name) + if not cond: + continue + draw.rounded_rectangle((1452, ly, 1452 + 26, ly + 10), radius=4, fill=color) + draw.text((1492, ly - 9), name, fill=PALETTE["ink"], font=ui_font(21, True)) + draw.text((1492, ly + 19), cond["note"], fill=PALETTE["muted"], font=ui_font(14)) + draw.text((1492, ly + 42), f"lock-on L{cond['lock_on_layer']} · peak p {cond['max_answer_p']:.2f} · {cond['chars_per_token']} chars/token", fill=color, font=mono_font(14)) + ly += 96 + + # Condition cards with real crops. + card_y = 938 + card_w = 660 + draw.text((64, card_y - 24), "what the model actually saw (answer region outlined)", fill=PALETTE["ink"], font=ui_font(22, True)) + positions = [(64, card_y + 10), (64 + card_w + 24, card_y + 10), (64 + 2 * (card_w + 24), card_y + 10)] + featured = ["base-8x13", "align-28x28", "repeat2-align-14x28"] + for (cx, cy), name in zip(positions, featured): + cond = conditions.get(name) + if not cond: + continue + color = dict(SERIES)[name] + draw.rounded_rectangle((cx, cy, cx + card_w, cy + 350), radius=20, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((cx + 22, cy + 14), name, fill=color, font=ui_font(23, True)) + draw.text((cx + 22, cy + 46), cond["note"], fill=PALETTE["muted"], font=ui_font(15)) + crop = crop_answer_region(result_dir / "images" / f"{name}.png", cond, q["answer_start"], q["answer_end"]) + scale = min((card_w - 44) / crop.width, 200 / crop.height) + crop_r = crop.resize((round(crop.width * scale), round(crop.height * scale)), Image.Resampling.NEAREST) + draw.rounded_rectangle((cx + 20, cy + 76, cx + card_w - 20, cy + 286), radius=12, fill=(244, 242, 230)) + canvas.paste(crop_r, (cx + 22 + (card_w - 44 - crop_r.width) // 2, cy + 78 + (206 - crop_r.height) // 2)) + draw.text((cx + 22, cy + 300), f"lock-on L{cond['lock_on_layer']} · peak p {cond['max_answer_p']:.2f} · {cond['chars_per_token']} chars/token · gen “{cond['generation']}”", fill=PALETTE["ink"], font=ui_font(16, True)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_pricing_viz.py b/packages/snapcompact/research/snapcompact_pricing_viz.py new file mode 100644 index 000000000..6e113791f --- /dev/null +++ b/packages/snapcompact/research/snapcompact_pricing_viz.py @@ -0,0 +1,137 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow"] +# /// +"""Pricing graphic for the snapcompact post: what a PNG bills vs what it carries.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +P = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "grid": (38, 49, 58), +} + +# (label, text-token equivalent, note). Billed image tokens are constant per canvas. +CARRY = [ + ("8x13 font · 1568²", 5219, "23,520 chars · measured BPE count", 3279), + ("6x10 font · 1568²", 10000, "40,716 chars · ~4 chars/token", 3279), + ("6x10 font · 2576²", 25000, "102,000 chars · one whole corpus", 4950), +] + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--out", default=str(HERE / "results" / "snapcompact-pricing.png")) + args = ap.parse_args() + + w, h = 2200, 1000 + canvas = Image.new("RGB", (w, h), P["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-260, -240, 800, 560), fill=(75, 220, 255, 26)) + gd.ellipse((1400, 400, 2460, 1240), fill=(255, 196, 68, 26)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(88))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 40), "THE BILLING MATH", fill=P["amber"], font=ui_font(24, True)) + draw.text((64, 80), "A flat fee per canvas, no matter what's inside", fill=P["ink"], font=ui_font(56, True)) + draw.text((66, 152), "Anthropic bills images at width × height ÷ 750 tokens. Text tokens scale with content; image tokens scale with pixels. Dense fonts exploit the gap.", fill=P["muted"], font=ui_font(22)) + + # Formula card. + card = (64, 224, 700, 420) + draw.rounded_rectangle(card, radius=22, fill=P["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 246), "flat fee per canvas", fill=P["cyan"], font=ui_font(21, True)) + draw.text((96, 286), "1568 × 1568 → 3,279 tokens", fill=P["ink"], font=mono_font(24)) + draw.text((96, 326), "2576 × 2576 → 4,950 tokens", fill=P["ink"], font=mono_font(24)) + draw.text((96, 372), "(2576 is silently downscaled 0.75x — still the best $/char)", fill=P["muted"], font=ui_font(15)) + + # Cache card. + card = (64, 452, 700, 660) + draw.rounded_rectangle(card, radius=22, fill=P["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 474), "with prompt caching", fill=P["green"], font=ui_font(21, True)) + draw.text((96, 514), "marginal re-ask ≈ 333 tokens/turn", fill=P["ink"], font=mono_font(22)) + draw.text((96, 554), "measured: 753 in · 3,330 cache-write", fill=P["muted"], font=mono_font(17)) + draw.text((96, 584), "16,650 cache-read over six calls", fill=P["muted"], font=mono_font(17)) + draw.text((96, 620), "$0.18 cached vs $0.33 uncached", fill=P["amber"], font=mono_font(18)) + + # Fine-print card. + card = (64, 692, 700, 920) + draw.rounded_rectangle(card, radius=22, fill=P["panel"], outline=(255, 112, 72), width=1) + draw.text((96, 714), "the decode tax", fill=P["orange"], font=ui_font(21, True)) + draw.text((96, 754), "Models reason their way through dense", fill=P["muted"], font=ui_font(18)) + draw.text((96, 782), "pixels: 5–10x more thinking tokens than", fill=P["muted"], font=ui_font(18)) + draw.text((96, 810), "text. Input savings are real; total cost", fill=P["muted"], font=ui_font(18)) + draw.text((96, 838), "depends on output pricing. Cache + re-ask", fill=P["muted"], font=ui_font(18)) + draw.text((96, 866), "is where it always wins.", fill=P["ink"], font=ui_font(18, True)) + + # Carry bars. + panel = (760, 224, 2136, 920) + draw.rounded_rectangle(panel, radius=26, fill=P["panel"], outline=(35, 49, 59), width=1) + draw.text((796, 250), "text-token equivalent carried vs image tokens billed", fill=P["ink"], font=ui_font(26, True)) + draw.text((796, 290), "same content, two meters — the orange bar is what you'd pay as text; the cyan bar is what the PNG bills", fill=P["muted"], font=ui_font(17)) + bx0, bx1 = 1100, 1860 + max_tokens = 25000 + y = 360 + for label, text_tokens, note, billed in CARRY: + draw.text((796, y + 6), label, fill=P["ink"], font=mono_font(17)) + draw.text((796, y + 32), note, fill=P["muted"], font=ui_font(13)) + tw = round((bx1 - bx0) * text_tokens / max_tokens) + bw = round((bx1 - bx0) * billed / max_tokens) + draw.rounded_rectangle((bx0, y, bx0 + tw, y + 26), radius=9, fill=P["orange"]) + draw.text((bx0 + tw + 12, y + 2), f"{text_tokens:,} as text", fill=P["orange"], font=mono_font(15)) + draw.rounded_rectangle((bx0, y + 34, bx0 + bw, y + 60), radius=9, fill=P["cyan"]) + ratio = text_tokens / billed + draw.text((bx0 + bw + 12, y + 36), f"{billed:,} billed · {ratio:.1f}x", fill=P["cyan"], font=mono_font(15)) + y += 130 + # Cached marginal bar. + draw.text((796, y + 6), "any font · cached re-ask", fill=P["ink"], font=mono_font(17)) + draw.text((796, y + 32), "image as cached prefix block", fill=P["muted"], font=ui_font(13)) + bw = max(6, round((bx1 - bx0) * 333 / max_tokens)) + draw.rounded_rectangle((bx0, y + 14, bx0 + bw, y + 40), radius=9, fill=P["green"]) + draw.text((bx0 + bw + 12, y + 16), "≈ 333 tokens/turn · 30x", fill=P["green"], font=mono_font(15)) + y += 110 + draw.line((796, y, 2100, y), fill=P["grid"], width=1) + draw.text((796, y + 16), "10,000 tokens of text, carried by 3,279 image tokens, amortizing to ~333 — that's the whole pitch.", fill=P["amber"], font=ui_font(19, True)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_qwen_control_intervention.py b/packages/snapcompact/research/snapcompact_qwen_control_intervention.py new file mode 100644 index 000000000..5c6f1376a --- /dev/null +++ b/packages/snapcompact/research/snapcompact_qwen_control_intervention.py @@ -0,0 +1,363 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers", "qwen-vl-utils"] +# /// +"""Qwen snapcompact controls: alternate prompt plus activation intervention.""" + +from __future__ import annotations + +import argparse +import json +import math +import random +import sys +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import sample_answer_questions # noqa: E402 +from snapcompact_text_image_compare import ( # noqa: E402 + cosine, + image_answer_token_indices, + normalize_heat, + run_image, + run_text, + to_device, +) + +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "purple": (188, 112, 255), + "red": (255, 76, 62), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def heat_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + stops = [ + (0.00, (4, 7, 20)), + (0.22, (24, 28, 88)), + (0.45, (49, 120, 190)), + (0.65, (54, 226, 195)), + (0.82, (188, 255, 120)), + (1.00, (255, 236, 128)), + ] + for (a, ca), (b, cb) in zip(stops, stops[1:]): + if t <= b: + u = (t - a) / (b - a) + return tuple(round(ca[i] + (cb[i] - ca[i]) * u) for i in range(3)) + return stops[-1][1] + + +def make_text_prompt(chunk: str, q: dict[str, Any]) -> str: + return ( + "Below is reference material. Answer the question using only it.\n\n" + f"{chunk}\n\nQuestion: {q['q']}\n" + "Answer with only the shortest extractive answer." + ) + + +def make_image_prompt(cols: int, rows: int, q: dict[str, Any]) -> str: + return load_prompt("qa-image.md").format(cols=cols, rows=rows) + f"\n\nQuestion: {q['q']}\nAnswer with only the shortest extractive answer." + + +def carrier_map(model: Any, processor: Any, img: Image.Image, chunk: str, q: dict[str, Any], cols: int, rows: int, device: Any) -> tuple[np.ndarray, np.ndarray, dict[str, Any]]: + text_layers, text_pos, _ = run_text(model, processor, make_text_prompt(chunk, q), chunk, q["answer_start"], q["answer_end"], device) + image_layers, image_positions, image_meta, _ = run_image(model, processor, img, make_image_prompt(cols, rows, q), device) + image_count = len(image_positions) + answer_indices = image_answer_token_indices(q["answer_start"], q["answer_end"], cols, 8, 13, img.width, img.height, image_count) + sims = [] + answer_cos = [] + for text_h, image_h in zip(text_layers, image_layers): + text_ans = text_h[text_pos["answer_start"] : text_pos["answer_end"]].mean(axis=0) + image_tokens = image_h[image_positions] + image_ans = image_tokens[answer_indices] if answer_indices else image_tokens + sims.append(cosine(np.repeat(text_ans[None, :], image_tokens.shape[0], axis=0), image_tokens).astype(np.float32, copy=False)) + answer_cos.append(float(cosine(text_ans[None, :], image_ans.mean(axis=0, keepdims=True))[0])) + raw = np.stack(sims, axis=0) + excess = raw - np.median(raw, axis=1, keepdims=True) + norm, lo, hi = normalize_heat(excess) + meta = { + "image_tokens": image_count, + "image_grid": round(math.sqrt(image_count)), + "image_meta": image_meta, + "answer_indices": answer_indices, + "answer_cosine": answer_cos, + "peak_layer": int(np.argmax(answer_cos)), + "peak_cosine": float(max(answer_cos)), + "final_cosine": float(answer_cos[-1]), + "heat_lo": lo, + "heat_hi": hi, + } + return raw, norm, meta + + +def generate_with_intervention( + model: Any, + processor: Any, + img: Image.Image, + prompt: str, + device: Any, + layer: int, + answer_indices: list[int], + mode: str, + seed: int, +) -> str: + import torch + + messages = [{"role": "user", "content": [{"type": "image", "image": img}, {"type": "text", "text": prompt}]}] + templated = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + batch = processor(images=img, text=templated, return_tensors="pt") + image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token) + image_positions = [i for i, token_id in enumerate(batch["input_ids"][0].tolist()) if token_id == image_token_id] + rng = random.Random(seed) + random_indices = sorted(rng.sample([i for i in range(len(image_positions)) if i not in set(answer_indices)], len(answer_indices))) if answer_indices else [] + target_indices = ( + answer_indices + if mode == "answer_mean_patch" + else random_indices + if mode == "random_mean_patch" + else list(range(len(image_positions))) + if mode == "all_image_zero" + else [] + ) + target_positions = [image_positions[i] for i in target_indices] + batch = to_device(batch, device) + + handle = None + if target_positions: + def hook(_module: Any, inputs: tuple[Any, ...]) -> tuple[Any, ...]: + hidden = inputs[0] + if hidden.ndim == 3 and hidden.shape[1] > max(target_positions): + patched = hidden.clone() + if mode == "all_image_zero": + patched[:, target_positions, :] = 0 + else: + source_positions = [p for p in image_positions if p not in target_positions] + mean_vec = hidden[:, source_positions, :].mean(dim=1, keepdim=True) + patched[:, target_positions, :] = mean_vec + return (patched, *inputs[1:]) + return inputs + + handle = model.model.language_model.layers[layer].register_forward_pre_hook(hook) + try: + with torch.no_grad(): + generated = model.generate(**batch, max_new_tokens=24, do_sample=False) + finally: + if handle is not None: + handle.remove() + new_tokens = generated[:, batch["input_ids"].shape[1] :] + return processor.batch_decode(new_tokens, skip_special_tokens=True)[0].strip() + + +def crop_answer(img: Image.Image, q: dict[str, Any], cols: int, adv: int = 8, pitch: int = 13) -> Image.Image: + start = q["answer_start"] + end = q["answer_end"] + row0 = max(0, start // cols - 5) + row1 = min(img.height // pitch, end // cols + 6) + col0 = max(0, start % cols - 42) + col1 = min(cols, end % cols + 42) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=PALETTE["orange"], width=3) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.NEAREST) + canvas.paste(resized, (x0 + (x1 - x0 - resized.width) // 2, y0 + (y1 - y0 - resized.height) // 2)) + + +def draw_grid(draw: ImageDraw.ImageDraw, grid_values: np.ndarray, answer_indices: list[int], box: tuple[int, int, int, int], title: str, subtitle: str, color: tuple[int, int, int]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=22, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((x0 + 20, y0 + 18), title, fill=color, font=ui_font(25, True)) + draw.text((x0 + 20, y0 + 50), subtitle, fill=PALETTE["muted"], font=ui_font(15)) + gx0, gy0, gx1, gy1 = x0 + 30, y0 + 84, x1 - 30, y1 - 28 + rows, cols = grid_values.shape + cw = (gx1 - gx0) / cols + ch = (gy1 - gy0) / rows + for r in range(rows): + for c in range(cols): + xa = round(gx0 + c * cw) + xb = round(gx0 + (c + 1) * cw) + ya = round(gy0 + r * ch) + yb = round(gy0 + (r + 1) * ch) + draw.rectangle((xa, ya, xb, yb), fill=heat_color(float(grid_values[r, c]))) + for idx in answer_indices: + r, c = divmod(idx, cols) + xa = round(gx0 + c * cw) + xb = round(gx0 + (c + 1) * cw) + ya = round(gy0 + r * ch) + yb = round(gy0 + (r + 1) * ch) + draw.rectangle((xa - 2, ya - 2, xb + 2, yb + 2), outline=PALETTE["orange"], width=2) + + +def render_figure(out_path: Path, img: Image.Image, primary: dict[str, Any], distractor: dict[str, Any], primary_norm: np.ndarray, distractor_norm: np.ndarray, primary_meta: dict[str, Any], distractor_meta: dict[str, Any], generations: dict[str, str], cols: int) -> None: + w, h = 2200, 1320 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-260, -240, 900, 760), fill=(75, 220, 255, 28)) + gd.ellipse((1240, 120, 2480, 1380), fill=(255, 112, 72, 27)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(86))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 42), "QWEN SNAPCOMPACT CONTROL + INTERVENTION", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((64, 84), "Ask a different thing; patch the hidden answer", fill=PALETTE["ink"], font=ui_font(62, True)) + draw.text((66, 166), "Same bitmap, two questions. Then patch answer-region image-token activations at the peak layer and watch generation change.", fill=PALETTE["muted"], font=ui_font(24)) + + draw.rounded_rectangle((64, 238, 616, 1234), radius=30, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 270), "same image carrier", fill=PALETTE["ink"], font=ui_font(32, True)) + draw.text((96, 312), "Qwen2.5-VL-7B, 1568px bitmap", fill=PALETTE["muted"], font=ui_font(18)) + for label, q, y, color in [("PRIMARY", primary, 374, PALETTE["orange"]), ("DISTRACTOR", distractor, 658, PALETTE["cyan"] )]: + draw.text((96, y), label, fill=color, font=ui_font(17, True)) + crop = crop_answer(img, q, cols) + draw.rounded_rectangle((96, y + 34, 584, y + 194), radius=14, fill=(244, 242, 230), outline=color, width=3) + paste_fit(canvas, crop, (112, y + 48, 568, y + 180)) + draw.text((96, y + 216), q["q"][:58], fill=PALETTE["ink"], font=ui_font(18)) + draw.text((96, y + 244), f"gold: {q['answer_text']}", fill=PALETTE["amber"], font=ui_font(22, True)) + draw.text((96, 1012), f"primary peak: L{primary_meta['peak_layer']} cosine {primary_meta['peak_cosine']:.3f}", fill=PALETTE["orange"], font=ui_font(20, True)) + draw.text((96, 1044), f"distractor peak: L{distractor_meta['peak_layer']} cosine {distractor_meta['peak_cosine']:.3f}", fill=PALETTE["cyan"], font=ui_font(20, True)) + draw.text((96, 1102), f"image tokens: {primary_meta['image_tokens']} ({primary_meta['image_grid']}×{primary_meta['image_grid']})", fill=PALETTE["muted"], font=ui_font(18)) + + grid = primary_meta["image_grid"] + draw_grid(draw, primary_norm[primary_meta["peak_layer"]].reshape(grid, grid), primary_meta["answer_indices"], (666, 238, 1386, 706), "primary question map", f"{primary['answer_text']} @ layer {primary_meta['peak_layer']} — orange box marks true answer", PALETTE["orange"]) + draw_grid(draw, distractor_norm[distractor_meta["peak_layer"]].reshape(grid, grid), distractor_meta["answer_indices"], (1420, 238, 2140, 706), "distractor question map", f"{distractor['answer_text']} @ layer {distractor_meta['peak_layer']} — map should move", PALETTE["cyan"]) + + draw.rounded_rectangle((666, 746, 2140, 1234), radius=30, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((704, 780), "activation patch test", fill=PALETTE["ink"], font=ui_font(34, True)) + draw.text((704, 822), "Before decoder layer 0, replace selected image-token residuals. Local answer patches test specificity; all-image zero is the sanity check.", fill=PALETTE["muted"], font=ui_font(20)) + rows = [ + ("normal", generations["normal"], PALETTE["green"]), + ("patch random region", generations["random_mean_patch"], PALETTE["cyan"]), + ("patch answer region", generations["answer_mean_patch"], PALETTE["red"]), + ("zero all image tokens", generations["all_image_zero"], PALETTE["purple"]), + ] + y = 878 + for label, text, color in rows: + draw.rounded_rectangle((704, y, 2078, y + 74), radius=18, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((730, y + 16), label.upper(), fill=color, font=ui_font(17, True)) + draw.text((1002, y + 15), text[:115], fill=PALETTE["ink"], font=ui_font(23, True)) + y += 86 + + out_path.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out_path) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default="Qwen/Qwen2.5-VL-7B-Instruct") + ap.add_argument("--font", default="8x13", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--limit-paras", type=int, default=80) + ap.add_argument("--qpc", type=int, default=24) + ap.add_argument("--question-index", type=int, default=12) + ap.add_argument("--distractor-index", type=int, default=3) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--out", default="qwen-control-intervention") + args = ap.parse_args() + + import torch + from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + chunk = flow[: min(len(flow), budget)] + questions = sample_answer_questions(paras, offsets, 0, len(chunk), args.qpc, args.seed) + if len(questions) < 2: + raise SystemExit("not enough questions fit in chunk") + primary = questions[min(args.question_index, len(questions) - 1)] + distractor = questions[min(args.distractor_index, len(questions) - 1)] + if distractor is primary: + distractor = questions[0 if args.question_index != 0 else 1] + img = render(chunk, cfg, CACHE, args.size, args.variant) + img.save(img_dir / "image-carrier.png") + + print(f"loading {args.model_dir}", flush=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + model = Qwen2_5_VLForConditionalGeneration.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto").eval() + device = next(model.parameters()).device + + primary_raw, primary_norm, primary_meta = carrier_map(model, processor, img, chunk, primary, cols, rows, device) + distractor_raw, distractor_norm, distractor_meta = carrier_map(model, processor, img, chunk, distractor, cols, rows, device) + + peak_layer = primary_meta["peak_layer"] + prompt = make_image_prompt(cols, rows, primary) + patch_layer = 0 + generations = { + "normal": generate_with_intervention(model, processor, img, prompt, device, patch_layer, primary_meta["answer_indices"], "none", args.seed), + "random_mean_patch": generate_with_intervention(model, processor, img, prompt, device, patch_layer, primary_meta["answer_indices"], "random_mean_patch", args.seed), + "answer_mean_patch": generate_with_intervention(model, processor, img, prompt, device, patch_layer, primary_meta["answer_indices"], "answer_mean_patch", args.seed), + "all_image_zero": generate_with_intervention(model, processor, img, prompt, device, patch_layer, primary_meta["answer_indices"], "all_image_zero", args.seed), + } + + summary = { + "args": vars(args), + "geometry": {"cols": cols, "rows": rows, "capacity": budget}, + "primary": primary, + "distractor": distractor, + "primary_meta": primary_meta, + "distractor_meta": distractor_meta, + "intervention_layer": patch_layer, + "generations": generations, + } + np.savez_compressed(out_dir / "control_intervention.npz", primary_raw=primary_raw, primary_norm=primary_norm, distractor_raw=distractor_raw, distractor_norm=distractor_norm) + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + render_figure(out_dir / "control-intervention.png", img, primary, distractor, primary_norm, distractor_norm, primary_meta, distractor_meta, generations, cols) + print(json.dumps(summary, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_qwen_spotlight_viz.py b/packages/snapcompact/research/snapcompact_qwen_spotlight_viz.py new file mode 100644 index 000000000..0e34ec882 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_qwen_spotlight_viz.py @@ -0,0 +1,278 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render a cleaner prompt-specific spotlight figure for Qwen snapcompact controls.""" + +from __future__ import annotations + +import argparse +import json +import random +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "orange": (255, 112, 72), + "cyan": (75, 220, 255), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "red": (255, 76, 62), + "purple": (188, 112, 255), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def heat_color(t: float, theme: str) -> tuple[int, int, int, int]: + t = max(0.0, min(1.0, t)) + if theme == "cyan": + rgb0, rgb1 = (10, 38, 62), (75, 220, 255) + else: + rgb0, rgb1 = (62, 25, 10), (255, 112, 72) + rgb = tuple(round(rgb0[i] + (rgb1[i] - rgb0[i]) * t) for i in range(3)) + alpha = round(20 + 210 * t) + return (*rgb, alpha) + + +def normalize_positive(arr: np.ndarray) -> np.ndarray: + arr = np.maximum(arr, 0) + hi = float(np.quantile(arr, 0.985)) if arr.size else 1.0 + if hi <= 0: + hi = 1.0 + return np.clip(arr / hi, 0, 1) + + +def smooth_map(grid: np.ndarray, radius: float = 1.45) -> np.ndarray: + g = normalize_positive(grid) + img = Image.fromarray(np.uint8(g * 255), mode="L").filter(ImageFilter.GaussianBlur(radius=radius)) + return np.asarray(img, dtype=np.float32) / 255.0 + + +def answer_bbox(indices: list[int], grid: int, image_w: int, image_h: int) -> tuple[int, int, int, int]: + rows = [idx // grid for idx in indices] + cols = [idx % grid for idx in indices] + x0 = int(min(cols) / grid * image_w) + x1 = int((max(cols) + 1) / grid * image_w) + y0 = int(min(rows) / grid * image_h) + y1 = int((max(rows) + 1) / grid * image_h) + return x0, y0, x1, y1 + + +def overlay_heat(base: Image.Image, heat: np.ndarray, bbox: tuple[int, int, int, int], theme: str) -> Image.Image: + base_rgba = base.convert("RGBA") + heat_img = Image.new("RGBA", base.size, (0, 0, 0, 0)) + # Upscale smoothed 56x56 field to bitmap size; threshold softens static. + up = Image.fromarray(np.uint8(heat * 255), mode="L").resize(base.size, Image.Resampling.BICUBIC) + vals = np.asarray(up, dtype=np.float32) / 255.0 + threshold = float(np.quantile(vals, 0.72)) + vals = np.clip((vals - threshold) / max(1e-6, 1.0 - threshold), 0, 1) + px = heat_img.load() + for y in range(0, heat_img.height, 2): + for x in range(0, heat_img.width, 2): + t = float(vals[y, x]) + if t <= 0: + continue + color = heat_color(t, theme) + px[x, y] = color + if x + 1 < heat_img.width: + px[x + 1, y] = color + if y + 1 < heat_img.height: + px[x, y + 1] = color + if x + 1 < heat_img.width and y + 1 < heat_img.height: + px[x + 1, y + 1] = color + out = Image.alpha_composite(base_rgba, heat_img).convert("RGB") + draw = ImageDraw.Draw(out) + color = PALETTE["cyan"] if theme == "cyan" else PALETTE["orange"] + draw.rounded_rectangle(bbox, radius=4, outline=color, width=6) + # spotlight ring around answer bbox + x0, y0, x1, y1 = bbox + pad = 24 + draw.rounded_rectangle((x0 - pad, y0 - pad, x1 + pad, y1 + pad), radius=16, outline=color, width=3) + return out + + +def crop_box(img: Image.Image, bbox: tuple[int, int, int, int], pad: int = 180) -> Image.Image: + x0, y0, x1, y1 = bbox + return img.crop((max(0, x0 - pad), max(0, y0 - pad), min(img.width, x1 + pad), min(img.height, y1 + pad))).convert("RGB") + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int], resample: int = Image.Resampling.LANCZOS) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), resample) + canvas.paste(resized, (x0 + (x1 - x0 - resized.width) // 2, y0 + (y1 - y0 - resized.height) // 2)) + + +def region_score(grid_map: np.ndarray, indices: list[int]) -> float: + if not indices: + return 0.0 + flat = grid_map.ravel() + return float(np.mean([flat[i] for i in indices if i < len(flat)])) + + +def random_region_scores(grid_map: np.ndarray, region_size: int, count: int = 600, seed: int = 7) -> np.ndarray: + rng = random.Random(seed) + flat = grid_map.ravel() + scores = [] + for _ in range(count): + picks = rng.sample(range(len(flat)), min(region_size, len(flat))) + scores.append(float(np.mean(flat[picks]))) + return np.array(scores, dtype=np.float32) + + +def draw_score_card(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], title: str, score: float, random_scores: np.ndarray, color: tuple[int, int, int]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=18, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((x0 + 20, y0 + 18), title, fill=color, font=ui_font(22, True)) + percentile = float((random_scores < score).mean() * 100) + draw.text((x0 + 20, y0 + 50), f"answer region beats {percentile:.0f}% of random same-size regions", fill=PALETTE["muted"], font=ui_font(16)) + gx0, gy0, gx1, gy1 = x0 + 28, y0 + 96, x1 - 28, y1 - 42 + lo = float(min(random_scores.min(), score)) + hi = float(max(random_scores.max(), score)) + if hi <= lo: + hi = lo + 1e-6 + bins = np.linspace(lo, hi, 30) + hist, _ = np.histogram(random_scores, bins=bins) + max_h = max(1, int(hist.max())) + bw = (gx1 - gx0) / len(hist) + for i, h in enumerate(hist): + xa = gx0 + i * bw + xb = gx0 + (i + 1) * bw - 1 + ya = gy1 - (gy1 - gy0) * int(h) / max_h + draw.rectangle((round(xa), round(ya), round(xb), gy1), fill=(37, 49, 58)) + sx = gx0 + (gx1 - gx0) * (score - lo) / (hi - lo) + draw.line((sx, gy0 - 8, sx, gy1 + 10), fill=color, width=5) + draw.text((round(sx) - 34, gy0 - 36), "answer", fill=color, font=ui_font(15, True)) + + +def draw_generation_rows(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], generations: dict[str, str]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=24, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((x0 + 28, y0 + 24), "causal patch check", fill=PALETTE["ink"], font=ui_font(30, True)) + draw.text((x0 + 28, y0 + 60), "Patch before decoder layer 0; only the true answer-region patch changes the answer.", fill=PALETTE["muted"], font=ui_font(18)) + rows = [ + ("normal", generations["normal"], PALETTE["green"]), + ("random region patch", generations["random_mean_patch"], PALETTE["cyan"]), + ("answer region patch", generations["answer_mean_patch"], PALETTE["red"]), + ("all image tokens zero", generations["all_image_zero"] or "∅", PALETTE["purple"]), + ] + y = y0 + 116 + for label, text, color in rows: + draw.rounded_rectangle((x0 + 28, y, x1 - 28, y + 62), radius=14, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((x0 + 48, y + 17), label.upper(), fill=color, font=ui_font(15, True)) + draw.text((x0 + 330, y + 13), text[:70], fill=PALETTE["ink"], font=ui_font(24, True)) + y += 78 + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-control-intervention-q3-d12-prehook")) + ap.add_argument("--out", default=str(HERE / "results" / "qwen-control-intervention-q3-d12-prehook" / "spotlight-control.png")) + args = ap.parse_args() + + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "control_intervention.npz") + img = Image.open(result_dir / "images" / "image-carrier.png").convert("RGB") + + primary = summary["primary"] + distractor = summary["distractor"] + primary_meta = summary["primary_meta"] + distractor_meta = summary["distractor_meta"] + grid = primary_meta["image_grid"] + primary_layer = primary_meta["peak_layer"] + distractor_layer = distractor_meta["peak_layer"] + primary_norm = data["primary_norm"][primary_layer].reshape(grid, grid) + distractor_norm = data["distractor_norm"][distractor_layer].reshape(grid, grid) + + # Prompt-specific contrast removes shared visual texture: what lights up more + # for this question than for the other question? + primary_contrast = smooth_map(primary_norm - distractor_norm) + distractor_contrast = smooth_map(distractor_norm - primary_norm) + primary_bbox = answer_bbox(primary_meta["answer_indices"], grid, img.width, img.height) + distractor_bbox = answer_bbox(distractor_meta["answer_indices"], grid, img.width, img.height) + primary_overlay = overlay_heat(img, primary_contrast, primary_bbox, "orange") + distractor_overlay = overlay_heat(img, distractor_contrast, distractor_bbox, "cyan") + + primary_score = region_score(primary_contrast, primary_meta["answer_indices"]) + distractor_score = region_score(distractor_contrast, distractor_meta["answer_indices"]) + primary_random = random_region_scores(primary_contrast, len(primary_meta["answer_indices"]), seed=11) + distractor_random = random_region_scores(distractor_contrast, len(distractor_meta["answer_indices"]), seed=13) + + w, h = 2200, 1320 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -200, 900, 700), fill=(255, 112, 72, 28)) + gd.ellipse((1160, 120, 2460, 1340), fill=(75, 220, 255, 25)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(86))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 42), "QWEN SNAPCOMPACT SPOTLIGHT", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((64, 84), "Subtract the other question; the signal stops looking like static", fill=PALETTE["ink"], font=ui_font(58, True)) + draw.text((66, 160), "These are not raw activation carpets. Each overlay is prompt-specific excess: this question’s map minus the other question’s map, smoothed and thresholded.", fill=PALETTE["muted"], font=ui_font(23)) + + # Overlay panels. + draw.rounded_rectangle((64, 230, 1068, 794), radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 260), "primary prompt spotlight", fill=PALETTE["orange"], font=ui_font(31, True)) + draw.text((96, 298), f"{primary['q']} → {primary['answer_text']}", fill=PALETTE["muted"], font=ui_font(18)) + primary_crop = crop_box(primary_overlay, primary_bbox, pad=300) + paste_fit(canvas, primary_crop, (96, 342, 694, 760), Image.Resampling.LANCZOS) + draw.rounded_rectangle((720, 342, 1036, 760), radius=18, fill=(244, 242, 230), outline=PALETTE["orange"], width=3) + paste_fit(canvas, crop_box(primary_overlay, primary_bbox, pad=90), (736, 358, 1020, 744), Image.Resampling.LANCZOS) + draw.text((736, 724), "zoom: answer region", fill=PALETTE["orange"], font=ui_font(15, True)) + + draw.rounded_rectangle((1132, 230, 2136, 794), radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((1164, 260), "distractor prompt spotlight", fill=PALETTE["cyan"], font=ui_font(31, True)) + draw.text((1164, 298), f"{distractor['q']} → {distractor['answer_text']}", fill=PALETTE["muted"], font=ui_font(18)) + distractor_crop = crop_box(distractor_overlay, distractor_bbox, pad=300) + paste_fit(canvas, distractor_crop, (1164, 342, 1762, 760), Image.Resampling.LANCZOS) + draw.rounded_rectangle((1788, 342, 2104, 760), radius=18, fill=(244, 242, 230), outline=PALETTE["cyan"], width=3) + paste_fit(canvas, crop_box(distractor_overlay, distractor_bbox, pad=90), (1804, 358, 2088, 744), Image.Resampling.LANCZOS) + draw.text((1804, 724), "zoom: answer region", fill=PALETTE["cyan"], font=ui_font(15, True)) + + draw_score_card(draw, (64, 836, 610, 1236), "primary answer-region score", primary_score, primary_random, PALETTE["orange"]) + draw_score_card(draw, (642, 836, 1188, 1236), "distractor answer-region score", distractor_score, distractor_random, PALETTE["cyan"]) + draw_generation_rows(draw, (1220, 836, 2136, 1236), summary["generations"]) + + # Save metrics alongside figure for the caption. + metrics = { + "primary_score": primary_score, + "primary_percentile": float((primary_random < primary_score).mean() * 100), + "distractor_score": distractor_score, + "distractor_percentile": float((distractor_random < distractor_score).mean() * 100), + "primary_peak_layer": primary_layer, + "distractor_peak_layer": distractor_layer, + "generations": summary["generations"], + } + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + (out.parent / "spotlight-metrics.json").write_text(json.dumps(metrics, indent=1)) + print(out) + print(json.dumps(metrics, indent=1)) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_r2_chord.py b/packages/snapcompact/research/snapcompact_r2_chord.py new file mode 100755 index 000000000..0c687c034 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_r2_chord.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python3 +"""Chord/ribbon diagram of cross-carrier matching in Qwen2.5-VL-7B. + +Left arc: 12 questions answered via the TEXT carrier (cyan). +Right arc: the same 12 questions answered via the IMAGE carrier (orange). +Ribbons between every (text_i, image_j) pair are sized by the REAL cosine +similarity of the carrier-centered hidden states at layer 19 (negatives +clipped to 0). Matched pairs (i == j) glow amber and visibly dominate. + +Data: results/qwen-carrier-convergence-n12/{carrier_convergence.npz,summary.json} +Output: results/agent-r2-chord/chord.png (~2200 px wide) +""" + +import json +import os + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.path import Path as MplPath +from matplotlib.patches import PathPatch + +HERE = os.path.dirname(os.path.abspath(__file__)) +SRC = os.path.join(HERE, "results", "qwen-carrier-convergence-n12") +OUT_DIR = os.path.join(HERE, "results", "agent-r2-chord") +LAYER = 19 # best layer per summary.json + +# ---------------------------------------------------------------- palette +BG = "#05070a" +PANEL = "#0c1117" +INK = "#f1efe0" +MUTED = "#8f9aa0" +AMBER = "#ffc444" +CYAN = "#4bdcff" +ORANGE = "#ff7048" +GREEN = "#94ff75" + +# ---------------------------------------------------------------- data +npz = np.load(os.path.join(SRC, "carrier_convergence.npz")) +cross = np.asarray(npz["cross_sim"][LAYER], dtype=np.float64) # [12,12] text x image +with open(os.path.join(SRC, "summary.json")) as fh: + summary = json.load(fh) +records = summary["records"] +best = summary["best"] +n = cross.shape[0] +assert len(records) == n == 12 + +labels = [r["gold"] for r in records] +matched_mean = float(np.trace(cross) / n) +mismatched_mean = float((cross.sum() - np.trace(cross)) / (n * n - n)) +# retrieval: for each text row, is the matched image column the argmax? +retrieved = int((cross.argmax(axis=1) == np.arange(n)).sum()) + +w = np.clip(cross, 0.0, None) # ribbon weights: clip negatives + +# ---------------------------------------------------------------- geometry +# Left arc (text carrier): 110deg -> 250deg, top-left to bottom-left. +# Right arc (image carrier): 70deg -> -70deg, mirrored so matched ribbons +# run roughly horizontally across the circle. +R_IN = 0.955 # inner radius where ribbons attach +R_OUT = 1.000 # outer radius of the node band +SEG_DEG = 8.6 # angular width of each node segment +ts = np.linspace(0.0, 1.0, n) +left_centers = 110.0 + ts * 140.0 +right_centers = 70.0 - ts * 140.0 + + +def seg_bounds(center_deg): + return center_deg - SEG_DEG / 2.0, center_deg + SEG_DEG / 2.0 + + +def pol(theta_deg, r): + a = np.deg2rad(theta_deg) + return np.array([r * np.cos(a), r * np.sin(a)]) + + +def allocate(centers, weights_per_node): + """Split each node's segment into sub-spans proportional to ribbon weight. + + weights_per_node: [n, n] -- weights_per_node[i, j] is the weight of the + ribbon to opposite-side node j, allocated within node i's segment. + Sub-spans are ordered by the opposite node index so ribbons fan smoothly. + Returns spans[i][j] = (a0, a1) in degrees (a0 < a1) or None if weight ~ 0. + """ + spans = [] + for i in range(n): + lo, hi = seg_bounds(centers[i]) + tot = weights_per_node[i].sum() + spans_i = [None] * n + if tot <= 1e-9: + spans.append(spans_i) + continue + cursor = lo + for j in range(n): + frac = weights_per_node[i, j] / tot + width = frac * (hi - lo) + if weights_per_node[i, j] > 1e-9: + spans_i[j] = (cursor, cursor + width) + cursor += width + spans.append(spans_i) + return spans + + +# Left node i sends ribbons to right nodes j with weight w[i, j]; +# right node j receives from left nodes i with weight w[i, j]. +left_spans = allocate(left_centers, w) +right_spans = allocate(right_centers, w.T) + + +def arc_points(a0, a1, r, steps=12): + th = np.linspace(a0, a1, steps) + return np.stack([r * np.cos(np.deg2rad(th)), r * np.sin(np.deg2rad(th))], axis=1) + + +def ribbon_path(la, lb, ra, rb, pull=0.18): + """Filled ribbon: arc(la->lb) on the left rim, cubic bezier to the right + rim, arc(ra->rb), bezier back. Control points pulled toward the center.""" + p_lb = pol(lb, R_IN) + p_ra = pol(ra, R_IN) + p_rb = pol(rb, R_IN) + p_la = pol(la, R_IN) + verts = [] + codes = [] + arc1 = arc_points(la, lb, R_IN) + verts.extend(arc1) + codes.extend([MplPath.MOVETO] + [MplPath.LINETO] * (len(arc1) - 1)) + # bezier left-edge-end -> right-edge-start + verts.extend([p_lb * pull, p_ra * pull, p_ra]) + codes.extend([MplPath.CURVE4] * 3) + arc2 = arc_points(ra, rb, R_IN)[1:] + verts.extend(arc2) + codes.extend([MplPath.LINETO] * len(arc2)) + verts.extend([p_rb * pull, p_la * pull, p_la]) + codes.extend([MplPath.CURVE4] * 3) + codes.append(MplPath.CLOSEPOLY) + verts.append(p_la) + return MplPath(verts, codes) + + +def center_bezier(a_deg, b_deg, pull=0.18, steps=60): + p0, p3 = pol(a_deg, R_IN), pol(b_deg, R_IN) + p1, p2 = p0 * pull, p3 * pull + t = np.linspace(0, 1, steps)[:, None] + return ((1 - t) ** 3 * p0 + 3 * (1 - t) ** 2 * t * p1 + + 3 * (1 - t) * t ** 2 * p2 + t ** 3 * p3) + + +# ---------------------------------------------------------------- figure +fig = plt.figure(figsize=(22, 16.5), dpi=100, facecolor=BG) +ax = fig.add_axes([0.02, 0.0, 0.96, 0.94]) +ax.set_facecolor(BG) +ax.set_xlim(-1.95, 1.95) +ax.set_ylim(-1.32, 1.30) +ax.set_aspect("equal") +ax.axis("off") + +w_max = w.max() + +# mismatched ribbons first (thin, dim), then matched (amber, glowing) on top +order = sorted(((i, j) for i in range(n) for j in range(n)), + key=lambda ij: (ij[0] == ij[1], w[ij])) +for i, j in order: + ls, rs = left_spans[i][j], right_spans[j][i] + if ls is None or rs is None: + continue + val = w[i, j] + matched = i == j + # right span traversed in reverse so the ribbon doesn't twist + path = ribbon_path(ls[0], ls[1], rs[1], rs[0]) + if matched: + # glow: layered strokes along the centerline under the fill + mid_l = 0.5 * (ls[0] + ls[1]) + mid_r = 0.5 * (rs[0] + rs[1]) + spine = center_bezier(mid_l, mid_r) + for lw, al in ((26, 0.045), (14, 0.075), (7, 0.12)): + ax.plot(spine[:, 0], spine[:, 1], color=AMBER, lw=lw, alpha=al, + solid_capstyle="round", zorder=4) + ax.add_patch(PathPatch(path, facecolor=AMBER, edgecolor=AMBER, + lw=0.7, alpha=0.78, zorder=5)) + else: + alpha = 0.10 + 0.45 * (val / w_max) + ax.add_patch(PathPatch(path, facecolor=MUTED, edgecolor="none", + alpha=alpha * 0.55, zorder=2)) + +# ---------------------------------------------------------------- node bands +for i in range(n): + for centers, color, side in ((left_centers, CYAN, "L"), + (right_centers, ORANGE, "R")): + a0, a1 = seg_bounds(centers[i]) + band = arc_points(a0, a1, R_OUT, 16) + band_in = arc_points(a1, a0, R_IN, 16) + poly = np.vstack([band, band_in]) + ax.add_patch(plt.Polygon(poly, closed=True, facecolor=color, + edgecolor="none", alpha=0.95, zorder=6)) + +# ---------------------------------------------------------------- labels +for i in range(n): + txt = labels[i] + for centers, color, ha in ((left_centers, CYAN, "right"), + (right_centers, ORANGE, "left")): + c = centers[i] + p = pol(c, 1.05) + ax.text(p[0], p[1], txt, color=INK, fontsize=15.5, ha=ha, va="center", + zorder=8, family="DejaVu Sans") + # small question index tick just inside the label + ax.text(p[0] + (0.018 if ha == "left" else -0.018), + p[1] - 0.052, f"Q{i + 1}", color=color, fontsize=10.5, + ha=ha, va="center", alpha=0.85, zorder=8) + +# arc side headers +ax.text(*pol(180, 1.62), "TEXT CARRIER", color=CYAN, fontsize=21, + ha="center", va="center", rotation=90, weight="bold", alpha=0.95) +ax.text(*pol(180, 1.69), "5,219 prose tokens", color=MUTED, fontsize=13, + ha="center", va="center", rotation=90) +ax.text(*pol(0, 1.62), "IMAGE CARRIER", color=ORANGE, fontsize=21, + ha="center", va="center", rotation=-90, weight="bold", alpha=0.95) +ax.text(*pol(0, 1.69), "same passage, rendered as pixels", color=MUTED, + fontsize=13, ha="center", va="center", rotation=-90) + +# ---------------------------------------------------------------- titles +fig.text(0.5, 0.965, "ONE MEMORY, TWO CARRIERS", color=INK, fontsize=34, + ha="center", va="center", weight="bold", family="DejaVu Sans") +fig.text(0.5, 0.932, + "Cross-carrier cosine of answer states at layer 19 -- every text" + " question finds its image twin (Qwen2.5-VL-7B)", + color=MUTED, fontsize=16.5, ha="center", va="center") + +# ---------------------------------------------------------------- stat plate +plate = fig.add_axes([0.035, 0.05, 0.215, 0.135]) +plate.set_facecolor(PANEL) +plate.set_xlim(0, 1) +plate.set_ylim(0, 1) +for s in plate.spines.values(): + s.set_color("#1c242e") +plate.set_xticks([]) +plate.set_yticks([]) +plate.text(0.5, 0.84, f"LAYER {LAYER} -- BEST SEPARATION", color=MUTED, + fontsize=12.5, ha="center", va="center") +stats = ( + (f"{matched_mean:+.2f}", "matched cosine", AMBER), + (f"{mismatched_mean:+.2f}", "mismatched", MUTED), + (f"{retrieved}/{n}", "retrieval", GREEN), +) +for k, (val, lab, color) in enumerate(stats): + x = 0.18 + 0.32 * k + plate.text(x, 0.48, val, color=color, fontsize=24, ha="center", + va="center", weight="bold") + plate.text(x, 0.18, lab, color=MUTED, fontsize=12, ha="center", + va="center") + +# footnote +fig.text(0.5, 0.012, + "Ribbon width/opacity = cosine(text_state_i, image_state_j)," + " negatives clipped; amber = matched pair (i = j)." + f" RSA r = {best['rsa_pearson']:.2f}.", + color=MUTED, fontsize=12.5, ha="center", va="center") + +os.makedirs(OUT_DIR, exist_ok=True) +out_path = os.path.join(OUT_DIR, "chord.png") +fig.savefig(out_path, dpi=100, facecolor=BG) +print(f"wrote {out_path}") +print(f"matched={matched_mean:.4f} mismatched={mismatched_mean:.4f} " + f"retrieval={retrieved}/{n}") diff --git a/packages/snapcompact/research/snapcompact_r2_crystal.py b/packages/snapcompact/research/snapcompact_r2_crystal.py new file mode 100755 index 000000000..fbe894bfc --- /dev/null +++ b/packages/snapcompact/research/snapcompact_r2_crystal.py @@ -0,0 +1,338 @@ +#!/usr/bin/env python3 +"""Crystallization: a pixel patch becomes a vocabulary token. + +Animated logit-lens GIF for Qwen2.5-VL-7B. Steps through layers L0..L28 for +visual token #310 of the image carrier (the 28x28 patch covering the rendered +word "spectacular"), showing the REAL top-5 decoded vocab tokens per layer and +the probability of the answer BPE 'acular', which stays ~0 until L24 then +climbs 0.14 -> 0.39 at L28. A control patch (token #1878) is shown alongside +and stays noise. All numbers are read from logit_lens.json; nothing is +fabricated. + +Output: results/agent-r2-crystal/crystal.gif (+ crystal_final.png) +""" + +import json +from pathlib import Path + +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +BASE = Path(__file__).resolve().parent +DATA = BASE / "results/qwen-logit-lens-q3/logit_lens.json" +CARRIER = BASE / "results/qwen-logit-lens-q3/images/image-carrier.png" +OUT_DIR = BASE / "results/agent-r2-crystal" + +# ---------------------------------------------------------------- palette +BG = (5, 7, 10) +PANEL = (12, 17, 23) +PANEL_EDGE = (28, 36, 46) +INK = (241, 239, 224) +MUTED = (143, 154, 160) +DIM = (80, 90, 98) +AMBER = (255, 196, 68) +CYAN = (75, 220, 255) +GREEN = (148, 255, 117) + +W, H = 1200, 720 +TOKEN_IDX = 310 +ANSWER_BPE = "acular" +ANSWER_SLOT = 1 # answer_token_p[1] == p('acular') + +# ---------------------------------------------------------------- fonts +MENLO = "/System/Library/Fonts/Menlo.ttc" +UNI = "/System/Library/Fonts/Supplemental/Arial Unicode.ttf" + + +def font(path: str, size: int, index: int = 0) -> ImageFont.FreeTypeFont: + try: + return ImageFont.truetype(path, size, index=index) + except OSError: + return ImageFont.truetype(MENLO, size) + + +F_TITLE = font(MENLO, 30, index=1) +F_SUB = font(MENLO, 15) +F_LABEL = font(MENLO, 13) +F_LABEL_B = font(MENLO, 13, index=1) +F_TINY = font(MENLO, 11) +F_LAYER = font(MENLO, 64, index=1) +F_STAGE = font(MENLO, 16, index=1) +F_NUM = font(MENLO, 14) +F_TOK = font(UNI, 18) +F_TOK_B = font(MENLO, 18, index=1) +F_TOK_S = font(UNI, 14) +F_BADGE = font(MENLO, 17, index=1) + + +def load() -> tuple[dict, list[dict], list[dict]]: + data = json.loads(DATA.read_text()) + assert data["layers"] == 29 and data["image_grid"] == 56 + assert data["answer_token_strs"][ANSWER_SLOT] == ANSWER_BPE + + def by_layer(idx: int) -> list[dict]: + return sorted( + (e for e in data["lens"] if e["token_index"] == idx), + key=lambda e: e["layer"], + ) + + target = by_layer(TOKEN_IDX) + control = by_layer(data["control_indices"][0]) + assert len(target) == 29 and len(control) == 29 + return data, target, control + + +def crop_cell(carrier: Image.Image, idx: int) -> Image.Image: + r, c = idx // 56, idx % 56 + return carrier.crop((c * 28, r * 28, (c + 1) * 28, (r + 1) * 28)) + + +def sanitize(s: str) -> str: + out = "".join(ch if ch.isprintable() else "\ufffd" for ch in s) + return out.replace(" ", "\u2423", 1) if s.startswith(" ") else out + + +def text_w(f: ImageFont.FreeTypeFont, s: str) -> float: + return f.getlength(s) + + +STAGES = [ + (0, "static"), + (8, "gibberish"), + (16, "morphemes"), + (23, "warming up"), + (24, "lock-on"), + (28, "crystallized"), +] + + +def stage_for(layer: int) -> str: + name = STAGES[0][1] + for lo, label in STAGES: + if layer >= lo: + name = label + return name + + +def render_frame( + layer: int, + data: dict, + target: list[dict], + control: list[dict], + patch: Image.Image, + ctrl_patch: Image.Image, + strip: Image.Image, + strip_cell: int, +) -> Image.Image: + e = target[layer] + ce = control[layer] + p_ans = e["answer_token_p"][ANSWER_SLOT] + p_ctrl = ce["answer_token_p"][ANSWER_SLOT] + locked = layer >= 24 + final = layer == 28 + heat = min(1.0, p_ans / 0.40) # glow scales with the real probability + + img = Image.new("RGB", (W, H), BG) + d = ImageDraw.Draw(img) + + # ---------------------------------------------------------- header + d.text((40, 22), "CRYSTALLIZATION", font=F_TITLE, fill=INK) + tw = text_w(F_TITLE, "CRYSTALLIZATION") + d.text( + (40 + tw + 18, 36), + "\u2014 28\u00d728 pixels become the token '" + ANSWER_BPE + "'", + font=F_SUB, + fill=MUTED, + ) + d.text( + (40, 60), + "logit lens on Qwen2.5-VL-7B \u00b7 image carrier 1568px \u00b7 visual token " + f"#{TOKEN_IDX} (grid {e['grid_rc'][0]},{e['grid_rc'][1]}) \u00b7 " + "answer \u201cspectacular\u201d = [spect][acular]", + font=F_SUB, + fill=MUTED, + ) + if final: # celebratory badge in the free top-right corner + badge = f"'{ANSWER_BPE}' \u00b7 p={p_ans:.2f} \u00b7 CRYSTALLIZED" + bw = text_w(F_BADGE, badge) + bx = W - bw - 76 + d.rounded_rectangle((bx, 22, bx + bw + 36, 56), radius=17, fill=(38, 30, 10), outline=AMBER, width=2) + d.text((bx + 18, 29), badge, font=F_BADGE, fill=AMBER) + + # ---------------------------------------------------------- context strip + sx, sy = (W - strip.width) // 2, 96 + img.paste(strip, (sx, sy)) + d.rectangle((sx, sy, sx + strip.width - 1, sy + strip.height - 1), outline=PANEL_EDGE) + col = AMBER if locked else CYAN + d.rectangle( + (sx + strip_cell, sy, sx + strip_cell + strip.height, sy + strip.height - 1), + outline=col, + width=3, + ) + cap = "carrier row 5 \u00b7 the model never sees glyphs \u2014 only these pixels" + d.text(((W - text_w(F_TINY, cap)) / 2, sy + strip.height + 6), cap, font=F_TINY, fill=DIM) + + # ---------------------------------------------------------- main row + top_y = 232 + # patch panel ------------------------------------------------------ + px, py, ps = 40, top_y, 280 + img.paste(patch.resize((ps, ps), Image.NEAREST), (px, py)) + d.rectangle((px - 1, py - 1, px + ps, py + ps), outline=col, width=2) + d.text((px, py + ps + 10), f"visual token #{TOKEN_IDX}", font=F_LABEL_B, fill=INK) + d.text((px, py + ps + 28), "28\u00d728 px \u00b7 reads: \u2018\"sp\u2019 / \u2018and\u2019", font=F_TINY, fill=MUTED) + + # layer counter ---------------------------------------------------- + cx = 392 + d.text((cx, top_y - 2), "LAYER", font=F_LABEL, fill=MUTED) + num = f"{layer:02d}" + d.text((cx, top_y + 16), num, font=F_LAYER, fill=AMBER if locked else INK) + d.text((cx + text_w(F_LAYER, num) + 8, top_y + 58), "/28", font=F_STAGE, fill=DIM) + stage = stage_for(layer) + d.text((cx, top_y + 98), stage.upper(), font=F_STAGE, fill=GREEN if final else (AMBER if locked else MUTED)) + ry = top_y + 134 # mini rail of 29 ticks + for i in range(29): + tx = cx + i * 5 + d.rectangle((tx, ry, tx + 3, ry + 10), fill=AMBER if i <= layer else (40, 48, 58)) + d.text((cx, ry + 18), f"p('{ANSWER_BPE}') = {p_ans:.4f}", font=F_NUM, fill=AMBER if p_ans > 0.01 else DIM) + + # top-5 panel ------------------------------------------------------ + tx0, ty0, tx1, ty1 = 580, top_y - 12, 1160, top_y + 318 + d.rounded_rectangle((tx0, ty0, tx1, ty1), radius=8, fill=PANEL, outline=PANEL_EDGE) + d.text((tx0 + 18, ty0 + 12), "TOP-5 DECODED VOCAB TOKENS \u00b7 what this patch \u201cmeans\u201d so far", font=F_LABEL, fill=MUTED) + bar_x = tx0 + 230 + bar_max = tx1 - bar_x - 86 + scale = 0.45 # fixed probability scale across all frames + for i, t in enumerate(e["top"]): + yy = ty0 + 48 + i * 54 + is_ans = t["id"] in data["answer_token_ids"] + tok_s = sanitize(t["str"]) + if len(tok_s) > 16: + tok_s = tok_s[:15] + "\u2026" + d.text((tx0 + 18, yy), f"'{tok_s}'", font=F_TOK_B if is_ans else F_TOK, fill=AMBER if is_ans else INK) + bw = max(2, int(min(t["p"] / scale, 1.0) * bar_max)) + d.rectangle((bar_x, yy + 4, bar_x + bw, yy + 18), fill=AMBER if is_ans else (58, 70, 84)) + if is_ans and heat > 0.3: + d.rectangle((bar_x, yy + 4, bar_x + bw, yy + 18), outline=INK) + d.text((bar_x + bw + 10, yy + 3), f"{t['p']:.3f}", font=F_NUM, fill=AMBER if is_ans else MUTED) + d.text((tx0 + 18, yy + 24), f"id {t['id']}", font=F_TINY, fill=DIM) + + # ---------------------------------------------------------- bottom row + by0, by1 = 596, 708 + # confidence meter for 'acular' + mx0, mx1 = 40, 730 + d.rounded_rectangle((mx0, by0, mx1, by1), radius=8, fill=PANEL, outline=PANEL_EDGE) + d.text((mx0 + 16, by0 + 8), f"CONFIDENCE \u00b7 p('{ANSWER_BPE}') across layers", font=F_LABEL, fill=MUTED) + leg_x = mx1 - 130 + d.rectangle((leg_x, by0 + 12, leg_x + 14, by0 + 15), fill=AMBER) + d.text((leg_x + 20, by0 + 6), "answer", font=F_TINY, fill=AMBER) + d.rectangle((leg_x + 74, by0 + 12, leg_x + 88, by0 + 15), fill=(60, 70, 80)) + d.text((leg_x + 94, by0 + 6), "ctrl", font=F_TINY, fill=(96, 108, 118)) + ch_x0, ch_x1 = mx0 + 52, mx1 - 64 + ch_y0, ch_y1 = by0 + 32, by1 - 22 + p_max = 0.45 + d.line((ch_x0, ch_y1, ch_x1, ch_y1), fill=PANEL_EDGE) + for gv in (0.2, 0.4): + gy = ch_y1 - gv / p_max * (ch_y1 - ch_y0) + d.line((ch_x0, gy, ch_x1, gy), fill=(22, 28, 36)) + d.text((mx0 + 16, gy - 6), f"{gv:.1f}", font=F_TINY, fill=DIM) + + def xs(l: int) -> float: + return ch_x0 + l / 28 * (ch_x1 - ch_x0) + + def ys(p: float) -> float: + return ch_y1 - min(p, p_max) / p_max * (ch_y1 - ch_y0) + + pts = [(xs(l), ys(target[l]["answer_token_p"][ANSWER_SLOT])) for l in range(layer + 1)] + cpts = [(xs(l), ys(control[l]["answer_token_p"][ANSWER_SLOT])) for l in range(layer + 1)] + if len(cpts) > 1: + d.line(cpts, fill=(60, 70, 80), width=2) + if len(pts) > 1: + poly = pts + [(pts[-1][0], ch_y1), (pts[0][0], ch_y1)] + d.polygon(poly, fill=(76, 56, 16) if locked else (46, 36, 14)) + d.line(pts, fill=AMBER, width=3) + hx, hy = pts[-1] + d.ellipse((hx - 5, hy - 5, hx + 5, hy + 5), fill=AMBER if p_ans > 0.01 else MUTED) + head = f"{p_ans:.2f}" if p_ans >= 0.005 else f"{p_ans:.4f}" + d.text((min(hx + 8, ch_x1 - 8), hy - 18), head, font=F_NUM, fill=AMBER if p_ans > 0.01 else MUTED) + for ml in (24, 28): + if layer >= ml: + mlx = xs(ml) + d.line((mlx, ch_y1, mlx, ys(target[ml]["answer_token_p"][ANSWER_SLOT])), fill=(90, 72, 30)) + d.text((mlx - 10, ch_y1 + 6), f"L{ml}", font=F_TINY, fill=AMBER) + d.text((ch_x0, ch_y1 + 6), "L0", font=F_TINY, fill=DIM) + + # control panel ---------------------------------------------------- + kx0, kx1 = 760, 1160 + d.rounded_rectangle((kx0, by0, kx1, by1), radius=8, fill=PANEL, outline=PANEL_EDGE) + d.text((kx0 + 16, by0 + 8), "CONTROL \u00b7 token #" + str(ce["token_index"]), font=F_LABEL, fill=MUTED) + cps = 60 + img.paste(ctrl_patch.resize((cps, cps), Image.NEAREST), (kx0 + 16, by0 + 30)) + d.rectangle((kx0 + 15, by0 + 29, kx0 + 16 + cps, by0 + 30 + cps), outline=PANEL_EDGE) + ct = ce["top"][0] + ct_s = sanitize(ct["str"]) + if len(ct_s) > 12: + ct_s = ct_s[:11] + "\u2026" + lx = kx0 + 16 + cps + 14 + d.text((lx, by0 + 30), "top-1: ", font=F_NUM, fill=INK) + tx = lx + text_w(F_NUM, "top-1: ") + d.text((tx, by0 + 30), f"'{ct_s}'", font=F_TOK_S, fill=INK) + d.text((tx + text_w(F_TOK_S, f"'{ct_s}'") + 10, by0 + 30), f"{ct['p']:.3f}", font=F_NUM, fill=INK) + d.text((lx, by0 + 52), f"p('{ANSWER_BPE}') = {p_ctrl:.5f}", font=F_NUM, fill=MUTED) + d.text((lx, by0 + 74), "still noise \u2713" if p_ctrl < 0.01 else "?!", font=F_LABEL_B, fill=GREEN) + d.text((kx0 + 224, by0 + 8), "never converges to the answer", font=F_TINY, fill=DIM) + + # ---------------------------------------------------------- glow + if locked and heat > 0: + glow = Image.new("RGB", (W, H), (0, 0, 0)) + gd = ImageDraw.Draw(glow) + a = int(70 + 110 * heat) + gd.rectangle((px - 6, py - 6, px + ps + 5, py + ps + 5), outline=(a, int(a * 0.77), int(a * 0.27)), width=10) + if final: + gd.rectangle((px - 14, py - 14, px + ps + 13, py + ps + 13), outline=(a, int(a * 0.77), int(a * 0.27)), width=8) + glow = glow.filter(ImageFilter.GaussianBlur(12 if final else 8)) + img = Image.composite(Image.new("RGB", (W, H), AMBER), img, glow.convert("L").point(lambda v: min(v, 140))) + return img + + +def main() -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + data, target, control = load() + carrier = Image.open(CARRIER).convert("RGB") + if carrier.size != (1568, 1568): + carrier = carrier.resize((1568, 1568), Image.LANCZOS) + patch = crop_cell(carrier, TOKEN_IDX) + ctrl_patch = crop_cell(carrier, data["control_indices"][0]) + # context strip: row 5, cols 26..39 (14 cells), scaled x3 -> 1176x84 + c0, c1 = 26, 40 + raw = carrier.crop((c0 * 28, 5 * 28, c1 * 28, 6 * 28)) + strip = raw.resize((raw.width * 3, raw.height * 3), Image.NEAREST) + strip_cell = (TOKEN_IDX % 56 - c0) * 84 + + frames, durations = [], [] + for layer in range(29): + fr = render_frame(layer, data, target, control, patch, ctrl_patch, strip, strip_cell) + frames.append(fr.quantize(colors=256, method=Image.MEDIANCUT, dither=Image.Dither.NONE)) + if layer < 23: + durations.append(220) + elif layer < 28: + durations.append(520) + else: + durations.append(2500) # celebratory hold on the final frame + + gif = OUT_DIR / "crystal.gif" + frames[0].save( + gif, + save_all=True, + append_images=frames[1:], + duration=durations, + loop=0, + disposal=2, + optimize=False, + ) + final_png = OUT_DIR / "crystal_final.png" + render_frame(28, data, target, control, patch, ctrl_patch, strip, strip_cell).save(final_png) + print(f"wrote {gif} ({gif.stat().st_size / 1024:.0f} KB, {len(frames)} frames)") + print(f"wrote {final_png}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_r2_filmstrip.py b/packages/snapcompact/research/snapcompact_r2_filmstrip.py new file mode 100755 index 000000000..5f88b31fc --- /dev/null +++ b/packages/snapcompact/research/snapcompact_r2_filmstrip.py @@ -0,0 +1,256 @@ +#!/usr/bin/env python3 +"""Twin reels: a filmstrip poster of text-carrier vs image-carrier similarity geometry. + +Renders ~7 layers of the carrier-centered 12x12 cosine matrices (`text_sim`, +`image_sim` from carrier_convergence.npz) as paired frames on two parallel +film reels, with the REAL per-layer RSA Pearson from summary.json as a match +meter under each frame, plus the cross-carrier matched cosine as a secondary +tick. Closes with a callout frame for the best layer (RSA 0.85 @ L19). + +Output: results/agent-r2-filmstrip/filmstrip.png (2200 px wide). +""" + +import json +import os + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import LinearSegmentedColormap +from matplotlib.patches import FancyBboxPatch, Rectangle + +HERE = os.path.dirname(os.path.abspath(__file__)) +DATA_DIR = os.path.join(HERE, "results", "qwen-carrier-convergence-n12") +OUT_DIR = os.path.join(HERE, "results", "agent-r2-filmstrip") + +# ---------------------------------------------------------------- palette +BG = "#05070a" +PANEL = "#0c1117" +FILM = "#0a0e14" +INK = "#f1efe0" +MUTED = "#8f9aa0" +AMBER = "#ffc444" +CYAN = "#4bdcff" +ORANGE = "#ff7048" +GREEN = "#94ff75" +EDGE = "#1d2630" + +DIVERGING = LinearSegmentedColormap.from_list( + "carrier_div", + [(0.0, CYAN), (0.30, "#16384a"), (0.50, "#0b1016"), (0.72, "#5c2c18"), (0.90, ORANGE), (1.0, AMBER)], +) + +LAYERS = [1, 5, 9, 13, 17, 19, 28] + +# ---------------------------------------------------------------- data +npz = np.load(os.path.join(DATA_DIR, "carrier_convergence.npz")) +text_sim = npz["text_sim"] # [29, 12, 12] +image_sim = npz["image_sim"] # [29, 12, 12] +assert text_sim.shape == image_sim.shape == (29, 12, 12) + +with open(os.path.join(DATA_DIR, "summary.json")) as fh: + summary = json.load(fh) +per_layer = {row["layer"]: row for row in summary["per_layer"]} +best = summary["best"] # layer 19: rsa 0.85, matched 0.66, mismatched -0.06 + +# ---------------------------------------------------------------- layout (pixel space) +W, H = 2200, 1000 +fig = plt.figure(figsize=(W / 100, H / 100), dpi=100) +fig.patch.set_facecolor(BG) +ax = fig.add_axes([0, 0, 1, 1]) +ax.set_xlim(0, W) +ax.set_ylim(H, 0) # y grows downward +ax.axis("off") +ax.set_facecolor(BG) + +X0, X1 = 150, 2136 +N_COLS = len(LAYERS) + 1 # 7 layer frames + closing callout +CW = (X1 - X0) / N_COLS +FS = 186 # matrix tile size + +TEXT_BAND_Y, BAND_H = 192, 262 +IMAGE_BAND_Y = TEXT_BAND_Y + BAND_H + 26 +METER_Y = IMAGE_BAND_Y + BAND_H + 36 +METER_H = 64 +FOOT_Y = METER_Y + METER_H + 52 + + +def col_cx(i: int) -> float: + return X0 + (i + 0.5) * CW + + +def sprockets(y: float, x_start: float, x_end: float) -> None: + x = x_start + 14 + while x + 20 < x_end: + ax.add_patch( + FancyBboxPatch( + (x, y), 20, 13, + boxstyle="round,pad=0,rounding_size=4", + facecolor=BG, edgecolor="#27313d", linewidth=1.0, zorder=6, + ) + ) + x += 49 + + +def film_band(y0: float, x_end: float) -> None: + ax.add_patch( + Rectangle((X0 - 26, y0), x_end - X0 + 26, BAND_H, + facecolor=FILM, edgecolor=EDGE, linewidth=1.2, zorder=2) + ) + sprockets(y0 + 11, X0 - 26, x_end) + sprockets(y0 + BAND_H - 24, X0 - 26, x_end) + + +BAND_X_END = X0 + (N_COLS - 1) * CW - 14 # bands stop before the callout column +film_band(TEXT_BAND_Y, BAND_X_END) +film_band(IMAGE_BAND_Y, BAND_X_END) + +# reel labels on the left edge +for y0, label, color in ( + (TEXT_BAND_Y, "TEXT REEL", CYAN), + (IMAGE_BAND_Y, "IMAGE REEL", ORANGE), +): + ax.text(X0 - 56, y0 + BAND_H / 2, label, color=color, fontsize=13, + fontweight="bold", rotation=90, ha="center", va="center", zorder=8) + ax.text(X0 - 84, y0 + BAND_H / 2, "12 \u00d7 12 carrier cosine", color=MUTED, + fontsize=8, rotation=90, ha="center", va="center", zorder=8) + +# ---------------------------------------------------------------- frames +VLIM = 0.75 # diagonal (cos=1) clips to amber, off-diagonal structure fills the range + + +def draw_matrix(mat: np.ndarray, cx: float, band_y: float) -> None: + x0m, y0m = cx - FS / 2, band_y + 36 + ax.imshow( + mat, cmap=DIVERGING, vmin=-VLIM, vmax=VLIM, + extent=(x0m, x0m + FS, y0m + FS, y0m), origin="upper", + interpolation="nearest", zorder=4, + ) + ax.add_patch(Rectangle((x0m, y0m), FS, FS, fill=False, + edgecolor=EDGE, linewidth=1.1, zorder=5)) + + +for i, layer in enumerate(LAYERS): + cx = col_cx(i) + draw_matrix(text_sim[layer], cx, TEXT_BAND_Y) + draw_matrix(image_sim[layer], cx, IMAGE_BAND_Y) + + # frame numbering, film style + ax.text(cx, TEXT_BAND_Y - 12, f"FRAME {i + 1:02d}", color=MUTED, + fontsize=8.5, ha="center", va="bottom", zorder=8) + for band_y in (TEXT_BAND_Y, IMAGE_BAND_Y): + ax.text(cx, band_y + 36 + FS + 14, f"LAYER {layer}", color=INK, + fontsize=10, fontweight="bold", ha="center", va="center", zorder=8) + + # dotted connector between the paired frames + ax.plot([cx, cx], [TEXT_BAND_Y + BAND_H + 3, IMAGE_BAND_Y - 3], + color="#3a4754", linewidth=1.2, linestyle=(0, (1, 3)), zorder=3) + +# ---------------------------------------------------------------- match meters +ax.text(X0 - 26, METER_Y - 14, "GEOMETRY MATCH", color=INK, fontsize=10, + fontweight="bold", ha="left", va="bottom", zorder=8) +ax.text(X0 + 152, METER_Y - 14, + "amber bar \u2014 RSA: Pearson r of the two reels' off-diagonal structure" + " cyan tick \u2014 matched cross-carrier cosine", + color=MUTED, fontsize=8.5, ha="left", va="bottom", zorder=8) + +BAR_W = FS +for i, layer in enumerate(LAYERS): + row = per_layer[layer] + rsa = row["rsa_pearson"] + matched = row["matched_cosine"] + cx = col_cx(i) + bx = cx - BAR_W / 2 + + ax.add_patch(Rectangle((bx, METER_Y), BAR_W, 12, facecolor=PANEL, + edgecolor=EDGE, linewidth=0.8, zorder=4)) + ax.add_patch(Rectangle((bx, METER_Y), BAR_W * rsa, 12, facecolor=AMBER, + edgecolor="none", zorder=5)) + ax.plot([bx + BAR_W * matched] * 2, [METER_Y - 4, METER_Y + 16], + color=CYAN, linewidth=2.0, zorder=6) + + ax.text(cx, METER_Y + 32, f"RSA {rsa:.2f}", color=AMBER, fontsize=10.5, + fontweight="bold", ha="center", va="center", zorder=8) + ax.text(cx, METER_Y + 50, f"matched cos {matched:.2f}", color=CYAN, + fontsize=8.5, ha="center", va="center", zorder=8) + +# ---------------------------------------------------------------- closing callout frame +cb_x = X0 + (N_COLS - 1) * CW + 2 +cb_w = X1 - cb_x +cb_y0, cb_y1 = TEXT_BAND_Y, METER_Y + METER_H +ax.add_patch( + FancyBboxPatch( + (cb_x, cb_y0), cb_w, cb_y1 - cb_y0, + boxstyle="round,pad=0,rounding_size=10", + facecolor=PANEL, edgecolor=AMBER, linewidth=1.6, zorder=4, + ) +) +ccx = cb_x + cb_w / 2 +ax.text(ccx, cb_y0 + 46, "THE SPLICE", color=MUTED, fontsize=10, + ha="center", va="center", zorder=8) +ax.text(ccx, cb_y0 + 122, f"RSA {best['rsa_pearson']:.2f}", color=AMBER, + fontsize=33, fontweight="bold", ha="center", va="center", zorder=8) +ax.text(ccx, cb_y0 + 168, f"@ LAYER {best['layer']}", color=INK, fontsize=14, + fontweight="bold", ha="center", va="center", zorder=8) + +ax.plot([cb_x + 28, cb_x + cb_w - 28], [cb_y0 + 206] * 2, + color=EDGE, linewidth=1.0, zorder=5) + +facts = [ + (f"matched cosine {best['matched_cosine']:.2f}", CYAN), + (f"mismatched {best['mismatched_cosine']:.2f}", MUTED), + (f"retrieval {int(round(best['match_rank_accuracy'] * 12))}/12", GREEN), +] +for j, (line, color) in enumerate(facts): + ax.text(ccx, cb_y0 + 244 + j * 34, line, color=color, fontsize=11.5, + fontweight="bold", ha="center", va="center", zorder=8) + +ax.text(ccx, cb_y0 + 380, "Read it as text or look at\nthe picture \u2014 by layer 19\nthe model files both under\nthe same geometry.", + color=INK, fontsize=10.5, ha="center", va="center", linespacing=1.6, zorder=8) + +# the actual L19 splice: the twin pair, miniaturized +MINI = 78 +for mat, mx, tag, tcol in ( + (text_sim[best["layer"]], ccx - MINI - 9, "text", CYAN), + (image_sim[best["layer"]], ccx + 9, "image", ORANGE), +): + ax.imshow(mat, cmap=DIVERGING, vmin=-VLIM, vmax=VLIM, + extent=(mx, mx + MINI, cb_y0 + 444 + MINI, cb_y0 + 444), + origin="upper", interpolation="nearest", zorder=6) + ax.add_patch(Rectangle((mx, cb_y0 + 444), MINI, MINI, fill=False, + edgecolor=EDGE, linewidth=1.0, zorder=7)) + ax.text(mx + MINI / 2, cb_y0 + 444 + MINI + 14, tag, color=tcol, + fontsize=9, ha="center", va="center", zorder=8) +ax.text(ccx, cb_y1 - 36, "two carriers,\none geometry", color=AMBER, fontsize=11, + fontweight="bold", fontstyle="italic", ha="center", va="center", zorder=8) + +# ---------------------------------------------------------------- title & footer +ax.text(X0 - 26, 64, "TWIN REELS", color=INK, fontsize=34, fontweight="bold", + ha="left", va="center", zorder=8) +ax.text(X0 + 318, 64, "\u2014 the same 12 facts, shot twice", color=AMBER, + fontsize=16, ha="left", va="center", zorder=8) +ax.text( + X0 - 26, 118, + "Twelve question\u2013answer pairs enter Qwen2.5-VL-7B twice: once as text, once rendered into pixels. " + "Each frame is the 12\u00d712 cosine similarity between carrier states at one layer \u2014 " + "the two reels print the same relational structure from the very first frames.", + color=MUTED, fontsize=11.5, ha="left", va="center", zorder=8, +) + +ax.text( + X0 - 26, FOOT_Y, + "data: results/qwen-carrier-convergence-n12 (carrier_convergence.npz \u00b7 summary.json) \u00b7 " + "carrier-centered cosine of hidden states, d = 3584, 29 layers \u00b7 " + "RSA = Pearson r over the 66 off-diagonal pairs \u00b7 " + "diverging scale \u2212%.2f \u2026 +%.2f (cyan \u2192 dark \u2192 orange)" % (VLIM, VLIM), + color=MUTED, fontsize=9, ha="left", va="center", zorder=8, +) + +# ---------------------------------------------------------------- save +os.makedirs(OUT_DIR, exist_ok=True) +out_path = os.path.join(OUT_DIR, "filmstrip.png") +fig.savefig(out_path, dpi=100, facecolor=BG) +print("wrote", out_path) diff --git a/packages/snapcompact/research/snapcompact_r2_hero.py b/packages/snapcompact/research/snapcompact_r2_hero.py new file mode 100755 index 000000000..49370267d --- /dev/null +++ b/packages/snapcompact/research/snapcompact_r2_hero.py @@ -0,0 +1,497 @@ +#!/usr/bin/env python3 +"""Hero / OG social card: "Two carriers, one thought." + +A single poster composition for the snapcompact blog post: a fragment of the +real text carrier (BPE tokens `spect`+`acular`) on the left, the real bitmap +patch pixels of the same word on the right, both flowing into one glowing +shared core annotated with the real layer-19 convergence stats. + +Outputs: + results/agent-r2-hero/hero-1200x630.png + results/agent-r2-hero/hero-2400x1260.png + +All numbers are read from: + results/qwen-carrier-convergence-n12/summary.json + results/qwen-token-entry-q3/token_entry.json + results/qwen-logit-lens-q3/images/image-carrier.png +""" + +from __future__ import annotations + +import json +import math +from pathlib import Path + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +ROOT = Path(__file__).resolve().parent +OUT_DIR = ROOT / "results" / "agent-r2-hero" + +# Layout is specified in 2400x1260 coordinates; the master canvas renders at +# MS x that for antialiasing, then is downsampled to both deliverables. +BASE_W, BASE_H = 2400, 1260 +MS = 2 +W, H = BASE_W * MS, BASE_H * MS + +# Palette (brief). +BG = (5, 7, 10) +PANEL = (12, 17, 23) +INK = (241, 239, 224) +MUTED = (143, 154, 160) +AMBER = (255, 196, 68) +CYAN = (75, 220, 255) +DIVIDER = (26, 34, 44) + +CORE_WORD = '"spectacular"' + + +def u(v: float) -> int: + return int(round(v * MS)) + + +def font_at(path: str, size: float, index: int = 0) -> ImageFont.FreeTypeFont | None: + p = Path(path) + if not p.exists(): + return None + try: + return ImageFont.truetype(str(p), u(size), index=index) + except OSError: + return None + + +def display_font(size: float) -> ImageFont.FreeTypeFont: + """Heavy display face for the title and big stats.""" + for path, index in [ + ("/System/Library/Fonts/Avenir Next.ttc", 8), # Heavy + ("/System/Library/Fonts/Supplemental/Arial Black.ttf", 0), + ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 0), + ]: + f = font_at(path, size, index) + if f is not None: + return f + return ImageFont.load_default() + + +def label_font(size: float) -> ImageFont.FreeTypeFont: + for path, index in [ + ("/System/Library/Fonts/Avenir Next.ttc", 2), # Demi Bold + ("/System/Library/Fonts/Supplemental/Arial Bold.ttf", 0), + ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 0), + ]: + f = font_at(path, size, index) + if f is not None: + return f + return ImageFont.load_default() + + +def body_font(size: float) -> ImageFont.FreeTypeFont: + for path, index in [ + ("/System/Library/Fonts/Avenir Next.ttc", 5), # Medium + ("/System/Library/Fonts/Supplemental/Arial.ttf", 0), + ("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 0), + ]: + f = font_at(path, size, index) + if f is not None: + return f + return ImageFont.load_default() + + +def mono_font(size: float) -> ImageFont.FreeTypeFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + f = font_at(path, size) + if f is not None: + return f + return ImageFont.load_default() + + +def tracked(draw: ImageDraw.ImageDraw, xy: tuple[int, int], text: str, font, fill, tracking: float = 0.0) -> int: + """Draw text with letterspacing; returns end x.""" + x, y = xy + t = u(tracking) + for ch in text: + draw.text((x, y), ch, font=font, fill=fill) + x += draw.textlength(ch, font=font) + t + return int(x) + + +def tracked_width(draw: ImageDraw.ImageDraw, text: str, font, tracking: float = 0.0) -> float: + t = u(tracking) + return sum(draw.textlength(ch, font=font) + t for ch in text) - (t if text else 0) + + +def bezier(p0, p1, p2, n=64): + pts = [] + for i in range(n + 1): + t = i / n + x = (1 - t) ** 2 * p0[0] + 2 * (1 - t) * t * p1[0] + t**2 * p2[0] + y = (1 - t) ** 2 * p0[1] + 2 * (1 - t) * t * p1[1] + t**2 * p2[1] + pts.append((x, y)) + return pts + + +def load_data(): + conv = json.loads((ROOT / "results" / "qwen-carrier-convergence-n12" / "summary.json").read_text()) + entry = json.loads((ROOT / "results" / "qwen-token-entry-q3" / "token_entry.json").read_text()) + carrier = Image.open(ROOT / "results" / "qwen-logit-lens-q3" / "images" / "image-carrier.png").convert("RGB") + if carrier.size != (1568, 1568): + carrier = carrier.resize((1568, 1568), Image.LANCZOS) + + best = conv["best"] + layer = conv["best_layer"] + n_q = conv["n_questions"] + stats = { + "layer": layer, + "n_layers": conv["layers"], + "matched": best["matched_cosine"], + "rsa": best["rsa_pearson"], + "retrieved": int(round(best["match_rank_accuracy"] * n_q)), + "n": n_q, + } + assert stats["layer"] == 19 and abs(stats["matched"] - 0.66) < 0.01 + assert abs(stats["rsa"] - 0.85) < 0.01 and stats["retrieved"] == 12 and stats["n"] == 12 + + toks = {t["i"]: t for t in entry["tokens"]} + answer = [t for t in entry["tokens"] if t["answer"]] + assert [t["str"] for t in answer] == ["spect", "acular"] + assert [t["id"] for t in answer] == [67082, 23006] + ctx_before = "…" + "".join(toks[i]["str"] for i in range(23, 32)) # " make the 50th Super Bowl \"" + ctx_after = "".join(toks[i]["str"] for i in range(34, 39)) + "…" # "\" and that it would" + + grid = entry["image_grid"] # 56 + word_idx = entry["image_answer_token_indices"][:4] # [310, 311, 312, 313] + assert word_idx == [310, 311, 312, 313] + assert word_idx[0] // grid == 5 and word_idx[0] % grid == 30 + + counts = { + "chars": entry["chunk_chars"], + "text_tokens": entry["chunk_text_tokens"], + "image_tokens": entry["image_tokens"], + "grid": grid, + "patch_px": entry["token_pixel_size"], # 28 + "embed_dim": entry["embed_dim"], + "visual_dim": entry["visual_out_dim"], + } + heads = { + "text": entry["text_entry"][0], # id 67082 "spect": 10-dim head + norm + "image": entry["image_entry"][0], # patch 310: 10-dim head + norm + } + assert heads["text"]["id"] == 67082 and heads["image"]["token_index"] == 310 + return stats, answer, (ctx_before, ctx_after), word_idx, counts, heads, carrier + + +# --------------------------------------------------------------------------- +# Composition geometry (2400x1260 space) +# --------------------------------------------------------------------------- +CORE = (1200, 660) +PANEL_TOP, PANEL_BOT = 332, 938 +LP = (92, PANEL_TOP, 782, PANEL_BOT) # left panel +RP = (1618, PANEL_TOP, 2308, PANEL_BOT) # right panel + + +def additive_base() -> np.ndarray: + """Background + ambient glow + orb fields + streamlines, all additive.""" + img = np.zeros((H, W, 3), dtype=np.float32) + img[:] = BG + + yy, xx = np.mgrid[0:H, 0:W].astype(np.float32) + cx, cy = u(CORE[0]), u(CORE[1]) + d = np.sqrt((xx - cx) ** 2 + (yy - cy) ** 2) + + # Ambient: amber wash on the left, cyan on the right, strongest near core. + side = np.clip((xx - cx) / u(900), -1.0, 1.0) + amber = np.array(AMBER, np.float32) / 255.0 + cyan = np.array(CYAN, np.float32) / 255.0 + mix = (1 - side[..., None]) / 2 * amber + (1 + side[..., None]) / 2 * cyan + img += 26.0 * mix * np.exp(-((d / u(760)) ** 2))[..., None] + + # Orb halo and warm nucleus (kept below blowout so the word stays legible). + img += 92.0 * mix * np.exp(-((d / u(250)) ** 2))[..., None] + warm = np.array((255, 240, 205), np.float32) / 255.0 + img += 96.0 * warm * np.exp(-((d / u(118)) ** 2))[..., None] + + # Subtle vignette. + ex = ((xx / W) - 0.5) ** 2 + ((yy / H) - 0.5) ** 2 + img *= (1.0 - 0.55 * ex)[..., None] + + # Streamlines: each carrier feeds the core. + glow = Image.new("RGB", (W, H), (0, 0, 0)) + gd = ImageDraw.Draw(glow) + sharp = Image.new("RGB", (W, H), (0, 0, 0)) + sd = ImageDraw.Draw(sharp) + rng = np.random.default_rng(19) + + def streams(x0: float, x_sign: float, color): + n = 9 + for k in range(n): + f = k / (n - 1) + y0 = PANEL_TOP + 120 + f * (PANEL_BOT - PANEL_TOP - 240) + rng.uniform(-14, 14) + ang = (f - 0.5) * 1.45 + rng.uniform(-0.07, 0.07) + r = 168 + x2 = CORE[0] - x_sign * r * math.cos(ang) + y2 = CORE[1] + r * math.sin(ang) + mid_x = (x0 + x2) / 2 + x_sign * 36 + mid_y = y0 + (CORE[1] - y0) * 0.62 + pts = bezier((x0, y0), (mid_x, mid_y), (x2, y2), 72) + spts = [(u(px), u(py)) for px, py in pts] + fade = 1.0 - abs(f - 0.5) * 0.9 + gd.line(spts, fill=tuple(int(c * 0.62 * fade) for c in color), width=u(7)) + sd.line(spts, fill=tuple(int(c * 0.5 * fade) for c in color), width=u(1.6)) + # Energy particles along the stream. + for t in (0.3, 0.55, 0.8): + i = int(t * len(pts)) + px, py = u(pts[i][0]), u(pts[i][1]) + rr = u(3.2) + gd.ellipse([px - rr, py - rr, px + rr, py + rr], fill=tuple(int(c * fade) for c in color)) + + streams(LP[2], 1.0, AMBER) + streams(RP[0], -1.0, CYAN) + + img += np.asarray(glow.filter(ImageFilter.GaussianBlur(u(11))), np.float32) * 0.9 + img += np.asarray(sharp.filter(ImageFilter.GaussianBlur(u(0.7))), np.float32) + + # Soft glow behind the core word (the sharp pass is drawn later, on top). + f = display_font(78) + layer = Image.new("RGB", (W, H), (0, 0, 0)) + ld = ImageDraw.Draw(layer) + tw = ld.textlength(CORE_WORD, font=f) + ld.text((u(CORE[0]) - tw / 2, u(CORE[1] - 54)), CORE_WORD, font=f, fill=(255, 232, 170)) + img += np.asarray(layer.filter(ImageFilter.GaussianBlur(u(9))), np.float32) * 0.8 + return img + + +def head_bars(ov: ImageDraw.ImageDraw, x: float, y_mid: float, values, color, label: str): + """Tiny bar strip of a real 10-dim vector head, centered on its axis.""" + vmax = max(abs(v) for v in values) + bw, gap, amp = 24, 11, 26 + total = len(values) * (bw + gap) - gap + ov.line([u(x), u(y_mid), u(x + total), u(y_mid)], fill=(*MUTED, 80), width=u(1)) + for i, v in enumerate(values): + bx = x + i * (bw + gap) + h = (v / vmax) * amp + y0, y1 = sorted((y_mid, y_mid - h)) + a = 120 + int(135 * abs(v) / vmax) + ov.rectangle([u(bx), u(y0), u(bx + bw), u(y1)], fill=(*color, a)) + f = mono_font(17) + tracked(ov, (u(x), u(y_mid - amp - 36)), label, f, (*MUTED, 255), tracking=0.5) + + +def draw_left_panel(ov: ImageDraw.ImageDraw, answer, ctx, counts, heads): + x0, y0, x1, _ = LP + pad = 44 + ctx_before, ctx_after = ctx + + tracked(ov, (u(x0 + pad), u(y0 + 34)), "TEXT CARRIER", label_font(30), AMBER, tracking=5) + sub = f"{counts['text_tokens']:,} BPE TOKENS" + f_sub = label_font(21) + tracked(ov, (int(u(x1 - pad) - tracked_width(ov, sub, f_sub, 2)), u(y0 + 42)), sub, f_sub, MUTED, tracking=2) + ov.line([u(x0 + pad), u(y0 + 88), u(x1 - pad), u(y0 + 88)], fill=(*DIVIDER, 255), width=u(1.2)) + + f_ctx = mono_font(23) + ov.text((u(x0 + pad), u(y0 + 116)), ctx_before, font=f_ctx, fill=(110, 118, 126)) + + # The two answer-token pills. + f_tok = mono_font(58) + f_id = mono_font(20) + px = x0 + pad + py = y0 + 184 + for t in answer: + s = t["str"] + wpx = ov.textlength(s, font=f_tok) / MS + ov.rounded_rectangle( + [u(px), u(py), u(px + wpx + 40), u(py + 96)], + radius=u(14), + fill=(38, 29, 10, 235), + outline=(*AMBER, 165), + width=u(1.6), + ) + ov.text((u(px + 20), u(py + 14)), s, font=f_tok, fill=(255, 224, 150)) + ov.text((u(px + 20), u(py + 110)), f"id {t['id']}", font=f_id, fill=(196, 156, 72)) + px += wpx + 40 + 22 + + ov.text((u(x0 + pad), u(y0 + 330)), ctx_after, font=f_ctx, fill=(110, 118, 126)) + + th = heads["text"] + head_bars( + ov, + x0 + pad, + y0 + 452, + th["vector_head"], + AMBER, + f"embedding row {th['id']} · dims 0-9 of {counts['embed_dim']:,} · norm {th['norm']:.2f}", + ) + + fy = y0 + 500 + f_fact = body_font(23) + ov.text((u(x0 + pad), u(fy)), f"{counts['chars']:,} characters of one SQuAD passage,", font=f_fact, fill=MUTED) + ov.text( + (u(x0 + pad), u(fy + 36)), + f"tokenized into {counts['text_tokens']:,} ids, each a {counts['embed_dim']:,}-dim row", + font=f_fact, + fill=MUTED, + ) + + +def draw_right_panel(base_img: Image.Image, ov: ImageDraw.ImageDraw, carrier: Image.Image, word_idx, counts, heads): + x0, y0, x1, _ = RP + pad = 44 + + tracked(ov, (u(x0 + pad), u(y0 + 34)), "IMAGE CARRIER", label_font(30), CYAN, tracking=5) + sub = f"{counts['image_tokens']:,} PATCHES" + f_sub = label_font(21) + tracked(ov, (int(u(x1 - pad) - tracked_width(ov, sub, f_sub, 2)), u(y0 + 42)), sub, f_sub, MUTED, tracking=2) + ov.line([u(x0 + pad), u(y0 + 88), u(x1 - pad), u(y0 + 88)], fill=(*DIVIDER, 255), width=u(1.2)) + + # Crop: patch rows 4..8, cols 27..38 of the 56x56 grid (28px cells). + pp = counts["patch_px"] + c0, c1, r0, r1 = 27, 38, 4, 8 + crop = carrier.crop((c0 * pp, r0 * pp, c1 * pp, r1 * pp)) # 308 x 112 + + scale = 2.0 # 28px cell -> 56px on the 2400 canvas + disp_w, disp_h = int(crop.width * scale), int(crop.height * scale) + big = crop.resize((u(disp_w), u(disp_h)), Image.NEAREST) + big = Image.eval(big, lambda v: int(v * 0.84)) # dim so the highlight pops + bx, by = x0 + pad, y0 + 122 + base_img.paste(big, (u(bx), u(by))) + + cell = pp * scale # 56 in 2400-space + grid_color = (CYAN[0], CYAN[1], CYAN[2], 46) + for c in range(c1 - c0 + 1): + ov.line([u(bx + c * cell), u(by), u(bx + c * cell), u(by + disp_h)], fill=grid_color, width=u(1)) + for r in range(r1 - r0 + 1): + ov.line([u(bx), u(by + r * cell), u(bx + disp_w), u(by + r * cell)], fill=grid_color, width=u(1)) + + # Highlight the four answer patches (grid row 5, cols 30..33) as one run. + grid = counts["grid"] + gr, gc = word_idx[0] // grid, word_idx[0] % grid + hx, hy = bx + (gc - c0) * cell, by + (gr - r0) * cell + hw = len(word_idx) * cell + ov.rectangle([u(hx), u(hy), u(hx + hw), u(hy + cell)], outline=(*CYAN, 240), width=u(2.4)) + for k in range(1, len(word_idx)): + ov.line([u(hx + k * cell), u(hy), u(hx + k * cell), u(hy + cell)], fill=(*CYAN, 130), width=u(1.2)) + + ov.text( + (u(bx), u(by + disp_h + 18)), + f"patches {word_idx[0]}-{word_idx[-1]} · grid row 5, cols 30-33 of {grid}×{grid}", + font=body_font(23), + fill=MUTED, + ) + + ih = heads["image"] + head_bars( + ov, + bx, + y0 + 460, + ih["vector_head"], + CYAN, + f"patch {ih['token_index']} vector · dims 0-9 of {counts['visual_dim']:,} · norm {ih['norm']:.1f}", + ) + + fy = y0 + 500 + f_fact = body_font(23) + ov.text((u(bx), u(fy)), "the same passage, rendered to a 1568 × 1568 px bitmap,", font=f_fact, fill=MUTED) + ov.text( + (u(bx), u(fy + 36)), + f"seen as {counts['image_tokens']:,} patches of {counts['patch_px']} px, each a {counts['visual_dim']:,}-dim vector", + font=f_fact, + fill=MUTED, + ) + + +def draw_title(ov: ImageDraw.ImageDraw): + kicker = "INSIDE QWEN2.5-VL · ONE FACT, TWO ENCODINGS" + f_k = label_font(24) + kw = tracked_width(ov, kicker, f_k, 7) + tracked(ov, (int((W - kw) / 2), u(64)), kicker, f_k, MUTED, tracking=7) + + f_t = display_font(96) + a, b = "TWO CARRIERS, ", "ONE THOUGHT." + wa, wb = ov.textlength(a, font=f_t), ov.textlength(b, font=f_t) + x = (W - wa - wb) / 2 + y = u(122) + ov.text((x, y), a, font=f_t, fill=INK) + ov.text((x + wa, y), b, font=f_t, fill=AMBER) + + +def draw_core(ov: ImageDraw.ImageDraw, stats): + f = display_font(78) + tw = ov.textlength(CORE_WORD, font=f) + ov.text( + (u(CORE[0]) - tw / 2, u(CORE[1] - 54)), + CORE_WORD, + font=f, + fill=(255, 248, 226), + stroke_width=u(1.4), + stroke_fill=(64, 44, 12, 160), + ) + cap = f"BY LAYER {stats['layer']} OF {stats['n_layers'] - 1}, ONE SHARED STATE" + f_c = label_font(23) + cw = tracked_width(ov, cap, f_c, 4) + tracked(ov, (int(u(CORE[0]) - cw / 2), u(CORE[1] + 96)), cap, f_c, (228, 222, 196), tracking=4) + + +def draw_stats_strip(ov: ImageDraw.ImageDraw, stats): + y_div = 992 + ov.line([u(92), u(y_div), u(2308), u(y_div)], fill=(*DIVIDER, 255), width=u(1.4)) + + groups = [ + (f"{stats['matched']:.2f}", "MATCHED COSINE · TEXT VS IMAGE", AMBER), + (f"{stats['rsa']:.2f}", "RSA · SAME RELATIONAL GEOMETRY", INK), + (f"{stats['retrieved']}/{stats['n']}", "CROSS-CARRIER RETRIEVAL", CYAN), + ] + centers = [500, 1200, 1900] + f_num = display_font(66) + f_cap = label_font(21) + for (num, cap, color), cx in zip(groups, centers): + nw = ov.textlength(num, font=f_num) + ov.text((u(cx) - nw / 2, u(1024)), num, font=f_num, fill=color) + cw = tracked_width(ov, cap, f_cap, 3) + tracked(ov, (int(u(cx) - cw / 2), u(1136)), cap, f_cap, MUTED, tracking=3) + for dx in (850, 1550): + ov.line([u(dx), u(1040), u(dx), u(1170)], fill=(*DIVIDER, 255), width=u(1.2)) + + foot = f"measured at layer {stats['layer']} · {stats['n']} SQuAD questions · carrier-convergence" + f_f = body_font(19) + fw = ov.textlength(foot, font=f_f) + ov.text(((W - fw) / 2, u(1206)), foot, font=f_f, fill=(92, 101, 108)) + + +def rounded_panel(overlay: ImageDraw.ImageDraw, box, accent, alpha_fill=216): + x0, y0, x1, y1 = (u(v) for v in box) + r = u(22) + overlay.rounded_rectangle([x0, y0, x1, y1], radius=r, fill=(*PANEL, alpha_fill)) + overlay.rounded_rectangle([x0, y0, x1, y1], radius=r, outline=(*accent, 70), width=u(1.4)) + + +def main() -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + stats, answer, ctx, word_idx, counts, heads, carrier = load_data() + + base = Image.fromarray(np.clip(additive_base(), 0, 255).astype(np.uint8), "RGB") + + overlay = Image.new("RGBA", (W, H), (0, 0, 0, 0)) + ov = ImageDraw.Draw(overlay) + rounded_panel(ov, LP, AMBER) + rounded_panel(ov, RP, CYAN) + base = Image.alpha_composite(base.convert("RGBA"), overlay).convert("RGB") + + # Panel contents drawn on a fresh overlay so the bitmap paste sits beneath grids. + overlay2 = Image.new("RGBA", (W, H), (0, 0, 0, 0)) + ov2 = ImageDraw.Draw(overlay2) + draw_title(ov2) + draw_left_panel(ov2, answer, ctx, counts, heads) + draw_right_panel(base, ov2, carrier, word_idx, counts, heads) + draw_core(ov2, stats) + draw_stats_strip(ov2, stats) + final = Image.alpha_composite(base.convert("RGBA"), overlay2).convert("RGB") + + retina = final.resize((2400, 1260), Image.LANCZOS) + og = final.resize((1200, 630), Image.LANCZOS) + retina.save(OUT_DIR / "hero-2400x1260.png") + og.save(OUT_DIR / "hero-1200x630.png") + print(f"wrote {OUT_DIR / 'hero-2400x1260.png'} {retina.size}") + print(f"wrote {OUT_DIR / 'hero-1200x630.png'} {og.size}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_r2_metro.py b/packages/snapcompact/research/snapcompact_r2_metro.py new file mode 100755 index 000000000..d7da1477a --- /dev/null +++ b/packages/snapcompact/research/snapcompact_r2_metro.py @@ -0,0 +1,348 @@ +#!/usr/bin/env python3 +"""Snapcompact R2 — "The Convergence Line". + +A transit/metro-map diagram of carrier convergence in Qwen2.5-VL-7B. +Two metro lines (cyan = text carrier, orange = image carrier) run through +29 stations (decoder layers L0..L28). The vertical gap between the lines at +each station is driven by real per-layer data: + + gap ~ 1 - matched_cosine (results/qwen-carrier-convergence-n12/summary.json) + +Named stations are grounded in the same summary.json plus +results/qwen-logit-lens-q3/logit_lens.json (visual tok[310] -> 'acular'). + +Output: results/agent-r2-metro/metro.png (~2200 px wide). +""" + +import json +import os + +import numpy as np +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib.patches import Circle # noqa: E402 + +HERE = os.path.dirname(os.path.abspath(__file__)) +SUMMARY_PATH = os.path.join(HERE, "results", "qwen-carrier-convergence-n12", "summary.json") +LENS_PATH = os.path.join(HERE, "results", "qwen-logit-lens-q3", "logit_lens.json") +OUT_DIR = os.path.join(HERE, "results", "agent-r2-metro") +OUT_PNG = os.path.join(OUT_DIR, "metro.png") + +# ---------------------------------------------------------------- palette +BG = "#05070a" +PANEL = "#0c1117" +INK = "#f1efe0" +MUTED = "#8f9aa0" +AMBER = "#ffc444" +CYAN = "#4bdcff" +ORANGE = "#ff7048" +GRID = "#0e141b" +MONO = "DejaVu Sans Mono" +SANS = "DejaVu Sans" + +MINUS = "\u2212" + + +def load_data(): + with open(SUMMARY_PATH) as f: + summary = json.load(f) + with open(LENS_PATH) as f: + lens = json.load(f) + + per = summary["per_layer"] + assert len(per) == summary["layers"] == 29 + assert summary["best_layer"] == 19 + + cos = np.array([p["matched_cosine"] for p in per]) + rsa = np.array([p["rsa_pearson"] for p in per]) + acc = np.array([p["match_rank_accuracy"] for p in per]) + mism = np.array([p["mismatched_cosine"] for p in per]) + + # logit lens: visual token 310, answer piece 'acular' + tok310 = {e["layer"]: e for e in lens["lens"] if e["token_index"] == 310} + assert tok310[24]["top"][0]["str"] == "acular" + p_acular_24 = tok310[24]["answer_token_p"][1] + p_acular_28 = tok310[28]["answer_token_p"][1] + answer = "".join(lens["answer_token_strs"]) # 'spectacular' + + return { + "summary": summary, + "cos": cos, + "rsa": rsa, + "acc": acc, + "mism": mism, + "p24": p_acular_24, + "p28": p_acular_28, + "answer": answer, + "n_q": summary["n_questions"], + "geometry": summary["geometry"], + "size_px": summary["args"]["size"], + "text_em": summary["text_em"], + "image_em": summary["image_em"], + } + + +def catmull_rom(xs, ys, samples=26): + """Centripetal-ish Catmull-Rom through all points (uniform parameter).""" + pts = np.column_stack([xs, ys]).astype(float) + ext = np.vstack([pts[0], pts, pts[-1]]) + out = [] + t = np.linspace(0.0, 1.0, samples, endpoint=False)[:, None] + for i in range(len(pts) - 1): + p0, p1, p2, p3 = ext[i], ext[i + 1], ext[i + 2], ext[i + 3] + a = 2.0 * p1 + b = p2 - p0 + c = 2.0 * p0 - 5.0 * p1 + 4.0 * p2 - p3 + d = -p0 + 3.0 * p1 - 3.0 * p2 + p3 + out.append(0.5 * (a + b * t + c * t**2 + d * t**3)) + out.append(pts[-1][None]) + return np.vstack(out) + + +def fmt2(v): + s = f"{v:.2f}" + return s.replace("-", MINUS) + + +def main(): + d = load_data() + cos = d["cos"] + n_layers = len(cos) + + # ---- track geometry: gap shrinks as matched cosine rises ------------- + c_min, c_max = float(cos.min()), float(cos.max()) # 0.0 (L0) .. 0.658 (L19) + GAP_MAX, GAP_MIN = 5.6, 0.78 + t = (cos - c_min) / (c_max - c_min) + gap = GAP_MAX + (GAP_MIN - GAP_MAX) * t + xs = np.arange(n_layers, dtype=float) + y_text = gap / 2.0 + y_img = -gap / 2.0 + + # depot stubs before L0 + xs_t = np.concatenate([[-1.5], xs]) + xs_i = np.concatenate([[-1.5], xs]) + yt = np.concatenate([[y_text[0]], y_text]) + yi = np.concatenate([[y_img[0]], y_img]) + + path_t = catmull_rom(xs_t, yt) + path_i = catmull_rom(xs_i, yi) + + # ---- figure ---------------------------------------------------------- + X0, X1 = -2.6, 32.2 + Y0, Y1 = -7.1, 8.3 + W_IN = 22.0 + H_IN = W_IN * (Y1 - Y0) / (X1 - X0) # equal data aspect + fig = plt.figure(figsize=(W_IN, H_IN), dpi=100) + fig.patch.set_facecolor(BG) + ax = fig.add_axes([0, 0, 1, 1]) + ax.set_facecolor(BG) + ax.set_xlim(X0, X1) + ax.set_ylim(Y0, Y1) + ax.set_aspect("equal", adjustable="box") + ax.axis("off") + + pt_per_unit = (W_IN / (X1 - X0)) * 72.0 # ~45.5 pt per data unit + + # faint vertical guides at every station + for i in range(n_layers): + ax.plot([i, i], [-3.55, 3.95], color=GRID, lw=1.0, zorder=1) + + # convergence axis + ax.plot([-1.5, 28.0], [0, 0], color=MUTED, lw=1.0, alpha=0.28, + linestyle=(0, (1, 3)), zorder=1) + ax.text(8.0, 0.16, "convergence axis", color=MUTED, alpha=0.55, + fontsize=8.5, style="italic", family=SANS, ha="left", zorder=2) + + # ---- metro lines: glow, casing, stroke ------------------------------- + for path, col in ((path_t, CYAN), (path_i, ORANGE)): + px, py = path[:, 0], path[:, 1] + ax.plot(px, py, color=col, lw=24, alpha=0.05, solid_capstyle="round", zorder=2) + ax.plot(px, py, color=col, lw=16, alpha=0.07, solid_capstyle="round", zorder=2) + ax.plot(px, py, color=BG, lw=14, solid_capstyle="round", + solid_joinstyle="round", zorder=3) + ax.plot(px, py, color=col, lw=9.5, solid_capstyle="round", + solid_joinstyle="round", zorder=4) + + # ---- stations --------------------------------------------------------- + named = {19, 24, 27, 28} + for i in range(n_layers): + for y, col in ((y_text[i], CYAN), (y_img[i], ORANGE)): + if i in named: + continue + ax.scatter([i], [y], s=115, facecolor=BG, edgecolor=col, + linewidths=2.1, zorder=6) + + # L19 interchange capsule (the two lines meet in one station) + cap_lw_outer = 0.56 * pt_per_unit + ax.plot([19, 19], [y_img[19], y_text[19]], color=INK, + lw=cap_lw_outer, solid_capstyle="round", zorder=5) + ax.plot([19, 19], [y_img[19], y_text[19]], color=PANEL, + lw=cap_lw_outer - 7.5, solid_capstyle="round", zorder=5) + ax.scatter([19, 19], [y_text[19], y_img[19]], s=92, + c=[CYAN, ORANGE], edgecolor=BG, linewidths=1.2, zorder=6) + + # L24 interchange ring on the image line (pixels decode to vocabulary) + ax.scatter([24], [y_img[24]], s=300, facecolor=PANEL, edgecolor=INK, + linewidths=2.8, zorder=6) + ax.scatter([24], [y_img[24]], s=58, facecolor=ORANGE, edgecolor="none", zorder=6) + ax.scatter([24], [y_text[24]], s=115, facecolor=BG, edgecolor=CYAN, + linewidths=2.1, zorder=6) + + # L27 white-ring stations on both lines (terminal approach) + for y, col in ((y_text[27], CYAN), (y_img[27], ORANGE)): + ax.scatter([27], [y], s=170, facecolor=PANEL, edgecolor=INK, + linewidths=2.3, zorder=6) + ax.scatter([27], [y], s=34, facecolor=col, edgecolor="none", zorder=6) + + # L28 terminus: double ring over both tracks + ax.add_patch(Circle((28, 0), 1.02, facecolor=PANEL, edgecolor=INK, + lw=3.2, zorder=5)) + ax.add_patch(Circle((28, 0), 0.66, facecolor="none", edgecolor=INK, + lw=1.3, alpha=0.85, zorder=5)) + ax.scatter([27.78, 28.22], [0, 0], s=120, c=[CYAN, ORANGE], + edgecolor=BG, linewidths=1.4, zorder=6) + ax.text(28, -0.42, "TERMINUS", color=MUTED, fontsize=6.8, family=MONO, + ha="center", va="center", zorder=7) + + # ---- carrier labels (depots) ------------------------------------------ + geo = d["geometry"] + ax.text(-1.55, y_text[0] + 0.95, "TEXT CARRIER", color=CYAN, fontsize=12.5, + family=SANS, fontweight="bold", ha="left", zorder=7) + ax.text(-1.55, y_text[0] + 0.48, + f"the page as typed tokens \u00b7 {geo['capacity']:,} chars", + color=MUTED, fontsize=9, family=SANS, ha="left", zorder=7) + ax.text(-1.55, y_img[0] - 0.62, "IMAGE CARRIER", color=ORANGE, fontsize=12.5, + family=SANS, fontweight="bold", ha="left", zorder=7) + ax.text(-1.55, y_img[0] - 1.09, + f"the same page as a {d['size_px']} px bitmap \u00b7 " + f"{geo['cols']}\u00d7{geo['rows']} cell grid", + color=MUTED, fontsize=9, family=SANS, ha="left", zorder=7) + + # ---- named-station callouts ------------------------------------------- + def leader(x, y_from, y_to, color=MUTED, alpha=0.65): + ax.plot([x, x], [y_from, y_to], color=color, lw=1.1, alpha=alpha, zorder=6) + + # L1: instant alignment + leader(1, y_img[1] - 0.18, -2.18) + ax.text(1.7, -2.35, "L1 \u00b7 INSTANT ALIGNMENT", color=INK, fontsize=10.5, + family=SANS, fontweight="bold", ha="left", zorder=7) + ax.text(1.7, -2.78, + f"matched cos {fmt2(cos[1])} \u00b7 RSA {fmt2(d['rsa'][1])}", + color=MUTED, fontsize=8.8, family=MONO, ha="left", zorder=7) + ax.text(1.7, -3.14, + f"retrieval {int(round(d['acc'][1] * 12))}/12 \u2014 12/12 from L2 onward", + color=MUTED, fontsize=8.8, family=MONO, ha="left", zorder=7) + + # L13: first close pass + leader(13, y_text[13] + 0.18, 1.62) + ax.text(13, 1.84, f"L13 \u00b7 first close pass \u00b7 cos {fmt2(cos[13])}", + color=MUTED, fontsize=8.8, family=MONO, ha="center", zorder=7) + + # L19: geometry locks (star station) + leader(19, y_text[19] + 0.62, 2.42, color=AMBER, alpha=0.8) + ax.text(19, 3.42, "L19 \u00b7 GEOMETRY LOCKS", color=AMBER, fontsize=14, + family=SANS, fontweight="bold", ha="center", zorder=7) + ax.text(19, 2.96, + f"matched cos {fmt2(cos[19])} \u00b7 mismatched {fmt2(d['mism'][19])}", + color=INK, fontsize=9.6, family=MONO, ha="center", zorder=7) + ax.text(19, 2.58, + f"RSA {fmt2(d['rsa'][19])} \u00b7 retrieval 12/12 \u2014 closest approach", + color=MUTED, fontsize=9.6, family=MONO, ha="center", zorder=7) + + # L23: small drift + leader(23, y_text[23] + 0.18, 1.30) + ax.text(23, 1.52, f"L23 \u00b7 small drift \u00b7 cos {fmt2(cos[23])}", + color=MUTED, fontsize=8.8, family=MONO, ha="center", zorder=7) + + # L24: pixels decode to vocabulary + leader(24, y_img[24] - 0.32, -1.92, color=ORANGE, alpha=0.8) + ax.text(24, -2.18, "L24 \u00b7 PIXELS DECODE TO VOCABULARY", color=ORANGE, + fontsize=12.5, family=SANS, fontweight="bold", ha="center", zorder=7) + ax.text(24, -2.62, + f"visual tok[310] top-1 \u2192 'acular' \u00b7 p {d['p24']:.2f}", + color=INK, fontsize=9.4, family=MONO, ha="center", zorder=7) + ax.text(24, -3.00, + f"rising to p {d['p28']:.2f} by L28 \u2014 " + "the answer's second BPE piece", + color=MUTED, fontsize=9.4, family=MONO, ha="center", zorder=7) + + # L27-L28 terminal (block above the terminus circle) + leader(28, 1.18, 1.86, color=AMBER, alpha=0.8) + ax.text(28, 3.00, "L27\u2013L28 \u00b7 TERMINAL", color=INK, fontsize=12.5, + family=SANS, fontweight="bold", ha="center", zorder=7) + ax.text(28, 2.56, f"SAME ANSWER: \u201c{d['answer']}\u201d", color=AMBER, + fontsize=10.5, family=SANS, fontweight="bold", ha="center", zorder=7) + ax.text(28, 2.18, + f"matched cos {fmt2(cos[27])} \u2192 {fmt2(cos[28])} \u00b7 retrieval 12/12", + color=MUTED, fontsize=8.8, family=MONO, ha="center", zorder=7) + + # ---- station index + matched-cosine gauge rows ------------------------- + hl = {19: AMBER, 24: ORANGE, 27: INK, 28: INK} + ax.text(-0.55, -4.45, "layer", color=MUTED, fontsize=8, style="italic", + family=SANS, ha="right", va="center", zorder=7) + ax.text(-0.55, -5.02, "matched cos", color=MUTED, fontsize=8, style="italic", + family=SANS, ha="right", va="center", zorder=7) + for i in range(n_layers): + col = hl.get(i, MUTED) + w = "bold" if i in hl else "normal" + ax.text(i, -4.45, f"L{i}", color=col, fontsize=7.6, family=MONO, + ha="center", va="center", fontweight=w, zorder=7) + val = f"{cos[i]:.2f}".lstrip("0") + ax.text(i, -5.02, val, color=col, fontsize=7.6, family=MONO, + ha="center", va="center", fontweight=w, zorder=7) + + # ---- title ------------------------------------------------------------- + ax.text(-1.9, 8.05, "THE CONVERGENCE LINE", color=INK, fontsize=29, + family=SANS, fontweight="bold", ha="left", va="top", zorder=7) + ax.text(-1.9, 6.92, + "One Wikipedia page, two carriers: typed tokens (cyan) and a " + f"{d['size_px']} px screenshot (orange) ride Qwen2.5-VL-7B's 29 decoder layers.", + color=MUTED, fontsize=12, family=SANS, ha="left", va="top", zorder=7) + ax.text(-1.9, 6.42, + "The closer the tracks, the more the two internal representations agree " + f"\u2014 track gap \u221d 1 {MINUS} matched cosine, n = {d['n_q']} questions.", + color=MUTED, fontsize=12, family=SANS, ha="left", va="top", zorder=7) + + # ---- legend (top right) ------------------------------------------------- + lx = 22.9 + ax.plot([lx, lx + 1.3], [7.95, 7.95], color=CYAN, lw=8, + solid_capstyle="round", zorder=7) + ax.text(lx + 1.65, 7.95, "TEXT CARRIER", color=INK, fontsize=10, + family=SANS, fontweight="bold", ha="left", va="center", zorder=7) + ax.plot([lx, lx + 1.3], [7.32, 7.32], color=ORANGE, lw=8, + solid_capstyle="round", zorder=7) + ax.text(lx + 1.65, 7.32, "IMAGE CARRIER", color=INK, fontsize=10, + family=SANS, fontweight="bold", ha="left", va="center", zorder=7) + ax.text(lx, 6.62, f"track gap \u221d 1 {MINUS} matched cosine(text, image)", + color=MUTED, fontsize=9, family=MONO, ha="left", va="center", zorder=7) + # wide pair = L0 + ax.plot([lx, lx + 1.0], [6.18, 6.18], color=CYAN, lw=4, solid_capstyle="round", zorder=7) + ax.plot([lx, lx + 1.0], [5.74, 5.74], color=ORANGE, lw=4, solid_capstyle="round", zorder=7) + ax.text(lx + 1.65, 5.96, f"cos {fmt2(cos[0])} \u2014 far apart (L0)", + color=MUTED, fontsize=9, family=MONO, ha="left", va="center", zorder=7) + # tight pair = L19 + ax.plot([lx, lx + 1.0], [5.28, 5.28], color=CYAN, lw=4, solid_capstyle="round", zorder=7) + ax.plot([lx, lx + 1.0], [5.14, 5.14], color=ORANGE, lw=4, solid_capstyle="round", zorder=7) + ax.text(lx + 1.65, 5.21, f"cos {fmt2(cos[19])} \u2014 almost touching (L19)", + color=MUTED, fontsize=9, family=MONO, ha="left", va="center", zorder=7) + + # ---- footer -------------------------------------------------------------- + ax.text(-1.9, -6.05, + "Across the same 12 questions the image carrier matches gold answers as often as the text carrier " + f"\u2014 image EM {d['image_em'] * 100:.1f}% vs text EM {d['text_em'] * 100:.1f}%.", + color=MUTED, fontsize=9.5, family=SANS, ha="left", zorder=7) + ax.text(-1.9, -6.58, + "Data: results/qwen-carrier-convergence-n12/summary.json (29 layers \u00b7 12 SQuAD questions) " + "+ results/qwen-logit-lens-q3/logit_lens.json \u00b7 Qwen2.5-VL-7B-Instruct \u00b7 agent r2-metro", + color=MUTED, alpha=0.7, fontsize=8.5, family=MONO, ha="left", zorder=7) + + os.makedirs(OUT_DIR, exist_ok=True) + fig.savefig(OUT_PNG, dpi=100, facecolor=BG) + plt.close(fig) + print(f"wrote {OUT_PNG}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_tensor_heatmap.py b/packages/snapcompact/research/snapcompact_tensor_heatmap.py new file mode 100644 index 000000000..0cc441e30 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_tensor_heatmap.py @@ -0,0 +1,354 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers", "sentencepiece", "protobuf", "einops"] +# /// +"""Capture and render token/layer hidden-state heatmaps for snapcompact masks. + +This is the blog-visual version of the white-box probe: it compares the same +prompt with the original bitmap, a gold-answer-region mask, and an equal random +mask. For every decoder layer and every image placeholder token, it plots +||hidden(original) - hidden(masked)||. +""" + +from __future__ import annotations + +import argparse +import json +import math +import random +import sys +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import mask_cells, random_span, sample_answer_questions # noqa: E402 + +DEFAULT_MODEL_DIR = ( + "/home/can/.cache/huggingface/hub/models--PaddlePaddle--PaddleOCR-VL/" + "snapshots/2b77538ef936207f60c16b45082841068987d08c" +) + +PALETTE = { + "bg": (5, 7, 10), + "panel": (13, 18, 23), + "ink": (239, 239, 224), + "muted": (132, 147, 154), + "cyan": (77, 218, 255), + "red": (255, 83, 62), + "green": (145, 255, 112), + "amber": (255, 194, 65), + "grid": (34, 45, 53), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if path and Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def heat_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + stops = [ + (0.00, (6, 8, 18)), + (0.20, (28, 20, 70)), + (0.43, (118, 29, 97)), + (0.67, (222, 72, 69)), + (0.85, (255, 164, 75)), + (1.00, (255, 243, 164)), + ] + for (a, ca), (b, cb) in zip(stops, stops[1:]): + if t <= b: + u = (t - a) / (b - a) + return tuple(round(ca[i] + (cb[i] - ca[i]) * u) for i in range(3)) + return stops[-1][1] + + +def downsample_cols(arr: np.ndarray, bins: int) -> np.ndarray: + if arr.shape[1] <= bins: + return arr + edges = np.linspace(0, arr.shape[1], bins + 1).round().astype(int) + out = np.zeros((arr.shape[0], bins), dtype=np.float32) + for i in range(bins): + lo, hi = edges[i], max(edges[i] + 1, edges[i + 1]) + out[:, i] = arr[:, lo:hi].mean(axis=1) + return out + + +def normalize(arr: np.ndarray, scale: float | None = None) -> tuple[np.ndarray, float]: + if scale is None: + scale = float(np.quantile(arr, 0.98)) if arr.size else 1.0 + if scale <= 0: + scale = 1.0 + return np.clip(arr / scale, 0, 1), scale + + +def draw_heatmap(draw: ImageDraw.ImageDraw, arr: np.ndarray, box: tuple[int, int, int, int], title: str, subtitle: str, color: tuple[int, int, int]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=22, fill=PALETTE["panel"], outline=(31, 42, 50), width=1) + draw.text((x0 + 24, y0 + 18), title, fill=color, font=ui_font(26, True)) + draw.text((x0 + 24, y0 + 50), subtitle, fill=PALETTE["muted"], font=ui_font(15)) + hx0, hy0, hx1, hy1 = x0 + 58, y0 + 84, x1 - 28, y1 - 44 + rows, cols = arr.shape + cw = (hx1 - hx0) / cols + ch = (hy1 - hy0) / rows + for r in range(rows): + y_a = round(hy0 + r * ch) + y_b = round(hy0 + (r + 1) * ch) + for c in range(cols): + x_a = round(hx0 + c * cw) + x_b = round(hx0 + (c + 1) * cw) + draw.rectangle((x_a, y_a, x_b, y_b), fill=heat_color(float(arr[r, c]))) + for r in range(0, rows, 4): + y = round(hy0 + (r + 0.5) * ch) + draw.text((x0 + 18, y - 8), str(r), fill=PALETTE["muted"], font=ui_font(12)) + draw.text((x0 + 16, hy0 - 4), "layer", fill=PALETTE["muted"], font=ui_font(12)) + draw.text((hx0, y1 - 31), "image token sequence →", fill=PALETTE["muted"], font=ui_font(13)) + + +def crop_with_box(img: Image.Image, start: int, end: int, cols: int, adv: int, pitch: int, pad_cells: int = 34) -> Image.Image: + row0 = max(0, start // cols - 5) + row1 = min(img.height // pitch, end // cols + 6) + col0 = max(0, start % cols - pad_cells) + col1 = min(cols, end % cols + pad_cells) + if col1 <= col0: + col1 = min(cols, col0 + 72) + x0, y0, x1, y1 = col0 * adv, row0 * pitch, col1 * adv, row1 * pitch + crop = img.crop((x0, y0, x1, y1)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=PALETTE["red"], width=3) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.NEAREST) + canvas.paste(resized, (x0 + (x1 - x0 - resized.width) // 2, y0 + (y1 - y0 - resized.height) // 2)) + + +def make_prompt(q: str, cols: int, rows: int) -> str: + return load_prompt("qa-image.md").format(cols=cols, rows=rows) + f"\n\nQuestion: {q}\nAnswer with only the shortest extractive answer." + + +def to_device(batch: dict[str, Any], device: Any) -> dict[str, Any]: + return {k: (v.to(device) if hasattr(v, "to") else v) for k, v in batch.items()} + + +def hidden_token_matrix(model: Any, processor: Any, image: Image.Image, prompt_text: str, device: Any) -> tuple[list[np.ndarray], list[int], dict[str, Any]]: + import torch + + messages = [{"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": prompt_text}]}] + templated = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + batch = processor(images=image, text=templated, return_tensors="pt") + image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token) + ids = batch["input_ids"][0].tolist() + image_positions = [i for i, token_id in enumerate(ids) if token_id == image_token_id] + meta = {k: (v.tolist() if hasattr(v, "tolist") else v) for k, v in batch.items() if k in ("image_grid_thw",)} + batch = to_device(batch, device) + with torch.no_grad(): + out = model(**batch, output_hidden_states=True, output_attentions=False, use_cache=False) + matrices: list[np.ndarray] = [] + for hidden in out.hidden_states: + token_hidden = hidden[0, image_positions, :].float().detach().cpu().numpy() + matrices.append(token_hidden.astype(np.float32, copy=False)) + return matrices, image_positions, meta + + +def render_tensor_card( + out_path: Path, + answer_heat: np.ndarray, + random_heat: np.ndarray, + ratio_heat: np.ndarray, + base_img: Image.Image, + answer_img: Image.Image, + record: dict[str, Any], + cols: int, + adv: int, + pitch: int, + summary: dict[str, Any], +) -> None: + w, h = 1900, 1180 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 14): + draw.line((0, y, w, y), fill=(8, 11 + (y % 9), 15 + (y % 13))) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-260, -180, 850, 640), fill=(255, 83, 62, 30)) + gd.ellipse((1080, 110, 2240, 1320), fill=(77, 218, 255, 30)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(80))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((58, 38), "SNAPCOMPACT WHITEBOX", fill=PALETTE["amber"], font=ui_font(22, True)) + draw.text((58, 76), "The hidden-state scar of a missing answer", fill=PALETTE["ink"], font=ui_font(58, True)) + draw.text((60, 148), "Each pixel below is a decoder layer × image-token bin. Bright = larger ||hidden(original) − hidden(masked)||.", fill=PALETTE["muted"], font=ui_font(24)) + + # Left evidence panel. + draw.rounded_rectangle((58, 205, 700, 1098), radius=28, fill=PALETTE["panel"], outline=(31, 42, 50), width=1) + draw.text((90, 236), "the visual intervention", fill=PALETTE["ink"], font=ui_font(30, True)) + draw.text((90, 274), "same prompt, same bitmap; only answer cells blanked", fill=PALETTE["muted"], font=ui_font(17)) + crop = crop_with_box(base_img, record["answer_start"], record["answer_end"], cols, adv, pitch) + masked_crop = crop_with_box(answer_img, record["answer_start"], record["answer_end"], cols, adv, pitch) + draw.text((90, 326), "ORIGINAL", fill=PALETTE["cyan"], font=ui_font(16, True)) + draw.rounded_rectangle((90, 352, 668, 528), radius=14, fill=(244, 242, 230), outline=PALETTE["cyan"], width=3) + paste_fit(canvas, crop, (108, 368, 650, 512)) + draw.text((90, 568), "ANSWER ERASED", fill=PALETTE["red"], font=ui_font(16, True)) + draw.rounded_rectangle((90, 594, 668, 770), radius=14, fill=(244, 242, 230), outline=PALETTE["red"], width=3) + paste_fit(canvas, masked_crop, (108, 610, 650, 754)) + question = record["q"] + if len(question) > 72: + question = question[:69] + "…" + draw.text((90, 828), "question", fill=PALETTE["muted"], font=ui_font(16, True)) + draw.text((90, 856), question, fill=PALETTE["ink"], font=ui_font(21)) + draw.text((90, 914), "gold answer", fill=PALETTE["muted"], font=ui_font(16, True)) + draw.text((90, 942), str(record["answer_text"]), fill=PALETTE["amber"], font=ui_font(32, True)) + draw.text((90, 1014), f"{summary['layers']} hidden layers × {summary['image_tokens']} image tokens", fill=PALETTE["muted"], font=ui_font(18)) + + draw_heatmap(draw, answer_heat, (742, 205, 1818, 488), "gold answer mask", "activation delta when the true answer is blanked", PALETTE["red"]) + draw_heatmap(draw, random_heat, (742, 520, 1818, 803), "random equal-size mask", "control: blank the same number of glyph cells elsewhere", PALETTE["green"]) + draw_heatmap(draw, ratio_heat, (742, 835, 1818, 1098), "answer / random ratio", "bright bands mark layers/tokens more sensitive to the answer region", PALETTE["amber"]) + + # Color scale. + for i in range(220): + draw.rectangle((1588 + i, 158, 1589 + i, 174), fill=heat_color(i / 219)) + draw.text((1588, 133), "low", fill=PALETTE["muted"], font=ui_font(13)) + draw.text((1758, 133), "high", fill=PALETTE["muted"], font=ui_font(13)) + + out_path.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out_path) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default=DEFAULT_MODEL_DIR) + ap.add_argument("--font", default="8x13", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=768) + ap.add_argument("--limit-paras", type=int, default=40) + ap.add_argument("--qpc", type=int, default=16) + ap.add_argument("--question-index", type=int, default=7) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--bins", type=int, default=180) + ap.add_argument("--out", default="tensor-heatmap-paddleocr") + args = ap.parse_args() + + import torch + from transformers import AutoModel, AutoProcessor + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + start, end = 0, min(len(flow), budget) + chunk = flow[start:end] + questions = sample_answer_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + raise SystemExit("no sampled questions fit in chunk") + q = questions[min(args.question_index, len(questions) - 1)] + + base_img = render(chunk, cfg, CACHE, args.size, args.variant) + fill = (255, 255, 255) if args.variant not in ("dark", "dark-sent") else (0, 0, 0) + span_len = max(1, q["answer_end"] - q["answer_start"]) + rng = random.Random(args.seed * 101 + args.question_index) + rand_start, rand_end = random_span(rng, len(chunk), span_len, q["answer_start"], q["answer_end"]) + answer_img = mask_cells(base_img, q["answer_start"], q["answer_end"], cols, cfg.adv, cfg.pitch, fill) + random_img = mask_cells(base_img, rand_start, rand_end, cols, cfg.adv, cfg.pitch, fill) + base_img.save(img_dir / "original.png") + answer_img.save(img_dir / "answer-mask.png") + random_img.save(img_dir / "random-mask.png") + + print(f"loading {args.model_dir}", flush=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 + model = AutoModel.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=dtype).to(device).eval() + + prompt = make_prompt(q["q"], cols, rows) + original, positions, meta = hidden_token_matrix(model, processor, base_img, prompt, device) + answer, answer_positions, _ = hidden_token_matrix(model, processor, answer_img, prompt, device) + random_mask, random_positions, _ = hidden_token_matrix(model, processor, random_img, prompt, device) + if positions != answer_positions or positions != random_positions: + raise SystemExit("image token positions changed across variants") + + answer_delta = np.stack([np.linalg.norm(a - b, axis=1) for a, b in zip(original, answer)], axis=0) + random_delta = np.stack([np.linalg.norm(a - b, axis=1) for a, b in zip(original, random_mask)], axis=0) + ratio = answer_delta / np.maximum(random_delta, 1e-6) + + answer_binned = downsample_cols(answer_delta, args.bins) + random_binned = downsample_cols(random_delta, args.bins) + ratio_binned = downsample_cols(ratio, args.bins) + common_scale = float(np.quantile(np.concatenate([answer_binned.ravel(), random_binned.ravel()]), 0.98)) + answer_norm, _ = normalize(answer_binned, common_scale) + random_norm, _ = normalize(random_binned, common_scale) + ratio_norm, ratio_scale = normalize(ratio_binned, float(np.quantile(ratio_binned, 0.98))) + + record = { + "q": q["q"], + "golds": q["golds"], + "answer_text": q["answer_text"], + "answer_start": q["answer_start"], + "answer_end": q["answer_end"], + "random_start": rand_start, + "random_end": rand_end, + } + summary = { + "args": vars(args), + "device": str(device), + "geometry": {"cols": cols, "rows": rows, "capacity": budget}, + "question": record, + "image_tokens": len(positions), + "layers": len(original), + "image_token_positions": {"first": positions[0], "last": positions[-1]}, + "processor_meta": meta, + "answer_delta_mean": float(answer_delta.mean()), + "random_delta_mean": float(random_delta.mean()), + "answer_over_random_delta": float(answer_delta.mean() / max(random_delta.mean(), 1e-6)), + "common_delta_scale_p98": common_scale, + "ratio_scale_p98": ratio_scale, + "max_ratio_layer": int(np.argmax(ratio.mean(axis=1))), + "mean_ratio_by_layer": [float(x) for x in ratio.mean(axis=1)], + } + + np.savez_compressed( + out_dir / "heatmaps.npz", + answer_delta=answer_delta, + random_delta=random_delta, + ratio=ratio, + answer_binned=answer_binned, + random_binned=random_binned, + ratio_binned=ratio_binned, + answer_norm=answer_norm, + random_norm=random_norm, + ratio_norm=ratio_norm, + ) + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + render_tensor_card(out_dir / "tensor-heatmap.png", answer_norm, random_norm, ratio_norm, base_img, answer_img, record, cols, cfg.adv, cfg.pitch, summary) + print(json.dumps(summary, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_text_image_3d_viz.py b/packages/snapcompact/research/snapcompact_text_image_3d_viz.py new file mode 100644 index 000000000..12723a4f4 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_text_image_3d_viz.py @@ -0,0 +1,211 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["matplotlib", "numpy", "pillow"] +# /// +"""Render a 3D text-vs-image activation comparison from paired carrier data.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +from matplotlib import cm +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +BG = (5, 7, 10) +PANEL = (12, 17, 23) +INK = (241, 239, 224) +MUTED = (143, 154, 160) +CYAN = (75, 220, 255) +ORANGE = (255, 112, 72) +AMBER = (255, 196, 68) +GREEN = (148, 255, 117) + + +def font(size: int, bold: bool = False): + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono(size: int): + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def downsample(arr: np.ndarray, cols: int) -> np.ndarray: + if arr.shape[1] <= cols: + return arr + edges = np.linspace(0, arr.shape[1], cols + 1).round().astype(int) + out = np.zeros((arr.shape[0], cols), dtype=np.float32) + for i in range(cols): + lo = edges[i] + hi = max(lo + 1, edges[i + 1]) + out[:, i] = arr[:, lo:hi].mean(axis=1) + return out + + +def normalize(arr: np.ndarray) -> np.ndarray: + lo = float(np.quantile(arr, 0.03)) + hi = float(np.quantile(arr, 0.985)) + if hi <= lo: + hi = lo + 1e-6 + return np.clip((arr - lo) / (hi - lo), 0, 1) + + +def render_surface(z: np.ndarray, answer_bins: list[int]) -> Image.Image: + fig = plt.figure(figsize=(14.5, 8.2), dpi=180) + fig.patch.set_facecolor("#05070a") + ax = fig.add_subplot(111, projection="3d") + ax.set_facecolor((0.02, 0.025, 0.035, 1)) + y = np.arange(z.shape[0]) + x = np.arange(z.shape[1]) + X, Y = np.meshgrid(x, y) + cmap = plt.colormaps["turbo"] + ax.plot_surface(X, Y, z, facecolors=cmap(z), linewidth=0, antialiased=True, shade=False, alpha=0.98) + ax.contour(X, Y, z, zdir="z", offset=-0.08, levels=12, cmap=cmap, linewidths=0.9, alpha=0.75) + for b in answer_bins: + if 0 <= b < z.shape[1]: + ax.plot([b, b], [0, z.shape[0] - 1], [1.08, 1.08], color="#ff7048", linewidth=2.6, alpha=0.78) + ax.plot([b, b], [0, z.shape[0] - 1], [-0.06, -0.06], color="#ff7048", linewidth=1.6, alpha=0.55) + ax.view_init(elev=32, azim=-58) + ax.set_box_aspect((3.2, 0.8, 0.72)) + ax.set_zlim(-0.08, 1.08) + ax.set_ylim(z.shape[0] - 1, 0) + ax.set_xlabel("image-token bins", color="#8f9aa0", labelpad=10) + ax.set_ylabel("decoder layer", color="#8f9aa0", labelpad=10) + ax.set_zlabel("excess cosine", color="#8f9aa0", labelpad=8) + ax.tick_params(colors="#8f9aa0", labelsize=8) + for axis in (ax.xaxis, ax.yaxis, ax.zaxis): + axis.pane.set_facecolor((0.02, 0.025, 0.035, 0.0)) + axis._axinfo["grid"]["color"] = (0.35, 0.45, 0.50, 0.18) + ax.set_title("text-answer vector ↔ image-token field", color="#efeede", fontsize=24, fontweight="bold", loc="left", pad=18) + tmp = HERE / "results" / ".text-image-3d-panel.png" + fig.savefig(tmp, facecolor=fig.get_facecolor(), transparent=False) + plt.close(fig) + img = Image.open(tmp).convert("RGB") + tmp.unlink(missing_ok=True) + return img + + +def crop_answer(img: Image.Image, q: dict, cols: int, adv: int = 8, pitch: int = 13) -> Image.Image: + start = q["answer_start"] + end = q["answer_end"] + row0 = max(0, start // cols - 5) + row1 = min(img.height // pitch, end // cols + 6) + col0 = max(0, start % cols - 34) + col1 = min(cols, end % cols + 34) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=ORANGE, width=3) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.NEAREST) + canvas.paste(resized, (x0 + (x1 - x0 - resized.width) // 2, y0 + (y1 - y0 - resized.height) // 2)) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "text-image-compare-paddleocr-q7")) + ap.add_argument("--out", default=str(HERE / "results" / "text-image-compare-paddleocr-q7" / "text-vs-image-3d.png")) + ap.add_argument("--bins", type=int, default=150) + args = ap.parse_args() + result_dir = Path(args.result_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "text_image_compare.npz") + raw = data["text_answer_to_image_excess"] if "text_answer_to_image_excess" in data else data["text_answer_to_image_cosine"] + z = normalize(downsample(raw, args.bins)) + token_count = summary["image_tokens"] + answer_bins = sorted({round(idx / max(1, token_count - 1) * (args.bins - 1)) for idx in summary["image_answer_token_indices"]}) + panel = render_surface(z, answer_bins) + + w, h = 2200, 1320 + canvas = Image.new("RGB", (w, h), BG) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-260, -220, 860, 680), fill=(75, 220, 255, 30)) + gd.ellipse((1160, 80, 2440, 1320), fill=(255, 112, 72, 28)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(84))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + q = summary["question"] + draw.text((64, 42), "TEXT ↔ IMAGE WHITEBOX", fill=AMBER, font=font(24, True)) + draw.text((64, 84), "Same input, different carrier, shared hidden space", fill=INK, font=font(61, True)) + draw.text((66, 166), "For every decoder layer, compare the raw-text answer state against all bitmap image-token states. Peaks = image regions whose hidden state becomes text-like.", fill=MUTED, font=font(24)) + + draw.rounded_rectangle((64, 238, 618, 1234), radius=30, fill=PANEL, outline=(35, 49, 59), width=1) + draw.text((96, 270), "two carriers", fill=INK, font=font(34, True)) + draw.text((96, 312), "same chunk + same question", fill=MUTED, font=font(18)) + draw.text((96, 366), "RAW TEXT", fill=CYAN, font=font(18, True)) + y = 402 + draw.text((96, y), "Question:", fill=MUTED, font=font(17, True)) + y += 30 + for line in [q["q"][i : i + 46] for i in range(0, len(q["q"]), 46)]: + draw.text((96, y), line, fill=INK, font=font(19)) + y += 26 + y += 24 + draw.text((96, y), "Gold answer token span:", fill=MUTED, font=font(17, True)) + y += 32 + answer_text = str(q["answer_text"]) + draw.rounded_rectangle((96, y, 108 + max(72, len(answer_text) * 24), y + 40), radius=7, fill=AMBER) + draw.text((108, y + 7), answer_text, fill=(5, 7, 10), font=mono(22)) + y += 66 + draw.text((96, y), "The raw-text run receives the same", fill=INK, font=font(18)) + draw.text((96, y + 28), "SQuAD passage as ordinary tokens;", fill=INK, font=font(18)) + draw.text((96, y + 56), "the image run receives the passage", fill=INK, font=font(18)) + draw.text((96, y + 84), "only through the bitmap carrier.", fill=INK, font=font(18)) + draw.text((96, 674), f"text reference: {summary['text_reference_tokens']} tokens", fill=MUTED, font=font(18)) + draw.text((96, 704), f"answer span: {summary['text_answer_tokens']} text tokens", fill=MUTED, font=font(18)) + + draw.text((96, 774), "SNAPCOMPACT IMAGE", fill=ORANGE, font=font(18, True)) + img = Image.open(result_dir / "images" / "image-carrier.png").convert("RGB") + crop = crop_answer(img, q, summary["geometry"]["cols"]) + draw.rounded_rectangle((96, 812, 586, 1052), radius=16, fill=(244, 242, 230), outline=ORANGE, width=3) + paste_fit(canvas, crop, (112, 828, 570, 1036)) + draw.text((96, 1092), f"image field: {summary['image_tokens']} tokens ({summary['image_grid']}×{summary['image_grid']})", fill=MUTED, font=font(18)) + draw.text((96, 1138), f"peak alignment: {summary['answer_region_cosine_max']:.3f} @ layer {summary['answer_region_cosine_argmax']}", fill=AMBER, font=font(22, True)) + draw.text((96, 1172), f"final alignment: {summary['answer_region_cosine_final']:.3f}", fill=MUTED, font=font(19)) + + draw.rounded_rectangle((650, 238, 2134, 1234), radius=30, fill=PANEL, outline=(35, 49, 59), width=1) + draw.text((686, 270), "3D cross-carrier resonance terrain", fill=INK, font=font(36, True)) + draw.text((686, 314), "z-axis = excess cosine after subtracting each layer's median image-token similarity; orange rails mark the bitmap answer region", fill=MUTED, font=font(20)) + panel = panel.resize((1408, 794), Image.Resampling.LANCZOS) + canvas.paste(panel, (692, 378)) + cmap = plt.colormaps["turbo"] + for i in range(280): + rgb = tuple(int(v * 255) for v in cmap(i / 279)[:3]) + draw.rectangle((1790 + i, 282, 1791 + i, 300), fill=rgb) + draw.text((1790, 254), "low excess", fill=MUTED, font=font(14)) + draw.text((1992, 254), "high excess", fill=MUTED, font=font(14)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_text_image_compare.py b/packages/snapcompact/research/snapcompact_text_image_compare.py new file mode 100644 index 000000000..000e46d81 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_text_image_compare.py @@ -0,0 +1,489 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers", "sentencepiece", "protobuf", "einops", "matplotlib"] +# /// +"""Compare raw-text vs snapcompact-image activations for the same input. + +The experiment feeds the same SQuAD chunk/question through a local VLM twice: +1. as ordinary raw text in a block +2. as a snapcompact bitmap plus the same question + +It then compares the text-carrier answer vector against every image-token vector +by decoder layer, producing a blog visual of cross-modal alignment. +""" + +from __future__ import annotations + +import argparse +import json +import math +import random +import sys +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS, load_prompt # noqa: E402 +from snapcompact_blackbox_occlusion import sample_answer_questions # noqa: E402 + +DEFAULT_MODEL_DIR = ( + "/home/can/.cache/huggingface/hub/models--PaddlePaddle--PaddleOCR-VL/" + "snapshots/2b77538ef936207f60c16b45082841068987d08c" +) + +PALETTE = { + "bg": (5, 7, 10), + "panel": (13, 18, 24), + "panel2": (9, 13, 18), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "purple": (180, 96, 255), + "grid": (38, 49, 58), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Monaco.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if path and Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Monaco.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf", + ] + for path in candidates: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def heat_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + stops = [ + (0.00, (4, 7, 20)), + (0.18, (22, 24, 80)), + (0.38, (62, 68, 168)), + (0.58, (38, 183, 208)), + (0.78, (160, 250, 145)), + (1.00, (255, 245, 166)), + ] + for (a, ca), (b, cb) in zip(stops, stops[1:]): + if t <= b: + u = (t - a) / (b - a) + return tuple(round(ca[i] + (cb[i] - ca[i]) * u) for i in range(3)) + return stops[-1][1] + + +def cosine(a: np.ndarray, b: np.ndarray) -> np.ndarray: + a_norm = np.linalg.norm(a, axis=-1, keepdims=True) + b_norm = np.linalg.norm(b, axis=-1, keepdims=True) + return (a * b).sum(axis=-1) / np.maximum((a_norm * b_norm).squeeze(-1), 1e-6) + + +def normalize_heat(arr: np.ndarray, lo: float | None = None, hi: float | None = None) -> tuple[np.ndarray, float, float]: + if lo is None: + lo = float(np.quantile(arr, 0.03)) + if hi is None: + hi = float(np.quantile(arr, 0.98)) + if hi <= lo: + hi = lo + 1e-6 + return np.clip((arr - lo) / (hi - lo), 0, 1), lo, hi + + +def apply_template(processor: Any, content: list[dict[str, Any]]) -> str: + return processor.apply_chat_template([{"role": "user", "content": content}], tokenize=False, add_generation_prompt=True) + + +def text_spans(processor: Any, templated: str, chunk: str, answer_start: int, answer_end: int) -> dict[str, int]: + tokenizer = processor.tokenizer + chunk_at = templated.index(chunk) + prefix = templated[:chunk_at] + def n_tokens(s: str) -> int: + return len(tokenizer(s, add_special_tokens=False)["input_ids"]) + ref_start = n_tokens(prefix) + ref_end = n_tokens(prefix + chunk) + answer_tok_start = n_tokens(prefix + chunk[:answer_start]) + answer_tok_end = max(answer_tok_start + 1, n_tokens(prefix + chunk[:answer_end])) + return { + "ref_start": ref_start, + "ref_end": ref_end, + "answer_start": answer_tok_start, + "answer_end": answer_tok_end, + } + + +def to_device(batch: dict[str, Any], device: Any) -> dict[str, Any]: + return {k: (v.to(device) if hasattr(v, "to") else v) for k, v in batch.items()} + + +def run_text(model: Any, processor: Any, text_prompt: str, chunk: str, answer_start: int, answer_end: int, device: Any) -> tuple[list[np.ndarray], dict[str, int], str]: + import torch + + templated = apply_template(processor, [{"type": "text", "text": text_prompt}]) + spans = text_spans(processor, templated, chunk, answer_start, answer_end) + batch = processor(text=templated, return_tensors="pt") + batch = to_device(batch, device) + with torch.no_grad(): + out = model(**batch, output_hidden_states=True, output_attentions=False, use_cache=False) + layers = [h[0].float().detach().cpu().numpy().astype(np.float32, copy=False) for h in out.hidden_states] + return layers, spans, templated + + +def run_image(model: Any, processor: Any, img: Image.Image, img_prompt: str, device: Any) -> tuple[list[np.ndarray], list[int], dict[str, Any], str]: + import torch + + templated = apply_template(processor, [{"type": "image", "image": img}, {"type": "text", "text": img_prompt}]) + batch = processor(images=img, text=templated, return_tensors="pt") + image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token) + image_positions = [i for i, token_id in enumerate(batch["input_ids"][0].tolist()) if token_id == image_token_id] + meta = {k: (v.tolist() if hasattr(v, "tolist") else v) for k, v in batch.items() if k in ("image_grid_thw",)} + batch = to_device(batch, device) + with torch.no_grad(): + out = model(**batch, output_hidden_states=True, output_attentions=False, use_cache=False) + layers = [h[0].float().detach().cpu().numpy().astype(np.float32, copy=False) for h in out.hidden_states] + return layers, image_positions, meta, templated + + +def image_answer_token_indices(answer_start: int, answer_end: int, text_cols: int, adv: int, pitch: int, image_w: int, image_h: int, image_token_count: int) -> list[int]: + grid = round(math.sqrt(image_token_count)) + if grid * grid != image_token_count: + return [] + row0 = max(0, answer_start // text_cols) + row1 = max(row0, (answer_end - 1) // text_cols) + col0 = max(0, answer_start % text_cols) + col1 = max(col0, (answer_end - 1) % text_cols) + x0 = max(0, col0 * adv - adv) + x1 = min(image_w, (col1 + 2) * adv) + y0 = max(0, row0 * pitch - 1) + y1 = min(image_h, (row1 + 1) * pitch + 1) + gx0 = max(0, min(grid - 1, int(x0 / image_w * grid))) + gx1 = max(0, min(grid - 1, int(math.ceil(x1 / image_w * grid)))) + gy0 = max(0, min(grid - 1, int(y0 / image_h * grid))) + gy1 = max(0, min(grid - 1, int(math.ceil(y1 / image_h * grid)))) + out: list[int] = [] + for gy in range(gy0, gy1 + 1): + for gx in range(gx0, gx1 + 1): + out.append(gy * grid + gx) + return sorted(set(out)) + + +def crop_answer(img: Image.Image, start: int, end: int, cols: int, adv: int, pitch: int, pad_cells: int = 34) -> Image.Image: + row0 = max(0, start // cols - 5) + row1 = min(img.height // pitch, end // cols + 6) + col0 = max(0, start % cols - pad_cells) + col1 = min(cols, end % cols + pad_cells) + if col1 <= col0: + col1 = min(cols, col0 + 72) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=PALETTE["orange"], width=3) + return crop + + +def draw_wrapped(draw: ImageDraw.ImageDraw, xy: tuple[int, int], text: str, width_chars: int, line_height: int, fill: tuple[int, int, int], fnt: ImageFont.ImageFont) -> int: + words = text.split() + lines: list[str] = [] + current = "" + for word in words: + cand = word if not current else current + " " + word + if len(cand) <= width_chars: + current = cand + else: + if current: + lines.append(current) + current = word + if current: + lines.append(current) + x, y = xy + for line in lines: + draw.text((x, y), line, fill=fill, font=fnt) + y += line_height + return y + + +def render_heat_grid(draw: ImageDraw.ImageDraw, grid: np.ndarray, box: tuple[int, int, int, int], title: str, layer: int, answer_indices: list[int], color: tuple[int, int, int]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=18, fill=PALETTE["panel2"], outline=(35, 49, 59), width=1) + draw.text((x0 + 18, y0 + 14), title, fill=color, font=ui_font(22, True)) + draw.text((x0 + 18, y0 + 43), f"decoder layer {layer}", fill=PALETTE["muted"], font=ui_font(15)) + gx0, gy0, gx1, gy1 = x0 + 26, y0 + 76, x1 - 26, y1 - 24 + rows, cols = grid.shape + cw = (gx1 - gx0) / cols + ch = (gy1 - gy0) / rows + for r in range(rows): + for c in range(cols): + xa = round(gx0 + c * cw) + xb = round(gx0 + (c + 1) * cw) + ya = round(gy0 + r * ch) + yb = round(gy0 + (r + 1) * ch) + draw.rectangle((xa, ya, xb, yb), fill=heat_color(float(grid[r, c]))) + for idx in answer_indices: + r, c = divmod(idx, cols) + xa = round(gx0 + c * cw) + xb = round(gx0 + (c + 1) * cw) + ya = round(gy0 + r * ch) + yb = round(gy0 + (r + 1) * ch) + draw.rectangle((xa - 2, ya - 2, xb + 2, yb + 2), outline=PALETTE["orange"], width=2) + + +def render_visual(out_path: Path, summary: dict[str, Any], arrays: dict[str, np.ndarray], original_img: Image.Image, chunk: str) -> None: + w, h = 2100, 1260 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -220, 860, 660), fill=(75, 220, 255, 28)) + gd.ellipse((1160, 80, 2420, 1320), fill=(255, 112, 72, 26)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(80))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + q = summary["question"] + draw.text((62, 42), "SNAPCOMPACT CARRIER COMPARISON", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((62, 82), "Same input, two internal languages", fill=PALETTE["ink"], font=ui_font(68, True)) + draw.text((64, 166), "Raw text tokens vs bitmap image tokens. Bright fields show where the text-carrier answer vector resonates with the image-carrier hidden state.", fill=PALETTE["muted"], font=ui_font(25)) + + # Carrier cards. + draw.rounded_rectangle((62, 236, 620, 760), radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((94, 270), "raw text carrier", fill=PALETTE["cyan"], font=ui_font(30, True)) + start = max(0, q["answer_start"] - 230) + end = min(len(chunk), q["answer_end"] + 230) + snippet = chunk[start:end].replace("\n", " ") + rel_a = q["answer_start"] - start + rel_b = q["answer_end"] - start + before = snippet[:rel_a] + answer = snippet[rel_a:rel_b] + after = snippet[rel_b:] + tx, ty = 94, 328 + ty = draw_wrapped(draw, (tx, ty), before[-260:], 52, 22, PALETTE["ink"], mono_font(15)) + draw.rounded_rectangle((tx, ty + 2, tx + 16 * max(3, len(answer)), ty + 27), radius=5, fill=(255, 196, 68)) + draw.text((tx + 4, ty + 5), answer, fill=(8, 10, 10), font=mono_font(16)) + ty += 36 + draw_wrapped(draw, (tx, ty), after[:260], 52, 22, PALETTE["ink"], mono_font(15)) + draw.text((94, 694), f"answer tokens: {summary['text_answer_tokens']}", fill=PALETTE["muted"], font=ui_font(18)) + draw.text((94, 724), f"reference tokens: {summary['text_reference_tokens']}", fill=PALETTE["muted"], font=ui_font(18)) + + draw.rounded_rectangle((62, 792, 620, 1192), radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((94, 826), "image carrier", fill=PALETTE["orange"], font=ui_font(30, True)) + crop = crop_answer(original_img, q["answer_start"], q["answer_end"], summary["geometry"]["cols"], 8, 13) + scale = min(478 / crop.width, 218 / crop.height) + crop_r = crop.resize((round(crop.width * scale), round(crop.height * scale)), Image.Resampling.NEAREST) + draw.rounded_rectangle((94, 888, 588, 1134), radius=16, fill=(244, 242, 230), outline=PALETTE["orange"], width=3) + canvas.paste(crop_r, (94 + (494 - crop_r.width) // 2, 888 + (246 - crop_r.height) // 2)) + draw.text((94, 1150), f"image tokens: {summary['image_tokens']} ({summary['image_grid']}×{summary['image_grid']})", fill=PALETTE["muted"], font=ui_font(18)) + + # Layer grids. + sim = arrays["text_answer_to_image_excess_norm"] + grid = summary["image_grid"] + answer_indices = summary["image_answer_token_indices"] + layers = summary["selected_layers"] + boxes = [(672, 236, 1088, 626), (1118, 236, 1534, 626), (1564, 236, 1980, 626)] + names = ["input layer", "middle layer", "peak alignment"] + colors = [PALETTE["cyan"], PALETTE["purple"], PALETTE["green"]] + for layer, box, name, color in zip(layers, boxes, names, colors): + render_heat_grid(draw, sim[layer].reshape(grid, grid), box, name, layer, answer_indices, color) + + # Cosine bridge panel. + draw.rounded_rectangle((672, 672, 1980, 1192), radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((704, 704), "cross-carrier convergence bridge", fill=PALETTE["ink"], font=ui_font(34, True)) + draw.text((704, 744), "Cosine similarity between pooled raw-text answer states and pooled bitmap answer-region states by layer", fill=PALETTE["muted"], font=ui_font(19)) + x0, y0, x1, y1 = 730, 820, 1908, 1096 + for i in range(5): + y = y0 + round((y1 - y0) * i / 4) + draw.line((x0, y, x1, y), fill=PALETTE["grid"], width=1) + local = arrays["answer_region_cosine"] + global_mean = arrays["global_mean_cosine"] + lo = float(min(local.min(), global_mean.min())) + hi = float(max(local.max(), global_mean.max())) + if hi <= lo: + hi = lo + 1e-6 + def pts(vals: np.ndarray) -> list[tuple[int, int]]: + out = [] + for i, v in enumerate(vals): + x = x0 + round((x1 - x0) * i / max(1, len(vals) - 1)) + y = y1 - round((y1 - y0) * (float(v) - lo) / (hi - lo)) + out.append((x, y)) + return out + p_local = pts(local) + p_global = pts(global_mean) + draw.line(p_global, fill=PALETTE["muted"], width=4) + draw.line(p_local, fill=PALETTE["amber"], width=6) + for x, y in p_local: + draw.ellipse((x - 5, y - 5, x + 5, y + 5), fill=PALETTE["amber"]) + draw.text((x0, y1 + 22), "layer 0", fill=PALETTE["muted"], font=ui_font(16)) + draw.text((x1 - 70, y1 + 22), f"layer {len(local) - 1}", fill=PALETTE["muted"], font=ui_font(16)) + peak_layer = int(np.argmax(local)) + draw.rounded_rectangle((1502, 790, 1938, 900), radius=18, fill=(9, 13, 18), outline=(38, 51, 60), width=1) + draw.text((1526, 812), f"answer cosine peaks: {local[peak_layer]:.3f} @L{peak_layer}", fill=PALETTE["amber"], font=ui_font(21, True)) + draw.text((1526, 842), f"final answer cosine: {local[-1]:.3f}", fill=PALETTE["muted"], font=ui_font(18)) + draw.text((1526, 868), f"final global carrier cosine: {global_mean[-1]:.3f}", fill=PALETTE["muted"], font=ui_font(18)) + draw.rounded_rectangle((704, 1120, 1238, 1168), radius=13, fill=(9, 13, 18), outline=(38, 51, 60), width=1) + draw.rectangle((724, 1138, 768, 1148), fill=PALETTE["amber"]) + draw.text((784, 1129), "answer region: text vector ↔ image region", fill=PALETTE["muted"], font=ui_font(17)) + draw.rectangle((1260, 1138, 1304, 1148), fill=PALETTE["muted"]) + draw.text((1320, 1129), "global carrier means", fill=PALETTE["muted"], font=ui_font(17)) + + out_path.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out_path) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default=DEFAULT_MODEL_DIR) + ap.add_argument("--font", default="8x13", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=768) + ap.add_argument("--limit-paras", type=int, default=40) + ap.add_argument("--qpc", type=int, default=16) + ap.add_argument("--question-index", type=int, default=7) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--out", default="text-image-compare-paddleocr-q7") + args = ap.parse_args() + + import torch + from transformers import AutoConfig, AutoModel, AutoProcessor + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + start, end = 0, min(len(flow), budget) + chunk = flow[start:end] + questions = sample_answer_questions(paras, offsets, start, end, args.qpc, args.seed) + if not questions: + raise SystemExit("no sampled questions fit in chunk") + q = questions[min(args.question_index, len(questions) - 1)] + img = render(chunk, cfg, CACHE, args.size, args.variant) + img.save(img_dir / "image-carrier.png") + + print(f"loading {args.model_dir}", flush=True) + config = AutoConfig.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True) + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + target_device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + dtype = torch.bfloat16 if target_device.type == "cuda" else torch.float32 + if getattr(config, "model_type", "") == "qwen2_5_vl": + from transformers import Qwen2_5_VLForConditionalGeneration + + model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + args.model_dir, + local_files_only=True, + trust_remote_code=True, + dtype=dtype, + device_map="auto" if target_device.type == "cuda" else None, + ).eval() + device = next(model.parameters()).device + else: + model = AutoModel.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=dtype).to(target_device).eval() + device = target_device + + text_prompt = ( + "Below is reference material. Answer the question using only it.\n\n" + f"{chunk}\n\nQuestion: {q['q']}\n" + "Answer with only the shortest extractive answer." + ) + img_prompt = load_prompt("qa-image.md").format(cols=cols, rows=rows) + f"\n\nQuestion: {q['q']}\nAnswer with only the shortest extractive answer." + + text_layers, text_pos, text_template = run_text(model, processor, text_prompt, chunk, q["answer_start"], q["answer_end"], device) + image_layers, image_positions, image_meta, image_template = run_image(model, processor, img, img_prompt, device) + image_token_count = len(image_positions) + image_grid = round(math.sqrt(image_token_count)) + answer_image_indices = image_answer_token_indices(q["answer_start"], q["answer_end"], cols, cfg.adv, cfg.pitch, img.width, img.height, image_token_count) + + answer_cos = [] + global_cos = [] + text_answer_to_image = [] + for text_h, image_h in zip(text_layers, image_layers): + text_ref = text_h[text_pos["ref_start"] : text_pos["ref_end"]] + text_ans = text_h[text_pos["answer_start"] : text_pos["answer_end"]] + image_tokens = image_h[image_positions] + image_ans = image_tokens[answer_image_indices] if answer_image_indices else image_tokens + text_ans_mean = text_ans.mean(axis=0) + image_ans_mean = image_ans.mean(axis=0) + text_ref_mean = text_ref.mean(axis=0) + image_mean = image_tokens.mean(axis=0) + answer_cos.append(float(cosine(text_ans_mean[None, :], image_ans_mean[None, :])[0])) + global_cos.append(float(cosine(text_ref_mean[None, :], image_mean[None, :])[0])) + sims = cosine(np.repeat(text_ans_mean[None, :], image_tokens.shape[0], axis=0), image_tokens) + text_answer_to_image.append(sims.astype(np.float32, copy=False)) + + text_answer_to_image_arr = np.stack(text_answer_to_image, axis=0) + layer_baseline = np.median(text_answer_to_image_arr, axis=1, keepdims=True) + text_answer_to_image_excess = text_answer_to_image_arr - layer_baseline + normed, heat_lo, heat_hi = normalize_heat(text_answer_to_image_excess) + answer_cos_arr = np.array(answer_cos, dtype=np.float32) + global_cos_arr = np.array(global_cos, dtype=np.float32) + selected_layers = [0, len(text_layers) // 2, int(answer_cos_arr.argmax())] + + summary = { + "args": vars(args), + "device": str(device), + "geometry": {"cols": cols, "rows": rows, "capacity": budget}, + "question": { + "q": q["q"], + "golds": q["golds"], + "answer_text": q["answer_text"], + "answer_start": q["answer_start"], + "answer_end": q["answer_end"], + }, + "layers": len(text_layers), + "image_tokens": image_token_count, + "image_grid": image_grid, + "image_answer_token_indices": answer_image_indices, + "image_meta": image_meta, + "text_positions": text_pos, + "text_reference_tokens": text_pos["ref_end"] - text_pos["ref_start"], + "text_answer_tokens": text_pos["answer_end"] - text_pos["answer_start"], + "selected_layers": selected_layers, + "answer_region_cosine_final": float(answer_cos_arr[-1]), + "global_mean_cosine_final": float(global_cos_arr[-1]), + "answer_region_cosine_max": float(answer_cos_arr.max()), + "answer_region_cosine_argmax": int(answer_cos_arr.argmax()), + "heat_normalization": {"lo_p03": heat_lo, "hi_p98": heat_hi}, + "text_template_prefix": text_template[:240], + "image_template_prefix": image_template[:240], + } + arrays = { + "text_answer_to_image_cosine": text_answer_to_image_arr, + "text_answer_to_image_excess": text_answer_to_image_excess, + "text_answer_to_image_excess_norm": normed, + "answer_region_cosine": answer_cos_arr, + "global_mean_cosine": global_cos_arr, + } + np.savez_compressed(out_dir / "text_image_compare.npz", **arrays) + (out_dir / "summary.json").write_text(json.dumps(summary, indent=1)) + render_visual(out_dir / "text-vs-image.png", summary, arrays, img, chunk) + print(json.dumps(summary, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_token_entry_dump.py b/packages/snapcompact/research/snapcompact_token_entry_dump.py new file mode 100644 index 000000000..5d9516192 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_token_entry_dump.py @@ -0,0 +1,158 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy", "torch", "transformers"] +# /// +"""Dump how the same content enters Qwen as text tokens vs visual tokens. + +Produces a JSON with real tokenizer output (token strings + ids), the real +embedding rows entering the decoder for the answer-word text tokens, and the +real visual-tower output vectors for the image tokens covering the same word. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +import numpy as np + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +import squad # noqa: E402 +from bdf import capacity, render # noqa: E402 +from run import CACHE, FONTS # noqa: E402 +from snapcompact_blackbox_occlusion import sample_answer_questions # noqa: E402 +from snapcompact_text_image_compare import image_answer_token_indices # noqa: E402 + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--model-dir", default="Qwen/Qwen2.5-VL-7B-Instruct") + ap.add_argument("--font", default="8x13", choices=sorted(FONTS)) + ap.add_argument("--variant", default="bw") + ap.add_argument("--size", type=int, default=1568) + ap.add_argument("--limit-paras", type=int, default=80) + ap.add_argument("--question-index", type=int, default=3) + ap.add_argument("--seed", type=int, default=42) + ap.add_argument("--context-chars", type=int, default=120) + ap.add_argument("--embed-dims", type=int, default=10) + ap.add_argument("--out", default="qwen-token-entry") + args = ap.parse_args() + + import torch + from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration + + out_dir = HERE / "results" / args.out + img_dir = out_dir / "images" + img_dir.mkdir(parents=True, exist_ok=True) + CACHE.mkdir(exist_ok=True) + + cfg = FONTS[args.font] + cols, rows, budget = capacity(cfg, args.size) + paras = squad.load_paragraphs(CACHE)[: args.limit_paras] + flow, offsets = squad.build_flow(paras) + chunk = flow[: min(len(flow), budget)] + questions = sample_answer_questions(paras, offsets, 0, len(chunk), 24, args.seed) + q = questions[min(args.question_index, len(questions) - 1)] + img = render(chunk, cfg, CACHE, args.size, args.variant) + img.save(img_dir / "image-carrier.png") + + print(f"loading {args.model_dir}", flush=True) + from transformers import AutoTokenizer + + processor = AutoProcessor.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, use_fast=False) + model = Qwen2_5_VLForConditionalGeneration.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True, dtype=torch.bfloat16, device_map="auto").eval() + device = next(model.parameters()).device + tokenizer = AutoTokenizer.from_pretrained(args.model_dir, local_files_only=True, trust_remote_code=True) # fast tokenizer for offsets + + # --- Text lane: real tokenization of the snippet around the answer. + snip_start = max(0, q["answer_start"] - args.context_chars) + snip_end = min(len(chunk), q["answer_end"] + args.context_chars) + snippet = chunk[snip_start:snip_end] + enc = tokenizer(snippet, add_special_tokens=False, return_offsets_mapping=True) + tokens = [] + answer_token_idx: list[int] = [] + rel_a = q["answer_start"] - snip_start + rel_b = q["answer_end"] - snip_start + for ti, (tok_id, (o0, o1)) in enumerate(zip(enc["input_ids"], enc["offset_mapping"])): + is_answer = o0 < rel_b and o1 > rel_a + if is_answer: + answer_token_idx.append(ti) + tokens.append({"i": ti, "id": int(tok_id), "str": tokenizer.decode([tok_id]), "answer": bool(is_answer)}) + + # Real embedding rows entering the decoder for the answer tokens. + embed = model.get_input_embeddings() + answer_ids = torch.tensor([tokens[i]["id"] for i in answer_token_idx], device=device) + with torch.no_grad(): + answer_embeds = embed(answer_ids).float().cpu().numpy() + text_entry = [ + { + "id": tokens[i]["id"], + "str": tokens[i]["str"], + "vector_head": [round(float(v), 4) for v in answer_embeds[k, : args.embed_dims]], + "norm": round(float(np.linalg.norm(answer_embeds[k])), 4), + } + for k, i in enumerate(answer_token_idx) + ] + chunk_token_count = len(tokenizer(chunk, add_special_tokens=False)["input_ids"]) + + # --- Image lane: real pixel patches and visual-tower output vectors. + batch = processor(images=img, text="<|vision_start|><|image_pad|><|vision_end|>", return_tensors="pt") + pixel_values = batch["pixel_values"] + grid_thw = batch["image_grid_thw"] + merge = int(getattr(processor.image_processor, "merge_size", 2)) + patch = int(getattr(processor.image_processor, "patch_size", 14)) + with torch.no_grad(): + visual_out = model.model.visual(pixel_values.to(device, dtype=torch.bfloat16), grid_thw=grid_thw.to(device)).float().cpu().numpy() + n_tokens = visual_out.shape[0] + grid = int(round(n_tokens**0.5)) + answer_img_indices = image_answer_token_indices(q["answer_start"], q["answer_end"], cols, cfg.adv, cfg.pitch, img.width, img.height, n_tokens) + image_entry = [ + { + "token_index": int(idx), + "grid_rc": [int(idx // grid), int(idx % grid)], + "vector_head": [round(float(v), 4) for v in visual_out[idx, : args.embed_dims]], + "norm": round(float(np.linalg.norm(visual_out[idx])), 4), + } + for idx in answer_img_indices + ] + # A few real normalized pixel values from the first answer patch (pre-visual-tower input). + patches_per_token = merge * merge + first_patch_row = answer_img_indices[0] * patches_per_token if answer_img_indices else 0 + pixel_head = [round(float(v), 4) for v in pixel_values[min(first_patch_row, pixel_values.shape[0] - 1), : args.embed_dims].tolist()] + + dump = { + "args": vars(args), + "question": {"q": q["q"], "answer_text": q["answer_text"], "answer_start": q["answer_start"], "answer_end": q["answer_end"]}, + "geometry": {"cols": cols, "rows": rows, "image_w": img.width, "image_h": img.height}, + "snippet": snippet, + "snippet_rel_answer": [rel_a, rel_b], + "tokens": tokens, + "text_entry": text_entry, + "chunk_chars": len(chunk), + "chunk_text_tokens": chunk_token_count, + "image_tokens": n_tokens, + "image_grid": grid, + "grid_thw": grid_thw.tolist(), + "patch_size": patch, + "merge_size": merge, + "token_pixel_size": patch * merge, + "processor_resized": [int(grid_thw[0][2]) * patch, int(grid_thw[0][1]) * patch], + "pixel_values_shape": list(pixel_values.shape), + "pixel_head_first_answer_patch": pixel_head, + "image_answer_token_indices": [int(i) for i in answer_img_indices], + "image_entry": image_entry, + "embed_dim": int(answer_embeds.shape[1]), + "visual_out_dim": int(visual_out.shape[1]), + } + (out_dir / "token_entry.json").write_text(json.dumps(dump, indent=1)) + print(json.dumps({k: v for k, v in dump.items() if k not in ("tokens", "snippet")}, indent=1)) + print(f"results -> {out_dir}") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_token_entry_viz.py b/packages/snapcompact/research/snapcompact_token_entry_viz.py new file mode 100644 index 000000000..3feaac0f6 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_token_entry_viz.py @@ -0,0 +1,225 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render how the same word enters Qwen as BPE tokens vs 28px visual patches.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +PALETTE = { + "bg": (5, 7, 10), + "panel": (12, 17, 23), + "panel2": (8, 12, 17), + "ink": (241, 239, 224), + "muted": (143, 154, 160), + "cyan": (75, 220, 255), + "orange": (255, 112, 72), + "green": (148, 255, 117), + "amber": (255, 196, 68), + "grid": (38, 49, 58), +} + + +def ui_font(size: int, bold: bool = False) -> ImageFont.ImageFont: + for path in [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mono_font(size: int) -> ImageFont.ImageFont: + for path in ["/System/Library/Fonts/Monaco.ttf", "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"]: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def vector_text(head: list[float]) -> str: + return "[" + ", ".join(f"{v:+.2f}" for v in head[:6]) + ", …]" + + +def draw_vector_bar(draw: ImageDraw.ImageDraw, xy: tuple[int, int], head: list[float], color: tuple[int, int, int], width: int = 330) -> None: + x, y = xy + n = len(head) + bw = width // n + hi = max(0.001, max(abs(v) for v in head)) + mid = y + 22 + for i, v in enumerate(head): + bh = round(20 * abs(v) / hi) + xa = x + i * bw + if v >= 0: + draw.rounded_rectangle((xa, mid - bh, xa + bw - 4, mid), radius=3, fill=color) + else: + draw.rounded_rectangle((xa, mid, xa + bw - 4, mid + bh), radius=3, fill=tuple(c // 2 for c in color)) + draw.line((x, mid, x + width, mid), fill=PALETTE["grid"], width=1) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(HERE / "results" / "qwen-token-entry-q3")) + ap.add_argument("--out", default=str(HERE / "results" / "qwen-token-entry-q3" / "token-entry.png")) + args = ap.parse_args() + result_dir = Path(args.result_dir) + dump = json.loads((result_dir / "token_entry.json").read_text()) + carrier = Image.open(result_dir / "images" / "image-carrier.png").convert("RGB") + + w, h = 2200, 1400 + canvas = Image.new("RGB", (w, h), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(7, 10 + y % 9, 15 + y % 11)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-240, -200, 940, 760), fill=(75, 220, 255, 27)) + gd.ellipse((1240, 540, 2460, 1480), fill=(255, 112, 72, 25)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(86))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + q = dump["question"] + answer = q["answer_text"] + draw.text((64, 42), "QWEN TOKEN ENTRY — SAME WORD, TWO ENCODINGS", fill=PALETTE["amber"], font=ui_font(24, True)) + draw.text((64, 84), f"How “{answer}” gets into the model", fill=PALETTE["ink"], font=ui_font(64, True)) + draw.text( + (66, 164), + "Real values, no schematic: actual BPE ids and embedding rows on the text path; actual 28×28 pixel patches and visual-tower output vectors on the image path.", + fill=PALETTE["muted"], + font=ui_font(23), + ) + + # ---- TEXT LANE ---- + lane = (64, 238, 2136, 700) + draw.rounded_rectangle(lane, radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 262), "text carrier — BPE tokens", fill=PALETTE["cyan"], font=ui_font(30, True)) + draw.text((96, 302), f"snippet around the answer · {dump['chunk_chars']:,} chars → {dump['chunk_text_tokens']:,} text tokens for the whole chunk", fill=PALETTE["muted"], font=ui_font(18)) + + # Token ribbon: show tokens around the answer. + tokens = dump["tokens"] + answer_positions = [t["i"] for t in tokens if t["answer"]] + mid_token = answer_positions[0] if answer_positions else len(tokens) // 2 + window = [t for t in tokens if mid_token - 7 <= t["i"] <= mid_token + 7] + fnt = mono_font(19) + fnt_id = mono_font(12) + x = 96 + y = 356 + for t in window: + label = t["str"].replace("\n", "⏎") + if not label.strip(): + label = "␣" * max(1, len(label)) + tw = max(int(draw.textlength(label, font=fnt)) + 22, 54) + if x + tw > 2100: + x = 96 + y += 96 + color = PALETTE["amber"] if t["answer"] else (30, 41, 50) + text_color = (8, 10, 12) if t["answer"] else PALETTE["ink"] + draw.rounded_rectangle((x, y, x + tw, y + 44), radius=9, fill=color, outline=(52, 68, 80), width=1) + draw.text((x + 11, y + 9), label, fill=text_color, font=fnt) + draw.text((x + 4, y + 50), f"id {t['id']}", fill=PALETTE["muted"], font=fnt_id) + x += tw + 8 + + draw.text((96, 500), "what actually enters the decoder (embedding row, first 6 of " + f"{dump['embed_dim']} dims):", fill=PALETTE["muted"], font=ui_font(18, True)) + ex = 96 + for entry in dump["text_entry"][:3]: + box = (ex, 536, ex + 470, 668) + draw.rounded_rectangle(box, radius=16, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((ex + 18, 548), f"“{entry['str']}” id {entry['id']}", fill=PALETTE["cyan"], font=ui_font(20, True)) + draw.text((ex + 18, 578), vector_text(entry["vector_head"]), fill=PALETTE["ink"], font=mono_font(15)) + draw_vector_bar(draw, (ex + 18, 606), entry["vector_head"], PALETTE["cyan"], width=430) + draw.text((ex + 360, 548), f"‖x‖={entry['norm']:.2f}", fill=PALETTE["muted"], font=ui_font(14)) + ex += 494 + + # ---- IMAGE LANE ---- + lane = (64, 736, 2136, 1336) + draw.rounded_rectangle(lane, radius=28, fill=PALETTE["panel"], outline=(35, 49, 59), width=1) + draw.text((96, 760), "image carrier — visual patch tokens", fill=PALETTE["orange"], font=ui_font(30, True)) + px = dump["token_pixel_size"] + draw.text( + (96, 800), + f"same word as pixels · bitmap resized to {dump['processor_resized'][0]}×{dump['processor_resized'][1]} → {dump['patch_size']}px patches, {dump['merge_size']}×{dump['merge_size']} merged → {dump['image_tokens']:,} tokens of {px}×{px}px", + fill=PALETTE["muted"], + font=ui_font(18), + ) + + # Zoomed answer region with the real patch grid. + grid = dump["image_grid"] + rw, rh = dump["processor_resized"] + resized = carrier.resize((rw, rh), Image.Resampling.LANCZOS) + indices = dump["image_answer_token_indices"] + rows = sorted({i // grid for i in indices}) + cols_ = sorted({i % grid for i in indices}) + pad = 3 + cx0 = max(0, (min(cols_) - pad) * px) + cx1 = min(rw, (max(cols_) + 1 + pad) * px) + cy0 = max(0, (min(rows) - pad) * px) + cy1 = min(rh, (max(rows) + 1 + pad) * px) + crop = resized.crop((cx0, cy0, cx1, cy1)) + scale = min(940 / crop.width, 225 / crop.height) + crop_big = crop.resize((round(crop.width * scale), round(crop.height * scale)), Image.Resampling.NEAREST) + ox, oy = 96, 852 + draw.rounded_rectangle((ox - 6, oy - 6, ox + crop_big.width + 6, oy + crop_big.height + 6), radius=10, fill=(244, 242, 230)) + canvas.paste(crop_big, (ox, oy)) + cd = ImageDraw.Draw(canvas) + for gx in range(cx0 // px, cx1 // px + 1): + lx = ox + (gx * px - cx0) * scale + cd.line((lx, oy, lx, oy + crop_big.height), fill=(150, 158, 162, 80), width=1) + for gy in range(cy0 // px, cy1 // px + 1): + ly = oy + (gy * px - cy0) * scale + cd.line((ox, ly, ox + crop_big.width, ly), fill=(150, 158, 162, 80), width=1) + for idx in indices: + r, c = idx // grid, idx % grid + xa = ox + (c * px - cx0) * scale + ya = oy + (r * px - cy0) * scale + cd.rectangle((xa, ya, xa + px * scale, ya + px * scale), outline=PALETTE["orange"], width=4) + draw.text((ox, oy + crop_big.height + 14), f"orange cells = the {len(indices)} visual tokens covering “{answer}” (token grid {grid}×{grid})", fill=PALETTE["muted"], font=ui_font(17)) + + # Magnified single patches. + sx = ox + crop_big.width + 60 + draw.text((sx, 852 - 26), "individual visual tokens (real input pixels):", fill=PALETTE["muted"], font=ui_font(18, True)) + for k, idx in enumerate(indices[:5]): + r, c = idx // grid, idx % grid + cell = resized.crop((c * px, r * px, (c + 1) * px, (r + 1) * px)).resize((132, 132), Image.Resampling.NEAREST) + bx = sx + k * 160 + draw.rounded_rectangle((bx - 4, 852 - 4, bx + 136, 852 + 136), radius=8, fill=(244, 242, 230), outline=PALETTE["orange"], width=3) + canvas.paste(cell, (bx, 852)) + draw.text((bx, 996), f"tok[{idx}]", fill=PALETTE["muted"], font=mono_font(13)) + draw.text((sx, 1030), f"pre-tower normalized pixels of first patch: {vector_text(dump['pixel_head_first_answer_patch'])}", fill=PALETTE["muted"], font=mono_font(14)) + + draw.text((96, 1106), f"what actually enters the decoder (visual-tower output, first 6 of {dump['visual_out_dim']} dims):", fill=PALETTE["muted"], font=ui_font(18, True)) + ex = 96 + for entry in dump["image_entry"][:4]: + box = (ex, 1142, ex + 470, 1274) + draw.rounded_rectangle(box, radius=16, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + r, c = entry["grid_rc"] + draw.text((ex + 18, 1154), f"visual tok[{entry['token_index']}] (row {r}, col {c})", fill=PALETTE["orange"], font=ui_font(20, True)) + draw.text((ex + 18, 1184), vector_text(entry["vector_head"]), fill=PALETTE["ink"], font=mono_font(15)) + draw_vector_bar(draw, (ex + 18, 1212), entry["vector_head"], PALETTE["orange"], width=430) + draw.text((ex + 360, 1154), f"‖x‖={entry['norm']:.2f}", fill=PALETTE["muted"], font=ui_font(14)) + ex += 494 + + # Comparison strip. + text_tok_for_word = len(dump["text_entry"]) + draw.rounded_rectangle((1100, 536, 2104, 668), radius=16, fill=PALETTE["panel2"], outline=(34, 48, 58), width=1) + draw.text((1128, 556), f"“{answer}” = {text_tok_for_word} text token(s) · {len(indices)} visual tokens", fill=PALETTE["ink"], font=ui_font(22, True)) + draw.text((1128, 592), f"both end up as {dump['embed_dim']}-dim rows in the same decoder", fill=PALETTE["ink"], font=ui_font(19)) + draw.text((1128, 626), "text path: lookup table row. image path: ViT forward over 4 raw patches → merger MLP.", fill=PALETTE["muted"], font=ui_font(16)) + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_atlas.py b/packages/snapcompact/research/snapcompact_viz_atlas.py new file mode 100644 index 000000000..29d05c57b --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_atlas.py @@ -0,0 +1,377 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["matplotlib", "numpy", "pillow"] +# /// +"""Render an Activation Atlas-style 2D geography from snapcompact activations.""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.patheffects as pe +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import LinearSegmentedColormap, Normalize +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +DEFAULT_DATA_DIR = HERE / "results" / "tensor-heatmap-paddleocr-q7" +DEFAULT_OUT_DIR = HERE / "results" / "agent-viz-atlas" + +BG = (4, 7, 12) +PANEL = (10, 15, 23) +INK = (241, 239, 224) +MUTED = (139, 151, 160) +CYAN = (70, 216, 255) +RED = (255, 75, 61) +AMBER = (255, 198, 68) +GREEN = (135, 255, 139) +PURPLE = (183, 108, 255) + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Supplemental/Avenir Next Condensed.ttc", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if path and Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def quantile_norm(x: np.ndarray, lo_q: float = 0.02, hi_q: float = 0.985) -> np.ndarray: + lo = float(np.quantile(x, lo_q)) + hi = float(np.quantile(x, hi_q)) + if hi <= lo: + return np.zeros_like(x, dtype=np.float32) + return np.clip((x - lo) / (hi - lo), 0.0, 1.0).astype(np.float32) + + +def pca2(features: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + x = features.astype(np.float64, copy=True) + x -= x.mean(axis=0, keepdims=True) + scale = x.std(axis=0, keepdims=True) + scale[scale < 1e-9] = 1.0 + x /= scale + _u, s, vt = np.linalg.svd(x, full_matrices=False) + coords = x @ vt[:2].T + explained = (s[:2] ** 2) / np.maximum(np.sum(s**2), 1e-12) + return coords.astype(np.float32), explained.astype(np.float32) + + +def normalize_coords(coords: np.ndarray) -> np.ndarray: + out = coords.copy() + for axis in range(2): + lo = float(np.quantile(out[:, axis], 0.01)) + hi = float(np.quantile(out[:, axis], 0.99)) + if hi <= lo: + out[:, axis] = 0.5 + else: + out[:, axis] = np.clip((out[:, axis] - lo) / (hi - lo), 0, 1) + out[:, 1] = 1.0 - out[:, 1] + return out + + +def kmeans(points: np.ndarray, k: int = 5, iters: int = 32) -> tuple[np.ndarray, np.ndarray]: + # Deterministic farthest-point seeding avoids random output drift. + centers = [points[np.argmax(points[:, 0] + points[:, 1])]] + for _ in range(1, k): + dist = np.min(np.sum((points[:, None, :] - np.asarray(centers)[None, :, :]) ** 2, axis=2), axis=1) + centers.append(points[int(np.argmax(dist))]) + c = np.asarray(centers, dtype=np.float32) + labels = np.zeros(points.shape[0], dtype=np.int32) + for _ in range(iters): + d = np.sum((points[:, None, :] - c[None, :, :]) ** 2, axis=2) + new_labels = np.argmin(d, axis=1).astype(np.int32) + if np.array_equal(new_labels, labels): + break + labels = new_labels + for i in range(k): + mask = labels == i + if np.any(mask): + c[i] = points[mask].mean(axis=0) + return labels, c + + +def crop_answer_strip(img: Image.Image, start: int, end: int, cols: int, adv: int = 8, pitch: int = 13) -> Image.Image: + row0 = max(0, start // cols - 4) + row1 = min(img.height // pitch, end // cols + 5) + col0 = max(0, start % cols - 32) + col1 = min(cols, end % cols + 34) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 2) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 2) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=4, outline=RED, width=3) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.NEAREST) + canvas.paste(resized, (x0 + (x1 - x0 - resized.width) // 2, y0 + (y1 - y0 - resized.height) // 2)) + + +def render_atlas_panel( + points: np.ndarray, + labels: np.ndarray, + centers: np.ndarray, + ratio_strength: np.ndarray, + answer_strength: np.ndarray, + peak_layers: np.ndarray, + explained: np.ndarray, + summary: dict, + out_dir: Path, +) -> Image.Image: + cmap = LinearSegmentedColormap.from_list("scar", ["#182132", "#245d7a", "#48d8ff", "#ffd04e", "#ff493d"]) + fig = plt.figure(figsize=(15.8, 10.6), dpi=170) + fig.patch.set_facecolor("#04070c") + ax = fig.add_axes((0.045, 0.06, 0.91, 0.88), facecolor="#07101a") + + x = points[:, 0] + y = points[:, 1] + hb = ax.hexbin(x, y, C=ratio_strength, gridsize=46, reduce_C_function=np.mean, cmap=cmap, mincnt=1, linewidths=0, alpha=0.64) + hb.set_clim(0.0, 1.0) + + cluster_colors = ["#46d8ff", "#ff4b3d", "#ffc644", "#87ff8b", "#b76cff"] + for i, color in enumerate(cluster_colors): + mask = labels == i + if np.count_nonzero(mask) < 4: + continue + ax.scatter(x[mask], y[mask], s=28 + answer_strength[mask] * 150, c=color, alpha=0.24, linewidths=0) + ax.scatter( + x[mask], + y[mask], + s=10 + answer_strength[mask] * 52, + c=ratio_strength[mask], + cmap=cmap, + norm=Normalize(0, 1), + alpha=0.93, + edgecolors=color, + linewidths=0.45, + ) + + hot = np.argsort(ratio_strength + answer_strength * 0.55)[-9:] + ax.scatter(x[hot], y[hot], s=210, facecolors="none", edgecolors="#fff0a8", linewidths=1.5, alpha=0.95) + for rank, idx in enumerate(hot[-5:][::-1], 1): + ax.text( + x[idx] + 0.012, + y[idx] + 0.010, + f"T{idx} · L{int(peak_layers[idx])}", + color="#fff3b0", + fontsize=8, + weight="bold", + path_effects=[pe.withStroke(linewidth=2.5, foreground="#05070a")], + ) + + names = ["answer ridge", "control basin", "early glyph shore", "late-context upland", "ratio reef"] + cluster_scores = [] + for i in range(len(centers)): + mask = labels == i + cluster_scores.append((float(ratio_strength[mask].mean()) if np.any(mask) else 0.0, i)) + order = {old: new for new, (_score, old) in enumerate(sorted(cluster_scores, reverse=True))} + for i, c in enumerate(centers): + mask = labels == i + if np.count_nonzero(mask) < 5: + continue + label = names[order[i] % len(names)] + ax.text( + c[0], + c[1], + label.upper(), + color=cluster_colors[i], + fontsize=11, + weight="bold", + ha="center", + va="center", + alpha=0.96, + path_effects=[pe.withStroke(linewidth=4, foreground="#05070a")], + ) + + ax.text( + 0.015, + 0.982, + "Activation Atlas projection", + transform=ax.transAxes, + color="#f1efe0", + fontsize=24, + weight="bold", + va="top", + ) + ax.text( + 0.017, + 0.942, + "Each island is one image token; geography = PCA of 19-layer answer-vs-random activation delta vectors.", + transform=ax.transAxes, + color="#8b97a0", + fontsize=11, + va="top", + ) + ax.text( + 0.017, + 0.905, + f"Question: {summary['question']['q']} · gold answer: {summary['question']['answer_text']} · answer/random mean Δ {summary['answer_over_random_delta']:.2f}×", + transform=ax.transAxes, + color="#ffc644", + fontsize=10, + weight="bold", + va="top", + ) + ax.text( + 0.99, + 0.02, + f"PCA variance: PC1 {explained[0] * 100:.1f}% · PC2 {explained[1] * 100:.1f}% color: answer/random scar size: answer-mask Δ labels: peak layer depth", + transform=ax.transAxes, + color="#8b97a0", + fontsize=9, + ha="right", + va="bottom", + ) + + cax = fig.add_axes((0.83, 0.865, 0.12, 0.014)) + cb = fig.colorbar(hb, cax=cax, orientation="horizontal") + cb.outline.set_visible(False) + cb.set_ticks([0, 1]) + cb.set_ticklabels(["random-like", "answer scar"]) + cb.ax.tick_params(colors="#cfd6d0", labelsize=8, length=0) + cax.set_facecolor("#07101a") + + for spine in ax.spines.values(): + spine.set_color("#223140") + spine.set_linewidth(1.0) + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_xlim(-0.04, 1.04) + ax.set_ylim(-0.04, 1.04) + ax.grid(color="#2d4255", alpha=0.14, linewidth=0.7) + + tmp = out_dir / ".atlas-panel.png" + fig.savefig(tmp, facecolor=fig.get_facecolor(), transparent=False) + plt.close(fig) + img = Image.open(tmp).convert("RGB") + tmp.unlink(missing_ok=True) + return img + + +def draw_shell(panel: Image.Image, summary: dict, data_dir: Path, out: Path) -> None: + w, h = 2400, 1500 + canvas = Image.new("RGB", (w, h), BG) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 18): + draw.line((0, y, w, y), fill=(6, 10 + y % 11, 16 + y % 9)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-360, -240, 1000, 760), fill=(70, 216, 255, 32)) + gd.ellipse((1270, 210, 2740, 1610), fill=(255, 75, 61, 32)) + gd.ellipse((690, 920, 1740, 1780), fill=(255, 198, 68, 18)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(90))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((74, 48), "SNAPCOMPACT WHITEBOX", fill=AMBER, font=font(24, True)) + draw.text((74, 88), "Activation Atlas of the missing answer", fill=INK, font=font(72, True)) + draw.text((78, 178), "A PCA geography of image-token residual scars: where blanking the gold answer ‘2003’ moves the model differently than a random blank.", fill=MUTED, font=font(27)) + + draw.rounded_rectangle((74, 252, 590, 1390), radius=32, fill=PANEL, outline=(32, 45, 58), width=1) + q = summary["question"] + cols = summary["geometry"]["cols"] + original = Image.open(data_dir / "images" / "original.png").convert("RGB") + answer_mask = Image.open(data_dir / "images" / "answer-mask.png").convert("RGB") + random_mask = Image.open(data_dir / "images" / "random-mask.png").convert("RGB") + strips = [ + ("ORIGINAL", original, CYAN), + ("ANSWER MASK", answer_mask, RED), + ("RANDOM MASK", random_mask, GREEN), + ] + y = 314 + for title, img, color in strips: + draw.text((110, y), title, fill=color, font=font(18, True)) + draw.rounded_rectangle((110, y + 28, 554, y + 166), radius=16, fill=(242, 241, 229), outline=color, width=3) + paste_fit(canvas, crop_answer_strip(img, q["answer_start"], q["answer_end"], cols), (124, y + 42, 540, y + 152)) + y += 226 + + draw.rounded_rectangle((110, 1002, 554, 1300), radius=24, fill=(7, 11, 18), outline=(35, 51, 66), width=1) + metrics = [ + ("answer", q["answer_text"], AMBER, 48), + ("layers", str(summary["layers"]), CYAN, 38), + ("image tokens", str(summary["image_tokens"]), GREEN, 38), + ("answer/random Δ", f"{summary['answer_over_random_delta']:.2f}×", RED, 38), + ] + yy = 1038 + for label, value, color, size in metrics: + draw.text((142, yy), label, fill=MUTED, font=font(16, True)) + draw.text((142, yy + 26), value, fill=color, font=font(size, True)) + yy += 68 + draw.text((110, 1336), "Actual heatmaps.npz + summary.json; no schematic points.", fill=MUTED, font=font(18)) + + draw.rounded_rectangle((622, 252, 2326, 1390), radius=32, fill=PANEL, outline=(32, 45, 58), width=1) + panel = panel.resize((1640, 1098), Image.Resampling.LANCZOS) + canvas.paste(panel, (654, 272)) + + out.parent.mkdir(parents=True, exist_ok=True) + canvas.save(out, quality=95) + + +def write_source_data(out_dir: Path, points: np.ndarray, labels: np.ndarray, ratio_strength: np.ndarray, answer_strength: np.ndarray, peak_layers: np.ndarray, explained: np.ndarray) -> None: + np.savez_compressed( + out_dir / "atlas_source.npz", + points=points, + cluster=labels, + ratio_strength=ratio_strength, + answer_strength=answer_strength, + peak_layer=peak_layers, + pca_explained=explained, + ) + with (out_dir / "atlas_points.csv").open("w", newline="") as f: + writer = csv.writer(f) + writer.writerow(["token", "atlas_x", "atlas_y", "cluster", "ratio_strength", "answer_strength", "peak_layer"]) + for i in range(points.shape[0]): + writer.writerow([i, f"{points[i, 0]:.6f}", f"{points[i, 1]:.6f}", int(labels[i]), f"{ratio_strength[i]:.6f}", f"{answer_strength[i]:.6f}", int(peak_layers[i])]) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--data-dir", default=str(DEFAULT_DATA_DIR)) + ap.add_argument("--out-dir", default=str(DEFAULT_OUT_DIR)) + args = ap.parse_args() + + data_dir = Path(args.data_dir) + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + summary = json.loads((data_dir / "summary.json").read_text()) + heatmaps = np.load(data_dir / "heatmaps.npz") + answer = heatmaps["answer_delta"].astype(np.float32) + random = heatmaps["random_delta"].astype(np.float32) + ratio = heatmaps["ratio"].astype(np.float32) + + contrast = np.log1p(answer) - np.log1p(random) + features = np.concatenate([contrast.T, np.log1p(ratio).T, np.log1p(answer).T], axis=1) + raw_coords, explained = pca2(features) + points = normalize_coords(raw_coords) + labels, centers = kmeans(points, k=5) + + ratio_strength = quantile_norm(np.log1p(ratio).mean(axis=0), 0.02, 0.99) + answer_strength = quantile_norm(np.log1p(answer).mean(axis=0), 0.02, 0.99) + peak_layers = np.argmax(ratio, axis=0).astype(np.int32) + + write_source_data(out_dir, points, labels, ratio_strength, answer_strength, peak_layers, explained) + panel = render_atlas_panel(points, labels, centers, ratio_strength, answer_strength, peak_layers, explained, summary, out_dir) + out = out_dir / "atlas.png" + draw_shell(panel, summary, data_dir, out) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_circuit.py b/packages/snapcompact/research/snapcompact_viz_circuit.py new file mode 100644 index 000000000..f78c70df3 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_circuit.py @@ -0,0 +1,427 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["numpy", "pillow"] +# /// +"""Render a circuit-graph visualization from snapcompact activation deltas.""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Iterable + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +DEFAULT_RESULT_DIR = HERE / "results" / "tensor-heatmap-paddleocr-q7" +DEFAULT_OUT_DIR = HERE / "results" / "agent-viz-circuit" + +BG = (4, 6, 10) +PANEL = (13, 17, 24) +PANEL_2 = (8, 12, 18) +INK = (242, 239, 225) +MUTED = (132, 146, 153) +BLUE = (83, 218, 255) +GOLD = (255, 199, 74) +ORANGE = (255, 122, 54) +RED = (255, 72, 82) +GREEN = (127, 245, 148) + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Supplemental/Helvetica Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Helvetica.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def clamp01(v: float) -> float: + return max(0.0, min(1.0, v)) + + +def mix(a: tuple[int, int, int], b: tuple[int, int, int], t: float) -> tuple[int, int, int]: + t = clamp01(t) + return tuple(round(x + (y - x) * t) for x, y in zip(a, b)) + + +def quantile_norm(values: np.ndarray, q: float = 0.97) -> np.ndarray: + scale = float(np.quantile(values, q)) + if not math.isfinite(scale) or scale <= 0: + scale = 1.0 + return np.clip(values / scale, 0, 1) + + +def rounded_panel(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], radius: int = 34) -> None: + draw.rounded_rectangle(box, radius=radius, fill=PANEL, outline=(31, 41, 51), width=1) + + +def multiline(draw: ImageDraw.ImageDraw, xy: tuple[int, int], text: str, *, fill: tuple[int, int, int], fnt: ImageFont.ImageFont, max_width: int, line_gap: int = 8) -> int: + words = text.split() + lines: list[str] = [] + cur = "" + for word in words: + trial = word if not cur else f"{cur} {word}" + if draw.textlength(trial, font=fnt) <= max_width: + cur = trial + else: + if cur: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + x, y = xy + step = fnt.size + line_gap if hasattr(fnt, "size") else 20 + for line in lines: + draw.text((x, y), line, fill=fill, font=fnt) + y += step + return y + + +def crop_answer_region(img: Image.Image, summary: dict, pad_cells: int = 42) -> Image.Image: + q = summary["question"] + cols = int(summary["geometry"]["cols"]) + rows = int(summary["geometry"]["rows"]) + adv = max(1, img.width // cols) + pitch = max(1, img.height // rows) + start = int(q["answer_start"]) + end = int(q["answer_end"]) + row0 = max(0, start // cols - 5) + row1 = min(rows, (end - 1) // cols + 7) + col0 = max(0, start % cols - pad_cells) + col1 = min(cols, (end - 1) % cols + pad_cells) + if col1 <= col0 + 8: + col1 = min(cols, col0 + 90) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 2) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 2) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=4, outline=RED, width=4) + return crop + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int], *, resample: int = Image.Resampling.LANCZOS) -> None: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + size = (max(1, round(img.width * scale)), max(1, round(img.height * scale))) + resized = img.resize(size, resample) + canvas.paste(resized, (x0 + (x1 - x0 - size[0]) // 2, y0 + (y1 - y0 - size[1]) // 2)) + + +def draw_bezier(draw: ImageDraw.ImageDraw, points: tuple[tuple[float, float], tuple[float, float], tuple[float, float], tuple[float, float]], *, fill: tuple[int, int, int, int], width: int) -> None: + p0, p1, p2, p3 = points + coords: list[tuple[float, float]] = [] + for i in range(46): + t = i / 45 + mt = 1 - t + x = mt**3 * p0[0] + 3 * mt**2 * t * p1[0] + 3 * mt * t**2 * p2[0] + t**3 * p3[0] + y = mt**3 * p0[1] + 3 * mt**2 * t * p1[1] + 3 * mt * t**2 * p2[1] + t**3 * p3[1] + coords.append((x, y)) + draw.line(coords, fill=fill, width=width, joint="curve") + + +def token_groups(grid_side: int = 27, tiles: int = 3) -> list[dict[str, int | str]]: + groups: list[dict[str, int | str]] = [] + names = ["upper-left", "upper", "upper-right", "left", "center", "right", "lower-left", "lower", "lower-right"] + idx = 0 + for gy in range(tiles): + y0 = round(gy * grid_side / tiles) + y1 = round((gy + 1) * grid_side / tiles) + for gx in range(tiles): + x0 = round(gx * grid_side / tiles) + x1 = round((gx + 1) * grid_side / tiles) + groups.append({"name": names[idx], "x0": x0, "x1": x1, "y0": y0, "y1": y1}) + idx += 1 + return groups + + +def group_indices(group: dict[str, int | str], grid_side: int = 27) -> np.ndarray: + ids: list[int] = [] + for y in range(int(group["y0"]), int(group["y1"])): + for x in range(int(group["x0"]), int(group["x1"])): + ids.append(y * grid_side + x) + return np.asarray(ids, dtype=np.int64) + + +def build_metrics(answer: np.ndarray, random: np.ndarray, ratio: np.ndarray) -> tuple[list[dict], list[dict], np.ndarray, np.ndarray]: + layers, tokens = answer.shape + grid_side = int(round(math.sqrt(tokens))) + if grid_side * grid_side != tokens: + raise ValueError(f"expected square image-token grid, got {tokens} tokens") + + groups = token_groups(grid_side) + layer_score = np.maximum(answer - random, 0.0) * np.log1p(np.maximum(ratio, 0.0)) + layer_norm = quantile_norm(layer_score, 0.975) + + group_rows: list[dict] = [] + for i, group in enumerate(groups): + ids = group_indices(group, grid_side) + a = answer[:, ids] + r = random[:, ids] + rr = ratio[:, ids] + raw = layer_score[:, ids] + group_rows.append( + { + "id": i, + "name": group["name"], + "x0": group["x0"], + "x1": group["x1"], + "y0": group["y0"], + "y1": group["y1"], + "answer_delta_mean": float(a.mean()), + "random_delta_mean": float(r.mean()), + "ratio_mean": float(rr.mean()), + "answer_minus_random_mean": float((a - r).mean()), + "edge_score": float(raw.mean()), + "layer_scores": [float(raw[j].mean()) for j in range(layers)], + "layer_ratios": [float(rr[j].mean()) for j in range(layers)], + } + ) + + layer_rows: list[dict] = [] + for layer in range(layers): + layer_rows.append( + { + "layer": layer, + "answer_delta_mean": float(answer[layer].mean()), + "random_delta_mean": float(random[layer].mean()), + "ratio_mean": float(ratio[layer].mean()), + "answer_minus_random_mean": float((answer[layer] - random[layer]).mean()), + "edge_score": float(layer_score[layer].mean()), + } + ) + return group_rows, layer_rows, layer_score, layer_norm + + +def draw_token_grid(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], token_strength: np.ndarray, group_rows: list[dict]) -> list[tuple[int, int]]: + x0, y0, x1, y1 = box + grid = token_strength.reshape(27, 27) + norm = quantile_norm(grid, 0.985) + cell = min((x1 - x0) // 27, (y1 - y0) // 27) + gx = x0 + ((x1 - x0) - 27 * cell) // 2 + gy = y0 + ((y1 - y0) - 27 * cell) // 2 + for y in range(27): + for x in range(27): + v = float(norm[y, x]) + color = mix((13, 24, 34), ORANGE, v) + if v > 0.72: + color = mix(color, GOLD, (v - 0.72) / 0.28) + draw.rectangle((gx + x * cell, gy + y * cell, gx + (x + 1) * cell - 1, gy + (y + 1) * cell - 1), fill=color) + draw.rectangle((gx - 1, gy - 1, gx + 27 * cell, gy + 27 * cell), outline=(70, 88, 101), width=2) + + centers: list[tuple[int, int]] = [] + scores = np.asarray([g["edge_score"] for g in group_rows], dtype=np.float32) + score_norm = quantile_norm(scores, 0.92) + for g, s in zip(group_rows, score_norm): + cx = gx + round((int(g["x0"]) + int(g["x1"])) * 0.5 * cell) + cy = gy + round((int(g["y0"]) + int(g["y1"])) * 0.5 * cell) + centers.append((cx, cy)) + rad = round(8 + 19 * float(s)) + draw.ellipse((cx - rad, cy - rad, cx + rad, cy + rad), outline=mix(BLUE, GOLD, float(s)), width=3) + return centers + + +def draw_layer_bands(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], layer_rows: list[dict]) -> list[tuple[int, int]]: + x0, y0, x1, y1 = box + scores = np.asarray([r["edge_score"] for r in layer_rows], dtype=np.float32) + ratios = np.asarray([r["ratio_mean"] for r in layer_rows], dtype=np.float32) + score_norm = quantile_norm(scores, 0.96) + ratio_norm = quantile_norm(ratios, 0.96) + centers: list[tuple[int, int]] = [] + gap = 7 + h = ((y1 - y0) - gap * (len(layer_rows) - 1)) / len(layer_rows) + for i, (row, s, rr) in enumerate(zip(layer_rows, score_norm, ratio_norm)): + yy0 = round(y0 + i * (h + gap)) + yy1 = round(yy0 + h) + inset = round(26 * (1 - float(s))) + color = mix((17, 27, 38), GOLD, float(rr) * 0.80) + outline = mix((54, 71, 83), RED, float(s)) + draw.rounded_rectangle((x0 + inset, yy0, x1 - inset, yy1), radius=8, fill=color, outline=outline, width=2) + draw.text((x0 - 74, yy0 + max(0, (yy1 - yy0 - 18) // 2)), f"L{int(row['layer']):02d}", fill=mix(MUTED, INK, float(s)), font=font(16, True)) + centers.append(((x0 + x1) // 2, (yy0 + yy1) // 2)) + return centers + + +def render(summary: dict, answer: np.ndarray, random: np.ndarray, ratio: np.ndarray, result_dir: Path, out_dir: Path) -> dict: + group_rows, layer_rows, layer_score, layer_norm = build_metrics(answer, random, ratio) + w, h = 2400, 1350 + canvas = Image.new("RGB", (w, h), BG) + draw = ImageDraw.Draw(canvas) + for y in range(0, h, 18): + shade = 8 + (y // 18) % 4 + draw.line((0, y, w, y), fill=(shade, shade + 2, shade + 7)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-380, -260, 980, 640), fill=(255, 72, 82, 36)) + gd.ellipse((780, 70, 2320, 1420), fill=(83, 218, 255, 22)) + gd.ellipse((1440, -120, 2760, 860), fill=(255, 199, 74, 26)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(90))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((70, 44), "SNAPCOMPACT CIRCUIT TRACE", fill=GOLD, font=font(25, True)) + draw.text((70, 82), "The answer glyphs light a decoder circuit", fill=INK, font=font(66, True)) + multiline( + draw, + (72, 166), + "Edges are computed from actual hidden-state deltas: max(answer − random, 0) × log(1 + answer/random ratio), averaged by image-token region and decoder layer.", + fill=MUTED, + fnt=font(24), + max_width=1470, + line_gap=7, + ) + + rounded_panel(draw, (62, 258, 520, 1238)) + rounded_panel(draw, (568, 258, 1002, 1238)) + rounded_panel(draw, (1126, 258, 1632, 1238)) + rounded_panel(draw, (1816, 258, 2338, 1238)) + + q = summary["question"] + base = Image.open(result_dir / "images" / "original.png").convert("RGB") + masked = Image.open(result_dir / "images" / "answer-mask.png").convert("RGB") + crop = crop_answer_region(base, summary) + masked_crop = crop_answer_region(masked, summary) + + draw.text((94, 294), "1. bitmap intervention", fill=INK, font=font(30, True)) + draw.text((94, 334), "question targets one visible year", fill=MUTED, font=font(18)) + draw.rounded_rectangle((94, 386, 488, 560), radius=16, fill=(240, 238, 224), outline=BLUE, width=3) + paste_fit(canvas, crop, (108, 400, 474, 546), resample=Image.Resampling.NEAREST) + draw.text((94, 574), "original answer region", fill=BLUE, font=font(18, True)) + draw.rounded_rectangle((94, 654, 488, 828), radius=16, fill=(240, 238, 224), outline=RED, width=3) + paste_fit(canvas, masked_crop, (108, 668, 474, 814), resample=Image.Resampling.NEAREST) + draw.text((94, 842), "blanked answer mask", fill=RED, font=font(18, True)) + draw.text((94, 930), "question", fill=MUTED, font=font(15, True)) + multiline(draw, (94, 956), str(q["q"]), fill=INK, fnt=font(23), max_width=370, line_gap=8) + draw.text((94, 1070), "gold answer", fill=MUTED, font=font(15, True)) + draw.text((94, 1098), str(q["answer_text"]), fill=GOLD, font=font(52, True)) + draw.text((94, 1172), f"global Δ ratio {summary['answer_over_random_delta']:.2f}×", fill=INK, font=font(22, True)) + + draw.text((600, 294), "2. image-token regions", fill=INK, font=font(30, True)) + draw.text((600, 334), "27×27 token lattice, colored by circuit score", fill=MUTED, font=font(18)) + token_strength = layer_score.mean(axis=0) + token_centers = draw_token_grid(draw, (616, 392, 954, 730), token_strength, group_rows) + top_groups = sorted(group_rows, key=lambda g: g["edge_score"], reverse=True)[:4] + draw.text((600, 794), "strongest token groups", fill=MUTED, font=font(16, True)) + y = 826 + group_score_norm = quantile_norm(np.asarray([g["edge_score"] for g in group_rows], dtype=np.float32), 0.92) + for g in top_groups: + s = float(group_score_norm[int(g["id"])]) + draw.rounded_rectangle((600, y, 970, y + 62), radius=14, fill=PANEL_2, outline=mix((44, 58, 68), GOLD, s), width=2) + draw.text((620, y + 12), str(g["name"]), fill=INK, font=font(20, True)) + draw.text((820, y + 12), f"{g['ratio_mean']:.2f}×", fill=mix(BLUE, GOLD, s), font=font(21, True)) + draw.text((620, y + 38), f"Δ {g['answer_delta_mean']:.2f} vs {g['random_delta_mean']:.2f}", fill=MUTED, font=font(14)) + y += 78 + + draw.text((1158, 294), "3. decoder layer bands", fill=INK, font=font(30, True)) + draw.text((1158, 334), "band width/color follows per-layer answer specificity", fill=MUTED, font=font(18)) + layer_centers = draw_layer_bands(draw, (1246, 394, 1566, 1122), layer_rows) + + draw.text((1848, 294), "4. output answer", fill=INK, font=font(30, True)) + draw.text((1848, 334), "residual stream converges on text", fill=MUTED, font=font(18)) + draw.rounded_rectangle((1880, 462, 2274, 730), radius=34, fill=(10, 13, 18), outline=(73, 82, 92), width=2) + draw.text((1918, 500), "PaddleOCR-VL", fill=MUTED, font=font(20, True)) + draw.text((1918, 558), "answers", fill=INK, font=font(32, True)) + draw.text((1918, 606), str(q["answer_text"]), fill=GOLD, font=font(82, True)) + draw.rounded_rectangle((1880, 820, 2274, 1034), radius=28, fill=PANEL_2, outline=(47, 62, 73), width=2) + draw.text((1918, 858), f"{summary['layers']} decoder layers", fill=INK, font=font(26, True)) + draw.text((1918, 900), f"{summary['image_tokens']} image tokens", fill=MUTED, font=font(21)) + draw.text((1918, 938), "edge thickness = grouped delta score", fill=MUTED, font=font(21)) + draw.text((1918, 976), "edge color = answer/random ratio", fill=MUTED, font=font(21)) + + # Edges live in a transparent layer so glow can sit behind node labels. + edges = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + ed = ImageDraw.Draw(edges) + all_group_layer = np.asarray([g["layer_scores"] for g in group_rows], dtype=np.float32) + group_layer_norm = quantile_norm(all_group_layer, 0.965) + group_layer_ratio = np.asarray([g["layer_ratios"] for g in group_rows], dtype=np.float32) + ratio_norm = quantile_norm(group_layer_ratio, 0.955) + + selected_groups = [int(g["id"]) for g in sorted(group_rows, key=lambda g: g["edge_score"], reverse=True)[:7]] + selected_layers = [0, 1, 2, 3, 4, 5, 7, 9, 12, 15, 18] + for gi in selected_groups: + sx, sy = token_centers[gi] + for li in selected_layers: + strength = float(group_layer_norm[gi, li]) + if strength < 0.10: + continue + ex, ey = layer_centers[li] + col = mix(BLUE, RED, float(ratio_norm[gi, li])) + alpha = round(54 + 156 * strength) + width = max(1, round(1 + 9 * strength)) + draw_bezier(ed, ((sx + 18, sy), (1046, sy), (1110, ey), (ex - 162, ey)), fill=(*col, alpha), width=width) + + layer_edge_norm = quantile_norm(np.asarray([r["edge_score"] for r in layer_rows], dtype=np.float32), 0.96) + layer_ratio_norm = quantile_norm(np.asarray([r["ratio_mean"] for r in layer_rows], dtype=np.float32), 0.96) + out_anchor = (1880, 596) + for li in selected_layers: + sx, sy = layer_centers[li] + strength = float(layer_edge_norm[li]) + col = mix(GOLD, RED, float(layer_ratio_norm[li])) + width = max(2, round(2 + 11 * strength)) + alpha = round(76 + 160 * strength) + draw_bezier(ed, ((sx + 162, sy), (1668, sy), (1748, out_anchor[1] + (sy - 760) * 0.18), out_anchor), fill=(*col, alpha), width=width) + + edges = edges.filter(ImageFilter.GaussianBlur(0.18)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), edges).convert("RGB") + draw = ImageDraw.Draw(canvas) + + # Connector arrows and legend drawn after glowing edges. + draw.line((520, 748, 568, 748), fill=(70, 84, 96), width=3) + draw.polygon([(568, 748), (550, 738), (550, 758)], fill=(70, 84, 96)) + draw.line((1002, 748, 1126, 748), fill=(70, 84, 96), width=3) + draw.polygon([(1126, 748), (1108, 738), (1108, 758)], fill=(70, 84, 96)) + draw.line((1632, 748, 1816, 748), fill=(70, 84, 96), width=3) + draw.polygon([(1816, 748), (1798, 738), (1798, 758)], fill=(70, 84, 96)) + + legend_x, legend_y = 590, 1168 + draw.text((legend_x, legend_y), "edge encoding", fill=INK, font=font(18, True)) + for i, (lab, val, col) in enumerate([("weak", 0.20, BLUE), ("medium", 0.55, GOLD), ("answer-specific", 0.95, RED)]): + yy = legend_y + 38 + i * 32 + draw.line((legend_x, yy, legend_x + 122, yy), fill=col, width=round(2 + 9 * val)) + draw.text((legend_x + 146, yy - 12), lab, fill=MUTED if i < 2 else INK, font=font(16)) + + draw.text((1158, 1164), "Data: heatmaps.npz answer_delta, random_delta, ratio. No schematic edges: every width/color is grouped from observed tensors.", fill=MUTED, font=font(17)) + + out_dir.mkdir(parents=True, exist_ok=True) + out_png = out_dir / "circuit.png" + canvas.save(out_png, optimize=True) + + source = { + "question": q, + "layers": summary["layers"], + "image_tokens": summary["image_tokens"], + "global_answer_over_random_delta": summary["answer_over_random_delta"], + "edge_formula": "max(answer_delta - random_delta, 0) * log1p(ratio)", + "token_groups": group_rows, + "layers_metrics": layer_rows, + "selected_token_groups": selected_groups, + "selected_layers": selected_layers, + } + (out_dir / "circuit-source-data.json").write_text(json.dumps(source, indent=2)) + return {"png": str(out_png), "source": str(out_dir / "circuit-source-data.json")} + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--result-dir", default=str(DEFAULT_RESULT_DIR)) + ap.add_argument("--out-dir", default=str(DEFAULT_OUT_DIR)) + args = ap.parse_args() + + result_dir = Path(args.result_dir) + out_dir = Path(args.out_dir) + summary = json.loads((result_dir / "summary.json").read_text()) + data = np.load(result_dir / "heatmaps.npz") + paths = render(summary, data["answer_delta"], data["random_delta"], data["ratio"], result_dir, out_dir) + print(paths["png"]) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_city.py b/packages/snapcompact/research/snapcompact_viz_city.py new file mode 100644 index 000000000..5b34f9fb9 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_city.py @@ -0,0 +1,237 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render an isometric activation city from snapcompact heatmap tensors.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Iterable + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +DATA_DIR = HERE / "results" / "tensor-heatmap-paddleocr-q7" +OUT_DIR = HERE / "results" / "agent-viz-city" + +BG = (8, 11, 24) +INK = (232, 238, 255) +MUTED = (136, 148, 184) +ANSWER = (255, 94, 117) +ANSWER_HI = (255, 198, 97) +RANDOM = (72, 201, 255) +RANDOM_HI = (127, 246, 213) +ROAD = (21, 28, 52) +GRID = (39, 48, 83) + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Supplemental/Helvetica Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Helvetica.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + try: + return ImageFont.truetype(path, size) + except OSError: + pass + return ImageFont.load_default() + + +def clamp255(v: float) -> int: + return max(0, min(255, int(round(v)))) + + +def mix(a: tuple[int, int, int], b: tuple[int, int, int], t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + return tuple(clamp255(x + (y - x) * t) for x, y in zip(a, b)) + + +def shade(c: tuple[int, int, int], factor: float) -> tuple[int, int, int]: + return tuple(clamp255(x * factor) for x in c) + + +def iso(x: float, y: float, origin: tuple[float, float], tile_w: float, tile_h: float) -> tuple[float, float]: + ox, oy = origin + return ox + (x - y) * tile_w * 0.5, oy + (x + y) * tile_h * 0.5 + + +def diamond(cx: float, cy: float, tile_w: float, tile_h: float) -> list[tuple[float, float]]: + return [ + (cx, cy - tile_h * 0.5), + (cx + tile_w * 0.5, cy), + (cx, cy + tile_h * 0.5), + (cx - tile_w * 0.5, cy), + ] + + +def building_faces(cx: float, cy: float, h: float, tile_w: float, tile_h: float) -> tuple[list[tuple[float, float]], list[tuple[float, float]], list[tuple[float, float]]]: + top = diamond(cx, cy - h, tile_w, tile_h) + right = [top[1], (cx + tile_w * 0.5, cy), (cx, cy + tile_h * 0.5), top[2]] + left = [top[3], top[2], (cx, cy + tile_h * 0.5), (cx - tile_w * 0.5, cy)] + return top, right, left + + +def draw_soft_line(draw: ImageDraw.ImageDraw, pts: Iterable[tuple[float, float]], fill: tuple[int, int, int], width: int = 1) -> None: + draw.line([(int(x), int(y)) for x, y in pts], fill=fill, width=width) + + +def draw_district( + draw: ImageDraw.ImageDraw, + values: np.ndarray, + ratios: np.ndarray, + origin: tuple[float, float], + base_color: tuple[int, int, int], + high_color: tuple[int, int, int], + label: str, + scale: float, + ratio_scale: float, +) -> None: + layers, bins = values.shape + tile_w = 9.0 + tile_h = 5.0 + max_h = 245.0 + + # Foundation grid and layer streets. + for layer in range(layers): + left = iso(0, layer, origin, tile_w, tile_h) + right = iso(bins - 1, layer, origin, tile_w, tile_h) + draw_soft_line(draw, [left, right], GRID if layer % 3 else (66, 76, 118), 1) + for token in range(0, bins, 10): + near = iso(token, 0, origin, tile_w, tile_h) + far = iso(token, layers - 1, origin, tile_w, tile_h) + draw_soft_line(draw, [near, far], (31, 39, 70), 1) + + # Draw far blocks first, near blocks last. + for layer in range(layers - 1, -1, -1): + for token in range(bins - 1, -1, -1): + v = float(values[layer, token]) + r = float(ratios[layer, token]) + intensity = min(1.0, np.log1p(v) / np.log1p(scale)) + h = 8.0 + (intensity**1.65) * max_h + if v <= 0.0: + h = 3.0 + cx, cy = iso(token, layer, origin, tile_w, tile_h) + top, right, left = building_faces(cx, cy, h, tile_w * 0.92, tile_h * 0.92) + ratio_t = min(1.0, np.log1p(max(r, 0.0)) / np.log1p(ratio_scale)) + c = mix(base_color, high_color, max(intensity * 0.55, ratio_t * 0.85)) + draw.polygon(left, fill=shade(c, 0.42)) + draw.polygon(right, fill=shade(c, 0.62)) + draw.polygon(top, fill=mix(shade(c, 0.95), (255, 255, 255), intensity * 0.20)) + if ratio_t > 0.80 or intensity > 0.90: + draw.line([(int(x), int(y)) for x, y in top + [top[0]]], fill=mix(c, (255, 255, 255), 0.25), width=1) + + # District label plaque. + x0, y0 = iso(-2, layers + 4, origin, tile_w, tile_h) + x1, y1 = iso(54, layers + 4, origin, tile_w, tile_h) + draw.rounded_rectangle((x0 - 26, y0 + 18, x1 + 26, y1 + 64), radius=14, fill=(13, 18, 37), outline=shade(base_color, 0.75), width=2) + draw.text((x0 - 8, y0 + 27), label, font=font(26, True), fill=mix(base_color, high_color, 0.55)) + + +def draw_legend(draw: ImageDraw.ImageDraw, summary: dict[str, object], scale: float) -> None: + draw.text((88, 70), "Snapcompact Activation City", font=font(54, True), fill=INK) + draw.text( + (92, 136), + "729 image tokens → 180 token-bin city blocks · 19 transformer layers → depth streets · building height = activation spike magnitude", + font=font(21), + fill=MUTED, + ) + question = str(summary.get("question", {}).get("q", "")) if isinstance(summary.get("question"), dict) else "" + answer = str(summary.get("question", {}).get("answer_text", "")) if isinstance(summary.get("question"), dict) else "" + draw.text((92, 173), f"Question: {question} Gold answer: {answer}", font=font(20), fill=(180, 190, 220)) + + ratio = float(summary.get("answer_over_random_delta", 0.0)) + draw.rounded_rectangle((1738, 72, 2286, 206), radius=24, fill=(12, 17, 36), outline=(50, 60, 99), width=2) + draw.text((1772, 96), "Answer-mask / random-mask mean delta", font=font(18), fill=MUTED) + draw.text((1772, 125), f"{ratio:.2f}×", font=font(52, True), fill=ANSWER_HI) + draw.text((1906, 149), f"common p98 height scale {scale:.1f}", font=font(17), fill=(176, 186, 218)) + + y = 1400 + draw.rounded_rectangle((88, y, 772, y + 96), radius=18, fill=(12, 17, 36), outline=(44, 54, 92), width=1) + draw.text((116, y + 18), "How to read it", font=font(22, True), fill=INK) + draw.text((116, y + 52), "Tall towers mark token/layer bins where masking changed hidden states most.", font=font(18), fill=MUTED) + draw.rounded_rectangle((836, y, 1520, y + 96), radius=18, fill=(12, 17, 36), outline=(44, 54, 92), width=1) + draw.text((864, y + 18), "Districts", font=font(22, True), fill=INK) + draw.text((864, y + 52), "Warm city = answer mask around “2003”; cool city = same-size random mask.", font=font(18), fill=MUTED) + draw.rounded_rectangle((1584, y, 2268, y + 96), radius=18, fill=(12, 17, 36), outline=(44, 54, 92), width=1) + draw.text((1612, y + 18), "Color halos", font=font(22, True), fill=INK) + draw.text((1612, y + 52), "Bright caps emphasize bins with high answer/random activation ratio.", font=font(18), fill=MUTED) + + +def render() -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + heatmaps = np.load(DATA_DIR / "heatmaps.npz") + summary = json.loads((DATA_DIR / "summary.json").read_text()) + + answer = np.asarray(heatmaps["answer_binned"], dtype=np.float32) + random = np.asarray(heatmaps["random_binned"], dtype=np.float32) + ratio = np.asarray(heatmaps["ratio_binned"], dtype=np.float32) + if answer.shape != random.shape or answer.shape != ratio.shape: + raise ValueError(f"expected matching binned shapes, got {answer.shape}, {random.shape}, {ratio.shape}") + + scale = float(summary.get("common_delta_scale_p98") or np.percentile(np.concatenate([answer.ravel(), random.ravel()]), 98)) + ratio_scale = float(summary.get("ratio_scale_p98") or np.percentile(ratio, 98)) + + w, h = 2400, 1600 + canvas = Image.new("RGB", (w, h), BG) + + # Atmospheric glow layer. + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((90, 250, 1090, 1260), fill=(255, 80, 95, 36)) + gd.ellipse((1220, 250, 2250, 1260), fill=(62, 190, 255, 34)) + gd.rectangle((0, 1240, w, h), fill=(2, 4, 11, 90)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(58))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + # Basemap plates. + draw.rounded_rectangle((52, 244, 1134, 1330), radius=42, fill=(9, 14, 30), outline=(42, 31, 55), width=2) + draw.rounded_rectangle((1234, 244, 2316, 1330), radius=42, fill=(8, 15, 31), outline=(24, 50, 70), width=2) + for y in range(312, 1300, 70): + draw.line((70, y, 1116, y), fill=ROAD, width=1) + draw.line((1252, y, 2298, y), fill=ROAD, width=1) + + draw_district(draw, answer, ratio, (252.0, 870.0), ANSWER, ANSWER_HI, "answer-mask district", scale, ratio_scale) + draw_district(draw, random, ratio, (1434.0, 870.0), RANDOM, RANDOM_HI, "random-mask district", scale, ratio_scale) + draw_legend(draw, summary, scale) + + # Fine vignette frame. + vignette = Image.new("L", (w, h), 0) + vd = ImageDraw.Draw(vignette) + vd.rectangle((0, 0, w, h), fill=255) + vignette = vignette.filter(ImageFilter.GaussianBlur(42)) + frame = Image.new("RGB", (w, h), (0, 0, 0)) + canvas = Image.composite(canvas, frame, vignette) + + np.savez_compressed( + OUT_DIR / "city_source_data.npz", + answer_binned=answer, + random_binned=random, + ratio_binned=ratio, + common_delta_scale_p98=np.array(scale, dtype=np.float32), + ratio_scale_p98=np.array(ratio_scale, dtype=np.float32), + ) + (OUT_DIR / "city_summary.json").write_text( + json.dumps( + { + "source_heatmaps": str(DATA_DIR / "heatmaps.npz"), + "source_summary": str(DATA_DIR / "summary.json"), + "shape": list(answer.shape), + "height_encoding": "log1p(delta) scaled by common_delta_scale_p98", + "districts": {"answer": "answer_binned", "random": "random_binned"}, + "ratio_encoding": "bright caps use ratio_binned / ratio_scale_p98", + "answer_over_random_delta": summary.get("answer_over_random_delta"), + }, + indent=2, + ) + + "\n" + ) + canvas.save(OUT_DIR / "city.png") + + +if __name__ == "__main__": + render() diff --git a/packages/snapcompact/research/snapcompact_viz_explainer.py b/packages/snapcompact/research/snapcompact_viz_explainer.py new file mode 100644 index 000000000..bc89e27f3 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_explainer.py @@ -0,0 +1,444 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render a single-frame snapcompact white-box explainer composite. + +The figure uses the saved tensor heatmaps from the PaddleOCR-VL run. It lays out +four linked stages: the source bitmap, the answer-region mask intervention, the +layer/token hidden-state tensor, and the interpretation stats. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Iterable + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +DEFAULT_DATA = HERE / "results" / "tensor-heatmap-paddleocr-q7" +DEFAULT_OUT = HERE / "results" / "agent-viz-explainer" + +BG = (4, 7, 13) +PANEL = (13, 18, 27) +PANEL_2 = (17, 24, 36) +INK = (244, 243, 231) +MUTED = (137, 153, 166) +DIM = (72, 84, 98) +CYAN = (78, 219, 255) +RED = (255, 82, 65) +AMBER = (255, 197, 78) +GREEN = (129, 255, 136) +VIOLET = (172, 116, 255) +LINE = (37, 50, 64) + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + names = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Supplemental/Helvetica Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Helvetica.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for name in names: + if name and Path(name).exists(): + return ImageFont.truetype(name, size) + return ImageFont.load_default() + + +F10 = font(10) +F12 = font(12) +F14 = font(14) +F16 = font(16) +F18 = font(18) +F20 = font(20) +F22 = font(22, True) +F26 = font(26, True) +F30 = font(30, True) +F38 = font(38, True) +F56 = font(56, True) +F72 = font(72, True) + + +def lerp(a: int, b: int, t: float) -> int: + return round(a + (b - a) * t) + + +def mix(a: tuple[int, int, int], b: tuple[int, int, int], t: float) -> tuple[int, int, int]: + return tuple(lerp(a[i], b[i], t) for i in range(3)) + + +def heat_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, float(t))) + stops = [ + (0.00, (5, 8, 20)), + (0.16, (22, 28, 75)), + (0.34, (81, 40, 125)), + (0.56, (208, 54, 101)), + (0.76, (255, 130, 69)), + (0.91, (255, 210, 94)), + (1.00, (255, 252, 200)), + ] + for (x0, c0), (x1, c1) in zip(stops, stops[1:]): + if t <= x1: + return mix(c0, c1, (t - x0) / (x1 - x0)) + return stops[-1][1] + + +def blue_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, float(t))) + return mix((8, 13, 28), CYAN, t**0.75) + + +def paste_round(base: Image.Image, img: Image.Image, box: tuple[int, int, int, int], radius: int = 24) -> None: + x0, y0, x1, y1 = box + img = img.convert("RGB") + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + resized = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.Resampling.LANCZOS) + px = x0 + (x1 - x0 - resized.width) // 2 + py = y0 + (y1 - y0 - resized.height) // 2 + mask = Image.new("L", resized.size, 0) + ImageDraw.Draw(mask).rounded_rectangle((0, 0, resized.width - 1, resized.height - 1), radius=radius, fill=255) + base.paste(resized, (px, py), mask) + + +def draw_panel(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], title: str, subtitle: str, accent: tuple[int, int, int]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=30, fill=PANEL, outline=LINE, width=2) + draw.rectangle((x0 + 28, y0 + 22, x0 + 84, y0 + 28), fill=accent) + draw.text((x0 + 28, y0 + 43), title, fill=INK, font=F26) + draw.text((x0 + 28, y0 + 78), subtitle, fill=MUTED, font=F16) + + +def draw_arrow(draw: ImageDraw.ImageDraw, start: tuple[int, int], end: tuple[int, int], color: tuple[int, int, int], label: str) -> None: + sx, sy = start + ex, ey = end + draw.line((sx, sy, ex - 18, ey), fill=color, width=5) + draw.polygon([(ex, ey), (ex - 22, ey - 14), (ex - 22, ey + 14)], fill=color) + if not label: + return + tw = round(draw.textlength(label, font=F14)) + draw.rounded_rectangle((sx + 20, sy - 31, sx + 42 + tw, sy - 6), radius=12, fill=(9, 14, 24), outline=mix(color, LINE, 0.35)) + draw.text((sx + 31, sy - 29), label, fill=color, font=F14) + + +def cell_box(summary: dict) -> tuple[int, int, int, int]: + q = summary["question"] + g = summary["geometry"] + cols = int(g["cols"]) + rows = int(g["rows"]) + start = int(q["answer_start"]) + end = int(q["answer_end"]) + cw = 768 / cols + ch = 768 / rows + r0, c0 = divmod(start, cols) + r1, c1 = divmod(max(start, end - 1), cols) + return (math.floor(c0 * cw), math.floor(r0 * ch), math.ceil((c1 + 1) * cw), math.ceil((r1 + 1) * ch)) + + +def answer_crop(img: Image.Image, summary: dict, pad_cells: int = 31) -> tuple[Image.Image, tuple[int, int, int, int]]: + g = summary["geometry"] + cols = int(g["cols"]) + rows = int(g["rows"]) + q = summary["question"] + start = int(q["answer_start"]) + end = int(q["answer_end"]) + row = start // cols + col0 = start % cols + col1 = (end - 1) % cols + 1 + cw = img.width / cols + ch = img.height / rows + x0 = max(0, math.floor((col0 - pad_cells) * cw)) + x1 = min(img.width, math.ceil((col1 + pad_cells) * cw)) + y0 = max(0, math.floor((row - 5) * ch)) + y1 = min(img.height, math.ceil((row + 6) * ch)) + crop = img.crop((x0, y0, x1, y1)).convert("RGB") + local = (round(col0 * cw - x0), round(row * ch - y0), round(col1 * cw - x0), round((row + 1) * ch - y0)) + return crop, local + + +def draw_crop_card(canvas: Image.Image, box: tuple[int, int, int, int], img: Image.Image, local_box: tuple[int, int, int, int], title: str, accent: tuple[int, int, int]) -> None: + draw = ImageDraw.Draw(canvas) + x0, y0, x1, y1 = box + draw.text((x0, y0 - 28), title, fill=accent, font=F16) + draw.rounded_rectangle(box, radius=18, fill=(236, 234, 219), outline=accent, width=3) + pad = 14 + scale = min((x1 - x0 - 2 * pad) / img.width, (y1 - y0 - 2 * pad) / img.height) + resized = img.resize((round(img.width * scale), round(img.height * scale)), Image.Resampling.NEAREST) + px = x0 + (x1 - x0 - resized.width) // 2 + py = y0 + (y1 - y0 - resized.height) // 2 + canvas.paste(resized, (px, py)) + bx = tuple(round(v * scale) for v in local_box) + draw.rounded_rectangle((px + bx[0] - 4, py + bx[1] - 4, px + bx[2] + 4, py + bx[3] + 4), radius=6, outline=accent, width=4) + + +def draw_heatmap(draw: ImageDraw.ImageDraw, arr: np.ndarray, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + rows, cols = arr.shape + cw = (x1 - x0) / cols + ch = (y1 - y0) / rows + for r in range(rows): + ya = round(y0 + r * ch) + yb = round(y0 + (r + 1) * ch) + for c in range(cols): + xa = round(x0 + c * cw) + xb = round(x0 + (c + 1) * cw) + draw.rectangle((xa, ya, xb, yb), fill=heat_color(float(arr[r, c]))) + for r in range(rows + 1): + y = round(y0 + r * ch) + draw.line((x0, y, x1, y), fill=(0, 0, 0, 90) if False else (20, 26, 36), width=1) + draw.rectangle(box, outline=(83, 101, 118), width=1) + + +def draw_tensor_ribbons(draw: ImageDraw.ImageDraw, answer: np.ndarray, random: np.ndarray, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + rows, cols = answer.shape + lane_h = (y1 - y0) / rows + for r in range(rows): + ya = y0 + r * lane_h + yb = y0 + (r + 0.68) * lane_h + for c in range(cols): + xa = x0 + c * (x1 - x0) / cols + xb = x0 + (c + 1) * (x1 - x0) / cols + a = float(answer[r, c]) + rr = float(random[r, c]) + color = heat_color(a) + if rr > a * 0.86: + color = mix(color, (45, 88, 77), min(0.50, rr * 0.45)) + draw.rectangle((round(xa), round(ya), round(xb), round(yb)), fill=color) + if r % 3 == 0: + draw.text((x0 - 33, round(ya + 2)), f"L{r}", fill=MUTED, font=F12) + draw.rectangle(box, outline=(91, 106, 122), width=1) + + +def draw_token_grid(draw: ImageDraw.ImageDraw, ratio: np.ndarray, box: tuple[int, int, int, int]) -> None: + x0, y0, x1, y1 = box + grid = ratio.mean(axis=0).reshape(27, 27) + q98 = float(np.quantile(grid, 0.98)) or 1.0 + norm = np.clip(grid / q98, 0, 1) + cell = min((x1 - x0) / 27, (y1 - y0) / 27) + ox = x0 + ((x1 - x0) - 27 * cell) / 2 + oy = y0 + ((y1 - y0) - 27 * cell) / 2 + for r in range(27): + for c in range(27): + xa = round(ox + c * cell) + ya = round(oy + r * cell) + xb = round(ox + (c + 1) * cell - 1) + yb = round(oy + (r + 1) * cell - 1) + draw.rounded_rectangle((xa, ya, xb, yb), radius=3, fill=blue_color(float(norm[r, c]))) + top = np.unravel_index(np.argsort(grid, axis=None)[-6:], grid.shape) + for r, c in zip(top[0], top[1]): + xa = round(ox + c * cell) + ya = round(oy + r * cell) + draw.rounded_rectangle((xa - 2, ya - 2, round(xa + cell + 1), round(ya + cell + 1)), radius=4, outline=AMBER, width=2) + + +def polyline(draw: ImageDraw.ImageDraw, values: Iterable[float], box: tuple[int, int, int, int], color: tuple[int, int, int], width: int = 4) -> None: + vals = list(values) + x0, y0, x1, y1 = box + lo = min(vals) + hi = max(vals) + span = hi - lo if hi > lo else 1.0 + points = [] + for i, v in enumerate(vals): + x = x0 + i * (x1 - x0) / max(1, len(vals) - 1) + y = y1 - ((v - lo) / span) * (y1 - y0) + points.append((round(x), round(y))) + for i in range(1, len(points)): + draw.line((points[i - 1], points[i]), fill=color, width=width) + for x, y in points: + draw.ellipse((x - 3, y - 3, x + 3, y + 3), fill=color) + + +def metric(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], label: str, value: str, sub: str, accent: tuple[int, int, int]) -> None: + x0, y0, x1, y1 = box + draw.rounded_rectangle(box, radius=20, fill=PANEL_2, outline=mix(accent, LINE, 0.35), width=2) + draw.text((x0 + 18, y0 + 16), label, fill=MUTED, font=F14) + draw.text((x0 + 18, y0 + 41), value, fill=accent, font=F38) + draw.text((x0 + 18, y1 - 32), sub, fill=INK, font=F14) + + +def wrap_text(draw: ImageDraw.ImageDraw, text: str, max_width: int, fnt: ImageFont.ImageFont) -> list[str]: + words = text.split() + lines: list[str] = [] + cur = "" + for word in words: + trial = word if not cur else f"{cur} {word}" + if draw.textlength(trial, font=fnt) <= max_width: + cur = trial + else: + if cur: + lines.append(cur) + cur = word + if cur: + lines.append(cur) + return lines + + +def save_source_metrics(out_dir: Path, summary: dict, arrays: dict[str, np.ndarray]) -> None: + ratio = arrays["ratio"] + answer = arrays["answer_delta"] + random = arrays["random_delta"] + top_cell = np.unravel_index(int(np.argmax(ratio)), ratio.shape) + metrics = { + "question": summary["question"]["q"], + "answer": summary["question"]["answer_text"], + "layers": int(summary["layers"]), + "image_tokens": int(summary["image_tokens"]), + "answer_delta_mean": float(summary["answer_delta_mean"]), + "random_delta_mean": float(summary["random_delta_mean"]), + "answer_over_random_delta": float(summary["answer_over_random_delta"]), + "max_ratio_layer": int(top_cell[0]), + "max_ratio_token": int(top_cell[1]), + "max_ratio": float(ratio[top_cell]), + "mean_answer_by_layer": [float(x) for x in answer.mean(axis=1)], + "mean_random_by_layer": [float(x) for x in random.mean(axis=1)], + "mean_ratio_by_layer": [float(x) for x in ratio.mean(axis=1)], + } + out_dir.mkdir(parents=True, exist_ok=True) + (out_dir / "explainer_metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") + + +def render(data_dir: Path, out_dir: Path) -> Path: + summary = json.loads((data_dir / "summary.json").read_text()) + npz = np.load(data_dir / "heatmaps.npz") + arrays = {name: npz[name] for name in npz.files} + original = Image.open(data_dir / "images" / "original.png").convert("RGB") + answer_mask = Image.open(data_dir / "images" / "answer-mask.png").convert("RGB") + + w, h = 2400, 1500 + canvas = Image.new("RGB", (w, h), BG) + draw = ImageDraw.Draw(canvas) + + # Subtle technical-paper background and activation glows. + for y in range(0, h, 16): + draw.line((0, y, w, y), fill=(6, 10 + (y % 5), 18 + (y % 7))) + for x in range(0, w, 32): + draw.line((x, 0, x, h), fill=(5, 8, 15)) + glow = Image.new("RGBA", (w, h), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-300, -210, 980, 760), fill=(78, 219, 255, 35)) + gd.ellipse((690, 200, 1850, 1370), fill=(255, 82, 65, 32)) + gd.ellipse((1550, -120, 2660, 980), fill=(255, 197, 78, 24)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(90))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((72, 52), "SNAPCOMPACT ACTIVATION EXPLAINER", fill=AMBER, font=F22) + draw.text((72, 86), "From erased pixels to a hidden-state scar", fill=INK, font=F72) + q = summary["question"]["q"] + subtitle = f"Question: {q} · gold answer: {summary['question']['answer_text']}" + draw.text((75, 172), subtitle, fill=MUTED, font=F22) + + # Stage panels. + p1 = (72, 240, 585, 1178) + p2 = (630, 240, 1143, 1178) + p3 = (1188, 240, 1770, 1178) + p4 = (1815, 240, 2328, 1178) + draw_panel(draw, p1, "1 · input bitmap", "Rendered context before intervention", CYAN) + draw_panel(draw, p2, "2 · mask intervention", "Only the true answer cells are blanked", RED) + draw_panel(draw, p3, "3 · hidden-state tensor", "Layer × image-token response", VIOLET) + draw_panel(draw, p4, "4 · interpretation", "Where the answer mattered most", AMBER) + draw_arrow(draw, (585, 715), (630, 715), CYAN, "") + draw_arrow(draw, (1143, 715), (1188, 715), RED, "") + draw_arrow(draw, (1770, 715), (1815, 715), AMBER, "") + + # Input / intervention panels. + paste_round(canvas, original, (108, 352, 548, 792), 24) + full_box = cell_box(summary) + scale = 440 / 768 + ox, oy = 108, 352 + draw.rounded_rectangle((ox + round(full_box[0] * scale), oy + round(full_box[1] * scale), ox + round(full_box[2] * scale), oy + round(full_box[3] * scale)), radius=5, outline=AMBER, width=4) + ocrop, local = answer_crop(original, summary) + mcrop, mlocal = answer_crop(answer_mask, summary) + draw_crop_card(canvas, (108, 885, 548, 1042), ocrop, local, "magnified answer glyphs", AMBER) + draw.text((108, 1083), "The OCR input is a fixed bitmap. The answer span", fill=MUTED, font=F16) + draw.text((108, 1108), f"occupies character cells {summary['question']['answer_start']}–{summary['question']['answer_end'] - 1}.", fill=MUTED, font=F16) + + paste_round(canvas, answer_mask, (666, 352, 1106, 792), 24) + draw.rounded_rectangle((666 + round(full_box[0] * scale), 352 + round(full_box[1] * scale), 666 + round(full_box[2] * scale), 352 + round(full_box[3] * scale)), radius=5, outline=RED, width=4) + draw_crop_card(canvas, (666, 885, 1106, 1042), mcrop, mlocal, "same crop after masking", RED) + draw.text((666, 1083), "Same prompt, same rendered page. Difference:", fill=MUTED, font=F16) + draw.text((666, 1108), "the four answer glyphs are removed before inference.", fill=MUTED, font=F16) + + # Tensor panel. + draw.text((1226, 332), "answer-mask delta", fill=RED, font=F18) + draw.text((1600, 332), "random control mixed in green", fill=GREEN, font=F14) + draw_tensor_ribbons(draw, arrays["answer_norm"], arrays["random_norm"], (1240, 372, 1718, 665)) + draw.text((1238, 686), "answer / random ratio", fill=AMBER, font=F18) + draw.text((1238, 712), "bright = answer-region deletion moves hidden states more than an equal random mask", fill=MUTED, font=F14) + draw_heatmap(draw, arrays["ratio_norm"], (1240, 748, 1718, 1000)) + max_layer = int(summary["max_ratio_layer"]) + draw.line((1240, 748 + round((max_layer + 0.5) * 252 / 19), 1718, 748 + round((max_layer + 0.5) * 252 / 19)), fill=AMBER, width=3) + for i in range(240): + draw.rectangle((1240 + i, 1046, 1241 + i, 1063), fill=heat_color(i / 239)) + draw.text((1240, 1022), "low", fill=MUTED, font=F12) + draw.text((1446, 1022), "high", fill=MUTED, font=F12) + draw.text((1240, 1094), f"{summary['layers']} decoder layers × {summary['image_tokens']} image tokens", fill=INK, font=F20) + draw.text((1240, 1124), "Each cell uses the saved heatmaps.npz tensor values.", fill=MUTED, font=F16) + + # Interpretation panel. + metric(draw, (1850, 344, 2075, 478), "mean delta ratio", f"{summary['answer_over_random_delta']:.2f}×", "answer mask vs control", AMBER) + metric(draw, (2086, 344, 2293, 478), "strongest layer", f"L{summary['max_ratio_layer']}", "mean ratio peak", VIOLET) + metric(draw, (1850, 500, 2075, 634), "answer delta", f"{summary['answer_delta_mean']:.2f}", "mean ||Δh||", RED) + metric(draw, (2086, 500, 2293, 634), "control delta", f"{summary['random_delta_mean']:.2f}", "mean ||Δh||", GREEN) + + draw.text((1852, 684), "layer sensitivity curve", fill=INK, font=F20) + curve_box = (1862, 725, 2290, 858) + draw.rounded_rectangle((1850, 704, 2304, 884), radius=20, fill=PANEL_2, outline=LINE, width=2) + for i in range(5): + y = curve_box[1] + i * (curve_box[3] - curve_box[1]) / 4 + draw.line((curve_box[0], round(y), curve_box[2], round(y)), fill=(31, 42, 54)) + polyline(draw, summary["mean_ratio_by_layer"], curve_box, AMBER, 4) + draw.text((1862, 862), "L0", fill=MUTED, font=F12) + draw.text((2262, 862), f"L{summary['layers'] - 1}", fill=MUTED, font=F12) + + draw.text((1852, 927), "image-token sensitivity field", fill=INK, font=F20) + draw.rounded_rectangle((1850, 955, 2067, 1150), radius=20, fill=PANEL_2, outline=LINE, width=2) + draw_token_grid(draw, arrays["ratio"], (1868, 970, 2049, 1132)) + explanation = "Answer deletion creates a high-ratio band in early layers; later layers diffuse it into surrounding context." + for i, line in enumerate(wrap_text(draw, explanation, 195, F14)): + draw.text((2092, 968 + i * 24), line, fill=INK if i == 0 else MUTED, font=F14) + draw.text((2092, 1090), "Interpretation:", fill=AMBER, font=F16) + draw.text((2092, 1118), "the answer glyphs are not just OCR text;", fill=MUTED, font=F14) + draw.text((2092, 1142), "they perturb the multimodal residual stream.", fill=MUTED, font=F14) + + # Footer with provenance. + footer = (72, 1228, 2328, 1422) + draw.rounded_rectangle(footer, radius=30, fill=(8, 12, 20), outline=LINE, width=2) + draw.text((108, 1266), "Reading the composite", fill=INK, font=F30) + bullets = [ + (CYAN, "Input bitmap", "is the rendered evidence page passed to PaddleOCR-VL."), + (RED, "Mask intervention", "removes only the gold answer span: 2003."), + (VIOLET, "Hidden-state tensor", "plots ||hidden(original) − hidden(masked)|| over saved layer/token arrays."), + (AMBER, "Interpretation", "compares that scar to an equal-size random mask: 2.52× stronger on average."), + ] + x = 108 + for color, head, text in bullets: + draw.rounded_rectangle((x, 1320, x + 500, 1384), radius=18, fill=PANEL_2, outline=mix(color, LINE, 0.35), width=2) + draw.ellipse((x + 18, 1343, x + 36, 1361), fill=color) + draw.text((x + 50, 1330), head, fill=color, font=F16) + draw.text((x + 50, 1355), text, fill=MUTED, font=F14) + x += 545 + + save_source_metrics(out_dir, summary, arrays) + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / "explainer.png" + canvas.save(out_path, optimize=True) + return out_path + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-dir", type=Path, default=DEFAULT_DATA) + parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT) + args = parser.parse_args() + out = render(args.data_dir, args.out_dir) + print(out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_glass_stack.py b/packages/snapcompact/research/snapcompact_viz_glass_stack.py new file mode 100644 index 000000000..29b206203 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_glass_stack.py @@ -0,0 +1,255 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render a perspective glass-stack view of snapcompact activation deltas.""" + +from __future__ import annotations + +import json +import math +from pathlib import Path +from typing import Iterable + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +DATA_DIR = HERE / "results" / "tensor-heatmap-paddleocr-q7" +OUT_DIR = HERE / "results" / "agent-viz-glass-stack" + +INK = (238, 248, 255) +MUTED = (133, 158, 174) +CYAN = (82, 226, 255) +GOLD = (255, 206, 93) +RED = (255, 78, 93) +PANEL = (9, 16, 26) + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Helvetica.ttc", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for candidate in candidates: + if candidate and Path(candidate).exists(): + return ImageFont.truetype(candidate, size) + return ImageFont.load_default() + + +def mix(a: tuple[int, int, int], b: tuple[int, int, int], t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + return tuple(round(a[i] + (b[i] - a[i]) * t) for i in range(3)) + + +def glass_heat(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + stops = [ + (0.00, (11, 24, 48)), + (0.24, (20, 67, 105)), + (0.50, (47, 183, 214)), + (0.72, (255, 91, 116)), + (0.88, (255, 178, 87)), + (1.00, (255, 252, 197)), + ] + for (ta, ca), (tb, cb) in zip(stops, stops[1:]): + if t <= tb: + return mix(ca, cb, (t - ta) / (tb - ta)) + return stops[-1][1] + + +def plane_corners(layer: int) -> tuple[tuple[float, float], tuple[float, float], tuple[float, float], tuple[float, float]]: + x = 225 + layer * 25.0 + y = 920 - layer * 34.0 + width = 930.0 + dx, dy = 190.0, -78.0 + return (x, y), (x + width, y), (x + width + dx, y + dy), (x + dx, y + dy) + + +def bilerp(corners: tuple[tuple[float, float], tuple[float, float], tuple[float, float], tuple[float, float]], u: float, v: float) -> tuple[float, float]: + fl, fr, br, bl = corners + ax = fl[0] + (fr[0] - fl[0]) * u + ay = fl[1] + (fr[1] - fl[1]) * u + bx = bl[0] + (br[0] - bl[0]) * u + by = bl[1] + (br[1] - bl[1]) * u + return ax + (bx - ax) * v, ay + (by - ay) * v + + +def poly(points: Iterable[tuple[float, float]]) -> list[tuple[int, int]]: + return [(round(x), round(y)) for x, y in points] + + +def select_scars(ratio_norm: np.ndarray, answer_binned: np.ndarray, random_binned: np.ndarray, count: int = 7) -> list[int]: + advantage = np.maximum(answer_binned - random_binned, 0.0) + if float(advantage.max(initial=0.0)) > 0: + advantage = advantage / float(np.quantile(advantage, 0.985)) + score = ratio_norm.mean(axis=0) * 0.68 + np.clip(advantage, 0, 1).mean(axis=0) * 0.32 + order = np.argsort(score)[::-1] + chosen: list[int] = [] + for idx in order: + i = int(idx) + if all(abs(i - old) >= 11 for old in chosen): + chosen.append(i) + if len(chosen) == count: + break + return sorted(chosen) + + +def draw_background(canvas: Image.Image) -> None: + draw = ImageDraw.Draw(canvas) + width, height = canvas.size + for y in range(height): + t = y / max(1, height - 1) + draw.line((0, y, width, y), fill=mix((2, 6, 13), (10, 18, 32), t)) + grid = Image.new("RGBA", canvas.size, (0, 0, 0, 0)) + gd = ImageDraw.Draw(grid) + for x in range(-240, width + 240, 48): + gd.line((x, height, x + 620, 0), fill=(65, 145, 190, 16), width=1) + for y in range(92, height, 46): + gd.line((0, y, width, y - 138), fill=(65, 145, 190, 12), width=1) + gd.ellipse((-220, -240, 850, 560), fill=(42, 188, 255, 34)) + gd.ellipse((1120, 420, 2040, 1380), fill=(255, 65, 112, 35)) + canvas.alpha_composite(grid.filter(ImageFilter.GaussianBlur(0.4))) + + +def draw_plane(canvas: Image.Image, values: np.ndarray, layer: int) -> None: + corners = plane_corners(layer) + overlay = Image.new("RGBA", canvas.size, (0, 0, 0, 0)) + draw = ImageDraw.Draw(overlay, "RGBA") + cols = values.shape[0] + + # Blue glass substrate. + draw.polygon(poly(corners), fill=(28, 85, 122, 24), outline=(108, 222, 255, 54)) + + for c, raw in enumerate(values): + u0 = c / cols + u1 = (c + 1) / cols + shade = float(np.clip(raw, 0, 1)) + rgb = glass_heat(shade) + alpha = round(26 + 96 * math.pow(shade, 0.82)) + draw.polygon( + poly((bilerp(corners, u0, 0.03), bilerp(corners, u1, 0.03), bilerp(corners, u1, 0.97), bilerp(corners, u0, 0.97))), + fill=(*rgb, alpha), + ) + + for u in np.linspace(0, 1, 13): + draw.line(poly((bilerp(corners, float(u), 0), bilerp(corners, float(u), 1))), fill=(190, 242, 255, 28), width=1) + for v in np.linspace(0, 1, 5): + draw.line(poly((bilerp(corners, 0, float(v)), bilerp(corners, 1, float(v)))), fill=(190, 242, 255, 24), width=1) + draw.line(poly((corners[0], corners[1], corners[2], corners[3], corners[0])), fill=(174, 241, 255, 70), width=2) + if layer in (0, 6, 12, 18): + x, y = corners[0] + draw.text((round(x - 64), round(y - 10)), f"L{layer:02d}", fill=(178, 226, 239, 150), font=font(16, True)) + canvas.alpha_composite(overlay) + + +def draw_scars(canvas: Image.Image, scar_bins: list[int], ratio_norm: np.ndarray) -> None: + glow = Image.new("RGBA", canvas.size, (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow, "RGBA") + cols = ratio_norm.shape[1] + scar_colors = [(255, 238, 164), (255, 102, 132), (87, 236, 255), (255, 190, 80), (205, 111, 255), (255, 255, 255), (72, 255, 190)] + for n, c in enumerate(scar_bins): + u = (c + 0.5) / cols + pts = [bilerp(plane_corners(layer), u, 0.46) for layer in range(ratio_norm.shape[0])] + color = scar_colors[n % len(scar_colors)] + gd.line(poly(pts), fill=(*color, 120), width=12) + for layer, pt in enumerate(pts): + r = 5 + 12 * float(ratio_norm[layer, c]) + x, y = pt + gd.ellipse((x - r, y - r, x + r, y + r), fill=(*color, 90)) + canvas.alpha_composite(glow.filter(ImageFilter.GaussianBlur(13))) + + draw = ImageDraw.Draw(canvas, "RGBA") + for n, c in enumerate(scar_bins): + u = (c + 0.5) / cols + pts = [bilerp(plane_corners(layer), u, 0.46) for layer in range(ratio_norm.shape[0])] + color = scar_colors[n % len(scar_colors)] + draw.line(poly(pts), fill=(*color, 235), width=3) + top = pts[-1] + draw.text((round(top[0] + 10), round(top[1] - 14)), f"bin {c}", fill=(*color, 220), font=font(13, True)) + for layer, pt in enumerate(pts): + r = 2.0 + 4.5 * float(ratio_norm[layer, c]) + x, y = pt + draw.ellipse((x - r, y - r, x + r, y + r), fill=(255, 255, 230, 225), outline=(*color, 255), width=1) + + +def draw_labels(canvas: Image.Image, summary: dict, scar_bins: list[int], ratio_binned: np.ndarray) -> None: + draw = ImageDraw.Draw(canvas, "RGBA") + q = summary["question"]["q"] + answer = summary["question"]["answer_text"] + draw.text((70, 54), "SNAPCOMPACT GLASS STACK", fill=GOLD, font=font(22, True)) + draw.text((70, 88), "Answer-mask scars through decoder depth", fill=INK, font=font(54, True)) + draw.text((73, 154), f"Question: {q} · gold answer: {answer}", fill=MUTED, font=font(22)) + + x0, y0, x1, y1 = 70, 960, 770, 1110 + draw.rounded_rectangle((x0, y0, x1, y1), radius=24, fill=(7, 13, 22, 205), outline=(115, 217, 255, 72), width=1) + ratio = summary["answer_over_random_delta"] + draw.text((x0 + 26, y0 + 22), f"{ratio:.2f}×", fill=GOLD, font=font(48, True)) + draw.text((x0 + 170, y0 + 31), "mean answer-mask / random-mask delta", fill=INK, font=font(22, True)) + draw.text((x0 + 28, y0 + 86), f"{summary['layers']} semi-transparent decoder planes · {summary['image_tokens']} image tokens binned to {ratio_binned.shape[1]} columns", fill=MUTED, font=font(18)) + + lx0, ly0 = 1240, 930 + draw.rounded_rectangle((lx0, ly0, lx0 + 475, ly0 + 182), radius=24, fill=(7, 13, 22, 210), outline=(115, 217, 255, 70), width=1) + draw.text((lx0 + 24, ly0 + 22), "encoding", fill=INK, font=font(25, True)) + draw.text((lx0 + 24, ly0 + 61), "plane color = answer/random ratio", fill=MUTED, font=font(18)) + draw.text((lx0 + 24, ly0 + 92), "vertical scar = high-ratio token bin", fill=MUTED, font=font(18)) + draw.text((lx0 + 24, ly0 + 124), "selected bins: " + ", ".join(map(str, scar_bins)), fill=(203, 231, 240), font=font(17)) + + # Color ramp. + for i in range(220): + draw.rectangle((lx0 + 230 + i, ly0 + 30, lx0 + 231 + i, ly0 + 49), fill=(*glass_heat(i / 219), 255)) + draw.text((lx0 + 230, ly0 + 54), "low", fill=MUTED, font=font(13)) + draw.text((lx0 + 417, ly0 + 54), "high", fill=MUTED, font=font(13)) + + +def main() -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + with (DATA_DIR / "summary.json").open("r", encoding="utf-8") as f: + summary = json.load(f) + heat = np.load(DATA_DIR / "heatmaps.npz") + ratio_norm = heat["ratio_norm"].astype(np.float32, copy=False) + ratio_binned = heat["ratio_binned"].astype(np.float32, copy=False) + answer_binned = heat["answer_binned"].astype(np.float32, copy=False) + random_binned = heat["random_binned"].astype(np.float32, copy=False) + scar_bins = select_scars(ratio_norm, answer_binned, random_binned) + + canvas = Image.new("RGBA", (1800, 1200), (0, 0, 0, 255)) + draw_background(canvas) + + # Paint upper layers first, lower layers last, so the stack reads as transparent sheets in perspective. + for layer in range(ratio_norm.shape[0] - 1, -1, -1): + draw_plane(canvas, ratio_norm[layer], layer) + draw_scars(canvas, scar_bins, ratio_norm) + draw_labels(canvas, summary, scar_bins, ratio_binned) + + ImageDraw.Draw(canvas).rounded_rectangle((42, 36, 1760, 1142), radius=38, outline=(128, 225, 255, 44), width=2) + out_path = OUT_DIR / "glass-stack.png" + canvas.convert("RGB").save(out_path, quality=95) + + np.savez_compressed( + OUT_DIR / "glass-stack-source-data.npz", + ratio_binned=ratio_binned, + ratio_norm=ratio_norm, + answer_minus_random_binned=answer_binned - random_binned, + scar_bins=np.array(scar_bins, dtype=np.int16), + ) + with (OUT_DIR / "glass-stack-source-summary.json").open("w", encoding="utf-8") as f: + json.dump( + { + "source_heatmaps": str(DATA_DIR / "heatmaps.npz"), + "source_summary": str(DATA_DIR / "summary.json"), + "output": str(out_path), + "scar_bins": scar_bins, + "answer_over_random_delta": summary["answer_over_random_delta"], + "layers": summary["layers"], + "image_tokens": summary["image_tokens"], + }, + f, + indent=2, + ) + print(out_path) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_glyph_matrix.py b/packages/snapcompact/research/snapcompact_viz_glyph_matrix.py new file mode 100644 index 000000000..77dd63f19 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_glyph_matrix.py @@ -0,0 +1,397 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render a glyph-to-activation matrix for the snapcompact answer scar. + +The figure keeps the original OCR bitmap visible, then projects the 27x27 image-token +activation field from heatmaps.npz back onto that bitmap. The answer glyphs are +outlined in text-cell coordinates; high-scar image tokens are outlined in model-token +coordinates; side bars show how the same region changes across decoder layers. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Iterable + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +DEFAULT_SOURCE = HERE / "results" / "tensor-heatmap-paddleocr-q7" +DEFAULT_OUT = HERE / "results" / "agent-viz-glyph-matrix" + +PALETTE = { + "bg": (3, 6, 12), + "panel": (10, 15, 22), + "panel2": (13, 21, 30), + "ink": (244, 245, 232), + "muted": (139, 154, 166), + "grid": (37, 53, 66), + "cyan": (86, 224, 255), + "red": (255, 78, 69), + "amber": (255, 202, 82), + "green": (135, 255, 159), + "violet": (173, 116, 255), +} + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Supplemental/Helvetica.ttc", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for path in candidates: + if path and Path(path).exists(): + return ImageFont.truetype(path, size) + return ImageFont.load_default() + + +def mix(a: tuple[int, int, int], b: tuple[int, int, int], t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + return tuple(round(a[i] + (b[i] - a[i]) * t) for i in range(3)) + + +def heat_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, t)) + stops = [ + (0.00, (8, 10, 26)), + (0.22, (33, 25, 81)), + (0.46, (117, 41, 117)), + (0.68, (224, 67, 79)), + (0.86, (255, 165, 73)), + (1.00, (255, 243, 174)), + ] + for (ta, ca), (tb, cb) in zip(stops, stops[1:]): + if t <= tb: + return mix(ca, cb, (t - ta) / (tb - ta)) + return stops[-1][1] + + +def quantile_norm(values: np.ndarray, q: float = 0.98) -> np.ndarray: + scale = float(np.quantile(values, q)) if values.size else 1.0 + if scale <= 0 or not math.isfinite(scale): + scale = 1.0 + return np.clip(values / scale, 0.0, 1.0) + + +def rounded_panel(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], title: str | None = None, subtitle: str | None = None) -> None: + draw.rounded_rectangle(box, radius=24, fill=PALETTE["panel"], outline=(34, 48, 61), width=1) + if title: + draw.text((box[0] + 24, box[1] + 18), title, fill=PALETTE["ink"], font=font(28, True)) + if subtitle: + draw.text((box[0] + 24, box[1] + 54), subtitle, fill=PALETTE["muted"], font=font(16)) + + +def token_boxes(side: int, grid: int) -> list[tuple[int, int, int, int]]: + boxes = [] + for idx in range(grid * grid): + r, c = divmod(idx, grid) + x0 = round(c * side / grid) + y0 = round(r * side / grid) + x1 = round((c + 1) * side / grid) + y1 = round((r + 1) * side / grid) + boxes.append((x0, y0, x1, y1)) + return boxes + + +def intersect_area(a: tuple[float, float, float, float], b: tuple[float, float, float, float]) -> float: + x0 = max(a[0], b[0]) + y0 = max(a[1], b[1]) + x1 = min(a[2], b[2]) + y1 = min(a[3], b[3]) + return max(0.0, x1 - x0) * max(0.0, y1 - y0) + + +def answer_bbox(start: int, end: int, cols: int, adv: int, pitch: int) -> tuple[int, int, int, int]: + row0, col0 = divmod(start, cols) + row1, col1 = divmod(max(start, end - 1), cols) + x0 = max(0, col0 * adv) + y0 = max(0, row0 * pitch) + x1 = (col1 + 1) * adv + y1 = (row1 + 1) * pitch + return x0, y0, x1, y1 + + +def draw_text_wrapped(draw: ImageDraw.ImageDraw, xy: tuple[int, int], text: str, width: int, fill: tuple[int, int, int], size: int, bold: bool = False, line_gap: int = 4) -> int: + words = text.split() + lines: list[str] = [] + current = "" + f = font(size, bold) + for word in words: + candidate = word if not current else f"{current} {word}" + if draw.textbbox((0, 0), candidate, font=f)[2] <= width or not current: + current = candidate + else: + lines.append(current) + current = word + if current: + lines.append(current) + x, y = xy + for line in lines: + draw.text((x, y), line, fill=fill, font=f) + y += size + line_gap + return y + + +def paste_shadowed(canvas: Image.Image, img: Image.Image, xy: tuple[int, int]) -> None: + shadow = Image.new("RGBA", img.size, (0, 0, 0, 0)) + alpha = Image.new("L", img.size, 180) + shadow.putalpha(alpha) + canvas.alpha_composite(shadow.filter(ImageFilter.GaussianBlur(12)), (xy[0] + 8, xy[1] + 10)) + canvas.alpha_composite(img, xy) + + +def draw_activation_overlay( + original: Image.Image, + answer_mask: Image.Image, + token_score: np.ndarray, + top_tokens: Iterable[int], + answer_tokens: Iterable[int], + bbox: tuple[int, int, int, int], +) -> Image.Image: + side = original.width + grid = int(round(math.sqrt(token_score.size))) + if grid * grid != token_score.size: + raise ValueError(f"expected square image token grid, got {token_score.size}") + base = original.convert("RGBA") + tint = Image.new("RGBA", base.size, (0, 0, 0, 0)) + td = ImageDraw.Draw(tint) + boxes = token_boxes(side, grid) + for idx, box in enumerate(boxes): + t = float(token_score[idx]) + if t <= 0: + continue + r, g, b = heat_color(t) + alpha = round(25 + 142 * t) + td.rectangle(box, fill=(r, g, b, alpha)) + composite = Image.alpha_composite(base, tint) + draw = ImageDraw.Draw(composite) + + for idx in top_tokens: + box = boxes[int(idx)] + draw.rounded_rectangle(box, radius=3, outline=PALETTE["amber"] + (235,), width=3) + for idx in answer_tokens: + box = boxes[int(idx)] + draw.rounded_rectangle(box, radius=4, outline=PALETTE["cyan"] + (245,), width=4) + + glow = Image.new("RGBA", composite.size, (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + for w, a in ((16, 46), (9, 80), (4, 235)): + gd.rounded_rectangle((bbox[0] - 8, bbox[1] - 7, bbox[2] + 8, bbox[3] + 8), radius=8, outline=PALETTE["red"] + (a,), width=w) + composite = Image.alpha_composite(composite, glow.filter(ImageFilter.GaussianBlur(4))) + draw = ImageDraw.Draw(composite) + draw.rounded_rectangle((bbox[0] - 8, bbox[1] - 7, bbox[2] + 8, bbox[3] + 8), radius=8, outline=PALETTE["red"] + (255,), width=3) + + mask_delta = Image.blend(original.convert("RGB"), answer_mask.convert("RGB"), 0.42).convert("RGBA") + crop = mask_delta.crop((max(0, bbox[0] - 76), max(0, bbox[1] - 42), min(side, bbox[2] + 154), min(side, bbox[3] + 48))) + crop = crop.resize((crop.width * 3, crop.height * 3), Image.Resampling.NEAREST) + crop_draw = ImageDraw.Draw(crop) + scale = 3 + cx0 = (bbox[0] - max(0, bbox[0] - 76)) * scale + cy0 = (bbox[1] - max(0, bbox[1] - 42)) * scale + cx1 = (bbox[2] - max(0, bbox[0] - 76)) * scale + cy1 = (bbox[3] - max(0, bbox[1] - 42)) * scale + crop_draw.rounded_rectangle((cx0 - 4, cy0 - 4, cx1 + 4, cy1 + 4), radius=8, outline=PALETTE["red"] + (255,), width=5) + composite.alpha_composite(crop, (side - crop.width - 20, 20)) + draw = ImageDraw.Draw(composite) + draw.text((side - crop.width - 16, 20 + crop.height + 8), "answer glyph crop: original → masked", fill=PALETTE["ink"] + (235,), font=font(18, True)) + return composite + + +def draw_layer_bars( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], + answer_layer: np.ndarray, + random_layer: np.ndarray, + answer_region_layer: np.ndarray, + ratio_layer: np.ndarray, +) -> None: + rounded_panel(draw, box, "layer-by-layer scar", "red = answer mask, green = equal random mask, cyan = answer glyph tokens") + x0, y0, x1, y1 = box + chart = (x0 + 74, y0 + 103, x1 - 34, y1 - 72) + rows = answer_layer.size + row_h = (chart[3] - chart[1]) / rows + scale = float(np.quantile(np.concatenate([answer_layer, random_layer, answer_region_layer]), 0.96)) + scale = max(scale, 1e-6) + for i in range(rows): + y = chart[1] + i * row_h + draw.text((x0 + 28, round(y + row_h * 0.18)), f"L{i:02d}", fill=PALETTE["muted"], font=font(12)) + max_w = chart[2] - chart[0] + aw = round(max_w * min(1.0, float(answer_layer[i]) / scale)) + rw = round(max_w * min(1.0, float(random_layer[i]) / scale)) + gw = round(max_w * min(1.0, float(answer_region_layer[i]) / scale)) + yy = round(y) + draw.rounded_rectangle((chart[0], yy + 2, chart[0] + aw, yy + 8), radius=3, fill=PALETTE["red"]) + draw.rounded_rectangle((chart[0], yy + 11, chart[0] + rw, yy + 17), radius=3, fill=PALETTE["green"]) + draw.rounded_rectangle((chart[0], yy + 20, chart[0] + gw, yy + 27), radius=3, fill=PALETTE["cyan"]) + ratio = float(ratio_layer[i]) + draw.text((chart[2] - 58, yy + 8), f"{ratio:4.1f}×", fill=PALETTE["amber"], font=font(13, True)) + draw.text((chart[0], y1 - 45), "Mean delta per decoder layer. Ratio labels compare answer-mask vs random-mask deltas.", fill=PALETTE["muted"], font=font(14)) + + +def draw_scar_strip(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], ratio_norm: np.ndarray, answer_tokens: list[int], top_tokens: list[int]) -> None: + rounded_panel(draw, box, "token scar matrix", "decoder layers × image tokens; vertical lines locate answer glyphs and top scar bins") + x0, y0, x1, y1 = box + hx0, hy0, hx1, hy1 = x0 + 58, y0 + 90, x1 - 28, y1 - 54 + rows, cols = ratio_norm.shape + cw = (hx1 - hx0) / cols + ch = (hy1 - hy0) / rows + for r in range(rows): + ya = round(hy0 + r * ch) + yb = round(hy0 + (r + 1) * ch) + for c in range(cols): + xa = round(hx0 + c * cw) + xb = round(hx0 + (c + 1) * cw) + draw.rectangle((xa, ya, xb, yb), fill=heat_color(float(ratio_norm[r, c]))) + for tok in answer_tokens: + x = round(hx0 + (tok + 0.5) * (hx1 - hx0) / 729) + draw.line((x, hy0 - 8, x, hy1 + 8), fill=PALETTE["cyan"], width=2) + for tok in top_tokens[:12]: + x = round(hx0 + (tok + 0.5) * (hx1 - hx0) / 729) + draw.line((x, hy0, x, hy1), fill=PALETTE["amber"], width=1) + for r in range(0, rows, 4): + y = round(hy0 + (r + 0.5) * ch) + draw.text((x0 + 22, y - 7), str(r), fill=PALETTE["muted"], font=font(12)) + draw.text((hx0, y1 - 32), "image-token sequence →", fill=PALETTE["muted"], font=font(13)) + + +def draw_top_token_table(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], top_tokens: list[int], ratio_mean: np.ndarray, answer_mean: np.ndarray, grid: int) -> None: + rounded_panel(draw, box, "highest-scar token bins", "actual heatmaps.npz token IDs") + x0, y0, _, y1 = box + y = y0 + 90 + row_gap = 26 + max_rows = max(1, min(8, (y1 - y - 18) // row_gap)) + bar_scale = max(1e-6, float(np.quantile(answer_mean, 0.98))) + for rank, tok in enumerate(top_tokens[:max_rows], start=1): + r, c = divmod(tok, grid) + draw.text((x0 + 28, y), f"{rank:02d}", fill=PALETTE["amber"], font=font(13, True)) + draw.text((x0 + 68, y), f"token {tok:03d}", fill=PALETTE["ink"], font=font(14, True)) + draw.text((x0 + 164, y), f"grid r{r:02d} c{c:02d}", fill=PALETTE["muted"], font=font(13)) + draw.text((x0 + 292, y), f"ratio {ratio_mean[tok]:.2f}×", fill=PALETTE["cyan"], font=font(13, True)) + bar_w = round(112 * min(1.0, float(answer_mean[tok]) / bar_scale)) + draw.rounded_rectangle((x0 + 408, y + 4, x0 + 408 + bar_w, y + 14), radius=4, fill=PALETTE["red"]) + y += row_gap + + +def render(source: Path, out_dir: Path) -> None: + summary = json.loads((source / "summary.json").read_text()) + data = np.load(source / "heatmaps.npz") + original = Image.open(source / "images" / "original.png").convert("RGB") + answer_mask = Image.open(source / "images" / "answer-mask.png").convert("RGB") + + answer_delta = data["answer_delta"].astype(np.float32) + random_delta = data["random_delta"].astype(np.float32) + ratio = data["ratio"].astype(np.float32) + ratio_binned = data["ratio_binned"].astype(np.float32) + ratio_norm_binned = data["ratio_norm"].astype(np.float32) + + q = summary["question"] + geom = summary["geometry"] + cols = int(geom["cols"]) + rows = int(geom["rows"]) + adv = original.width // cols + pitch = max(1, original.height // rows) + bbox = answer_bbox(int(q["answer_start"]), int(q["answer_end"]), cols, adv, pitch) + + token_count = answer_delta.shape[1] + grid = int(round(math.sqrt(token_count))) + boxes = token_boxes(original.width, grid) + answer_area = (bbox[0], bbox[1], bbox[2], bbox[3]) + answer_tokens = [i for i, b in enumerate(boxes) if intersect_area(answer_area, b) > 0] + if not answer_tokens: + center_x = (bbox[0] + bbox[2]) / 2 + center_y = (bbox[1] + bbox[3]) / 2 + answer_tokens = [min(token_count - 1, max(0, int(center_y / original.height * grid) * grid + int(center_x / original.width * grid)))] + + ratio_mean = ratio.mean(axis=0) + answer_mean = answer_delta.mean(axis=0) + token_score = quantile_norm(ratio_mean, 0.985) + answer_set = set(answer_tokens) + top_tokens = [int(i) for i in np.argsort(ratio_mean)[::-1] if int(i) not in answer_set][:24] + answer_region_layer = answer_delta[:, answer_tokens].mean(axis=1) + answer_layer = answer_delta.mean(axis=1) + random_layer = random_delta.mean(axis=1) + ratio_layer = answer_layer / np.maximum(random_layer, 1e-6) + + out_dir.mkdir(parents=True, exist_ok=True) + overlay = draw_activation_overlay(original, answer_mask, token_score, top_tokens[:18], answer_tokens, bbox) + overlay = overlay.resize((760, 760), Image.Resampling.LANCZOS) + + W, H = 1900, 1260 + canvas = Image.new("RGBA", (W, H), PALETTE["bg"] + (255,)) + draw = ImageDraw.Draw(canvas) + for y in range(0, H, 16): + draw.line((0, y, W, y), fill=(6, 11 + y % 17, 18 + y % 11, 255)) + glow = Image.new("RGBA", (W, H), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-220, 120, 900, 1040), fill=PALETTE["red"] + (28,)) + gd.ellipse((820, -260, 2100, 820), fill=PALETTE["cyan"] + (26,)) + gd.ellipse((980, 650, 2050, 1510), fill=PALETTE["violet"] + (20,)) + canvas = Image.alpha_composite(canvas, glow.filter(ImageFilter.GaussianBlur(78))) + draw = ImageDraw.Draw(canvas) + + draw.text((64, 42), "SNAPCOMPACT GLYPH MATRIX", fill=PALETTE["amber"], font=font(22, True)) + draw.text((64, 78), "The answer glyphs leave a hidden activation scar", fill=PALETTE["ink"], font=font(56, True)) + subtitle = "Original dense text bitmap, overlaid with answer/random activation ratios from 19 decoder layers × 729 image tokens." + draw.text((68, 145), subtitle, fill=PALETTE["muted"], font=font(22)) + + rounded_panel(draw, (52, 205, 862, 1066), "visible glyphs ↔ hidden tokens", "red box = actual answer cells; cyan = intersecting image tokens; amber = top scar bins") + paste_shadowed(canvas, overlay, (78, 282)) + draw.text((82, 1085), f"Question: {q['q']}", fill=PALETTE["ink"], font=font(21, True)) + draw.text((82, 1120), f"Gold answer: {q['answer_text']} · cells {q['answer_start']}–{q['answer_end'] - 1}", fill=PALETTE["amber"], font=font(24, True)) + draw.text((82, 1160), f"Answer/random mean delta: {summary['answer_over_random_delta']:.2f}×", fill=PALETTE["cyan"], font=font(22, True)) + + draw_layer_bars(draw, (900, 205, 1838, 628), answer_layer, random_layer, answer_region_layer, ratio_layer) + draw_scar_strip(draw, (900, 662, 1838, 930), ratio_norm_binned, answer_tokens, top_tokens) + draw_top_token_table(draw, (900, 964, 1838, 1196), top_tokens, ratio_mean, answer_mean, grid) + + for i in range(240): + draw.rectangle((1568 + i, 156, 1569 + i, 174), fill=heat_color(i / 239)) + draw.text((1568, 132), "activation ratio", fill=PALETTE["muted"], font=font(13)) + draw.text((1568, 179), "low", fill=PALETTE["muted"], font=font(12)) + draw.text((1776, 179), "high", fill=PALETTE["muted"], font=font(12)) + + png = out_dir / "glyph-matrix.png" + canvas.convert("RGB").save(png, quality=96) + + source_data = { + "source": str(source), + "question": q, + "geometry": {"text_cols": cols, "text_rows": rows, "glyph_adv": adv, "glyph_pitch": pitch, "image_token_grid": [grid, grid]}, + "answer_bbox_pixels": list(map(int, bbox)), + "answer_image_tokens": [int(x) for x in answer_tokens], + "top_scar_tokens": [ + { + "token": int(tok), + "row": int(tok // grid), + "col": int(tok % grid), + "mean_ratio": float(ratio_mean[tok]), + "mean_answer_delta": float(answer_mean[tok]), + } + for tok in top_tokens[:24] + ], + "mean_answer_delta": float(answer_delta.mean()), + "mean_random_delta": float(random_delta.mean()), + "answer_over_random_delta": float(summary["answer_over_random_delta"]), + "max_ratio_binned": float(ratio_binned.max()), + } + (out_dir / "glyph-matrix-data.json").write_text(json.dumps(source_data, indent=2)) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--source", type=Path, default=DEFAULT_SOURCE) + parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT) + args = parser.parse_args() + render(args.source, args.out_dir) + print(args.out_dir / "glyph-matrix.png") + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_radial.py b/packages/snapcompact/research/snapcompact_viz_radial.py new file mode 100644 index 000000000..a5f477d8c --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_radial.py @@ -0,0 +1,219 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["matplotlib", "numpy"] +# /// +"""Render a radial sonar view of snapcompact answer/random activation echoes.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.colors as mcolors +import matplotlib.patheffects as pe +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.patches import Wedge + +HERE = Path(__file__).resolve().parent +DATA_DIR = HERE / "results" / "tensor-heatmap-paddleocr-q7" +OUT_DIR = HERE / "results" / "agent-viz-radial" + +BG = "#02060a" +GRID = "#2cf5d044" +CYAN = "#38f4ff" +GREEN = "#81ffb4" +AMBER = "#ffc247" +RED = "#ff4d42" +INK = "#f3f1dd" +MUTED = "#8da0a8" + + +def robust_norm(values: np.ndarray, q: float = 0.975) -> np.ndarray: + scale = float(np.quantile(values, q)) + if not np.isfinite(scale) or scale <= 0: + scale = float(np.max(values)) or 1.0 + return np.clip(values / scale, 0.0, 1.0) + + +def polar_edges(cols: int, layers: int) -> tuple[np.ndarray, np.ndarray]: + theta = np.linspace(0.0, 2.0 * np.pi, cols + 1) + radius = np.arange(layers + 1, dtype=np.float32) + 1.0 + return theta, radius + + +def radar_cmap() -> mcolors.LinearSegmentedColormap: + colors = [ + (0.00, "#02060a"), + (0.10, "#03241f"), + (0.32, "#08705e"), + (0.55, "#16f0be"), + (0.76, "#fff06a"), + (1.00, "#fff8e0"), + ] + return mcolors.LinearSegmentedColormap.from_list("snapcompact_radar", colors) + + +def top_echoes(ratio: np.ndarray, answer: np.ndarray, random: np.ndarray, limit: int = 18) -> list[dict[str, float | int]]: + flat = np.argpartition(ratio.ravel(), -limit)[-limit:] + flat = flat[np.argsort(ratio.ravel()[flat])[::-1]] + rows: list[dict[str, float | int]] = [] + for idx in flat: + layer, bin_idx = np.unravel_index(int(idx), ratio.shape) + rows.append( + { + "rank": len(rows) + 1, + "layer": int(layer), + "bin": int(bin_idx), + "angle_degrees": round(float((bin_idx + 0.5) * 360.0 / ratio.shape[1]), 2), + "answer_delta": round(float(answer[layer, bin_idx]), 4), + "random_delta": round(float(random[layer, bin_idx]), 4), + "answer_random_ratio": round(float(ratio[layer, bin_idx]), 4), + } + ) + return rows + + +def add_glow_spikes(ax: plt.Axes, ratio: np.ndarray, norm_ratio: np.ndarray) -> None: + layers, bins = ratio.shape + theta_centers = (np.arange(bins) + 0.5) * 2.0 * np.pi / bins + threshold = float(np.quantile(norm_ratio, 0.91)) + for layer in range(layers): + active = np.flatnonzero(norm_ratio[layer] >= threshold) + if active.size == 0: + active = np.argpartition(norm_ratio[layer], -3)[-3:] + for idx in active: + v = float(norm_ratio[layer, idx]) + base_r = layer + 1.18 + tip_r = base_r + 0.12 + 0.58 * v + theta = float(theta_centers[idx]) + color = AMBER if v > 0.78 else CYAN + ax.plot([theta, theta], [base_r, tip_r], color=color, linewidth=0.7 + 1.8 * v, alpha=0.30 + 0.55 * v) + ax.scatter([theta], [tip_r], s=5 + 28 * v, color=color, alpha=0.26 + 0.55 * v, linewidths=0) + + +def draw_radial(summary: dict, answer: np.ndarray, random: np.ndarray, ratio: np.ndarray) -> plt.Figure: + layers, bins = ratio.shape + norm_ratio = robust_norm(ratio, 0.972) + theta_edges, radius_edges = polar_edges(bins, layers) + theta_grid, radius_grid = np.meshgrid(theta_edges, radius_edges) + + fig = plt.figure(figsize=(16, 10), dpi=180, facecolor=BG) + ax = fig.add_axes([0.04, 0.04, 0.68, 0.90], projection="polar", facecolor=BG) + ax.set_theta_zero_location("N") + ax.set_theta_direction(-1) + ax.set_ylim(0, layers + 2.05) + ax.set_xticks(np.deg2rad(np.arange(0, 360, 30))) + ax.set_xticklabels([f"{d}°" for d in range(0, 360, 30)], color=MUTED, fontsize=8) + ax.set_yticks(np.arange(1, layers + 1) + 0.5) + ax.set_yticklabels([str(i) for i in range(layers)], color="#8da0a888", fontsize=7) + ax.grid(color=GRID, linewidth=0.6, alpha=0.55) + ax.spines["polar"].set_color("#38f4ff66") + ax.spines["polar"].set_linewidth(1.2) + + ax.pcolormesh(theta_grid, radius_grid, norm_ratio, cmap=radar_cmap(), shading="flat", alpha=0.96) + + # Soft trace underneath the hottest angular bearings, like phosphor persistence. + bearing_strength = norm_ratio.mean(axis=0) + norm_ratio.max(axis=0) * 0.42 + sweep_bin = int(np.argmax(bearing_strength)) + sweep_angle = float((sweep_bin + 0.5) * 360.0 / bins) + sweep_theta = np.deg2rad(sweep_angle) + for width, alpha in ((38, 0.055), (22, 0.075), (8, 0.14)): + half_width = np.deg2rad(width / 2) + theta = np.linspace(sweep_theta - half_width, sweep_theta + half_width, 80) + ax.fill_between(theta, 0.0, layers + 1.75, color=GREEN, alpha=alpha, linewidth=0) + + add_glow_spikes(ax, ratio, norm_ratio) + + for r in range(1, layers + 2): + ax.plot(np.linspace(0, 2 * np.pi, 360), np.full(360, r), color="#6fffe522", linewidth=0.55) + for deg in range(0, 360, 15): + th = np.deg2rad(deg) + ax.plot([th, th], [1, layers + 1.4], color="#6fffe516", linewidth=0.45) + + ax.text(0.5, 0.5, "ECHO\nCORE", color="#dff", fontsize=13, fontweight="bold", ha="center", va="center", transform=ax.transAxes) + ax.text(np.deg2rad(sweep_angle), layers + 1.35, "strongest bearing", color=GREEN, fontsize=8, ha="center", va="center") + + side = fig.add_axes([0.70, 0.06, 0.27, 0.86], facecolor=BG) + side.axis("off") + side.set_xlim(0, 1) + side.set_ylim(0, 1) + q = summary["question"] + ratio_mean = float(summary["answer_over_random_delta"]) + max_layer = int(summary.get("max_ratio_layer", int(np.argmax(ratio.mean(axis=1))))) + top = top_echoes(ratio, answer, random, 7) + max_echo = top[0] + + title_fx = [pe.withStroke(linewidth=4, foreground="#0b1918")] + side.text(0.00, 0.98, "SNAPCOMPACT RADAR", color=GREEN, fontsize=12, fontweight="bold", va="top") + side.text(0.00, 0.925, "Where the missing\nanswer echoes", color=INK, fontsize=27, fontweight="bold", va="top", linespacing=0.92, path_effects=title_fx) + side.text(0.00, 0.765, "Concentric rings are decoder layers. Angles are image-token bins. Bright spikes are answer-mask residuals divided by the random-mask control.", color=MUTED, fontsize=9.5, va="top", wrap=True) + + metrics = [ + ("gold answer", str(q["answer_text"]), AMBER), + ("question", q["q"], INK), + ("image tokens", f"{summary['image_tokens']:,}", CYAN), + ("layers", f"{summary['layers']}", CYAN), + ("mean answer/random Δ", f"{ratio_mean:.2f}×", AMBER), + ("max-ratio layer", f"L{max_layer}", GREEN), + ("loudest echo", f"L{max_echo['layer']} · bin {max_echo['bin']} · {max_echo['answer_random_ratio']:.1f}×", RED), + ] + y = 0.655 + for label, value, color in metrics: + side.text(0.00, y, label.upper(), color=MUTED, fontsize=7.2, fontweight="bold", va="top") + value_size = 12.6 if len(value) < 34 else 8.7 + side.text(0.00, y - 0.026, value, color=color, fontsize=value_size, fontweight="bold" if label != "question" else "normal", va="top", wrap=True) + y -= 0.075 if label != "question" else 0.105 + + side.text(0.00, y - 0.006, "TOP ECHOES", color=GREEN, fontsize=7.6, fontweight="bold", va="top") + y -= 0.040 + for row in top[:4]: + intensity = min(1.0, float(row["answer_random_ratio"]) / float(max_echo["answer_random_ratio"])) + side.plot([0.00, 0.36 * intensity], [y - 0.004, y - 0.004], color=AMBER, linewidth=3.2, alpha=0.35 + 0.55 * intensity, solid_capstyle="round") + side.text(0.40, y - 0.014, f"L{row['layer']:02d} bin {row['bin']:03d} {row['answer_random_ratio']:>5.1f}×", color=INK, fontsize=7.4, va="bottom", family="monospace") + y -= 0.032 + + # Tiny color scale and data provenance line. + grad_ax = fig.add_axes([0.708, 0.048, 0.19, 0.014], facecolor=BG) + grad_ax.imshow(np.linspace(0, 1, 512)[None, :], cmap=radar_cmap(), aspect="auto") + grad_ax.set_axis_off() + side.text(0.00, 0.006, "low ratio", color=MUTED, fontsize=7, va="bottom") + side.text(0.59, 0.006, "high answer echo", color=MUTED, fontsize=7, va="bottom") + fig.text(0.045, 0.018, "Actual heatmaps.npz arrays: ratio_binned, answer_binned, random_binned", color="#8da0a888", fontsize=8) + return fig + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", default=str(DATA_DIR)) + parser.add_argument("--out-dir", default=str(OUT_DIR)) + args = parser.parse_args() + + data_dir = Path(args.data_dir) + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + summary = json.loads((data_dir / "summary.json").read_text()) + heatmaps = np.load(data_dir / "heatmaps.npz") + answer = np.asarray(heatmaps["answer_binned"], dtype=np.float32) + random = np.asarray(heatmaps["random_binned"], dtype=np.float32) + ratio = np.asarray(heatmaps["ratio_binned"], dtype=np.float32) + + fig = draw_radial(summary, answer, random, ratio) + out_png = out_dir / "radial.png" + fig.savefig(out_png, facecolor=BG) + plt.close(fig) + + echoes = top_echoes(ratio, answer, random, 24) + (out_dir / "radial_top_echoes.json").write_text(json.dumps({"source": str(data_dir / "heatmaps.npz"), "top_echoes": echoes}, indent=2) + "\n") + np.savez_compressed(out_dir / "radial_source.npz", answer_binned=answer, random_binned=random, ratio_binned=ratio, ratio_norm=robust_norm(ratio, 0.972)) + print(out_png) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_token_grid.py b/packages/snapcompact/research/snapcompact_viz_token_grid.py new file mode 100644 index 000000000..7087390ad --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_token_grid.py @@ -0,0 +1,396 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["pillow", "numpy"] +# /// +"""Render a spatial image-token map for the snapcompact white-box run. + +The PaddleOCR-VL processor reports a 1 x 54 x 54 visual patch grid, while the +recorded hidden states contain 729 image tokens. This script folds the token +axis back to 27 x 27 (2 x 2 patch merge) and projects answer-mask delta / ratio +onto the original bitmap so the scar is visible in image space. +""" + +from __future__ import annotations + +import json +import math +from pathlib import Path +from typing import Iterable + +import numpy as np +from PIL import Image, ImageDraw, ImageFilter, ImageFont + +HERE = Path(__file__).resolve().parent +SOURCE = HERE / "results" / "tensor-heatmap-paddleocr-q7" +OUT_DIR = HERE / "results" / "agent-viz-token-grid" +OUT = OUT_DIR / "token-grid.png" + +PALETTE = { + "bg": (4, 6, 10), + "panel": (12, 17, 24), + "panel2": (17, 23, 31), + "ink": (244, 241, 225), + "muted": (139, 153, 163), + "grid": (49, 64, 75), + "cyan": (75, 218, 255), + "red": (255, 80, 66), + "amber": (255, 194, 72), + "green": (148, 255, 126), +} + + +def font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + candidates = [ + "/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Arial.ttf", + "/System/Library/Fonts/Supplemental/Helvetica Bold.ttf" if bold else "/System/Library/Fonts/Supplemental/Helvetica.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + for candidate in candidates: + if candidate and Path(candidate).exists(): + return ImageFont.truetype(candidate, size) + return ImageFont.load_default() + + +def lerp(a: int, b: int, t: float) -> int: + return round(a + (b - a) * t) + + +def mix(a: tuple[int, int, int], b: tuple[int, int, int], t: float) -> tuple[int, int, int]: + return (lerp(a[0], b[0], t), lerp(a[1], b[1], t), lerp(a[2], b[2], t)) + + +def heat_color(t: float) -> tuple[int, int, int]: + t = max(0.0, min(1.0, float(t))) + stops = [ + (0.00, (7, 12, 25)), + (0.18, (28, 24, 91)), + (0.38, (113, 31, 112)), + (0.62, (220, 61, 72)), + (0.82, (255, 152, 67)), + (1.00, (255, 242, 158)), + ] + for (pa, ca), (pb, cb) in zip(stops, stops[1:]): + if t <= pb: + return mix(ca, cb, (t - pa) / (pb - pa)) + return stops[-1][1] + + +def normalize(arr: np.ndarray, q: float = 0.985) -> tuple[np.ndarray, float]: + scale = float(np.quantile(arr, q)) if arr.size else 1.0 + if not math.isfinite(scale) or scale <= 0: + scale = 1.0 + return np.clip(arr / scale, 0, 1), scale + + +def token_side(summary: dict, token_count: int) -> int: + side = math.isqrt(token_count) + if side * side == token_count: + return side + grid = summary.get("processor_meta", {}).get("image_grid_thw", [[1, 0, 0]])[0] + _, gh, gw = grid + merge = math.isqrt(max(1, (gh * gw) // token_count)) + if merge and gh % merge == 0 and gw % merge == 0 and (gh // merge) * (gw // merge) == token_count: + return gh // merge + raise ValueError(f"cannot fold {token_count} image tokens into a square grid") + + +def fold_tokens(arr: np.ndarray, side: int) -> np.ndarray: + if arr.ndim == 1: + return arr.reshape(side, side) + return arr.reshape(arr.shape[0], side, side) + + +def heat_overlay(base: Image.Image, heat: np.ndarray, alpha_floor: int = 28, alpha_peak: int = 220) -> Image.Image: + norm, _ = normalize(heat) + small = Image.new("RGBA", (heat.shape[1], heat.shape[0]), (0, 0, 0, 0)) + pix = small.load() + for y in range(heat.shape[0]): + for x in range(heat.shape[1]): + t = float(norm[y, x]) + r, g, b = heat_color(t) + pix[x, y] = (r, g, b, round(alpha_floor + (alpha_peak - alpha_floor) * (t ** 0.85))) + overlay = small.resize(base.size, Image.Resampling.BICUBIC).filter(ImageFilter.GaussianBlur(1.0)) + dim = Image.blend(base.convert("RGB"), Image.new("RGB", base.size, (5, 8, 13)), 0.28).convert("RGBA") + return Image.alpha_composite(dim, overlay).convert("RGB") + + +def draw_token_grid(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], side: int, color: tuple[int, int, int] = (255, 255, 255)) -> None: + x0, y0, x1, y1 = box + for i in range(side + 1): + x = round(x0 + (x1 - x0) * i / side) + y = round(y0 + (y1 - y0) * i / side) + fill = (*color, 36) if hasattr(draw, "mode") else color + draw.line((x, y0, x, y1), fill=fill, width=1) + draw.line((x0, y, x1, y), fill=fill, width=1) + + +def paste_fit(canvas: Image.Image, img: Image.Image, box: tuple[int, int, int, int], resample: int = Image.Resampling.LANCZOS) -> tuple[int, int, int, int]: + x0, y0, x1, y1 = box + scale = min((x1 - x0) / img.width, (y1 - y0) / img.height) + w = max(1, round(img.width * scale)) + h = max(1, round(img.height * scale)) + resized = img.resize((w, h), resample) + px = x0 + (x1 - x0 - w) // 2 + py = y0 + (y1 - y0 - h) // 2 + canvas.paste(resized, (px, py)) + return (px, py, px + w, py + h) + + +def crop_answer(img: Image.Image, start: int, end: int, cols: int, adv: int, pitch: int, pad_cells: int = 34) -> Image.Image: + rows = img.height // pitch + row0 = max(0, start // cols - 5) + row1 = min(rows, end // cols + 6) + col0 = max(0, start % cols - pad_cells) + col1 = min(cols, end % cols + pad_cells) + if col1 <= col0: + col1 = min(cols, col0 + 72) + crop = img.crop((col0 * adv, row0 * pitch, col1 * adv, row1 * pitch)).convert("RGB") + d = ImageDraw.Draw(crop) + bx0 = max(0, (start % cols - col0) * adv - adv) + bx1 = min(crop.width - 1, ((end - 1) % cols - col0 + 2) * adv) + by0 = max(0, (start // cols - row0) * pitch - 1) + by1 = min(crop.height - 1, ((end - 1) // cols - row0 + 1) * pitch + 1) + d.rounded_rectangle((bx0, by0, bx1, by1), radius=3, outline=PALETTE["red"], width=3) + return crop + + +def answer_bbox(start: int, end: int, cols: int, adv: int, pitch: int) -> tuple[int, int, int, int]: + return ( + max(0, (start % cols) * adv - adv), + max(0, (start // cols) * pitch - 2), + min(cols * adv, ((end - 1) % cols + 2) * adv), + ((end - 1) // cols + 1) * pitch + 2, + ) + + +def draw_panel(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], title: str, subtitle: str | None = None) -> None: + draw.rounded_rectangle(box, radius=26, fill=PALETTE["panel"], outline=(32, 43, 55), width=1) + x0, y0, _, _ = box + draw.text((x0 + 24, y0 + 20), title, fill=PALETTE["ink"], font=font(28, True)) + if subtitle: + draw.text((x0 + 24, y0 + 56), subtitle, fill=PALETTE["muted"], font=font(17)) + + +def draw_micro_grid(canvas: Image.Image, heat: np.ndarray, box: tuple[int, int, int, int], title: str, subtitle: str) -> None: + draw = ImageDraw.Draw(canvas) + draw_panel(draw, box, title, subtitle) + x0, y0, x1, y1 = box + gx0, gy0, gx1, gy1 = x0 + 32, y0 + 96, x1 - 32, y1 - 42 + side = heat.shape[0] + norm, _ = normalize(heat) + cw = (gx1 - gx0) / side + ch = (gy1 - gy0) / side + for r in range(side): + for c in range(side): + xa = round(gx0 + c * cw) + ya = round(gy0 + r * ch) + xb = round(gx0 + (c + 1) * cw) + yb = round(gy0 + (r + 1) * ch) + draw.rectangle((xa, ya, xb, yb), fill=heat_color(float(norm[r, c]))) + for i in range(0, side + 1, 3): + x = round(gx0 + (gx1 - gx0) * i / side) + y = round(gy0 + (gy1 - gy0) * i / side) + draw.line((x, gy0, x, gy1), fill=(255, 255, 255, 34)) + draw.line((gx0, y, gx1, y), fill=(255, 255, 255, 34)) + + +def label(draw: ImageDraw.ImageDraw, xy: tuple[int, int], text: str, color: tuple[int, int, int], size: int = 18, bold: bool = True) -> None: + x, y = xy + pad = 8 + f = font(size, bold) + box = draw.textbbox((x, y), text, font=f) + draw.rounded_rectangle((box[0] - pad, box[1] - 4, box[2] + pad, box[3] + 5), radius=9, fill=(4, 6, 10), outline=color, width=1) + draw.text((x, y), text, fill=color, font=f) + + +def draw_hotspots(draw: ImageDraw.ImageDraw, box: tuple[int, int, int, int], heat: np.ndarray, count: int = 9) -> None: + x0, y0, x1, y1 = box + side = heat.shape[0] + flat = heat.ravel() + # Suppress immediate duplicates by greedily keeping separated cells. + chosen: list[int] = [] + for idx in np.argsort(flat)[::-1]: + r, c = divmod(int(idx), side) + if all(abs(r - divmod(j, side)[0]) + abs(c - divmod(j, side)[1]) >= 3 for j in chosen): + chosen.append(int(idx)) + if len(chosen) == count: + break + for rank, idx in enumerate(chosen, start=1): + r, c = divmod(idx, side) + cx = round(x0 + (c + 0.5) * (x1 - x0) / side) + cy = round(y0 + (r + 0.5) * (y1 - y0) / side) + rad = 11 if rank <= 3 else 8 + draw.ellipse((cx - rad, cy - rad, cx + rad, cy + rad), outline=PALETTE["amber"], width=3) + if rank <= 5: + draw.text((cx + 10, cy - 16), str(rank), fill=PALETTE["amber"], font=font(16, True)) + + +def text_block(draw: ImageDraw.ImageDraw, xy: tuple[int, int], lines: Iterable[str], fill: tuple[int, int, int], size: int = 20, gap: int = 8) -> None: + x, y = xy + f = font(size) + for line in lines: + draw.text((x, y), line, fill=fill, font=f) + y += size + gap + + +def render() -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + summary = json.loads((SOURCE / "summary.json").read_text()) + arrays = np.load(SOURCE / "heatmaps.npz") + original = Image.open(SOURCE / "images" / "original.png").convert("RGB") + masked = Image.open(SOURCE / "images" / "answer-mask.png").convert("RGB") + + answer_delta = arrays["answer_delta"].astype(np.float32, copy=False) + random_delta = arrays["random_delta"].astype(np.float32, copy=False) + ratio = arrays["ratio"].astype(np.float32, copy=False) + side = token_side(summary, answer_delta.shape[1]) + answer_grid = fold_tokens(answer_delta, side) + random_grid = fold_tokens(random_delta, side) + ratio_grid = fold_tokens(ratio, side) + + answer_mean = answer_grid.mean(axis=0) + random_mean = random_grid.mean(axis=0) + ratio_mean = ratio_grid.mean(axis=0) + early_ratio = ratio_grid[:4].mean(axis=0) + mid_delta = answer_grid[6:13].mean(axis=0) + late_delta = answer_grid[-4:].mean(axis=0) + + np.savez_compressed( + OUT_DIR / "token_grid_source.npz", + answer_mean=answer_mean, + random_mean=random_mean, + ratio_mean=ratio_mean, + early_ratio=early_ratio, + mid_answer_delta=mid_delta, + late_answer_delta=late_delta, + image_grid_thw=np.array(summary["processor_meta"]["image_grid_thw"][0], dtype=np.int32), + ) + (OUT_DIR / "token_grid_summary.json").write_text( + json.dumps( + { + "source": str(SOURCE), + "image_grid_thw": summary["processor_meta"]["image_grid_thw"][0], + "image_tokens": int(summary["image_tokens"]), + "rendered_token_grid": [side, side], + "patch_merge": int(summary["processor_meta"]["image_grid_thw"][0][1] // side), + "answer_over_random_delta": float(summary["answer_over_random_delta"]), + "question": summary["question"]["q"], + "answer_text": summary["question"]["answer_text"], + }, + indent=2, + ) + ) + + W, H = 2200, 1500 + canvas = Image.new("RGB", (W, H), PALETTE["bg"]) + draw = ImageDraw.Draw(canvas) + for y in range(0, H, 18): + draw.line((0, y, W, y), fill=(7, 10 + (y % 11), 17 + (y % 13))) + glow = Image.new("RGBA", (W, H), (0, 0, 0, 0)) + gd = ImageDraw.Draw(glow) + gd.ellipse((-320, -240, 960, 780), fill=(255, 80, 66, 34)) + gd.ellipse((920, -120, 2350, 1100), fill=(75, 218, 255, 30)) + gd.ellipse((760, 860, 1810, 1760), fill=(255, 194, 72, 18)) + canvas = Image.alpha_composite(canvas.convert("RGBA"), glow.filter(ImageFilter.GaussianBlur(95))).convert("RGB") + draw = ImageDraw.Draw(canvas) + + draw.text((64, 48), "SNAPCOMPACT TOKEN FIELD", fill=PALETTE["amber"], font=font(24, True)) + draw.text((64, 86), "Where the hidden-state scar lands on the bitmap", fill=PALETTE["ink"], font=font(62, True)) + draw.text( + (66, 164), + "PaddleOCR-VL reports a 1×54×54 image patch grid; 729 hidden-state image tokens fold back to 27×27 spatial cells.", + fill=PALETTE["muted"], + font=font(24), + ) + + # Main spatial map. + main_panel = (545, 225, 1455, 1340) + draw_panel(draw, main_panel, "answer-mask delta projected onto image tokens", "mean ||hidden(original) − hidden(answer-mask)|| across 19 layers") + map_box = (610, 330, 1390, 1110) + projected = heat_overlay(original, answer_mean) + pasted = paste_fit(canvas, projected, map_box, Image.Resampling.LANCZOS) + # Grid + answer box sit over the pasted square. + overlay = Image.new("RGBA", canvas.size, (0, 0, 0, 0)) + od = ImageDraw.Draw(overlay) + draw_token_grid(od, pasted, side, (255, 255, 255)) + bbox = answer_bbox(summary["question"]["answer_start"], summary["question"]["answer_end"], summary["geometry"]["cols"], 8, 13) + sx = (pasted[2] - pasted[0]) / original.width + sy = (pasted[3] - pasted[1]) / original.height + answer_rect = ( + round(pasted[0] + bbox[0] * sx), + round(pasted[1] + bbox[1] * sy), + round(pasted[0] + bbox[2] * sx), + round(pasted[1] + bbox[3] * sy), + ) + od.rounded_rectangle(answer_rect, radius=8, outline=(*PALETTE["red"], 255), width=5) + draw_hotspots(od, pasted, answer_mean) + canvas = Image.alpha_composite(canvas.convert("RGBA"), overlay).convert("RGB") + draw = ImageDraw.Draw(canvas) + label(draw, (pasted[0] + 18, pasted[1] + 18), "27×27 reconstructed image-token grid", PALETTE["cyan"], 19) + label(draw, (answer_rect[2] + 14, answer_rect[1] - 5), "erased answer text", PALETTE["red"], 18) + text_block( + draw, + (620, 1162), + [ + f"Q: {summary['question']['q']}", + f"Gold answer: {summary['question']['answer_text']} answer-mask mean delta: {summary['answer_delta_mean']:.2f} random-mask mean delta: {summary['random_delta_mean']:.2f}", + f"Answer/random delta ratio: {summary['answer_over_random_delta']:.2f}×. Bright cells are token locations most perturbed by hiding the answer span.", + ], + PALETTE["muted"], + 20, + 7, + ) + + # Evidence crops. + left = (64, 225, 505, 1340) + draw_panel(draw, left, "bitmap intervention", "original crop vs. answer erased") + crop = crop_answer(original, summary["question"]["answer_start"], summary["question"]["answer_end"], summary["geometry"]["cols"], 8, 13) + mcrop = crop_answer(masked, summary["question"]["answer_start"], summary["question"]["answer_end"], summary["geometry"]["cols"], 8, 13) + draw.text((96, 332), "ORIGINAL", fill=PALETTE["cyan"], font=font(17, True)) + draw.rounded_rectangle((94, 360, 475, 525), radius=16, fill=(240, 238, 226), outline=PALETTE["cyan"], width=3) + paste_fit(canvas, crop, (108, 374, 461, 511), Image.Resampling.NEAREST) + draw.text((96, 572), "ANSWER MASK", fill=PALETTE["red"], font=font(17, True)) + draw.rounded_rectangle((94, 600, 475, 765), radius=16, fill=(240, 238, 226), outline=PALETTE["red"], width=3) + paste_fit(canvas, mcrop, (108, 614, 461, 751), Image.Resampling.NEAREST) + draw.text((96, 822), "source arrays", fill=PALETTE["muted"], font=font(17, True)) + text_block( + draw, + (96, 858), + [ + "heatmaps.npz:", + "answer_delta[19,729]", + "random_delta[19,729]", + "ratio[19,729]", + "", + "fold rule:", + "54×54 patches / 2×2 merge", + "→ 27×27 visual tokens", + ], + PALETTE["ink"], + 21, + 8, + ) + draw.rounded_rectangle((96, 1110, 472, 1268), radius=18, fill=PALETTE["panel2"], outline=(38, 51, 64), width=1) + draw.text((118, 1132), "scar strength", fill=PALETTE["amber"], font=font(18, True)) + draw.text((118, 1170), f"{summary['answer_over_random_delta']:.2f}×", fill=PALETTE["ink"], font=font(54, True)) + draw.text((120, 1232), "answer-mask / random-mask mean delta", fill=PALETTE["muted"], font=font(17)) + + # Right analytical small multiples. + draw_micro_grid(canvas, ratio_mean, (1495, 225, 2136, 590), "ratio field", "mean answer_delta / random_delta") + draw_micro_grid(canvas, early_ratio, (1495, 620, 1810, 975), "early layers", "ratio, layers 0–3") + draw_micro_grid(canvas, mid_delta, (1820, 620, 2136, 975), "middle layers", "answer delta, layers 6–12") + draw_micro_grid(canvas, late_delta, (1495, 1005, 1810, 1340), "late layers", "answer delta, last 4") + draw_micro_grid(canvas, random_mean, (1820, 1005, 2136, 1340), "random control", "random-mask delta") + + # Color legend. + lx0, ly0, lx1, ly1 = 1530, 530, 2100, 552 + for x in range(lx0, lx1): + draw.line((x, ly0, x, ly1), fill=heat_color((x - lx0) / (lx1 - lx0))) + draw.text((lx0, ly1 + 10), "low", fill=PALETTE["muted"], font=font(15)) + draw.text((lx1 - 34, ly1 + 10), "high", fill=PALETTE["muted"], font=font(15)) + + canvas.save(OUT) + + +if __name__ == "__main__": + render() diff --git a/packages/snapcompact/research/snapcompact_viz_volume.py b/packages/snapcompact/research/snapcompact_viz_volume.py new file mode 100644 index 000000000..7089b3608 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_volume.py @@ -0,0 +1,235 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["matplotlib", "numpy"] +# /// +"""Render a volumetric tensor-cube visualization for snapcompact activations.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib import cm +from matplotlib.colors import LinearSegmentedColormap +from mpl_toolkits.mplot3d.art3d import Line3DCollection + +HERE = Path(__file__).resolve().parent +DATA_DIR = HERE / "results" / "tensor-heatmap-paddleocr-q7" +OUT_DIR = HERE / "results" / "agent-viz-volume" + +BG = "#05070b" +PANEL = "#0b1017" +GRID = "#33424c" +INK = "#efe9d5" +MUTED = "#87959b" +CYAN = "#44d9ff" +RED = "#ff5d4c" +AMBER = "#ffc84a" +GREEN = "#87ff80" + + +def robust01(values: np.ndarray, q: float = 0.985) -> np.ndarray: + """Quantile-normalize positive activation magnitudes without copying when possible.""" + scale = float(np.nanquantile(values, q)) + if not np.isfinite(scale) or scale <= 0.0: + scale = 1.0 + return np.clip(values / scale, 0.0, 1.0).astype(np.float32, copy=False) + + +def tinted_cmap(name: str, low: str, high: str) -> LinearSegmentedColormap: + return LinearSegmentedColormap.from_list(name, [(0.0, BG), (0.24, low), (1.0, high)], N=256) + + +def load_volume(data_dir: Path) -> tuple[np.ndarray, dict, list[str]]: + npz = np.load(data_dir / "heatmaps.npz") + summary = json.loads((data_dir / "summary.json").read_text()) + + answer = robust01(npz["answer_binned"]) + random = robust01(npz["random_binned"]) + ratio = robust01(npz["ratio_binned"]) + volume = np.stack([answer, random, ratio], axis=0) + labels = ["ANSWER Δ", "RANDOM Δ", "ANSWER/RANDOM"] + return volume, summary, labels + + +def cube_edges(x0: float, x1: float, y0: float, y1: float, z0: float, z1: float) -> list[list[tuple[float, float, float]]]: + p = { + "000": (x0, y0, z0), + "100": (x1, y0, z0), + "010": (x0, y1, z0), + "110": (x1, y1, z0), + "001": (x0, y0, z1), + "101": (x1, y0, z1), + "011": (x0, y1, z1), + "111": (x1, y1, z1), + } + return [ + [p["000"], p["100"]], [p["010"], p["110"]], [p["001"], p["101"]], [p["011"], p["111"]], + [p["000"], p["010"]], [p["100"], p["110"]], [p["001"], p["011"]], [p["101"], p["111"]], + [p["000"], p["001"]], [p["100"], p["101"]], [p["010"], p["011"]], [p["110"], p["111"]], + ] + + +def style_3d(ax) -> None: + ax.set_facecolor(BG) + for axis in (ax.xaxis, ax.yaxis, ax.zaxis): + axis.pane.set_facecolor((0.02, 0.03, 0.04, 0.0)) + axis._axinfo["grid"]["color"] = (0.35, 0.48, 0.56, 0.12) + axis._axinfo["tick"]["color"] = (0.75, 0.82, 0.82, 0.55) + ax.tick_params(colors=MUTED, labelsize=8, pad=0) + ax.set_xlabel("image-token bins (729 → 180)", color=MUTED, labelpad=9) + ax.set_ylabel("condition", color=MUTED, labelpad=7) + ax.set_zlabel("decoder layer", color=MUTED, labelpad=7) + ax.set_xlim(0, 179) + ax.set_ylim(-0.38, 2.38) + ax.set_zlim(0, 18) + ax.set_yticks([0, 1, 2]) + ax.set_yticklabels(["answer", "random", "ratio"], color=INK) + ax.set_xticks([0, 45, 90, 135, 179]) + ax.set_zticks([0, 4, 9, 14, 18]) + ax.view_init(elev=23, azim=-61) + ax.set_box_aspect((3.9, 1.0, 1.15)) + + +def add_volume(ax, volume: np.ndarray) -> None: + cmaps = [tinted_cmap("answer_ct", "#063842", CYAN), tinted_cmap("random_ct", "#461813", RED), tinted_cmap("ratio_ct", "#3c2b05", AMBER)] + edge_colors = [CYAN, RED, AMBER] + layers = np.arange(volume.shape[1]) + bins = np.arange(volume.shape[2]) + x, z = np.meshgrid(bins, layers) + + # Translucent CT slices: condition is depth, layer is vertical, image-token bin is horizontal. + for cond, cmap in enumerate(cmaps): + vals = volume[cond] + rgba = cmap(vals) + rgba[..., 3] = 0.08 + 0.68 * np.power(vals, 1.55) + y = np.full_like(x, cond, dtype=np.float32) + ax.plot_surface(x, y, z, facecolors=rgba, rstride=1, cstride=1, linewidth=0, antialiased=False, shade=False) + + # Bright activation voxels above each condition's 98th percentile. + threshold = float(np.quantile(vals, 0.982)) + zz, xx = np.where(vals >= threshold) + yy = np.full(xx.shape, cond, dtype=np.float32) + strength = vals[zz, xx] + ax.scatter(xx, yy, zz, s=10 + 90 * strength, c=edge_colors[cond], marker="s", alpha=0.58, depthshade=False, linewidths=0) + + ax.add_collection3d(Line3DCollection(cube_edges(0, 179, -0.23, 2.23, 0, 18), colors=(0.42, 0.72, 0.82, 0.22), linewidths=0.9)) + + # Crosshair slices through the strongest answer/random separation. + ratio = volume[2] + layer_profile = ratio.mean(axis=1) + bin_profile = ratio.mean(axis=0) + peak_layer = int(layer_profile.argmax()) + peak_bin = int(bin_profile.argmax()) + ax.plot([peak_bin, peak_bin], [-0.28, 2.28], [peak_layer, peak_layer], color=GREEN, alpha=0.9, linewidth=1.5) + ax.plot([0, 179], [2.28, 2.28], [peak_layer, peak_layer], color=GREEN, alpha=0.45, linewidth=1.1) + ax.text(peak_bin + 3, 2.35, peak_layer + 0.2, "hottest ratio slice", color=GREEN, fontsize=8) + + +def add_projection_panel(ax, volume: np.ndarray, labels: list[str]) -> None: + ax.set_facecolor(PANEL) + cmap = tinted_cmap("small_ct", "#10252f", "#f2d87b") + strip = np.vstack([volume[0], np.full((2, volume.shape[2]), np.nan), volume[1], np.full((2, volume.shape[2]), np.nan), volume[2]]) + masked = np.ma.masked_invalid(strip) + cmap.set_bad(PANEL) + ax.imshow(masked, aspect="auto", interpolation="nearest", cmap=cmap, vmin=0, vmax=1) + ax.set_xticks([0, 45, 90, 135, 179]) + ax.set_yticks([9, 30, 51]) + ax.set_yticklabels(labels, color=INK, fontsize=8) + ax.tick_params(colors=MUTED, labelsize=8, length=0) + ax.set_title("unwrapped tensor volume", color=INK, fontsize=12, loc="left", pad=8) + for spine in ax.spines.values(): + spine.set_color("#27323a") + + +def add_layer_panel(ax, volume: np.ndarray) -> None: + ax.set_facecolor(PANEL) + colors = [CYAN, RED, AMBER] + names = ["answer", "random", "ratio"] + for cond, color in enumerate(colors): + profile = volume[cond].mean(axis=1) + ax.plot(np.arange(profile.size), profile, color=color, linewidth=2.0, label=names[cond]) + ax.fill_between(np.arange(profile.size), profile, 0, color=color, alpha=0.08) + ax.set_xlim(0, 18) + ax.set_ylim(0, 1.0) + ax.set_xlabel("layer", color=MUTED, fontsize=8) + ax.set_ylabel("mean normalized intensity", color=MUTED, fontsize=8) + ax.tick_params(colors=MUTED, labelsize=8) + ax.grid(color=GRID, alpha=0.18, linewidth=0.7) + ax.legend(frameon=False, labelcolor=INK, fontsize=8, loc="upper right") + ax.set_title("layer dose curve", color=INK, fontsize=12, loc="left", pad=8) + for spine in ax.spines.values(): + spine.set_color("#27323a") + + +def render(volume: np.ndarray, summary: dict, labels: list[str], out_path: Path) -> None: + fig = plt.figure(figsize=(18, 11), dpi=180, facecolor=BG) + gs = fig.add_gridspec(3, 5, width_ratios=[1.35, 1.35, 1.35, 0.95, 0.95], height_ratios=[0.12, 1.0, 0.42], wspace=0.22, hspace=0.24) + + title_ax = fig.add_subplot(gs[0, :]) + title_ax.axis("off") + title_ax.text(0.0, 0.70, "SNAPCOMPACT ACTIVATION CT", color=INK, fontsize=27, fontweight="bold", transform=title_ax.transAxes) + title_ax.text(0.0, 0.24, "volumetric tensor cube: 19 layers × 180 image-token bins × 3 conditions", color=MUTED, fontsize=11, transform=title_ax.transAxes) + title_ax.text(0.985, 0.58, f"PaddleOCR-VL · Q: {summary['question']['q']}", color=MUTED, fontsize=9, ha="right", transform=title_ax.transAxes) + title_ax.text(0.985, 0.24, f"gold answer {summary['question']['answer_text']} · answer/random mean Δ {summary['answer_over_random_delta']:.2f}×", color=AMBER, fontsize=10, ha="right", transform=title_ax.transAxes) + + ax3d = fig.add_subplot(gs[1:, :3], projection="3d") + style_3d(ax3d) + add_volume(ax3d, volume) + ax3d.set_title("MRI-style scan of hidden-state deltas", color=INK, fontsize=15, loc="left", pad=12) + + ax_proj = fig.add_subplot(gs[1, 3:]) + add_projection_panel(ax_proj, volume, labels) + + ax_layer = fig.add_subplot(gs[2, 3:]) + add_layer_panel(ax_layer, volume) + + fig.text(0.055, 0.055, "source: heatmaps.npz arrays answer_binned, random_binned, ratio_binned · quantile normalized per condition", color="#617078", fontsize=8) + fig.text(0.055, 0.033, "cyan=answer evidence · red=random control · gold=answer/random amplification · green=crosshair at peak ratio slice", color="#617078", fontsize=8) + + out_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(out_path, facecolor=BG, bbox_inches="tight", pad_inches=0.22) + plt.close(fig) + + +def write_source_data(volume: np.ndarray, summary: dict, out_dir: Path) -> None: + out_dir.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + out_dir / "volume_source.npz", + volume=volume, + answer_norm=volume[0], + random_norm=volume[1], + ratio_norm=volume[2], + ) + payload = { + "description": "Quantile-normalized tensor cube used by snapcompact_viz_volume.py.", + "shape": {"condition": 3, "layers": int(volume.shape[1]), "image_token_bins": int(volume.shape[2])}, + "conditions": ["answer_delta", "random_delta", "answer_over_random_ratio"], + "question": summary["question"], + "answer_over_random_delta": summary["answer_over_random_delta"], + "source": str(DATA_DIR / "heatmaps.npz"), + } + (out_dir / "volume_source.json").write_text(json.dumps(payload, indent=2) + "\n") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-dir", type=Path, default=DATA_DIR) + parser.add_argument("--out-dir", type=Path, default=OUT_DIR) + parser.add_argument("--out", type=str, default="volume.png") + args = parser.parse_args() + + volume, summary, labels = load_volume(args.data_dir) + write_source_data(volume, summary, args.out_dir) + render(volume, summary, labels, args.out_dir / args.out) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/snapcompact_viz_waterfall.py b/packages/snapcompact/research/snapcompact_viz_waterfall.py new file mode 100755 index 000000000..de4bf49f4 --- /dev/null +++ b/packages/snapcompact/research/snapcompact_viz_waterfall.py @@ -0,0 +1,301 @@ +#!/usr/bin/env python3 +"""Layered snapcompact activation waterfall. + +Builds a seismic/ridgeline rendering from the PaddleOCR-VL white-box +activation deltas in results/tensor-heatmap-paddleocr-q7/heatmaps.npz. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import matplotlib.patheffects as pe +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import LinearSegmentedColormap +from matplotlib.patches import Rectangle + + +SCRIPT_DIR = Path(__file__).resolve().parent +DATA_DIR = SCRIPT_DIR / "results" / "tensor-heatmap-paddleocr-q7" +OUT_DIR = SCRIPT_DIR / "results" / "agent-viz-waterfall" + + +def smooth_rows(values: np.ndarray, radius: int = 3) -> np.ndarray: + """Small separable bin smoother; preserves shape and avoids scipy.""" + if radius <= 0: + return values.copy() + x = np.arange(-radius, radius + 1, dtype=np.float32) + kernel = np.exp(-(x * x) / (2.0 * (radius / 1.8) ** 2)) + kernel /= kernel.sum() + padded = np.pad(values, ((0, 0), (radius, radius)), mode="edge") + out = np.empty_like(values, dtype=np.float32) + for row in range(values.shape[0]): + out[row] = np.convolve(padded[row], kernel, mode="valid") + return out + + +def robust_unit(values: np.ndarray, high: float) -> np.ndarray: + scaled = np.log1p(np.maximum(values, 0.0)) / np.log1p(high) + return np.clip(scaled, 0.0, 1.0).astype(np.float32) + + +def load_source() -> tuple[dict, dict[str, np.ndarray]]: + with (DATA_DIR / "summary.json").open("r", encoding="utf-8") as handle: + summary = json.load(handle) + with np.load(DATA_DIR / "heatmaps.npz") as npz: + arrays = {name: npz[name].astype(np.float32) for name in npz.files} + return summary, arrays + + +def build_waterfall_data(summary: dict, arrays: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + answer = smooth_rows(arrays["answer_binned"], radius=3) + random = smooth_rows(arrays["random_binned"], radius=3) + ratio = smooth_rows(arrays["ratio_binned"], radius=2) + + # Use one common robust scale so answer/random amplitudes are visually comparable. + common_high = float(summary.get("common_delta_scale_p98") or np.percentile(np.r_[answer, random], 98)) + answer_u = robust_unit(answer, common_high) + random_u = robust_unit(random, common_high) + contrast = np.tanh((answer_u - random_u) * 2.8).astype(np.float32) + ratio_u = robust_unit(ratio, float(summary.get("ratio_scale_p98") or np.percentile(ratio, 98))) + + layers, bins = answer.shape + x = np.linspace(0.0, 1.0, bins, dtype=np.float32) + baselines = np.arange(layers, dtype=np.float32)[::-1] + return { + "x": x, + "baselines": baselines, + "answer": answer, + "random": random, + "ratio": ratio, + "answer_unit": answer_u, + "random_unit": random_u, + "contrast": contrast, + "ratio_unit": ratio_u, + } + + +def draw_glow_line(ax, x, y, color, lw=1.4, z=5, alpha=1.0): + line, = ax.plot(x, y, color=color, lw=lw, alpha=alpha, zorder=z, solid_joinstyle="round") + line.set_path_effects( + [ + pe.Stroke(linewidth=lw + 8.5, foreground=color, alpha=0.055), + pe.Stroke(linewidth=lw + 4.5, foreground=color, alpha=0.12), + pe.Normal(), + ] + ) + return line + + +def render(summary: dict, data: dict[str, np.ndarray]) -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + out_png = OUT_DIR / "waterfall.png" + + x = data["x"] + baselines = data["baselines"] + answer_u = data["answer_unit"] + random_u = data["random_unit"] + contrast = data["contrast"] + ratio_u = data["ratio_unit"] + layers, bins = answer_u.shape + + fig = plt.figure(figsize=(15.5, 10.5), dpi=210, facecolor="#05070d") + ax = fig.add_axes([0.055, 0.09, 0.89, 0.80], facecolor="#05070d") + + # Background ratio field: a dim spectrogram behind the ridges. + cmap = LinearSegmentedColormap.from_list( + "snap_seismic", + ["#05070d", "#0b1831", "#133f59", "#2d6f75", "#f3a44c", "#fff0bd"], + ) + extent = (0.0, 1.0, -0.78, layers - 0.22) + ax.imshow( + ratio_u[::-1], + extent=extent, + aspect="auto", + cmap=cmap, + interpolation="bicubic", + alpha=0.38, + zorder=0, + ) + + # Seismic paper grid and scanlines. + for gx in np.linspace(0, 1, 13): + ax.axvline(gx, color="#7cc8ff", lw=0.45, alpha=0.10, zorder=1) + for y in range(layers): + ax.axhline(y, color="#d4eaff", lw=0.38, alpha=0.085, zorder=1) + for yy in np.linspace(-0.6, layers - 0.35, 78): + ax.axhline(yy, color="#ffffff", lw=0.2, alpha=0.018, zorder=1) + + answer_color = "#ffc05a" + random_color = "#33d7ff" + gain_color = "#ff4d8d" + amplitude = 0.72 + + for layer_index, base in enumerate(baselines): + ans = base + answer_u[layer_index] * amplitude + rnd = base - random_u[layer_index] * amplitude * 0.82 + mid = base + contrast[layer_index] * amplitude * 0.62 + + ax.fill_between(x, base, ans, color=answer_color, alpha=0.075, zorder=2) + ax.fill_between(x, base, rnd, color=random_color, alpha=0.045, zorder=2) + ax.fill_between( + x, + rnd, + ans, + where=answer_u[layer_index] >= random_u[layer_index], + interpolate=True, + color=gain_color, + alpha=0.055, + zorder=2, + ) + + # Double-trace each layer: random-mask lower trace, answer-mask upper trace, + # plus a magenta differential tremor to make the comparison readable. + draw_glow_line(ax, x, rnd, random_color, lw=0.9, z=4, alpha=0.74) + draw_glow_line(ax, x, ans, answer_color, lw=1.16, z=5, alpha=0.92) + ax.plot(x, mid, color=gain_color, lw=0.48, alpha=0.55, zorder=3) + + if layer_index in {0, 4, 9, 14, layers - 1}: + ax.text( + -0.018, + base, + f"L{layer_index:02d}", + ha="right", + va="center", + color="#b9dfff", + fontsize=9, + family="monospace", + alpha=0.82, + ) + + # Highlight the strongest answer-vs-random bin per layer with tiny hot pips. + gain = answer_u - random_u + strongest = np.argmax(gain, axis=1) + ax.scatter( + x[strongest], + baselines + answer_u[np.arange(layers), strongest] * amplitude + 0.045, + s=8 + 32 * np.clip(gain[np.arange(layers), strongest], 0, 1), + c="#fff6cf", + alpha=0.72, + edgecolors="none", + zorder=6, + ) + + # Framing labels. + answer_mean = float(summary["answer_delta_mean"]) + random_mean = float(summary["random_delta_mean"]) + ratio = float(summary["answer_over_random_delta"]) + question = summary["question"]["q"] + answer_text = summary["question"]["answer_text"] + image_tokens = int(summary["image_tokens"]) + + ax.text( + 0.0, + layers + 0.62, + "SNAPCOMPACT ACTIVATION WATERFALL", + color="#f7fbff", + fontsize=22, + weight="bold", + family="monospace", + ha="left", + va="bottom", + ) + ax.text( + 0.0, + layers + 0.20, + f"PaddleOCR-VL · {layers} decoder layers · {image_tokens} image tokens binned into {bins} traces · answer '{answer_text}' vs random mask", + color="#9cc7e5", + fontsize=10.5, + family="monospace", + ha="left", + va="bottom", + ) + ax.text( + 0.0, + layers - 0.19, + f"Q: {question}", + color="#d8ecff", + fontsize=9.5, + family="monospace", + ha="left", + va="top", + alpha=0.84, + ) + ax.text( + 1.0, + layers + 0.27, + f"Δmean {answer_mean:.2f} / {random_mean:.2f} = {ratio:.2f}×", + color="#ffd37a", + fontsize=13, + family="monospace", + weight="bold", + ha="right", + va="bottom", + ) + + # Legend built as luminous calibration bars. + legend_y = -1.35 + ax.plot([0.02, 0.10], [legend_y, legend_y], color=answer_color, lw=2.2) + ax.text(0.112, legend_y, "answer-mask ridge", color="#ffdca0", fontsize=9, va="center", family="monospace") + ax.plot([0.32, 0.40], [legend_y, legend_y], color=random_color, lw=2.2) + ax.text(0.412, legend_y, "random-mask ridge", color="#9ff0ff", fontsize=9, va="center", family="monospace") + ax.plot([0.62, 0.70], [legend_y, legend_y], color=gain_color, lw=1.4) + ax.text(0.712, legend_y, "answer excess tremor", color="#ff9bbb", fontsize=9, va="center", family="monospace") + + # Outer phosphor frame. + ax.add_patch(Rectangle((0, -0.78), 1, layers - 0.44, fill=False, lw=0.9, edgecolor="#5fb7ff", alpha=0.34, zorder=10)) + ax.set_xlim(-0.055, 1.02) + ax.set_ylim(-1.62, layers + 0.98) + ax.set_xticks(np.linspace(0, 1, 7)) + ax.set_xticklabels([f"{int(t * image_tokens):03d}" for t in np.linspace(0, 1, 7)], color="#8fbede", fontsize=8, family="monospace") + ax.set_yticks([]) + ax.set_xlabel("image-token bin →", color="#9cc7e5", fontsize=10, family="monospace", labelpad=12) + for spine in ax.spines.values(): + spine.set_visible(False) + ax.tick_params(axis="x", length=0) + + # Save source arrays used by this rendering for reproducibility. + np.savez_compressed( + OUT_DIR / "waterfall_source.npz", + x=x, + baselines=baselines, + answer_binned=data["answer"], + random_binned=data["random"], + ratio_binned=data["ratio"], + answer_unit=answer_u, + random_unit=random_u, + ratio_unit=ratio_u, + contrast=contrast, + ) + with (OUT_DIR / "waterfall_source.json").open("w", encoding="utf-8") as handle: + json.dump( + { + "question": question, + "answer_text": answer_text, + "layers": layers, + "bins": bins, + "image_tokens": image_tokens, + "answer_delta_mean": answer_mean, + "random_delta_mean": random_mean, + "answer_over_random_delta": ratio, + "source_npz": str(DATA_DIR / "heatmaps.npz"), + }, + handle, + indent=2, + ) + + fig.savefig(out_png, facecolor=fig.get_facecolor(), bbox_inches="tight", pad_inches=0.14) + plt.close(fig) + print(out_png) + + +def main() -> None: + summary, arrays = load_source() + data = build_waterfall_data(summary, arrays) + render(summary, data) + + +if __name__ == "__main__": + main() diff --git a/packages/snapcompact/research/squad.py b/packages/snapcompact/research/squad.py new file mode 100644 index 000000000..39d206f4d --- /dev/null +++ b/packages/snapcompact/research/squad.py @@ -0,0 +1,120 @@ +"""SQuAD v1.1 dev: passage flow, question sampling, official EM/F1 scoring.""" + +import json +import random +import re +import string +import urllib.request +from collections import Counter +from pathlib import Path + +SQUAD_URL = "https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json" + + +def load_paragraphs(cache: Path) -> list[dict]: + """Flattened [{ctx, qas, title}] in deterministic dataset order.""" + path = cache / "squad-dev-v1.1.json" + if not path.exists(): + urllib.request.urlretrieve(SQUAD_URL, path) + data = json.loads(path.read_text())["data"] + out = [] + for art in data: + for p in art["paragraphs"]: + out.append({"ctx": " ".join(p["context"].split()), "qas": p["qas"], "title": art["title"]}) + return out + + +def build_flow(paras: list[dict], max_chars: int | None = None) -> tuple[str, list[int]]: + """Space-joined passage stream + start offset of each passage.""" + flow, offsets = "", [] + for p in paras: + offsets.append(len(flow)) + flow += p["ctx"] + " " + if max_chars is not None and len(flow) >= max_chars: + break + return flow, offsets + + +def sample_chunk_questions( + paras: list[dict], offsets: list[int], start: int, end: int, n: int, seed: int +) -> list[dict]: + """Up to n questions from passages fully inside [start, end), evenly spread. + + Passages straddling a chunk boundary are skipped (their answers may be cut). + Each question records pos_rel: passage start relative to the chunk, 0..1. + """ + rng = random.Random(seed * 1_000_003 + start) + eligible = [ + i + for i in range(len(offsets)) + if offsets[i] >= start and offsets[i] + len(paras[i]["ctx"]) <= end + ] + if not eligible: + return [] + n = min(n, len(eligible)) + step = len(eligible) / n + picked = [] + for k in range(n): + pi = eligible[int(k * step)] + qa = rng.choice(paras[pi]["qas"]) + picked.append( + { + "q": " ".join(qa["question"].split()), + "golds": sorted({a["text"] for a in qa["answers"]}), + "pos_rel": (offsets[pi] - start) / (end - start), + } + ) + return picked + + +# --- official SQuAD normalization / metrics --- + + +def _normalize(s: str) -> str: + s = s.lower() + s = "".join(ch for ch in s if ch not in string.punctuation) + s = re.sub(r"\b(a|an|the)\b", " ", s) + return " ".join(s.split()) + + +def exact_match(pred: str, golds: list[str]) -> float: + return float(any(_normalize(pred) == _normalize(g) for g in golds)) + + +def f1(pred: str, golds: list[str]) -> float: + best = 0.0 + for g in golds: + p_tok, g_tok = _normalize(pred).split(), _normalize(g).split() + common = Counter(p_tok) & Counter(g_tok) + overlap = sum(common.values()) + if overlap == 0: + continue + prec, rec = overlap / len(p_tok), overlap / len(g_tok) + best = max(best, 2 * prec * rec / (prec + rec)) + return best + + +def parse_numbered(text: str, n: int) -> list[str]: + """Extract answers from a numbered list; missing entries become ''. """ + answers = [""] * n + for line in text.splitlines(): + m = re.match(r"\s*(\d+)[.):]\s*(.*\S)?\s*$", line) + if m and m.group(2): + idx = int(m.group(1)) - 1 + if 0 <= idx < n and not answers[idx]: + answers[idx] = m.group(2).strip() + return answers + + +def score(answers: list[str], questions: list[dict]) -> dict: + ems = [exact_match(a, q["golds"]) for a, q in zip(answers, questions)] + f1s = [f1(a, q["golds"]) for a, q in zip(answers, questions)] + return { + "em": sum(ems) / len(ems), + "f1": sum(f1s) / len(f1s), + "abstained": sum("unreadable" in a.lower() for a in answers), + "per_question": [ + {"answer": a, "golds": q["golds"], "em": e, "f1": f} + for a, q, e, f in zip(answers, questions, ems, f1s) + ], + } diff --git a/packages/snapcompact/src/index.ts b/packages/snapcompact/src/index.ts new file mode 100644 index 000000000..efd6a6990 --- /dev/null +++ b/packages/snapcompact/src/index.ts @@ -0,0 +1 @@ +export * from "./snapcompact"; diff --git a/packages/snapcompact/src/prompts/file-operations.md b/packages/snapcompact/src/prompts/file-operations.md new file mode 100644 index 000000000..053b76839 --- /dev/null +++ b/packages/snapcompact/src/prompts/file-operations.md @@ -0,0 +1,10 @@ +{{#if readFiles.length}} +{{#xml "read-files"}} +{{join readFiles "\n"}} +{{/xml}} +{{/if}} +{{#if modifiedFiles.length}} +{{#xml "modified-files"}} +{{join modifiedFiles "\n"}} +{{/xml}} +{{/if}} diff --git a/packages/snapcompact/src/prompts/snapcompact-summary.md b/packages/snapcompact/src/prompts/snapcompact-summary.md new file mode 100644 index 000000000..0d8d22323 --- /dev/null +++ b/packages/snapcompact/src/prompts/snapcompact-summary.md @@ -0,0 +1,17 @@ +Prior conversation history has been archived verbatim onto {{frameCount}} snapcompact frame{{#if multipleFrames}}s{{/if}} — the bitmap image{{#if multipleFrames}}s{{/if}} attached below{{#if multipleFrames}}, ordered oldest to newest{{/if}}. + +Reading a frame: monospace {{fontCell}} pixel font on a white background, {{cols}} characters per row, {{rows}} text rows per frame; read left to right, top to bottom. Text flows continuously with no word wrap, so words may break across row ends. Whitespace runs (including newlines) were collapsed to single spaces. {{#if sentenceInk}}Ink color cycles through six colors, advancing at sentence boundaries — a color change marks a new sentence.{{else}}Glyphs are plain black ink.{{/if}}{{#if lineRepeated}} Every text line is printed twice in a row — first on the white background, then repeated on a pale yellow band. The copies are identical: read each line once and use the duplicate only to double-check hard glyphs.{{/if}} Roles are tagged inline as [User]:, [Assistant]:, [Assistant thinking]:, [Assistant tool calls]:, and [Tool result]:. +{{#if mixedShapes}} + +Older frames may use a different font, grid, or ink coloring than described above; the reading order is always the same (left to right, top to bottom, oldest frame first). +{{/if}} +{{#if includedPreviousSummary}} + +The earliest frame begins with "[Summary of earlier history]" — a condensed digest of context that predates the archived conversation. +{{/if}} +{{#if truncatedChars}} + +{{truncatedChars}} characters of older history were dropped to respect the frame budget. The first frame (session start) is always kept, so the missing span sits between the first frame and the next. +{{/if}} + +Total archived: {{totalChars}} characters. Consult the frames whenever you need exact earlier details (user wording, decisions, file paths, tool output). If a region is hard to read, re-derive the fact from the workspace (re-read files, re-run commands) rather than guessing. diff --git a/packages/snapcompact/src/snapcompact.ts b/packages/snapcompact/src/snapcompact.ts new file mode 100644 index 000000000..8e688bc8e --- /dev/null +++ b/packages/snapcompact/src/snapcompact.ts @@ -0,0 +1,638 @@ +/** + * Snapcompact compaction: archive conversation history as dense bitmap images. + * + * Instead of asking an LLM to summarize discarded history, the serialized + * conversation is rendered into square PNG frames of pixel-font text that + * vision models read back directly, like an archivist at a snapcompact frame + * reader. + * + * The frame shape is provider-aware, following the snapcompact SQuAD evals + * (`packages/snapcompact`, 200k-token monolithic runs): + * + * - **Anthropic** (`8x8r-bw`): unscii-8 square cells, black ink, every line + * printed twice with the copy on a pale highlight band. Read at F1 parity + * with raw text at ~2x lower cost; the colored variants drew refusals at + * scale, the repeated plain shape did not. + * - **Google** (`8x8r-sent`): same repeated grid with six-hue sentence + * coloring (0.90 F1 at ~2.9x lower cost on gemini-3.5-flash). + * - **OpenAI** (`6x6u-sent`): OpenAI bills a flat ~2.9k tokens per image, so + * image count is the only cost lever — unscii-8 Lanczos-stretched to 6x6 + * cells packs the most readable chars per frame. Frames request + * `detail: "original"`; the default `auto` downscale destroys 6px glyphs. + * - **Unknown providers** default to the Anthropic shape (most + * refusal-robust). Gateways that resize images (e.g. OpenRouter normalizes + * visual payloads to a fixed token budget) defeat any shape — optical + * context fails silently there. + * + * The whole pass is local and deterministic — no LLM call, no API key, no + * latency beyond rendering. Rasterization and PNG encoding happen in native + * code (`renderSnapcompactPng` in `crates/pi-natives/src/snapcompact.rs`). + * Frames persist in the compaction entry's `preserveData` and are + * re-attached to the compaction summary message on every context rebuild. + */ + +import type { Api, ImageContent, Message, Model } from "@oh-my-pi/pi-ai"; +import { renderSnapcompactPng } from "@oh-my-pi/pi-natives"; +import { prompt } from "@oh-my-pi/pi-utils"; +import fileOperationsTemplate from "./prompts/file-operations.md" with { type: "text" }; +import snapcompactSummaryPrompt from "./prompts/snapcompact-summary.md" with { type: "text" }; + +// ============================================================================ +// Shapes +// ============================================================================ + +/** One eval-validated frame shape: font, cell, ink, repetition, and size. */ +export interface SnapcompactShape { + /** Bundled font in the native renderer. */ + font: "5x8" | "8x8"; + /** Target cell advance in pixels; differing from the font's natural cell + * renders via Lanczos stretch (anti-aliased RGB frame). */ + cellWidth: number; + /** Target cell pitch in pixels. */ + cellHeight: number; + /** Ink: `sent` cycles six hues at sentence boundaries; `bw` is black. */ + variant: "sent" | "bw"; + /** Each text line is printed this many times; copies after the first sit + * on a pale highlight band (redundancy coding). */ + lineRepeat: number; + /** Frame edge in pixels. */ + frameSize: number; + /** Per-frame billed-token estimate for the shape's target provider. */ + frameTokenEstimate: number; + /** Resolution hint attached to frame images (OpenAI-only). */ + imageDetail?: ImageContent["detail"]; +} + +/** Eval-validated shapes, keyed by the provider family they won on. */ +export const SNAPCOMPACT_SHAPES = { + /** `8x8r-bw`: unscii square, black ink, lines doubled on highlight bands. */ + anthropic: { + font: "8x8", + cellWidth: 8, + cellHeight: 8, + variant: "bw", + lineRepeat: 2, + frameSize: 1568, + frameTokenEstimate: 3300, + }, + /** `8x8r-sent`: the repeated grid with sentence-hue ink. */ + google: { + font: "8x8", + cellWidth: 8, + cellHeight: 8, + variant: "sent", + lineRepeat: 2, + frameSize: 1568, + frameTokenEstimate: 1100, + }, + /** `6x6u-sent`: unscii stretched to 6x6 — densest readable cell, fewest + * frames (OpenAI bills per image, ~2.9k tokens flat). */ + openaiDense: { + font: "8x8", + cellWidth: 6, + cellHeight: 6, + variant: "sent", + lineRepeat: 1, + frameSize: 1568, + frameTokenEstimate: 2900, + imageDetail: "original", + }, + /** Original 5x8 X.org shape (pre-shape-table sessions rendered this). */ + legacy: { + font: "5x8", + cellWidth: 5, + cellHeight: 8, + variant: "sent", + lineRepeat: 1, + frameSize: 2576, + frameTokenEstimate: 3300, + }, +} as const satisfies Record; + +/** Pick the eval-optimal frame shape for a provider API. */ +export function resolveSnapcompactShape(api?: Api): SnapcompactShape { + switch (api) { + case "openai-completions": + case "openai-responses": + case "openai-codex-responses": + case "azure-openai-responses": + return SNAPCOMPACT_SHAPES.openaiDense; + case "google-generative-ai": + case "google-gemini-cli": + case "google-vertex": + return SNAPCOMPACT_SHAPES.google; + default: + // anthropic-messages, bedrock-converse-stream, and anything unknown: + // the plain repeated grid is the most refusal-robust reader shape. + return SNAPCOMPACT_SHAPES.anthropic; + } +} + +// ============================================================================ +// Constants +// ============================================================================ + +/** Legacy frame edge in pixels (the 5x8 shape's eval-validated size). New + * shapes carry their own `frameSize`. */ +export const SNAPCOMPACT_FRAME_SIZE = 2576; + +/** Maximum frames carried on a compaction entry. Oldest frames are dropped + * first once the budget is exceeded (mirrors how iterative text summaries + * fade the oldest detail). */ +export const SNAPCOMPACT_MAX_FRAMES = 8; + +/** Conservative per-frame token estimate used for context budgeting + * (upper bound across shapes: Anthropic bills 1568*1568/750 ≈ 3,278). */ +export const SNAPCOMPACT_FRAME_TOKEN_ESTIMATE = 3300; + +/** Key under `CompactionEntry.preserveData` holding the frame archive. */ +export const SNAPCOMPACT_PRESERVE_KEY = "snapcompact"; + +// ============================================================================ +// Types +// ============================================================================ + +/** One developed snapcompact frame: a base64 PNG plus its reading geometry. */ +export interface SnapcompactFrame { + /** Base64-encoded PNG. */ + data: string; + mimeType: string; + /** Characters per row in the frame grid. */ + cols: number; + /** Text rows in the frame grid (unique lines, not repeated copies). */ + rows: number; + /** Characters actually printed onto this frame. */ + chars: number; + /** Shape metadata (absent on legacy frames, which are 5x8 `sent`). */ + font?: SnapcompactShape["font"]; + variant?: SnapcompactShape["variant"]; + lineRepeat?: number; + /** Resolution hint forwarded to the provider when re-attaching. */ + detail?: ImageContent["detail"]; +} + +/** Frame archive persisted under `preserveData[SNAPCOMPACT_PRESERVE_KEY]`. */ +export interface SnapcompactArchive { + /** Frames ordered oldest to newest. */ + frames: SnapcompactFrame[]; + /** Characters currently readable across all frames. */ + totalChars: number; + /** Characters dropped so far to respect the frame budget. */ + truncatedChars: number; +} + +export interface SnapcompactGeometry { + cols: number; + rows: number; + /** Characters that fit one frame (cols * rows). */ + capacity: number; +} + +export interface SnapcompactOptions { + /** App-level message transformer (same contract as agent-core's `SummaryOptions.convertToLlm`). */ + convertToLlm?: SnapcompactConvertToLlm; + /** Model whose provider API selects the frame shape. */ + model?: Pick; + /** Explicit shape override; wins over `model`. */ + shape?: SnapcompactShape; + /** Frame edge in pixels. Defaults to the shape's `frameSize`. */ + frameSize?: number; + /** Frame budget. Defaults to {@link SNAPCOMPACT_MAX_FRAMES}. */ + maxFrames?: number; +} + +/** Result of rendering one frame, before base64 packing. */ +export interface RenderedFrame { + png: Uint8Array; + cols: number; + rows: number; + /** Characters printed (input may be shorter than capacity). */ + chars: number; +} + +// ============================================================================ +// Compaction data contracts +// ============================================================================ + +export interface SnapcompactFileOperations { + read: Set; + written: Set; + edited: Set; +} + +export interface SnapcompactCompactionDetails { + readFiles: string[]; + modifiedFiles: string[]; +} + +export interface SnapcompactCompactionPreparation { + /** UUID of first entry to keep. */ + firstKeptEntryId: string; + /** Messages that will be archived and discarded. */ + messagesToSummarize: TMessage[]; + /** Messages that will be archived as the split-turn prefix, if any. */ + turnPrefixMessages: TMessage[]; + tokensBefore: number; + /** Summary from previous compaction, for continuity when no prior snapcompact archive exists. */ + previousSummary?: string; + /** Preserved opaque compaction payload from the previous compaction, if any. */ + previousPreserveData?: Record; + /** File operations extracted by the host agent. */ + fileOps: SnapcompactFileOperations; +} + +export interface SnapcompactCompactionResult { + summary: string; + shortSummary?: string; + firstKeptEntryId: string; + tokensBefore: number; + details?: T; + preserveData?: Record; +} + +export type SnapcompactConvertToLlm = (messages: TMessage[]) => Message[]; + +function defaultConvertToLlm(messages: TMessage[]): Message[] { + return messages as unknown as Message[]; +} + +// ============================================================================ +// File operation helpers +// ============================================================================ + +export function createSnapcompactFileOps(): SnapcompactFileOperations { + return { + read: new Set(), + written: new Set(), + edited: new Set(), + }; +} + +export function computeSnapcompactFileLists(fileOps: SnapcompactFileOperations): SnapcompactCompactionDetails { + const modified = new Set([...fileOps.edited, ...fileOps.written]); + const readFiles = [...fileOps.read].filter(file => !modified.has(file)).sort(); + const modifiedFiles = [...modified].sort(); + return { readFiles, modifiedFiles }; +} + +const FILE_OPERATION_SUMMARY_LIMIT = 20; + +function truncateFileList(files: string[]): string[] { + if (files.length <= FILE_OPERATION_SUMMARY_LIMIT) return files; + const omitted = files.length - FILE_OPERATION_SUMMARY_LIMIT; + return [...files.slice(0, FILE_OPERATION_SUMMARY_LIMIT), `… (${omitted} more files omitted)`]; +} + +function stripFileOperationTags(summary: string): string { + const withoutReadFiles = summary.replace(/[\s\S]*?<\/read-files>\s*/g, ""); + const withoutModifiedFiles = withoutReadFiles.replace(/[\s\S]*?<\/modified-files>\s*/g, ""); + return withoutModifiedFiles.trimEnd(); +} + +function formatFileOperations(readFiles: string[], modifiedFiles: string[]): string { + if (readFiles.length === 0 && modifiedFiles.length === 0) return ""; + return prompt.render(fileOperationsTemplate, { + readFiles: truncateFileList(readFiles), + modifiedFiles: truncateFileList(modifiedFiles), + }); +} + +export function upsertSnapcompactFileOperations(summary: string, readFiles: string[], modifiedFiles: string[]): string { + const baseSummary = stripFileOperationTags(summary); + const fileOperations = formatFileOperations(readFiles, modifiedFiles); + if (!fileOperations) return baseSummary; + if (!baseSummary) return fileOperations; + return `${baseSummary}\n\n${fileOperations}`; +} + +// ============================================================================ +// Message serialization +// ============================================================================ + +const TOOL_RESULT_MAX_CHARS = 2000; + +function truncateForSummary(text: string, maxChars: number): string { + if (text.length <= maxChars) return text; + const truncatedChars = text.length - maxChars; + return `${text.slice(0, maxChars)}\n\n[... ${truncatedChars} more characters truncated]`; +} + +export function serializeSnapcompactConversation(messages: Message[]): string { + const parts: string[] = []; + + for (const msg of messages) { + if (msg.role === "user") { + const content = + typeof msg.content === "string" + ? msg.content + : msg.content + .filter((content): content is { type: "text"; text: string } => content.type === "text") + .map(content => content.text) + .join(""); + if (content) parts.push(`[User]: ${content}`); + } else if (msg.role === "assistant") { + const textParts: string[] = []; + const thinkingParts: string[] = []; + const toolCalls: string[] = []; + + for (const block of msg.content) { + if (block.type === "text") { + textParts.push(block.text); + } else if (block.type === "thinking") { + thinkingParts.push(block.thinking); + } else if (block.type === "toolCall") { + const args = block.arguments as Record; + const argsStr = Object.entries(args) + .map(([key, value]) => `${key}=${JSON.stringify(value)}`) + .join(", "); + toolCalls.push(`${block.name}(${argsStr})`); + } + } + + if (thinkingParts.length > 0) { + parts.push(`[Assistant thinking]: ${thinkingParts.join("\n")}`); + } + if (textParts.length > 0) { + parts.push(`[Assistant]: ${textParts.join("\n")}`); + } + if (toolCalls.length > 0) { + parts.push(`[Assistant tool calls]: ${toolCalls.join("; ")}`); + } + } else if (msg.role === "toolResult") { + const content = msg.content + .filter((block): block is { type: "text"; text: string } => block.type === "text") + .map(block => block.text) + .join(""); + if (content) { + parts.push(`[Tool result]: ${truncateForSummary(content, TOOL_RESULT_MAX_CHARS)}`); + } + } + } + + return parts.join("\n\n"); +} + +// ============================================================================ +// Preserve-data helpers +// ============================================================================ + +const OPENAI_REMOTE_COMPACTION_PRESERVE_KEY = "openaiRemoteCompaction"; + +function stripOpenAiRemoteCompactionPreserveData( + preserveData: Record | undefined, +): Record | undefined { + if (!preserveData || !(OPENAI_REMOTE_COMPACTION_PRESERVE_KEY in preserveData)) { + return preserveData; + } + const { [OPENAI_REMOTE_COMPACTION_PRESERVE_KEY]: _removed, ...rest } = preserveData; + return Object.keys(rest).length > 0 ? rest : undefined; +} + +// ============================================================================ +// Text normalization +// ============================================================================ + +/** Folds for common non-Latin-1 characters the bundled fonts cannot draw. */ +const CHAR_FOLD: Record = { + "\u2018": "'", + "\u2019": "'", + "\u201a": "'", + "\u201b": "'", + "\u201c": '"', + "\u201d": '"', + "\u201e": '"', + "\u2013": "-", + "\u2014": "-", + "\u2015": "-", + "\u2212": "-", + "\u2026": "...", + "\u2022": "*", + "\u25cf": "*", + "\u25a0": "*", + "\u25aa": "*", + "\u2190": "<-", + "\u2192": "->", + "\u21d2": "=>", + "\u2713": "v", + "\u2714": "v", + "\u2717": "x", + "\u2718": "x", +}; + +/** + * Prepare text for printing: collapse whitespace runs (incl. newlines) to + * single spaces — the eval's "paragraph breaks collapsed to spaces" format — + * then fold everything outside the fonts' ASCII + Latin-1 coverage to ASCII + * approximations (`?` as the last resort). + */ +export function normalizeForSnapcompact(text: string): string { + const collapsed = text.replace(/\s+/g, " ").trim(); + let out = ""; + for (const ch of collapsed) { + const cp = ch.codePointAt(0) as number; + if (cp < 0x7f || (cp >= 0xa0 && cp <= 0xff)) { + out += ch; + continue; + } + const fold = CHAR_FOLD[ch]; + if (fold !== undefined) { + out += fold; + } else if (cp >= 0x2500 && cp <= 0x257f) { + // Box drawing: keep table skeletons legible. + out += cp === 0x2502 || cp === 0x2503 ? "|" : cp === 0x2500 || cp === 0x2501 ? "-" : "+"; + } else { + out += "?"; + } + } + return out; +} + +// ============================================================================ +// Rendering +// ============================================================================ + +export function snapcompactGeometry(shape: SnapcompactShape, size: number = shape.frameSize): SnapcompactGeometry { + const cols = Math.floor(size / shape.cellWidth); + const rows = Math.floor(size / shape.cellHeight / shape.lineRepeat); + return { cols, rows, capacity: cols * rows }; +} + +/** Render one snapcompact frame from already-normalized text. */ +export function renderSnapcompactFrame( + text: string, + shape: SnapcompactShape, + size: number = shape.frameSize, +): RenderedFrame { + const { cols, rows, capacity } = snapcompactGeometry(shape, size); + const chars = Math.min(text.length, capacity); + const png = renderSnapcompactPng(text, { + size, + font: shape.font, + cellWidth: shape.cellWidth, + cellHeight: shape.cellHeight, + variant: shape.variant, + lineRepeat: shape.lineRepeat, + }); + return { png, cols, rows, chars }; +} + +// ============================================================================ +// Archive helpers +// ============================================================================ + +/** Validate and extract a persisted frame archive from `preserveData`. */ +export function getPreservedSnapcompactArchive( + preserveData: Record | undefined, +): SnapcompactArchive | undefined { + const candidate = preserveData?.[SNAPCOMPACT_PRESERVE_KEY]; + if (!candidate || typeof candidate !== "object") return undefined; + const archive = candidate as SnapcompactArchive; + if (!Array.isArray(archive.frames)) return undefined; + const frames = archive.frames.filter( + frame => + !!frame && + typeof frame.data === "string" && + frame.data.length > 0 && + typeof frame.mimeType === "string" && + typeof frame.cols === "number" && + typeof frame.rows === "number" && + typeof frame.chars === "number", + ); + if (frames.length === 0) return undefined; + return { + frames, + totalChars: typeof archive.totalChars === "number" ? archive.totalChars : 0, + truncatedChars: typeof archive.truncatedChars === "number" ? archive.truncatedChars : 0, + }; +} + +/** Convert archive frames into LLM image blocks (oldest first). */ +export function snapcompactImages(archive: SnapcompactArchive): ImageContent[] { + return archive.frames.map(frame => ({ + type: "image", + data: frame.data, + mimeType: frame.mimeType, + ...(frame.detail ? { detail: frame.detail } : {}), + })); +} + +// ============================================================================ +// Compaction entry point +// ============================================================================ + +/** + * Run a snapcompact compaction over prepared messages. Fully local: serializes + * the discarded history, prints it onto PNG frames in the provider-optimal + * shape, merges previously archived frames (oldest dropped beyond the + * budget), and produces a deterministic summary explaining how to read the + * frames. + * + * Frames archived under a different shape (provider switches, legacy 5x8 + * sessions) are kept as-is — each frame carries its own geometry, and the + * summary describes the newest shape while noting that older frames may + * differ. + * + * If the previous compaction was text-based, its summary is printed at the + * head of the frame archive as `[Summary of earlier history]` so no continuity is lost. + */ +export async function snapcompactCompact( + preparation: SnapcompactCompactionPreparation, + options?: SnapcompactOptions, +): Promise { + const { firstKeptEntryId, tokensBefore, previousSummary, previousPreserveData, fileOps } = preparation; + if (!firstKeptEntryId) { + throw new Error("First kept entry has no ID - session may need migration"); + } + const shape = options?.shape ?? resolveSnapcompactShape(options?.model?.api); + const frameSize = options?.frameSize ?? shape.frameSize; + const maxFrames = Math.max(1, options?.maxFrames ?? SNAPCOMPACT_MAX_FRAMES); + const geometry = snapcompactGeometry(shape, frameSize); + + const messages = preparation.messagesToSummarize.concat(preparation.turnPrefixMessages); + const llmMessages = (options?.convertToLlm ?? defaultConvertToLlm)(messages); + let archiveText = normalizeForSnapcompact(serializeSnapcompactConversation(llmMessages)); + + const previousArchive = getPreservedSnapcompactArchive(previousPreserveData); + const includedPreviousSummary = !previousArchive && !!previousSummary; + if (includedPreviousSummary && previousSummary) { + const head = `[Summary of earlier history] ${normalizeForSnapcompact(previousSummary)}`; + archiveText = archiveText.length > 0 ? `${head} [Recent conversation] ${archiveText}` : head; + } + + let truncatedChars = previousArchive?.truncatedChars ?? 0; + + const newFrames: SnapcompactFrame[] = []; + for (let offset = 0; offset < archiveText.length; offset += geometry.capacity) { + const chunk = archiveText.slice(offset, offset + geometry.capacity); + const rendered = renderSnapcompactFrame(chunk, shape, frameSize); + newFrames.push({ + data: Buffer.from(rendered.png).toBase64(), + mimeType: "image/png", + cols: rendered.cols, + rows: rendered.rows, + chars: rendered.chars, + font: shape.font, + variant: shape.variant, + lineRepeat: shape.lineRepeat, + ...(shape.imageDetail ? { detail: shape.imageDetail } : {}), + }); + // Keep the event loop responsive between native render passes. + await Bun.sleep(0); + } + + const frames = [...(previousArchive?.frames ?? []), ...newFrames]; + if (frames.length > maxFrames) { + // Pin the earliest frame: it anchors the session head (the original + // request, or the filmed summary of even older history) the way the + // LLM-summary strategies keep the original goal alive across rounds. + // Eviction removes the oldest *unpinned* frames, so the archive fades + // from the middle out — head and tail survive. With a budget of one + // frame the pin is moot; keep the newest frame instead. + const evictStart = maxFrames >= 2 ? 1 : 0; + const dropped = frames.splice(evictStart, frames.length - maxFrames); + for (const frame of dropped) truncatedChars += frame.chars; + } + const totalChars = frames.reduce((sum, frame) => sum + frame.chars, 0); + const mixedShapes = frames.some( + frame => + frame.cols !== geometry.cols || + frame.rows !== geometry.rows || + (frame.variant ?? "sent") !== shape.variant || + (frame.lineRepeat ?? 1) !== shape.lineRepeat, + ); + + let summary: string; + if (frames.length === 0) { + summary = "No prior history."; + } else { + summary = prompt.render(snapcompactSummaryPrompt, { + frameCount: frames.length, + multipleFrames: frames.length > 1, + fontCell: `${shape.cellWidth}x${shape.cellHeight}`, + cols: geometry.cols, + rows: geometry.rows, + sentenceInk: shape.variant === "sent", + lineRepeated: shape.lineRepeat > 1, + mixedShapes, + totalChars, + truncatedChars, + includedPreviousSummary, + }); + } + const { readFiles, modifiedFiles } = computeSnapcompactFileLists(fileOps); + summary = upsertSnapcompactFileOperations(summary, readFiles, modifiedFiles); + + // A snapcompact pass replaces any provider-side replacement history; strip the + // OpenAI remote-compaction payload like the default summarizer path does. + const basePreserve = stripOpenAiRemoteCompactionPreserveData(previousPreserveData) ?? {}; + const archive: SnapcompactArchive = { frames, totalChars, truncatedChars }; + + return { + summary, + shortSummary: `Archived ${totalChars.toLocaleString()} chars of history onto ${frames.length} snapcompact frame${frames.length === 1 ? "" : "s"}`, + firstKeptEntryId, + tokensBefore, + details: { readFiles, modifiedFiles }, + preserveData: { ...basePreserve, [SNAPCOMPACT_PRESERVE_KEY]: archive }, + }; +} diff --git a/packages/agent/test/snapcompact.test.ts b/packages/snapcompact/test/snapcompact.test.ts similarity index 64% rename from packages/agent/test/snapcompact.test.ts rename to packages/snapcompact/test/snapcompact.test.ts index fe224304c..316389ebb 100644 --- a/packages/agent/test/snapcompact.test.ts +++ b/packages/snapcompact/test/snapcompact.test.ts @@ -1,24 +1,60 @@ import { describe, expect, it } from "bun:test"; -import type { ImageContent } from "@oh-my-pi/pi-ai"; -import { type CompactionPreparation, type CompactionResult, estimateTokens } from "../src/compaction/compaction"; -import { createCompactionSummaryMessage, defaultConvertToLlm } from "../src/compaction/messages"; +import type { AssistantMessage, Message, Usage } from "@oh-my-pi/pi-ai"; import { + createSnapcompactFileOps, getPreservedSnapcompactArchive, normalizeForSnapcompact, renderSnapcompactFrame, - SNAPCOMPACT_FRAME_TOKEN_ESTIMATE, + resolveSnapcompactShape, SNAPCOMPACT_PRESERVE_KEY, + SNAPCOMPACT_SHAPES, type SnapcompactArchive, + type SnapcompactCompactionPreparation, + type SnapcompactCompactionResult, snapcompactCompact, snapcompactGeometry, -} from "../src/compaction/snapcompact"; -import { createFileOps } from "../src/compaction/utils"; -import { createAssistantMessage, createUserMessage } from "./helpers"; + snapcompactImages, +} from "../src"; -// Small frames keep render time negligible: 320px → 64 cols x 40 rows = 2560 chars. +// Small frames keep render time negligible. Legacy 5x8 shape: 320px → 64 cols +// x 40 rows = 2560 chars. Default (anthropic 8x8r-bw): 40 cols x 20 rows = 800. const TEST_FRAME_SIZE = 320; -function makePreparation(overrides: Partial = {}): CompactionPreparation { +function createUserMessage(content: string): Message { + return { role: "user", content, timestamp: 0 }; +} + +const ZERO_USAGE: Usage = { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + total: 0, + }, +}; + +function createAssistantMessage(content: AssistantMessage["content"]): Message { + return { + role: "assistant", + content, + api: "mock", + provider: "mock", + model: "mock", + usage: ZERO_USAGE, + stopReason: "stop", + timestamp: 0, + }; +} + +function makePreparation( + overrides: Partial> = {}, +): SnapcompactCompactionPreparation { return { firstKeptEntryId: "kept-1", messagesToSummarize: [ @@ -26,13 +62,10 @@ function makePreparation(overrides: Partial = {}): Compac createAssistantMessage([{ type: "text", text: "Fixed the TTL comparison in src/login.ts." }]), ], turnPrefixMessages: [], - recentMessages: [], - isSplitTurn: false, tokensBefore: 99000, previousSummary: undefined, previousPreserveData: undefined, - fileOps: createFileOps(), - settings: { enabled: true, reserveTokens: 16384, keepRecentTokens: 20000 }, + fileOps: createSnapcompactFileOps(), ...overrides, }; } @@ -106,12 +139,42 @@ describe("normalizeForSnapcompact", () => { }); }); +describe("shape resolution", () => { + it("maps provider APIs to their eval-winning shapes", () => { + expect(resolveSnapcompactShape("anthropic-messages")).toBe(SNAPCOMPACT_SHAPES.anthropic); + expect(resolveSnapcompactShape("openai-responses")).toBe(SNAPCOMPACT_SHAPES.openaiDense); + expect(resolveSnapcompactShape("azure-openai-responses")).toBe(SNAPCOMPACT_SHAPES.openaiDense); + expect(resolveSnapcompactShape("google-generative-ai")).toBe(SNAPCOMPACT_SHAPES.google); + // Unknown and absent APIs fall back to the refusal-robust plain shape. + expect(resolveSnapcompactShape("some-future-api")).toBe(SNAPCOMPACT_SHAPES.anthropic); + expect(resolveSnapcompactShape(undefined)).toBe(SNAPCOMPACT_SHAPES.anthropic); + }); + + it("snapcompactImages forwards the per-frame detail hint", () => { + const archive: SnapcompactArchive = { + frames: [ + { data: "ZmFrZQ==", mimeType: "image/png", cols: 10, rows: 10, chars: 5, detail: "original" }, + { data: "ZmFrZTI=", mimeType: "image/png", cols: 10, rows: 10, chars: 5 }, + ], + totalChars: 10, + truncatedChars: 0, + }; + const [withDetail, without] = snapcompactImages(archive); + expect(withDetail.detail).toBe("original"); + expect("detail" in without).toBe(false); + }); +}); + describe("renderSnapcompactFrame", () => { - it("produces an indexed PNG of the declared geometry with sentence-cycled ink", () => { - const geometry = snapcompactGeometry(TEST_FRAME_SIZE); + it("produces an indexed PNG of the declared geometry with sentence-cycled ink (legacy 5x8)", () => { + const geometry = snapcompactGeometry(SNAPCOMPACT_SHAPES.legacy, TEST_FRAME_SIZE); expect(geometry).toEqual({ cols: 64, rows: 40, capacity: 2560 }); - const frame = renderSnapcompactFrame("First sentence here. Second one differs.", TEST_FRAME_SIZE); + const frame = renderSnapcompactFrame( + "First sentence here. Second one differs.", + SNAPCOMPACT_SHAPES.legacy, + TEST_FRAME_SIZE, + ); expect(frame.cols).toBe(64); expect(frame.rows).toBe(40); expect(frame.chars).toBe(40); @@ -128,24 +191,46 @@ describe("renderSnapcompactFrame", () => { expect(used.has(3)).toBe(false); }); + it("renders the anthropic shape with doubled lines, black ink, and highlight bands", () => { + const geometry = snapcompactGeometry(SNAPCOMPACT_SHAPES.anthropic, TEST_FRAME_SIZE); + expect(geometry).toEqual({ cols: 40, rows: 20, capacity: 800 }); + + const frame = renderSnapcompactFrame("Hello world. Again.", SNAPCOMPACT_SHAPES.anthropic, TEST_FRAME_SIZE); + const decoded = decodePng(frame.png); + expect(decoded.colorType).toBe(3); + const used = new Set(decoded.pixels); + expect(used.has(7)).toBe(true); // black bw ink + expect(used.has(8)).toBe(true); // repeat highlight band + expect(used.has(1)).toBe(false); // no sentence hues in bw + }); + + it("renders the openai stretch shape as truecolor RGB", () => { + const frame = renderSnapcompactFrame("Hello world.", SNAPCOMPACT_SHAPES.openaiDense, TEST_FRAME_SIZE); + // IHDR color type byte: 2 = truecolor RGB (anti-aliased stretch output). + expect(frame.png[25]).toBe(2); + expect(frame.cols).toBe(Math.floor(TEST_FRAME_SIZE / 6)); + }); + it("caps printed characters at frame capacity", () => { - const { capacity } = snapcompactGeometry(TEST_FRAME_SIZE); - const frame = renderSnapcompactFrame("x".repeat(capacity + 500), TEST_FRAME_SIZE); + const { capacity } = snapcompactGeometry(SNAPCOMPACT_SHAPES.legacy, TEST_FRAME_SIZE); + const frame = renderSnapcompactFrame("x".repeat(capacity + 500), SNAPCOMPACT_SHAPES.legacy, TEST_FRAME_SIZE); expect(frame.chars).toBe(capacity); }); }); describe("snapcompactCompact", () => { it("archives history onto frames with a self-describing summary", async () => { - const fileOps = createFileOps(); + const fileOps = createSnapcompactFileOps(); fileOps.read.add("src/auth.ts"); fileOps.edited.add("src/login.ts"); const result = await snapcompactCompact(makePreparation({ fileOps }), { frameSize: TEST_FRAME_SIZE }); expect(result.firstKeptEntryId).toBe("kept-1"); expect(result.tokensBefore).toBe(99000); - // Reading instructions reflect the actual grid geometry. - expect(result.summary).toContain("64 characters per row"); + // Reading instructions reflect the default (anthropic 8x8r-bw) shape. + expect(result.summary).toContain("40 characters per row"); + expect(result.summary).toContain("printed twice"); + expect(result.summary).toContain("plain black ink"); expect(result.summary).toContain("snapcompact frame"); // File operations are upserted like every other compaction summary. expect(result.summary).toContain(""); @@ -157,6 +242,9 @@ describe("snapcompactCompact", () => { expect(archive?.frames.length).toBe(1); expect(archive?.frames[0].mimeType).toBe("image/png"); expect(archive?.frames[0].chars).toBe(archive?.totalChars); + expect(archive?.frames[0].font).toBe("8x8"); + expect(archive?.frames[0].variant).toBe("bw"); + expect(archive?.frames[0].lineRepeat).toBe(2); expect(archive?.truncatedChars).toBe(0); // Frame data round-trips as a decodable PNG. const decoded = decodePng(Buffer.from(archive?.frames[0].data ?? "", "base64")); @@ -164,7 +252,7 @@ describe("snapcompactCompact", () => { }); it("splits oversized history across frames and evicts beyond the budget", async () => { - const { capacity } = snapcompactGeometry(TEST_FRAME_SIZE); + const { capacity } = snapcompactGeometry(SNAPCOMPACT_SHAPES.anthropic, TEST_FRAME_SIZE); // Sentences avoid whitespace collapse shrinking the payload below 2.5 frames. const longText = "Important fact number one. ".repeat(Math.ceil((capacity * 2.5) / 28)); const result = await snapcompactCompact(makePreparation({ messagesToSummarize: [createUserMessage(longText)] }), { @@ -178,7 +266,7 @@ describe("snapcompactCompact", () => { }); it("evicts the oldest unpinned frames, keeping the session-head frame alive", async () => { - let previous: CompactionResult | undefined; + let previous: SnapcompactCompactionResult | undefined; let headFrameData = ""; let secondFrameData = ""; for (let pass = 1; pass <= 4; pass++) { @@ -237,46 +325,27 @@ describe("snapcompactCompact", () => { expect(second.preserveData?.openaiRemoteCompaction).toBeUndefined(); expect(second.preserveData?.appKey).toBe("kept"); }); + + it("flags mixed shapes when merged frames disagree with the active shape", async () => { + const first = await snapcompactCompact(makePreparation(), { + frameSize: TEST_FRAME_SIZE, + shape: SNAPCOMPACT_SHAPES.legacy, + }); + const second = await snapcompactCompact( + makePreparation({ + messagesToSummarize: [createUserMessage("Another turn after a provider switch.")], + previousSummary: first.summary, + previousPreserveData: first.preserveData, + }), + { frameSize: TEST_FRAME_SIZE, model: { api: "anthropic-messages" } }, + ); + expect(second.summary).toContain("Older frames may use a different font"); + // Same-shape merges stay silent. + expect(first.summary).not.toContain("Older frames may use a different font"); + }); }); -describe("compaction summary message with snapcompact frames", () => { - const images: ImageContent[] = [ - { type: "image", data: "ZmFrZQ==", mimeType: "image/png" }, - { type: "image", data: "ZmFrZTI=", mimeType: "image/png" }, - ]; - - it("estimateTokens charges per attached frame", () => { - const bare = createCompactionSummaryMessage("summary text", 1000, new Date().toISOString()); - const withFrames = createCompactionSummaryMessage( - "summary text", - 1000, - new Date().toISOString(), - undefined, - undefined, - images, - ); - expect(estimateTokens(withFrames) - estimateTokens(bare)).toBe(2 * SNAPCOMPACT_FRAME_TOKEN_ESTIMATE); - }); - - it("defaultConvertToLlm appends frames as image blocks after the summary text", () => { - const message = createCompactionSummaryMessage( - "the snapcompact archive", - 1000, - new Date().toISOString(), - undefined, - undefined, - images, - ); - const [converted] = defaultConvertToLlm([message]); - expect(converted.role).toBe("user"); - const content = converted.content as Array<{ type: string; text?: string; data?: string }>; - expect(content.length).toBe(3); - expect(content[0].type).toBe("text"); - expect(content[0].text).toContain("the snapcompact archive"); - expect(content[1]).toEqual(images[0]); - expect(content[2]).toEqual(images[1]); - }); - +describe("archive helpers", () => { it("getPreservedSnapcompactArchive rejects malformed payloads", () => { expect(getPreservedSnapcompactArchive(undefined)).toBeUndefined(); expect(getPreservedSnapcompactArchive({ [SNAPCOMPACT_PRESERVE_KEY]: "nope" })).toBeUndefined(); diff --git a/packages/snapcompact/tsconfig.json b/packages/snapcompact/tsconfig.json new file mode 100644 index 000000000..08130e07c --- /dev/null +++ b/packages/snapcompact/tsconfig.json @@ -0,0 +1,7 @@ +{ + "extends": "../tsconfig.workspace.json", + "include": [ + "src", + "test" + ] +} diff --git a/packages/snapcompact/tsconfig.publish.json b/packages/snapcompact/tsconfig.publish.json new file mode 100644 index 000000000..c216426de --- /dev/null +++ b/packages/snapcompact/tsconfig.publish.json @@ -0,0 +1,25 @@ +{ + "extends": "./tsconfig.json", + "compilerOptions": { + "noEmit": false, + "emitDeclarationOnly": true, + "declaration": true, + "declarationMap": false, + "sourceMap": false, + "inlineSources": false, + "rootDir": "src", + "outDir": "dist/types", + "noCheck": true + }, + "include": [ + "src" + ], + "exclude": [ + "dist", + "node_modules", + "test", + "tests", + "scripts", + "bench" + ] +} diff --git a/packages/typescript-edit-benchmark/package.json b/packages/typescript-edit-benchmark/package.json index 9ad9028fe..21c8553b7 100644 --- a/packages/typescript-edit-benchmark/package.json +++ b/packages/typescript-edit-benchmark/package.json @@ -4,6 +4,7 @@ "name": "@oh-my-pi/typescript-edit-benchmark", "version": "0.0.1", "description": "Edit benchmark suite using TypeScript source code mutations", + "homepage": "https://omp.sh", "author": "Can Boluk", "license": "MIT", "repository": { diff --git a/packages/utils/test/prompt.test.ts b/packages/utils/test/prompt.test.ts index 7dd5e2126..7b0700bf3 100644 --- a/packages/utils/test/prompt.test.ts +++ b/packages/utils/test/prompt.test.ts @@ -96,7 +96,7 @@ describe("compile cache", () => { }); describe("helpers: join", () => { - it('unescapes \\n and \\t in the separator (Handlebars string literals carry no escapes)', () => { + it("unescapes \\n and \\t in the separator (Handlebars string literals carry no escapes)", () => { // Regression: `{{join files "\n"}}` used to emit the literal two-char `\n` // between entries (visible in compaction lists). expect(prompt.render('{{join files "\\n"}}', { files: ["a.ts", "b.ts"] })).toBe("a.ts\nb.ts"); diff --git a/python/robomp/web/package.json b/python/robomp/web/package.json index feb6caeaa..681d6471c 100644 --- a/python/robomp/web/package.json +++ b/python/robomp/web/package.json @@ -4,6 +4,7 @@ "version": "0.1.0", "type": "module", "description": "Glassmorphic SolidJS dashboard bundled by Vite and served by robomp's FastAPI app.", + "homepage": "https://omp.sh", "scripts": { "dev": "vite", "build": "vite build", diff --git a/scripts/ci-release-publish.ts b/scripts/ci-release-publish.ts index b5b9af951..87b7d36e1 100644 --- a/scripts/ci-release-publish.ts +++ b/scripts/ci-release-publish.ts @@ -91,6 +91,7 @@ export const packages: PublishPackage[] = [ { dir: "packages/tui", kind: "typescript" }, { dir: "packages/hashline", kind: "typescript" }, { dir: "packages/mnemopi", kind: "typescript" }, + { dir: "packages/snapcompact", kind: "typescript" }, { dir: "packages/stats", kind: "typescript", diff --git a/scripts/install-tests/run-ci.sh b/scripts/install-tests/run-ci.sh index 62f50b140..f401b1895 100755 --- a/scripts/install-tests/run-ci.sh +++ b/scripts/install-tests/run-ci.sh @@ -92,7 +92,7 @@ cp "$natives_pkg_backup" "$ROOT_DIR/packages/natives/package.json" # 3. Pack the remaining workspace packages (natives core and coding-agent # handled separately). -for pkg in utils hashline catalog ai mnemopi agent tui stats; do +for pkg in utils hashline catalog ai mnemopi snapcompact agent tui stats; do ( cd "$ROOT_DIR/packages/$pkg" bun pm pack --destination "$TARBALL_DIR" --quiet >/dev/null @@ -122,6 +122,7 @@ hashline_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-hashline-*.tgz)" catalog_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-pi-catalog-*.tgz)" ai_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-pi-ai-*.tgz)" mnemopi_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-pi-mnemopi-*.tgz)" +snapcompact_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-snapcompact-*.tgz)" agent_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-pi-agent-core-*.tgz)" tui_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-pi-tui-*.tgz)" stats_tgz="$(find_tarball "$TARBALL_DIR"/oh-my-pi-omp-stats-*.tgz)" @@ -145,6 +146,7 @@ mkdir -p "$TARBALL_APP_DIR" '@oh-my-pi/pi-ai': '$ai_tgz', '@oh-my-pi/pi-catalog': '$catalog_tgz', '@oh-my-pi/pi-mnemopi': '$mnemopi_tgz', + '@oh-my-pi/snapcompact': '$snapcompact_tgz', '@oh-my-pi/pi-agent-core': '$agent_tgz', '@oh-my-pi/pi-tui': '$tui_tgz', '@oh-my-pi/omp-stats': '$stats_tgz', @@ -153,7 +155,7 @@ mkdir -p "$TARBALL_APP_DIR" require('fs').writeFileSync('package.json', JSON.stringify(pkg, null, 2)); " - bun add "$utils_tgz" "$natives_tgz" "$hashline_tgz" "$catalog_tgz" "$ai_tgz" "$mnemopi_tgz" "$agent_tgz" "$tui_tgz" "$stats_tgz" "$coding_agent_tgz" + bun add "$utils_tgz" "$natives_tgz" "$hashline_tgz" "$catalog_tgz" "$ai_tgz" "$mnemopi_tgz" "$snapcompact_tgz" "$agent_tgz" "$tui_tgz" "$stats_tgz" "$coding_agent_tgz" # The platform leaf must arrive through the core's optionalDependencies + # override, not as a direct dependency — assert it landed before smoking so a # resolution regression is distinguishable from a runtime loader bug.