Nadir 1.0: 15.9M on-device dark-mode model, int8/int4 quants, runtime, extension
Browse files- README.md +218 -0
- config.json +36 -0
- model.py +141 -0
- model.safetensors +3 -0
- nadir-extension.zip +3 -0
- onnx/nadir-fp32.onnx +3 -0
- onnx/nadir-int4.onnx +3 -0
- onnx/nadir-int8.onnx +3 -0
- reports/eval_judge.json +88 -0
- reports/export_report.json +20 -0
- reports/ortweb_timing_m4max.json +67 -0
- reports/quant_report.json +63 -0
- reports/train_log.json +501 -0
- runtime/DARKREADER-LICENSE +23 -0
- runtime/darkreader-nadir.js +0 -0
- runtime/extract.js +825 -0
- runtime/nadir-infer.js +193 -0
- runtime/nadir-map.js +244 -0
- training/build_ds.py +281 -0
- training/export.py +211 -0
- training/haiku_gen.py +621 -0
- training/judge.py +110 -0
- training/model.py +141 -0
- training/predict.py +93 -0
- training/quantize.py +330 -0
- training/safety.py +188 -0
- training/train.py +292 -0
README.md
ADDED
|
@@ -0,0 +1,218 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: apache-2.0
|
| 3 |
+
library_name: onnx
|
| 4 |
+
pipeline_tag: other
|
| 5 |
+
tags:
|
| 6 |
+
- dark-mode
|
| 7 |
+
- browser-extension
|
| 8 |
+
- onnx
|
| 9 |
+
- onnxruntime-web
|
| 10 |
+
- accessibility
|
| 11 |
+
- design
|
| 12 |
+
- distillation
|
| 13 |
+
---
|
| 14 |
+
|
| 15 |
+
# Nadir
|
| 16 |
+
|
| 17 |
+
**Nadir** is a small on-device model that gives light websites a tasteful, readable dark mode, with the
|
| 18 |
+
care of a real dark theme rather than a filter. It runs entirely in your browser (onnxruntime-web, WASM),
|
| 19 |
+
in about 35 ms of model time per page, and ships with a ready-to-load Chromium extension.
|
| 20 |
+
|
| 21 |
+
Nadir follows three rules:
|
| 22 |
+
|
| 23 |
+
1. **Already dark? Leave it alone.** Pages that are already dark are detected and never touched.
|
| 24 |
+
2. **Has its own dark mode? Use it.** Nadir switches on the site's own dark theme (`prefers-color-scheme`
|
| 25 |
+
media rules, or `dark`/`theme-dark` class and attribute toggles it detects in the page's CSS).
|
| 26 |
+
3. **No dark mode? The model designs one.** Nadir reads the page's colors and images and predicts a dark
|
| 27 |
+
palette for that specific site: brand colors stay recognizable, surfaces keep their hierarchy, text stays
|
| 28 |
+
readable, and logos stay visible: a logo on its own white rectangle is adapted to sit cleanly on the dark
|
| 29 |
+
surface. Photos, videos and ads are never color-shifted.
|
| 30 |
+
|
| 31 |
+
In a blind pairwise test on 146 sites it never trained on, judged by Claude Haiku 5.5, the preference
|
| 32 |
+
went to Nadir's dark mode over Dark Reader's on **93 sites vs 50** (3 ties): 65% of the decided pages.
|
| 33 |
+
|
| 34 |
+
## How it works
|
| 35 |
+
|
| 36 |
+
Nadir doesn't paint pixels. It redesigns the page's **color system**:
|
| 37 |
+
|
| 38 |
+
1. **Extraction (`runtime/extract.js`).** One synchronous pass over the DOM collects every distinct color
|
| 39 |
+
*in its role*: background, text, border, shadow, gradient, SVG fill, SVG stroke. Each token gets 49
|
| 40 |
+
features: lightness, chroma, hue, area, text size and weight, where it appears (header, nav, link, form,
|
| 41 |
+
code, footer), what it sits on and the original contrast. It also records four relations between tokens
|
| 42 |
+
(text-on-surface, surface-on-surface, border-on-surface, same-element) and every visible image or
|
| 43 |
+
inline-SVG logo with 39 features plus a 32×32 thumbnail. Ad slots are excluded.
|
| 44 |
+
2. **Model (`NadirNet`).** A relation-biased transformer over `[page] + tokens + images`. Each attention head
|
| 45 |
+
gets a learned bias from the 11 relation channels, so the model knows which text sits on which surface.
|
| 46 |
+
It predicts for every token an OKLab color as a parameterized edit of the input (lightness, chroma scale,
|
| 47 |
+
hue rotation, tint, alpha). This keeps brand hues stable and makes impossible colors rare. For every
|
| 48 |
+
image it predicts a treatment: `keep`, `dim`, `invert`, `light` (monochrome dark glyph → near-white) or
|
| 49 |
+
`plate` (dark multicolor logo that needs help; shown with a hue-preserving inversion at runtime). A small
|
| 50 |
+
CNN reads the thumbnails.
|
| 51 |
+
3. **Design rules (`runtime/nadir-infer.js`, `runtime/nadir-map.js`).** Deterministic guards on top of the
|
| 52 |
+
model:
|
| 53 |
+
- The page canvas is always dark.
|
| 54 |
+
- Every text and icon color keeps WCAG contrast on *every* surface it actually sits on. It never needs
|
| 55 |
+
more contrast than the original design had.
|
| 56 |
+
- Brand colors keep their hue (±45°).
|
| 57 |
+
- Things that were invisible in the light design stay invisible: text-fade gradients and borders that
|
| 58 |
+
matched their surface follow the surface.
|
| 59 |
+
- Colors drawn over video and photos are left alone when they already read as dark mode (player
|
| 60 |
+
controls, scrims, captions).
|
| 61 |
+
- Photos and video are never recolored.
|
| 62 |
+
- Light, flat background images behind content (paper textures, pale washes, SVG gradients) are cleared so the
|
| 63 |
+
designed dark surface shows through. Images with real detail are left alone.
|
| 64 |
+
- Inline SVG logos stay readable: their shapes are drawn as ink, and SVG mask and clip-path shapes (knockout
|
| 65 |
+
text, cut-outs) are never recolored.
|
| 66 |
+
- **Logos on their own white rectangle are adapted.** The white background is removed so Nadir's surface
|
| 67 |
+
shows through, black or gray ink becomes light, and brand colors keep their hue (lifted only as far as
|
| 68 |
+
needed to read on dark). This applies only to images the page itself marks as its logo.
|
| 69 |
+
- Dark logos on transparency get a hue-preserving inversion, and already-light ones are kept.
|
| 70 |
+
4. **Engine.** The predicted palette is applied through a patched
|
| 71 |
+
[Dark Reader](https://github.com/darkreader/darkreader) 4.9.133 dynamic engine (MIT), which handles the
|
| 72 |
+
hard plumbing: stylesheets, inline styles, shadow DOM and dynamically inserted content. It asks Nadir for
|
| 73 |
+
every color it rewrites. Colors the model didn't see (new content) are mapped by nearest-neighbor
|
| 74 |
+
lookup in the predicted palette.
|
| 75 |
+
|
| 76 |
+
## Model
|
| 77 |
+
|
| 78 |
+
| | |
|
| 79 |
+
|---|---|
|
| 80 |
+
| Parameters | 15.9M (15,895,995) |
|
| 81 |
+
| Width / depth / heads | 384 / 8 layers / 6 heads |
|
| 82 |
+
| Inputs | ≤256 color tokens, ≤32 images, 11 relation channels |
|
| 83 |
+
| Outputs | OKLab color + alpha per token, 5-way image treatment |
|
| 84 |
+
| Runtime | onnxruntime-web 1.29 (WASM SIMD, multi-threaded) |
|
| 85 |
+
|
| 86 |
+
### Files
|
| 87 |
+
|
| 88 |
+
| File | Precision | Size | Δ color vs fp32 (mean / max, OKLab) | Image agreement |
|
| 89 |
+
|---|---|---|---|---|
|
| 90 |
+
| `onnx/nadir-fp32.onnx` | fp32 | 63.7 MB | 0 / 0 | 100% |
|
| 91 |
+
| `onnx/nadir-int8.onnx` | 8-bit weights | 20.6 MB | 0.00015 / 0.025 | 99.99% |
|
| 92 |
+
| `onnx/nadir-int4.onnx` | 4-bit weights (8-bit input/output layers) | 13.3 MB | 0.00036 / 0.071 | 99.76% |
|
| 93 |
+
|
| 94 |
+
Fidelity is measured on 400 held-out pages. A difference of 0.02 in OKLab is about the smallest one you can see.
|
| 95 |
+
|
| 96 |
+
Speed in onnxruntime-web 1.29 (Chrome, Apple M4 Max, 4 WASM threads), per page:
|
| 97 |
+
|
| 98 |
+
| File | Load | Inference |
|
| 99 |
+
|---|---|---|
|
| 100 |
+
| fp32 | ~7 s | 34 ms |
|
| 101 |
+
| int8 | ~1.5 s | 33–35 ms |
|
| 102 |
+
| int4 | ~1.0 s | 38 ms |
|
| 103 |
+
|
| 104 |
+
The extension's own end-to-end time for a first visit is 0.7–3 s, mostly reading the page and its images. Revisits apply
|
| 105 |
+
from cache.
|
| 106 |
+
|
| 107 |
+
- `onnx/nadir-int8.onnx` is the extension's default.
|
| 108 |
+
- The 8-bit and 4-bit models are **calibrated weight-only quants**:
|
| 109 |
+
- Layer-sequential GPTQ, with Hessians collected from 256 real crawled pages flowing through the
|
| 110 |
+
already-quantized earlier layers.
|
| 111 |
+
- Asymmetric codes in groups of 32, with each group's clipping range chosen by an MSE search.
|
| 112 |
+
- The 4-bit model keeps its input and output layers at 8 bits.
|
| 113 |
+
- Storage format:
|
| 114 |
+
- 4-bit layers use ONNX `MatMulNBits` (block 32).
|
| 115 |
+
- 8-bit layers store uint8 codes with a per-block scale and zero point, dequantized in the graph by standard
|
| 116 |
+
ops (`Cast`/`Sub`/`Mul`/`Reshape`). onnxruntime folds these into the weights when it loads the model, so
|
| 117 |
+
the 8-bit layers run as exact fp32 math with no special kernels. Both files are tested on onnxruntime 1.22
|
| 118 |
+
(CPU) and onnxruntime-web 1.29.
|
| 119 |
+
- Each quant is verified three ways: packing parity against a PyTorch fake-quant model (≤2e-6), fidelity
|
| 120 |
+
against fp32, and quality against the training targets.
|
| 121 |
+
- They run in onnxruntime-web with outputs matching onnxruntime CPU to ~1e-5.
|
| 122 |
+
- Activation-quantized int8 (ORT dynamic/static) was rejected: worst-case color error up to 0.99.
|
| 123 |
+
|
| 124 |
+
`model.safetensors` + `model.py` hold the PyTorch weights and architecture.
|
| 125 |
+
|
| 126 |
+
## Training data
|
| 127 |
+
|
| 128 |
+
- **Pages.** A crawl of 47,569 sites from the [Tranco](https://tranco-list.eu/) top list at 1280×800.
|
| 129 |
+
Each page was captured in light mode, in emulated dark mode, and with dark-class toggles applied.
|
| 130 |
+
- **Teacher.** 19,563 light pages were redesigned by **Claude Haiku 5.5** (Anthropic Batch API,
|
| 131 |
+
high effort) acting as the site's own product designer. It was given the screenshot, a contact sheet of
|
| 132 |
+
the images, and the full color inventory with roles and locations, and returned a dark palette plus image
|
| 133 |
+
treatments in a strict JSON schema. The design brief is in `training/haiku_gen.py`.
|
| 134 |
+
- **Native pairs.** 4,142 pages with a real built-in dark theme give human-designed light→dark
|
| 135 |
+
pairs, aligned token by token.
|
| 136 |
+
- **Quality control.**
|
| 137 |
+
- Designs that missed ids, kept a light canvas, or left >15% of text below 3:1 contrast were dropped.
|
| 138 |
+
- The rest were passed through the same contrast rule the runtime uses (14,338 of 19,464 designs adjusted).
|
| 139 |
+
- The train/validation split is by site (hash of the hostname).
|
| 140 |
+
- **Augmentation.** Random token dropout (it also simulates ad blockers and partial loads) and whole-design
|
| 141 |
+
hue rotation.
|
| 142 |
+
|
| 143 |
+
## Evaluation
|
| 144 |
+
|
| 145 |
+
The judge was a blind pairwise one: Claude Haiku 5.5 saw the original light page plus two dark versions in
|
| 146 |
+
random order, A and B. It judged readability first, then breakage (glaring light patches, lost logos,
|
| 147 |
+
distorted photos), then design quality. All pages are held-out sites the model never trained on.
|
| 148 |
+
|
| 149 |
+
| Nadir vs … (146 held-out sites) | Nadir preferred | Other preferred | Ties |
|
| 150 |
+
|---|---|---|---|
|
| 151 |
+
| Dark Reader 4.9 (dynamic mode) | **93** | 50 | 3 |
|
| 152 |
+
| Claude Haiku 5.5's own designs (the teacher)¹ | 76 | 61 | 9 |
|
| 153 |
+
| Nadir before its final design rules² | 45 | 29 | 72 |
|
| 154 |
+
|
| 155 |
+
¹ The teacher's designs were rendered with the earlier runtime, so part of this margin comes from the runtime's newer
|
| 156 |
+
design rules rather than from the model itself. Before those rules, the model and its teacher were even (63 vs 67).
|
| 157 |
+
² The rules for light background textures, SVG logos and SVG masks.
|
| 158 |
+
|
| 159 |
+
**Model size.** A 34.0M-parameter version (512 / 10 layers / 8 heads) trained on the same data was no better: 49 vs 53
|
| 160 |
+
with 45 ties against this one. It reached the same color error and runs about 2× slower in the browser, so the 15.9M
|
| 161 |
+
model is the release.
|
| 162 |
+
|
| 163 |
+
Held-out color error vs the teacher designs (OKLab ΔE): 0.040 (Haiku pages), 0.103
|
| 164 |
+
(native dark themes). Image-treatment accuracy: 79.4% over 5 classes.
|
| 165 |
+
|
| 166 |
+
## Use it
|
| 167 |
+
|
| 168 |
+
### Browser extension (Chrome, Edge, Brave, Arc, Aside and other Chromium browsers)
|
| 169 |
+
|
| 170 |
+
1. Download `nadir-extension.zip` and unzip it.
|
| 171 |
+
2. Open `chrome://extensions`, enable **Developer mode**, click **Load unpacked**, and choose the folder.
|
| 172 |
+
3. Browse. The toolbar badge shows what Nadir did on each site:
|
| 173 |
+
- `AI`: model-designed dark mode.
|
| 174 |
+
- `ON`: the site's own dark mode.
|
| 175 |
+
- No badge: already dark.
|
| 176 |
+
|
| 177 |
+
The popup can turn Nadir off for a site.
|
| 178 |
+
|
| 179 |
+
The extension works alongside content blockers such as uBlock Origin Lite: removed ads simply aren't part
|
| 180 |
+
of the palette, and ads that remain are treated as images (never recolored).
|
| 181 |
+
|
| 182 |
+
### In your own code
|
| 183 |
+
|
| 184 |
+
```js
|
| 185 |
+
import * as ort from 'onnxruntime-web';
|
| 186 |
+
// runtime/extract.js, runtime/nadir-infer.js and runtime/nadir-map.js define the globals used below
|
| 187 |
+
const page = await __nadirExtract({describe: false, thumbs: false}); // in the page: colors, roles, relations, images
|
| 188 |
+
// pixel stats + thumbnails for the images (any context that can fetch them, e.g. an extension page)
|
| 189 |
+
const st = await __nadirImageStats(page.images, {dpr: devicePixelRatio});
|
| 190 |
+
page.images.forEach((im, i) => { Object.assign(im, st[i]); im.f = __nadirImageFeatures(im, page.vw, page.vh, page.docH); });
|
| 191 |
+
const built = NadirInfer.buildInputs(page);
|
| 192 |
+
const session = await ort.InferenceSession.create('nadir-int8.onnx', {executionProviders: ['wasm']});
|
| 193 |
+
const feeds = Object.fromEntries(Object.entries(built.feeds).map(([k, [data, dims]]) =>
|
| 194 |
+
[k, new ort.Tensor(k === 'tok_role' || k === 'src' ? 'int64' : k === 'mask' ? 'bool' : 'float32', data, dims)]));
|
| 195 |
+
const out = await session.run(feeds);
|
| 196 |
+
const {colors, keep, treatments} = NadirInfer.decode(page, built,
|
| 197 |
+
{lab: out.lab.data, alpha_logit: out.alpha_logit.data, img_logits: out.img_logits.data});
|
| 198 |
+
// colors[i] is the dark replacement for page.tokens[i] (null = keep); treatments[j] applies to page.images[j]
|
| 199 |
+
```
|
| 200 |
+
|
| 201 |
+
## Limitations
|
| 202 |
+
|
| 203 |
+
- Nadir recolors CSS. Text and UI baked into images, canvas drawings and cross-origin iframes keep their
|
| 204 |
+
original colors, though the iframes get their own Nadir pass in the extension.
|
| 205 |
+
- Colors are mapped per (color, role). If one exact color is used both as a light card and as a light
|
| 206 |
+
button, both get the same dark color.
|
| 207 |
+
- Background images that mix light areas with strong color (for example a colored hero shape with a white base) are
|
| 208 |
+
kept as they are, so their light parts stay light.
|
| 209 |
+
- The model reads at most 256 colors per page. On very colorful pages the rarest colors are mapped by nearest
|
| 210 |
+
neighbor in the predicted palette.
|
| 211 |
+
- The teacher is an LLM. Its taste is good but not perfect, and the model inherits its preferences.
|
| 212 |
+
- The judge is also an LLM, so the evaluation is automated, not a human study.
|
| 213 |
+
|
| 214 |
+
## License
|
| 215 |
+
|
| 216 |
+
Apache-2.0 for the model, weights and Nadir code. The bundled page engine is a modified
|
| 217 |
+
[Dark Reader](https://github.com/darkreader/darkreader) (MIT License, © Dark Reader Ltd.); its license is
|
| 218 |
+
included in the extension (`lib/DARKREADER-LICENSE`).
|
config.json
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"d": 384,
|
| 3 |
+
"layers": 8,
|
| 4 |
+
"heads": 6,
|
| 5 |
+
"ff": 4,
|
| 6 |
+
"drop": 0.1,
|
| 7 |
+
"architecture": "NadirNet",
|
| 8 |
+
"token_features": 49,
|
| 9 |
+
"image_features": 39,
|
| 10 |
+
"relation_channels": 11,
|
| 11 |
+
"max_tokens": 256,
|
| 12 |
+
"max_images": 32,
|
| 13 |
+
"image_treatments": [
|
| 14 |
+
"keep",
|
| 15 |
+
"dim",
|
| 16 |
+
"invert",
|
| 17 |
+
"light",
|
| 18 |
+
"plate"
|
| 19 |
+
],
|
| 20 |
+
"inputs": [
|
| 21 |
+
"tok_f",
|
| 22 |
+
"tok_role",
|
| 23 |
+
"tok_lab",
|
| 24 |
+
"img_f",
|
| 25 |
+
"img_thumb",
|
| 26 |
+
"page_f",
|
| 27 |
+
"src",
|
| 28 |
+
"rel",
|
| 29 |
+
"mask"
|
| 30 |
+
],
|
| 31 |
+
"outputs": [
|
| 32 |
+
"lab",
|
| 33 |
+
"alpha_logit",
|
| 34 |
+
"img_logits"
|
| 35 |
+
]
|
| 36 |
+
}
|
model.py
ADDED
|
@@ -0,0 +1,141 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""NadirNet: a relation-biased set transformer over a page's colour palette and images.
|
| 2 |
+
|
| 3 |
+
Tokens: [PAGE] + one token per (colour, role) the page uses + one token per image.
|
| 4 |
+
Attention logits get a learned per-head bias from the page's structure graph (text-on-surface,
|
| 5 |
+
surface-on-surface, border-on-surface, same-element, image-on-surface, same-colour), so the model
|
| 6 |
+
reasons about hierarchy and contrast the way a designer would. Outputs a dark colour per token
|
| 7 |
+
(OKLab, parameterised relative to the input so hue is kept unless the model chooses otherwise),
|
| 8 |
+
an alpha, and a treatment class per image.
|
| 9 |
+
"""
|
| 10 |
+
import math
|
| 11 |
+
import torch
|
| 12 |
+
import torch.nn as nn
|
| 13 |
+
import torch.nn.functional as F
|
| 14 |
+
|
| 15 |
+
FT, FI, FP = 49, 39, 8
|
| 16 |
+
N_ROLES, N_TREAT = 7, 5
|
| 17 |
+
K_REL = 11
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def fourier_lab(lab, n=6):
|
| 21 |
+
x = torch.stack([lab[..., 0], lab[..., 1] * 4, lab[..., 2] * 4], -1) # (..., 3)
|
| 22 |
+
freqs = (2.0 ** torch.arange(n, device=lab.device, dtype=lab.dtype)) * math.pi
|
| 23 |
+
ang = x[..., None] * freqs # (...,3,n)
|
| 24 |
+
return torch.cat([torch.sin(ang), torch.cos(ang)], -1).flatten(-2) # (..., 6n)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
class MLP(nn.Module):
|
| 28 |
+
def __init__(self, i, h, o, drop=0.0):
|
| 29 |
+
super().__init__()
|
| 30 |
+
self.net = nn.Sequential(nn.Linear(i, h), nn.GELU(), nn.Dropout(drop), nn.Linear(h, o))
|
| 31 |
+
|
| 32 |
+
def forward(self, x):
|
| 33 |
+
return self.net(x)
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
class ThumbCNN(nn.Module):
|
| 37 |
+
"""32x32 RGBA thumbnail -> embedding. Alpha matters: transparent logos behave differently from photos."""
|
| 38 |
+
def __init__(self, o=128):
|
| 39 |
+
super().__init__()
|
| 40 |
+
c = [5, 32, 64, 96, 128]
|
| 41 |
+
layers = []
|
| 42 |
+
for a, b in zip(c[:-1], c[1:]):
|
| 43 |
+
layers += [nn.Conv2d(a, b, 3, padding=1), nn.GELU(), nn.Conv2d(b, b, 3, padding=1, stride=2), nn.GELU()]
|
| 44 |
+
self.net = nn.Sequential(*layers)
|
| 45 |
+
self.proj = nn.Linear(128 * 2 * 2, o)
|
| 46 |
+
|
| 47 |
+
def forward(self, x): # (B, M, 4, 32, 32) in [0,1]
|
| 48 |
+
B, M = x.shape[:2]
|
| 49 |
+
x = x.reshape(B * M, 4, 32, 32)
|
| 50 |
+
a = x[:, 3:4]
|
| 51 |
+
lum = (0.2126 * x[:, 0:1] + 0.7152 * x[:, 1:2] + 0.0722 * x[:, 2:3]) * a # premultiplied luminance
|
| 52 |
+
x = torch.cat([x[:, :3] * a, a, lum], 1)
|
| 53 |
+
x = self.net(x) # (BM,128,2,2)
|
| 54 |
+
return self.proj(x.flatten(1)).reshape(B, M, -1)
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
class Block(nn.Module):
|
| 58 |
+
def __init__(self, d, h, ff, drop):
|
| 59 |
+
super().__init__()
|
| 60 |
+
self.h = h
|
| 61 |
+
self.ln1 = nn.LayerNorm(d)
|
| 62 |
+
self.qkv = nn.Linear(d, 3 * d)
|
| 63 |
+
self.proj = nn.Linear(d, d)
|
| 64 |
+
self.rel = nn.Linear(K_REL, h, bias=False)
|
| 65 |
+
nn.init.normal_(self.rel.weight, std=0.5)
|
| 66 |
+
self.ln2 = nn.LayerNorm(d)
|
| 67 |
+
self.ff = nn.Sequential(nn.Linear(d, ff * d), nn.GELU(), nn.Dropout(drop), nn.Linear(ff * d, d))
|
| 68 |
+
self.drop = nn.Dropout(drop)
|
| 69 |
+
|
| 70 |
+
def forward(self, x, rel, pad_bias):
|
| 71 |
+
B, T, D = x.shape
|
| 72 |
+
q, k, v = self.qkv(self.ln1(x)).reshape(B, T, 3, self.h, D // self.h).permute(2, 0, 3, 1, 4)
|
| 73 |
+
bias = self.rel(rel).permute(0, 3, 1, 2) + pad_bias # (B,H,T,T)
|
| 74 |
+
a = F.scaled_dot_product_attention(q, k, v, attn_mask=bias.to(q.dtype))
|
| 75 |
+
x = x + self.drop(self.proj(a.transpose(1, 2).reshape(B, T, D)))
|
| 76 |
+
return x + self.drop(self.ff(self.ln2(x)))
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
class NadirNet(nn.Module):
|
| 80 |
+
def __init__(self, d=384, layers=8, heads=6, ff=4, drop=0.1):
|
| 81 |
+
super().__init__()
|
| 82 |
+
self.cfg = dict(d=d, layers=layers, heads=heads, ff=ff, drop=drop)
|
| 83 |
+
self.tok_in = MLP(FT + N_ROLES + 36, 2 * d, d)
|
| 84 |
+
self.cnn = ThumbCNN(128)
|
| 85 |
+
self.img_in = MLP(FI + 128, 2 * d, d)
|
| 86 |
+
self.page_in = MLP(FP, d, d)
|
| 87 |
+
self.src = nn.Embedding(2, d)
|
| 88 |
+
self.typ = nn.Embedding(3, d)
|
| 89 |
+
self.blocks = nn.ModuleList([Block(d, heads, ff, drop) for _ in range(layers)])
|
| 90 |
+
self.ln = nn.LayerNorm(d)
|
| 91 |
+
self.color_head = MLP(d + 3, d, 6)
|
| 92 |
+
self.img_head = MLP(d, d // 2, N_TREAT)
|
| 93 |
+
|
| 94 |
+
def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask):
|
| 95 |
+
"""tok_f (B,N,FT) tok_role (B,N) long, tok_lab (B,N,3), img_f (B,M,FI), img_thumb (B,M,4,32,32) float 0-1,
|
| 96 |
+
page_f (B,FP), src (B,) long, rel (B,T,T,K) with T=1+N+M, mask (B,T) bool (True=real)."""
|
| 97 |
+
B, N = tok_role.shape
|
| 98 |
+
M = img_f.shape[1]
|
| 99 |
+
tok = self.tok_in(torch.cat([tok_f, F.one_hot(tok_role, N_ROLES).to(tok_f.dtype), fourier_lab(tok_lab)], -1)) + self.typ.weight[1]
|
| 100 |
+
if M > 0:
|
| 101 |
+
img = self.img_in(torch.cat([img_f, self.cnn(img_thumb)], -1)) + self.typ.weight[2]
|
| 102 |
+
else:
|
| 103 |
+
img = tok.new_zeros(B, 0, tok.shape[-1])
|
| 104 |
+
cls = self.page_in(page_f)[:, None] + self.typ.weight[0]
|
| 105 |
+
x = torch.cat([cls, tok, img], 1) + self.src(src)[:, None]
|
| 106 |
+
pad_bias = torch.zeros(mask.shape, dtype=x.dtype, device=x.device).masked_fill(~mask, -1e4)[:, None, None, :]
|
| 107 |
+
for blk in self.blocks:
|
| 108 |
+
x = blk(x, rel, pad_bias)
|
| 109 |
+
x = self.ln(x)
|
| 110 |
+
ht = x[:, 1:1 + N]
|
| 111 |
+
z = self.color_head(torch.cat([ht, tok_lab], -1))
|
| 112 |
+
L = torch.sigmoid(z[..., 0])
|
| 113 |
+
s = 1.6 * torch.sigmoid(z[..., 1])
|
| 114 |
+
r = 0.6 * torch.tanh(z[..., 2])
|
| 115 |
+
a_in, b_in = tok_lab[..., 1], tok_lab[..., 2]
|
| 116 |
+
a = s * a_in - r * b_in + 0.12 * torch.tanh(z[..., 3])
|
| 117 |
+
b = r * a_in + s * b_in + 0.12 * torch.tanh(z[..., 4])
|
| 118 |
+
lab = torch.stack([L, a, b], -1)
|
| 119 |
+
alpha_logit = z[..., 5]
|
| 120 |
+
img_logits = self.img_head(x[:, 1 + N:]) if M > 0 else x.new_zeros(B, 0, N_TREAT)
|
| 121 |
+
return lab, alpha_logit, img_logits
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
# ---------- colour maths (differentiable) ----------
|
| 125 |
+
OKLAB_TO_LMS = torch.tensor([[1.0, 0.3963377774, 0.2158037573], [1.0, -0.1055613458, -0.0638541728], [1.0, -0.0894841775, -1.2914855480]])
|
| 126 |
+
LMS_TO_RGB = torch.tensor([[4.0767416621, -3.3077115913, 0.2309699292], [-1.2684380046, 2.6097574011, -0.3413193965], [-0.0041960863, -0.7034186147, 1.7076147010]])
|
| 127 |
+
|
| 128 |
+
|
| 129 |
+
def oklab_to_linrgb(lab):
|
| 130 |
+
lms = lab @ OKLAB_TO_LMS.to(lab).T
|
| 131 |
+
return (lms ** 3) @ LMS_TO_RGB.to(lab).T
|
| 132 |
+
|
| 133 |
+
|
| 134 |
+
def luminance(lab):
|
| 135 |
+
rgb = oklab_to_linrgb(lab).clamp(0, 1)
|
| 136 |
+
return rgb @ torch.tensor([0.2126, 0.7152, 0.0722]).to(lab)
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
def contrast_ratio(lab1, lab2):
|
| 140 |
+
l1, l2 = luminance(lab1), luminance(lab2)
|
| 141 |
+
return (torch.maximum(l1, l2) + 0.05) / (torch.minimum(l1, l2) + 0.05)
|
model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e7acc89ba8a039a4889d89ec9f01afce49b2f765d1ac57a489eb4c604aef2ceb
|
| 3 |
+
size 63596788
|
nadir-extension.zip
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:4a5da0a2671306af61a162f81c8cfa276ab2dd70fb05f7e927a401d9e1f08757
|
| 3 |
+
size 23421769
|
onnx/nadir-fp32.onnx
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:83c54558313415596d4d3b3e753ff30e3a72cdc65e27a5c8eefcb132b9d5c934
|
| 3 |
+
size 63699911
|
onnx/nadir-int4.onnx
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:5e11864ad3d7f1edb7bc8b7006e5d7a1c12401bbc17578f590182c0ca764802a
|
| 3 |
+
size 13313059
|
onnx/nadir-int8.onnx
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:bbdb0a2b13bfca0c05f8d19d54d3b721fb4fcd459d40e8519a7a5f77af15d634
|
| 3 |
+
size 20626640
|
reports/eval_judge.json
ADDED
|
@@ -0,0 +1,88 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"judge": "Claude Haiku 5.5 (Batch API, effort medium), blind pairwise: original light page + two dark versions in random order; criteria readability > breakage > design",
|
| 3 |
+
"pages": "held-out sites (hash split by hostname), rendered live at 1280x800",
|
| 4 |
+
"final_release_vs_others": {
|
| 5 |
+
"final_m3:dr": {
|
| 6 |
+
"n": 146,
|
| 7 |
+
"final_m3_wins": 93,
|
| 8 |
+
"dr_wins": 50,
|
| 9 |
+
"ties": 3,
|
| 10 |
+
"final_m3_winrate_excl_ties": 0.65,
|
| 11 |
+
"mean_score_final_m3": 6.85,
|
| 12 |
+
"mean_score_dr": 6.1
|
| 13 |
+
},
|
| 14 |
+
"final_m3:final_m": {
|
| 15 |
+
"n": 146,
|
| 16 |
+
"final_m3_wins": 45,
|
| 17 |
+
"final_m_wins": 29,
|
| 18 |
+
"ties": 72,
|
| 19 |
+
"final_m3_winrate_excl_ties": 0.608,
|
| 20 |
+
"mean_score_final_m3": 6.62,
|
| 21 |
+
"mean_score_final_m": 6.16
|
| 22 |
+
},
|
| 23 |
+
"final_m3:haiku": {
|
| 24 |
+
"n": 146,
|
| 25 |
+
"final_m3_wins": 76,
|
| 26 |
+
"haiku_wins": 61,
|
| 27 |
+
"ties": 9,
|
| 28 |
+
"final_m3_winrate_excl_ties": 0.555,
|
| 29 |
+
"mean_score_final_m3": 6.72,
|
| 30 |
+
"mean_score_haiku": 6.31
|
| 31 |
+
}
|
| 32 |
+
},
|
| 33 |
+
"before_final_design_rules": {
|
| 34 |
+
"final_m:dr": {
|
| 35 |
+
"n": 147,
|
| 36 |
+
"final_m_wins": 83,
|
| 37 |
+
"dr_wins": 57,
|
| 38 |
+
"ties": 7,
|
| 39 |
+
"final_m_winrate_excl_ties": 0.593,
|
| 40 |
+
"mean_score_final_m": 6.5,
|
| 41 |
+
"mean_score_dr": 6.28
|
| 42 |
+
},
|
| 43 |
+
"final_l:dr": {
|
| 44 |
+
"n": 147,
|
| 45 |
+
"final_l_wins": 81,
|
| 46 |
+
"dr_wins": 60,
|
| 47 |
+
"ties": 6,
|
| 48 |
+
"final_l_winrate_excl_ties": 0.574,
|
| 49 |
+
"mean_score_final_l": 6.61,
|
| 50 |
+
"mean_score_dr": 6.31
|
| 51 |
+
},
|
| 52 |
+
"final_m:final_l": {
|
| 53 |
+
"n": 147,
|
| 54 |
+
"final_m_wins": 49,
|
| 55 |
+
"final_l_wins": 53,
|
| 56 |
+
"ties": 45,
|
| 57 |
+
"final_m_winrate_excl_ties": 0.48,
|
| 58 |
+
"mean_score_final_m": 6.43,
|
| 59 |
+
"mean_score_final_l": 6.54
|
| 60 |
+
},
|
| 61 |
+
"final_m:haiku": {
|
| 62 |
+
"n": 146,
|
| 63 |
+
"final_m_wins": 63,
|
| 64 |
+
"haiku_wins": 67,
|
| 65 |
+
"ties": 16,
|
| 66 |
+
"final_m_winrate_excl_ties": 0.485,
|
| 67 |
+
"mean_score_final_m": 6.32,
|
| 68 |
+
"mean_score_haiku": 6.38
|
| 69 |
+
},
|
| 70 |
+
"final_l:haiku": {
|
| 71 |
+
"n": 146,
|
| 72 |
+
"final_l_wins": 64,
|
| 73 |
+
"haiku_wins": 70,
|
| 74 |
+
"ties": 12,
|
| 75 |
+
"final_l_winrate_excl_ties": 0.478,
|
| 76 |
+
"mean_score_final_l": 6.47,
|
| 77 |
+
"mean_score_haiku": 6.44
|
| 78 |
+
}
|
| 79 |
+
},
|
| 80 |
+
"notes": {
|
| 81 |
+
"final_m3": "released model (15.9M) with final runtime rules",
|
| 82 |
+
"final_m": "same model, runtime before the backdrop/SVG rules",
|
| 83 |
+
"final_l": "34.0M model, earlier runtime",
|
| 84 |
+
"dr": "Dark Reader 4.9.133 dynamic mode",
|
| 85 |
+
"haiku": "teacher designs from Claude Haiku 5.5, earlier runtime"
|
| 86 |
+
},
|
| 87 |
+
"cost_usd": 0.368
|
| 88 |
+
}
|
reports/export_report.json
ADDED
|
@@ -0,0 +1,20 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"params": 15895995,
|
| 3 |
+
"parity_max_abs": 1.2516975402832031e-06,
|
| 4 |
+
"variants": {
|
| 5 |
+
"fp32": {
|
| 6 |
+
"dE_vs_fp32": 0.0,
|
| 7 |
+
"dE_max": 0.0,
|
| 8 |
+
"img_agree": 1.0,
|
| 9 |
+
"size_mb": 63.7,
|
| 10 |
+
"cpu_ms_median": 7.48
|
| 11 |
+
},
|
| 12 |
+
"fp16": {
|
| 13 |
+
"dE_vs_fp32": 0.0002421401860071152,
|
| 14 |
+
"dE_max": 0.23341728746891022,
|
| 15 |
+
"img_agree": 1.0,
|
| 16 |
+
"size_mb": 31.91,
|
| 17 |
+
"cpu_ms_median": 10.17
|
| 18 |
+
}
|
| 19 |
+
}
|
| 20 |
+
}
|
reports/ortweb_timing_m4max.json
ADDED
|
@@ -0,0 +1,67 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"runtime": "onnxruntime-web 1.29 wasm, Chrome for Testing 141 headless, Apple M4 Max; t1/t4 = 1/4 threads; 12 held-out pages, median of 3 runs each",
|
| 3 |
+
"fp32": {
|
| 4 |
+
"t1": {
|
| 5 |
+
"isolated": false,
|
| 6 |
+
"load_ms": 7246.9,
|
| 7 |
+
"ms_median": 117.08,
|
| 8 |
+
"max_abs_lab_vs_ort_cpu": 6.16908073425293e-06,
|
| 9 |
+
"max_abs_img_vs_ort_cpu": 4.172325134277344e-06
|
| 10 |
+
},
|
| 11 |
+
"t4": {
|
| 12 |
+
"isolated": false,
|
| 13 |
+
"load_ms": 7363.2,
|
| 14 |
+
"ms_median": 33.93,
|
| 15 |
+
"max_abs_lab_vs_ort_cpu": 6.16908073425293e-06,
|
| 16 |
+
"max_abs_img_vs_ort_cpu": 4.172325134277344e-06
|
| 17 |
+
}
|
| 18 |
+
},
|
| 19 |
+
"int8": {
|
| 20 |
+
"t1": {
|
| 21 |
+
"isolated": false,
|
| 22 |
+
"load_ms": 1356.6,
|
| 23 |
+
"ms_median": 116.78,
|
| 24 |
+
"max_abs_lab_vs_ort_cpu": 1.9669532775878906e-06,
|
| 25 |
+
"max_abs_img_vs_ort_cpu": 2.9802322387695312e-06
|
| 26 |
+
},
|
| 27 |
+
"t4": {
|
| 28 |
+
"isolated": false,
|
| 29 |
+
"load_ms": 1628.3,
|
| 30 |
+
"ms_median": 34.78,
|
| 31 |
+
"max_abs_lab_vs_ort_cpu": 1.9669532775878906e-06,
|
| 32 |
+
"max_abs_img_vs_ort_cpu": 2.9802322387695312e-06
|
| 33 |
+
}
|
| 34 |
+
},
|
| 35 |
+
"int8_repeat": {
|
| 36 |
+
"t1": {
|
| 37 |
+
"isolated": false,
|
| 38 |
+
"load_ms": 1523.5,
|
| 39 |
+
"ms_median": 115.82,
|
| 40 |
+
"max_abs_lab_vs_ort_cpu": 1.9669532775878906e-06,
|
| 41 |
+
"max_abs_img_vs_ort_cpu": 2.9802322387695312e-06
|
| 42 |
+
},
|
| 43 |
+
"t4": {
|
| 44 |
+
"isolated": false,
|
| 45 |
+
"load_ms": 1519,
|
| 46 |
+
"ms_median": 32.68,
|
| 47 |
+
"max_abs_lab_vs_ort_cpu": 1.9669532775878906e-06,
|
| 48 |
+
"max_abs_img_vs_ort_cpu": 2.9802322387695312e-06
|
| 49 |
+
}
|
| 50 |
+
},
|
| 51 |
+
"int4": {
|
| 52 |
+
"t1": {
|
| 53 |
+
"isolated": false,
|
| 54 |
+
"load_ms": 1074.8,
|
| 55 |
+
"ms_median": 137.33,
|
| 56 |
+
"max_abs_lab_vs_ort_cpu": 1.3709068298339844e-06,
|
| 57 |
+
"max_abs_img_vs_ort_cpu": 7.659196853637695e-06
|
| 58 |
+
},
|
| 59 |
+
"t4": {
|
| 60 |
+
"isolated": false,
|
| 61 |
+
"load_ms": 1047,
|
| 62 |
+
"ms_median": 38.47,
|
| 63 |
+
"max_abs_lab_vs_ort_cpu": 1.3709068298339844e-06,
|
| 64 |
+
"max_abs_img_vs_ort_cpu": 7.659196853637695e-06
|
| 65 |
+
}
|
| 66 |
+
}
|
| 67 |
+
}
|
reports/quant_report.json
ADDED
|
@@ -0,0 +1,63 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"format": {
|
| 3 |
+
"8-bit layers": "uint8 blocks of 32 with fp32 scale + uint8 zero point, dequantized in-graph (Cast/Sub/Mul, folded at load)",
|
| 4 |
+
"4-bit layers": "uint4 blocks of 32, MatMulNBits"
|
| 5 |
+
},
|
| 6 |
+
"group_size": 32,
|
| 7 |
+
"calib_pages": 256,
|
| 8 |
+
"eval_pages": 400,
|
| 9 |
+
"fp32_layers": [
|
| 10 |
+
"cnn.proj",
|
| 11 |
+
"page_in.net.0",
|
| 12 |
+
"page_in.net.3",
|
| 13 |
+
"blocks.0.rel",
|
| 14 |
+
"blocks.1.rel",
|
| 15 |
+
"blocks.2.rel",
|
| 16 |
+
"blocks.3.rel",
|
| 17 |
+
"blocks.4.rel",
|
| 18 |
+
"blocks.5.rel",
|
| 19 |
+
"blocks.6.rel",
|
| 20 |
+
"blocks.7.rel"
|
| 21 |
+
],
|
| 22 |
+
"int4_8bit_edge_layers": [
|
| 23 |
+
"tok_in.net.0",
|
| 24 |
+
"tok_in.net.3",
|
| 25 |
+
"img_in.net.0",
|
| 26 |
+
"img_in.net.3",
|
| 27 |
+
"color_head.net.0",
|
| 28 |
+
"color_head.net.3",
|
| 29 |
+
"img_head.net.0",
|
| 30 |
+
"img_head.net.3"
|
| 31 |
+
],
|
| 32 |
+
"variants": {
|
| 33 |
+
"fp32": {
|
| 34 |
+
"haiku_dE": 0.039349029815445344,
|
| 35 |
+
"native_dE": 0.08639045081450604,
|
| 36 |
+
"img_acc": 0.7966292134831461,
|
| 37 |
+
"size_mb": 63.7,
|
| 38 |
+
"cpu_ms_median": 11.42
|
| 39 |
+
},
|
| 40 |
+
"int8": {
|
| 41 |
+
"dE_vs_fp32": 0.00015194275738394936,
|
| 42 |
+
"dE_max": 0.02521813102066517,
|
| 43 |
+
"img_agree": 0.9998743718592965,
|
| 44 |
+
"haiku_dE": 0.03935358844697476,
|
| 45 |
+
"native_dE": 0.0863548206735868,
|
| 46 |
+
"img_acc": 0.7968164794007491,
|
| 47 |
+
"pack_parity_max_abs": 2.5033950805664062e-06,
|
| 48 |
+
"size_mb": 20.63,
|
| 49 |
+
"cpu_ms_median": 9.19
|
| 50 |
+
},
|
| 51 |
+
"int4": {
|
| 52 |
+
"dE_vs_fp32": 0.00036296586506068707,
|
| 53 |
+
"dE_max": 0.0709674283862114,
|
| 54 |
+
"img_agree": 0.9976130653266332,
|
| 55 |
+
"haiku_dE": 0.039340872559696434,
|
| 56 |
+
"native_dE": 0.08632776886341162,
|
| 57 |
+
"img_acc": 0.797003745318352,
|
| 58 |
+
"pack_parity_max_abs": 1.5497207641601562e-06,
|
| 59 |
+
"size_mb": 13.31,
|
| 60 |
+
"cpu_ms_median": 9.76
|
| 61 |
+
}
|
| 62 |
+
}
|
| 63 |
+
}
|
reports/train_log.json
ADDED
|
@@ -0,0 +1,501 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"args": {
|
| 3 |
+
"data": "/home/user/.local/share/rtx-temporary-jobs/tj-nadir-0fe76575c0f3/payload/work/data/ds_final.pkl",
|
| 4 |
+
"out": "/job/export/run_m",
|
| 5 |
+
"d": 384,
|
| 6 |
+
"layers": 8,
|
| 7 |
+
"heads": 6,
|
| 8 |
+
"epochs": 30,
|
| 9 |
+
"bs": 48,
|
| 10 |
+
"lr": 0.0004,
|
| 11 |
+
"wd": 0.05,
|
| 12 |
+
"drop": 0.1,
|
| 13 |
+
"native_weight": 0.5,
|
| 14 |
+
"seed": 0,
|
| 15 |
+
"workers": 4
|
| 16 |
+
},
|
| 17 |
+
"params": 15895995,
|
| 18 |
+
"log": [
|
| 19 |
+
{
|
| 20 |
+
"epoch": 0,
|
| 21 |
+
"step": 468,
|
| 22 |
+
"secs": 81,
|
| 23 |
+
"loss": 0.3405,
|
| 24 |
+
"dE": 0.12132,
|
| 25 |
+
"con": 0.50822,
|
| 26 |
+
"img": 1.29224,
|
| 27 |
+
"s0_dE": 0.06306,
|
| 28 |
+
"s0_dL": 0.06024,
|
| 29 |
+
"s0_con": 0.04883,
|
| 30 |
+
"img_acc": 0.73416,
|
| 31 |
+
"s1_dE": 0.16366,
|
| 32 |
+
"s1_dL": 0.15909,
|
| 33 |
+
"s1_con": 0.11938
|
| 34 |
+
},
|
| 35 |
+
{
|
| 36 |
+
"epoch": 1,
|
| 37 |
+
"step": 936,
|
| 38 |
+
"secs": 148,
|
| 39 |
+
"loss": 0.18873,
|
| 40 |
+
"dE": 0.06161,
|
| 41 |
+
"con": 0.03771,
|
| 42 |
+
"img": 0.99138,
|
| 43 |
+
"s0_dE": 0.05432,
|
| 44 |
+
"s0_dL": 0.05155,
|
| 45 |
+
"s0_con": 0.03668,
|
| 46 |
+
"img_acc": 0.78603,
|
| 47 |
+
"s1_dE": 0.1523,
|
| 48 |
+
"s1_dL": 0.14838,
|
| 49 |
+
"s1_con": 0.08049
|
| 50 |
+
},
|
| 51 |
+
{
|
| 52 |
+
"epoch": 2,
|
| 53 |
+
"step": 1404,
|
| 54 |
+
"secs": 215,
|
| 55 |
+
"loss": 0.17375,
|
| 56 |
+
"dE": 0.05668,
|
| 57 |
+
"con": 0.03498,
|
| 58 |
+
"img": 0.90371,
|
| 59 |
+
"s0_dE": 0.04977,
|
| 60 |
+
"s0_dL": 0.04725,
|
| 61 |
+
"s0_con": 0.02243,
|
| 62 |
+
"img_acc": 0.78005,
|
| 63 |
+
"s1_dE": 0.14476,
|
| 64 |
+
"s1_dL": 0.14153,
|
| 65 |
+
"s1_con": 0.26806
|
| 66 |
+
},
|
| 67 |
+
{
|
| 68 |
+
"epoch": 3,
|
| 69 |
+
"step": 1872,
|
| 70 |
+
"secs": 281,
|
| 71 |
+
"loss": 0.16779,
|
| 72 |
+
"dE": 0.05488,
|
| 73 |
+
"con": 0.04139,
|
| 74 |
+
"img": 0.86775,
|
| 75 |
+
"s0_dE": 0.0506,
|
| 76 |
+
"s0_dL": 0.04808,
|
| 77 |
+
"s0_con": 0.0331,
|
| 78 |
+
"img_acc": 0.79611,
|
| 79 |
+
"s1_dE": 0.13697,
|
| 80 |
+
"s1_dL": 0.13372,
|
| 81 |
+
"s1_con": 0.1237
|
| 82 |
+
},
|
| 83 |
+
{
|
| 84 |
+
"epoch": 4,
|
| 85 |
+
"step": 2340,
|
| 86 |
+
"secs": 349,
|
| 87 |
+
"loss": 0.16278,
|
| 88 |
+
"dE": 0.05274,
|
| 89 |
+
"con": 0.0414,
|
| 90 |
+
"img": 0.85412,
|
| 91 |
+
"s0_dE": 0.04849,
|
| 92 |
+
"s0_dL": 0.04604,
|
| 93 |
+
"s0_con": 0.02439,
|
| 94 |
+
"img_acc": 0.78982,
|
| 95 |
+
"s1_dE": 0.13195,
|
| 96 |
+
"s1_dL": 0.12841,
|
| 97 |
+
"s1_con": 0.16648
|
| 98 |
+
},
|
| 99 |
+
{
|
| 100 |
+
"epoch": 5,
|
| 101 |
+
"step": 2808,
|
| 102 |
+
"secs": 416,
|
| 103 |
+
"loss": 0.15903,
|
| 104 |
+
"dE": 0.05166,
|
| 105 |
+
"con": 0.03833,
|
| 106 |
+
"img": 0.84086,
|
| 107 |
+
"s0_dE": 0.0472,
|
| 108 |
+
"s0_dL": 0.04477,
|
| 109 |
+
"s0_con": 0.03465,
|
| 110 |
+
"img_acc": 0.80467,
|
| 111 |
+
"s1_dE": 0.12882,
|
| 112 |
+
"s1_dL": 0.12538,
|
| 113 |
+
"s1_con": 0.12963
|
| 114 |
+
},
|
| 115 |
+
{
|
| 116 |
+
"epoch": 6,
|
| 117 |
+
"step": 3276,
|
| 118 |
+
"secs": 484,
|
| 119 |
+
"loss": 0.15729,
|
| 120 |
+
"dE": 0.05076,
|
| 121 |
+
"con": 0.04035,
|
| 122 |
+
"img": 0.83185,
|
| 123 |
+
"s0_dE": 0.0457,
|
| 124 |
+
"s0_dL": 0.04338,
|
| 125 |
+
"s0_con": 0.02999,
|
| 126 |
+
"img_acc": 0.80345,
|
| 127 |
+
"s1_dE": 0.15936,
|
| 128 |
+
"s1_dL": 0.15619,
|
| 129 |
+
"s1_con": 0.51491
|
| 130 |
+
},
|
| 131 |
+
{
|
| 132 |
+
"epoch": 7,
|
| 133 |
+
"step": 3744,
|
| 134 |
+
"secs": 552,
|
| 135 |
+
"loss": 0.15392,
|
| 136 |
+
"dE": 0.04978,
|
| 137 |
+
"con": 0.0403,
|
| 138 |
+
"img": 0.81841,
|
| 139 |
+
"s0_dE": 0.04612,
|
| 140 |
+
"s0_dL": 0.04395,
|
| 141 |
+
"s0_con": 0.01962,
|
| 142 |
+
"img_acc": 0.79899,
|
| 143 |
+
"s1_dE": 0.14854,
|
| 144 |
+
"s1_dL": 0.14578,
|
| 145 |
+
"s1_con": 0.42056
|
| 146 |
+
},
|
| 147 |
+
{
|
| 148 |
+
"epoch": 8,
|
| 149 |
+
"step": 4212,
|
| 150 |
+
"secs": 619,
|
| 151 |
+
"loss": 0.15244,
|
| 152 |
+
"dE": 0.04894,
|
| 153 |
+
"con": 0.04062,
|
| 154 |
+
"img": 0.81195,
|
| 155 |
+
"s0_dE": 0.04632,
|
| 156 |
+
"s0_dL": 0.04425,
|
| 157 |
+
"s0_con": 0.02413,
|
| 158 |
+
"img_acc": 0.79917,
|
| 159 |
+
"s1_dE": 0.12446,
|
| 160 |
+
"s1_dL": 0.12131,
|
| 161 |
+
"s1_con": 0.23072
|
| 162 |
+
},
|
| 163 |
+
{
|
| 164 |
+
"epoch": 9,
|
| 165 |
+
"step": 4680,
|
| 166 |
+
"secs": 687,
|
| 167 |
+
"loss": 0.14931,
|
| 168 |
+
"dE": 0.0481,
|
| 169 |
+
"con": 0.03785,
|
| 170 |
+
"img": 0.79837,
|
| 171 |
+
"s0_dE": 0.04452,
|
| 172 |
+
"s0_dL": 0.04238,
|
| 173 |
+
"s0_con": 0.02434,
|
| 174 |
+
"img_acc": 0.80094,
|
| 175 |
+
"s1_dE": 0.12096,
|
| 176 |
+
"s1_dL": 0.11779,
|
| 177 |
+
"s1_con": 0.21844
|
| 178 |
+
},
|
| 179 |
+
{
|
| 180 |
+
"epoch": 10,
|
| 181 |
+
"step": 5148,
|
| 182 |
+
"secs": 755,
|
| 183 |
+
"loss": 0.1484,
|
| 184 |
+
"dE": 0.04786,
|
| 185 |
+
"con": 0.03854,
|
| 186 |
+
"img": 0.79393,
|
| 187 |
+
"s0_dE": 0.04549,
|
| 188 |
+
"s0_dL": 0.04344,
|
| 189 |
+
"s0_con": 0.02666,
|
| 190 |
+
"img_acc": 0.79123,
|
| 191 |
+
"s1_dE": 0.11822,
|
| 192 |
+
"s1_dL": 0.11502,
|
| 193 |
+
"s1_con": 0.25328
|
| 194 |
+
},
|
| 195 |
+
{
|
| 196 |
+
"epoch": 11,
|
| 197 |
+
"step": 5616,
|
| 198 |
+
"secs": 822,
|
| 199 |
+
"loss": 0.14645,
|
| 200 |
+
"dE": 0.04703,
|
| 201 |
+
"con": 0.03983,
|
| 202 |
+
"img": 0.78304,
|
| 203 |
+
"s0_dE": 0.04345,
|
| 204 |
+
"s0_dL": 0.04139,
|
| 205 |
+
"s0_con": 0.01948,
|
| 206 |
+
"img_acc": 0.81053,
|
| 207 |
+
"s1_dE": 0.11486,
|
| 208 |
+
"s1_dL": 0.11153,
|
| 209 |
+
"s1_con": 0.20122
|
| 210 |
+
},
|
| 211 |
+
{
|
| 212 |
+
"epoch": 12,
|
| 213 |
+
"step": 6084,
|
| 214 |
+
"secs": 891,
|
| 215 |
+
"loss": 0.13803,
|
| 216 |
+
"dE": 0.04635,
|
| 217 |
+
"con": 0.03723,
|
| 218 |
+
"img": 0.76961,
|
| 219 |
+
"s0_dE": 0.04328,
|
| 220 |
+
"s0_dL": 0.04125,
|
| 221 |
+
"s0_con": 0.01831,
|
| 222 |
+
"img_acc": 0.8051,
|
| 223 |
+
"s1_dE": 0.11435,
|
| 224 |
+
"s1_dL": 0.11091,
|
| 225 |
+
"s1_con": 0.16656
|
| 226 |
+
},
|
| 227 |
+
{
|
| 228 |
+
"epoch": 13,
|
| 229 |
+
"step": 6552,
|
| 230 |
+
"secs": 958,
|
| 231 |
+
"loss": 0.13627,
|
| 232 |
+
"dE": 0.04589,
|
| 233 |
+
"con": 0.03889,
|
| 234 |
+
"img": 0.76454,
|
| 235 |
+
"s0_dE": 0.04247,
|
| 236 |
+
"s0_dL": 0.04047,
|
| 237 |
+
"s0_con": 0.01875,
|
| 238 |
+
"img_acc": 0.80161,
|
| 239 |
+
"s1_dE": 0.11794,
|
| 240 |
+
"s1_dL": 0.11469,
|
| 241 |
+
"s1_con": 0.23754
|
| 242 |
+
},
|
| 243 |
+
{
|
| 244 |
+
"epoch": 14,
|
| 245 |
+
"step": 7020,
|
| 246 |
+
"secs": 1025,
|
| 247 |
+
"loss": 0.13324,
|
| 248 |
+
"dE": 0.045,
|
| 249 |
+
"con": 0.0383,
|
| 250 |
+
"img": 0.75126,
|
| 251 |
+
"s0_dE": 0.04194,
|
| 252 |
+
"s0_dL": 0.03998,
|
| 253 |
+
"s0_con": 0.02424,
|
| 254 |
+
"img_acc": 0.80332,
|
| 255 |
+
"s1_dE": 0.11805,
|
| 256 |
+
"s1_dL": 0.11474,
|
| 257 |
+
"s1_con": 0.22873
|
| 258 |
+
},
|
| 259 |
+
{
|
| 260 |
+
"epoch": 15,
|
| 261 |
+
"step": 7488,
|
| 262 |
+
"secs": 1093,
|
| 263 |
+
"loss": 0.13267,
|
| 264 |
+
"dE": 0.04502,
|
| 265 |
+
"con": 0.0394,
|
| 266 |
+
"img": 0.74724,
|
| 267 |
+
"s0_dE": 0.04183,
|
| 268 |
+
"s0_dL": 0.03983,
|
| 269 |
+
"s0_con": 0.01954,
|
| 270 |
+
"img_acc": 0.81353,
|
| 271 |
+
"s1_dE": 0.11306,
|
| 272 |
+
"s1_dL": 0.10954,
|
| 273 |
+
"s1_con": 0.20409
|
| 274 |
+
},
|
| 275 |
+
{
|
| 276 |
+
"epoch": 16,
|
| 277 |
+
"step": 7956,
|
| 278 |
+
"secs": 1161,
|
| 279 |
+
"loss": 0.1305,
|
| 280 |
+
"dE": 0.04465,
|
| 281 |
+
"con": 0.03933,
|
| 282 |
+
"img": 0.72731,
|
| 283 |
+
"s0_dE": 0.04195,
|
| 284 |
+
"s0_dL": 0.04006,
|
| 285 |
+
"s0_con": 0.02284,
|
| 286 |
+
"img_acc": 0.80485,
|
| 287 |
+
"s1_dE": 0.11437,
|
| 288 |
+
"s1_dL": 0.11122,
|
| 289 |
+
"s1_con": 0.27215
|
| 290 |
+
},
|
| 291 |
+
{
|
| 292 |
+
"epoch": 17,
|
| 293 |
+
"step": 8424,
|
| 294 |
+
"secs": 1229,
|
| 295 |
+
"loss": 0.12864,
|
| 296 |
+
"dE": 0.04415,
|
| 297 |
+
"con": 0.03796,
|
| 298 |
+
"img": 0.71419,
|
| 299 |
+
"s0_dE": 0.04177,
|
| 300 |
+
"s0_dL": 0.0399,
|
| 301 |
+
"s0_con": 0.0207,
|
| 302 |
+
"img_acc": 0.79611,
|
| 303 |
+
"s1_dE": 0.10924,
|
| 304 |
+
"s1_dL": 0.10592,
|
| 305 |
+
"s1_con": 0.22063
|
| 306 |
+
},
|
| 307 |
+
{
|
| 308 |
+
"epoch": 18,
|
| 309 |
+
"step": 8892,
|
| 310 |
+
"secs": 1296,
|
| 311 |
+
"loss": 0.12691,
|
| 312 |
+
"dE": 0.04379,
|
| 313 |
+
"con": 0.03646,
|
| 314 |
+
"img": 0.69736,
|
| 315 |
+
"s0_dE": 0.04175,
|
| 316 |
+
"s0_dL": 0.03993,
|
| 317 |
+
"s0_con": 0.02416,
|
| 318 |
+
"img_acc": 0.801,
|
| 319 |
+
"s1_dE": 0.10774,
|
| 320 |
+
"s1_dL": 0.10449,
|
| 321 |
+
"s1_con": 0.24435
|
| 322 |
+
},
|
| 323 |
+
{
|
| 324 |
+
"epoch": 19,
|
| 325 |
+
"step": 9360,
|
| 326 |
+
"secs": 1364,
|
| 327 |
+
"loss": 0.12499,
|
| 328 |
+
"dE": 0.04319,
|
| 329 |
+
"con": 0.036,
|
| 330 |
+
"img": 0.68816,
|
| 331 |
+
"s0_dE": 0.0413,
|
| 332 |
+
"s0_dL": 0.03945,
|
| 333 |
+
"s0_con": 0.02047,
|
| 334 |
+
"img_acc": 0.79813,
|
| 335 |
+
"s1_dE": 0.11372,
|
| 336 |
+
"s1_dL": 0.11037,
|
| 337 |
+
"s1_con": 0.20421
|
| 338 |
+
},
|
| 339 |
+
{
|
| 340 |
+
"epoch": 20,
|
| 341 |
+
"step": 9828,
|
| 342 |
+
"secs": 1432,
|
| 343 |
+
"loss": 0.12434,
|
| 344 |
+
"dE": 0.04334,
|
| 345 |
+
"con": 0.03581,
|
| 346 |
+
"img": 0.67546,
|
| 347 |
+
"s0_dE": 0.0409,
|
| 348 |
+
"s0_dL": 0.03905,
|
| 349 |
+
"s0_con": 0.01907,
|
| 350 |
+
"img_acc": 0.79483,
|
| 351 |
+
"s1_dE": 0.10692,
|
| 352 |
+
"s1_dL": 0.1036,
|
| 353 |
+
"s1_con": 0.22495
|
| 354 |
+
},
|
| 355 |
+
{
|
| 356 |
+
"epoch": 21,
|
| 357 |
+
"step": 10296,
|
| 358 |
+
"secs": 1499,
|
| 359 |
+
"loss": 0.12248,
|
| 360 |
+
"dE": 0.04287,
|
| 361 |
+
"con": 0.03581,
|
| 362 |
+
"img": 0.65995,
|
| 363 |
+
"s0_dE": 0.04105,
|
| 364 |
+
"s0_dL": 0.0392,
|
| 365 |
+
"s0_con": 0.02163,
|
| 366 |
+
"img_acc": 0.7999,
|
| 367 |
+
"s1_dE": 0.10461,
|
| 368 |
+
"s1_dL": 0.10127,
|
| 369 |
+
"s1_con": 0.18494
|
| 370 |
+
},
|
| 371 |
+
{
|
| 372 |
+
"epoch": 22,
|
| 373 |
+
"step": 10764,
|
| 374 |
+
"secs": 1567,
|
| 375 |
+
"loss": 0.12033,
|
| 376 |
+
"dE": 0.04222,
|
| 377 |
+
"con": 0.03395,
|
| 378 |
+
"img": 0.64253,
|
| 379 |
+
"s0_dE": 0.04075,
|
| 380 |
+
"s0_dL": 0.03893,
|
| 381 |
+
"s0_con": 0.02266,
|
| 382 |
+
"img_acc": 0.7999,
|
| 383 |
+
"s1_dE": 0.10178,
|
| 384 |
+
"s1_dL": 0.09855,
|
| 385 |
+
"s1_con": 0.23005
|
| 386 |
+
},
|
| 387 |
+
{
|
| 388 |
+
"epoch": 23,
|
| 389 |
+
"step": 11232,
|
| 390 |
+
"secs": 1635,
|
| 391 |
+
"loss": 0.11905,
|
| 392 |
+
"dE": 0.04192,
|
| 393 |
+
"con": 0.03691,
|
| 394 |
+
"img": 0.63201,
|
| 395 |
+
"s0_dE": 0.04024,
|
| 396 |
+
"s0_dL": 0.03843,
|
| 397 |
+
"s0_con": 0.01664,
|
| 398 |
+
"img_acc": 0.79049,
|
| 399 |
+
"s1_dE": 0.10458,
|
| 400 |
+
"s1_dL": 0.10125,
|
| 401 |
+
"s1_con": 0.1973
|
| 402 |
+
},
|
| 403 |
+
{
|
| 404 |
+
"epoch": 24,
|
| 405 |
+
"step": 11700,
|
| 406 |
+
"secs": 1703,
|
| 407 |
+
"loss": 0.11767,
|
| 408 |
+
"dE": 0.04169,
|
| 409 |
+
"con": 0.03462,
|
| 410 |
+
"img": 0.61761,
|
| 411 |
+
"s0_dE": 0.04043,
|
| 412 |
+
"s0_dL": 0.03861,
|
| 413 |
+
"s0_con": 0.02038,
|
| 414 |
+
"img_acc": 0.79379,
|
| 415 |
+
"s1_dE": 0.1025,
|
| 416 |
+
"s1_dL": 0.09926,
|
| 417 |
+
"s1_con": 0.2135
|
| 418 |
+
},
|
| 419 |
+
{
|
| 420 |
+
"epoch": 25,
|
| 421 |
+
"step": 12168,
|
| 422 |
+
"secs": 1770,
|
| 423 |
+
"loss": 0.11655,
|
| 424 |
+
"dE": 0.04131,
|
| 425 |
+
"con": 0.03564,
|
| 426 |
+
"img": 0.60858,
|
| 427 |
+
"s0_dE": 0.04068,
|
| 428 |
+
"s0_dL": 0.0389,
|
| 429 |
+
"s0_con": 0.01981,
|
| 430 |
+
"img_acc": 0.79715,
|
| 431 |
+
"s1_dE": 0.10372,
|
| 432 |
+
"s1_dL": 0.10038,
|
| 433 |
+
"s1_con": 0.20402
|
| 434 |
+
},
|
| 435 |
+
{
|
| 436 |
+
"epoch": 26,
|
| 437 |
+
"step": 12636,
|
| 438 |
+
"secs": 1838,
|
| 439 |
+
"loss": 0.11502,
|
| 440 |
+
"dE": 0.04106,
|
| 441 |
+
"con": 0.03415,
|
| 442 |
+
"img": 0.59825,
|
| 443 |
+
"s0_dE": 0.04032,
|
| 444 |
+
"s0_dL": 0.03852,
|
| 445 |
+
"s0_con": 0.01906,
|
| 446 |
+
"img_acc": 0.79691,
|
| 447 |
+
"s1_dE": 0.10293,
|
| 448 |
+
"s1_dL": 0.09969,
|
| 449 |
+
"s1_con": 0.22009
|
| 450 |
+
},
|
| 451 |
+
{
|
| 452 |
+
"epoch": 27,
|
| 453 |
+
"step": 13104,
|
| 454 |
+
"secs": 1906,
|
| 455 |
+
"loss": 0.11521,
|
| 456 |
+
"dE": 0.04109,
|
| 457 |
+
"con": 0.03603,
|
| 458 |
+
"img": 0.59315,
|
| 459 |
+
"s0_dE": 0.0404,
|
| 460 |
+
"s0_dL": 0.03861,
|
| 461 |
+
"s0_con": 0.02053,
|
| 462 |
+
"img_acc": 0.79721,
|
| 463 |
+
"s1_dE": 0.10302,
|
| 464 |
+
"s1_dL": 0.09978,
|
| 465 |
+
"s1_con": 0.2078
|
| 466 |
+
},
|
| 467 |
+
{
|
| 468 |
+
"epoch": 28,
|
| 469 |
+
"step": 13572,
|
| 470 |
+
"secs": 1973,
|
| 471 |
+
"loss": 0.11402,
|
| 472 |
+
"dE": 0.04082,
|
| 473 |
+
"con": 0.03683,
|
| 474 |
+
"img": 0.58722,
|
| 475 |
+
"s0_dE": 0.04034,
|
| 476 |
+
"s0_dL": 0.03856,
|
| 477 |
+
"s0_con": 0.01831,
|
| 478 |
+
"img_acc": 0.79098,
|
| 479 |
+
"s1_dE": 0.10383,
|
| 480 |
+
"s1_dL": 0.10053,
|
| 481 |
+
"s1_con": 0.181
|
| 482 |
+
},
|
| 483 |
+
{
|
| 484 |
+
"epoch": 29,
|
| 485 |
+
"step": 14040,
|
| 486 |
+
"secs": 2041,
|
| 487 |
+
"loss": 0.11321,
|
| 488 |
+
"dE": 0.04059,
|
| 489 |
+
"con": 0.03495,
|
| 490 |
+
"img": 0.5803,
|
| 491 |
+
"s0_dE": 0.04018,
|
| 492 |
+
"s0_dL": 0.03841,
|
| 493 |
+
"s0_con": 0.01978,
|
| 494 |
+
"img_acc": 0.79416,
|
| 495 |
+
"s1_dE": 0.10315,
|
| 496 |
+
"s1_dL": 0.09986,
|
| 497 |
+
"s1_con": 0.19971
|
| 498 |
+
}
|
| 499 |
+
],
|
| 500 |
+
"best_score": 0.07564117624840601
|
| 501 |
+
}
|
runtime/DARKREADER-LICENSE
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
MIT License
|
| 2 |
+
|
| 3 |
+
Copyright (c) 2026 Dark Reader Ltd.
|
| 4 |
+
|
| 5 |
+
All rights reserved.
|
| 6 |
+
|
| 7 |
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
| 8 |
+
of this software and associated documentation files (the "Software"), to deal
|
| 9 |
+
in the Software without restriction, including without limitation the rights
|
| 10 |
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
| 11 |
+
copies of the Software, and to permit persons to whom the Software is
|
| 12 |
+
furnished to do so, subject to the following conditions:
|
| 13 |
+
|
| 14 |
+
The above copyright notice and this permission notice shall be included in all
|
| 15 |
+
copies or substantial portions of the Software.
|
| 16 |
+
|
| 17 |
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
| 18 |
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
| 19 |
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
| 20 |
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
| 21 |
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
| 22 |
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
| 23 |
+
SOFTWARE.
|
runtime/darkreader-nadir.js
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
runtime/extract.js
ADDED
|
@@ -0,0 +1,825 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
// Nadir page analyzer. Shared by the training crawler (Playwright) and the browser extension.
|
| 2 |
+
// Defines globalThis.__nadirExtract(opts) -> Promise<PageDescription>.
|
| 3 |
+
// Everything the model sees is computed here, so training and inference features match exactly.
|
| 4 |
+
(function () {
|
| 5 |
+
'use strict';
|
| 6 |
+
if (globalThis.__nadirExtract) return;
|
| 7 |
+
|
| 8 |
+
const ROLES = ['bg', 'text', 'border', 'shadow', 'gradient', 'fill', 'stroke'];
|
| 9 |
+
const R = Object.fromEntries(ROLES.map((r, i) => [r, i]));
|
| 10 |
+
const SKIP_TAGS = new Set(['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEMPLATE', 'HEAD', 'META', 'LINK', 'TITLE', 'BR', 'WBR', 'SOURCE', 'TRACK', 'PARAM']);
|
| 11 |
+
const INTERACTIVE = new Set(['A', 'BUTTON', 'INPUT', 'SELECT', 'TEXTAREA', 'SUMMARY', 'LABEL']);
|
| 12 |
+
const FORM = new Set(['INPUT', 'SELECT', 'TEXTAREA']);
|
| 13 |
+
const HEADINGS = new Set(['H1', 'H2', 'H3', 'H4', 'H5', 'H6']);
|
| 14 |
+
const CODE = new Set(['CODE', 'PRE', 'KBD', 'SAMP']);
|
| 15 |
+
const SVG_SHAPES = new Set(['path', 'circle', 'rect', 'ellipse', 'line', 'polyline', 'polygon', 'text', 'tspan', 'use', 'g', 'svg']);
|
| 16 |
+
|
| 17 |
+
// ---------- color parsing ----------
|
| 18 |
+
const colorCache = new Map();
|
| 19 |
+
let cctx = null;
|
| 20 |
+
const BAD_COLOR = /var\(|env\(|attr\(|calc\(|^(inherit|initial|unset|revert|revert-layer|currentcolor|auto|none|transparent|-webkit-[\w-]+|-moz-[\w-]+)$/i;
|
| 21 |
+
function splitTop(s) { // split on top-level commas
|
| 22 |
+
const out = []; let depth = 0, cur = '';
|
| 23 |
+
for (const ch of s) { if (ch === '(') depth++; else if (ch === ')') depth--; if (ch === ',' && depth === 0) { out.push(cur); cur = ''; } else cur += ch; }
|
| 24 |
+
out.push(cur); return out.map(x => x.trim());
|
| 25 |
+
}
|
| 26 |
+
function canvasParse(str) {
|
| 27 |
+
const ld = str.match(/^light-dark\((.*)\)$/i);
|
| 28 |
+
if (ld) { const parts = splitTop(ld[1]); return parts.length === 2 ? parseColor(parts[0]) : null; }
|
| 29 |
+
if (BAD_COLOR.test(str)) return null;
|
| 30 |
+
if (typeof CSS !== 'undefined' && CSS.supports && !CSS.supports('color', str)) return null;
|
| 31 |
+
if (!cctx) {
|
| 32 |
+
const c = (typeof OffscreenCanvas !== 'undefined') ? new OffscreenCanvas(1, 1) : Object.assign(document.createElement('canvas'), {width: 1, height: 1});
|
| 33 |
+
cctx = c.getContext('2d', {willReadFrequently: true});
|
| 34 |
+
}
|
| 35 |
+
cctx.fillStyle = '#010203';
|
| 36 |
+
const before = cctx.fillStyle;
|
| 37 |
+
cctx.fillStyle = str;
|
| 38 |
+
if (cctx.fillStyle === before && !/^#?010203/i.test(str)) return null; // canvas rejected it
|
| 39 |
+
cctx.clearRect(0, 0, 1, 1);
|
| 40 |
+
cctx.fillRect(0, 0, 1, 1);
|
| 41 |
+
const d = cctx.getImageData(0, 0, 1, 1).data;
|
| 42 |
+
return [d[0], d[1], d[2], d[3] / 255];
|
| 43 |
+
}
|
| 44 |
+
function parseColor(str) {
|
| 45 |
+
if (!str) return null;
|
| 46 |
+
let c = colorCache.get(str);
|
| 47 |
+
if (c !== undefined) return c;
|
| 48 |
+
c = null;
|
| 49 |
+
const s = str.trim();
|
| 50 |
+
if (s === 'transparent' || s === 'none' || s === '' || s.startsWith('url(') || s === 'currentcolor') c = null;
|
| 51 |
+
else {
|
| 52 |
+
const m = s.match(/^rgba?\(\s*([\d.]+)[\s,]+([\d.]+)[\s,]+([\d.]+)(?:\s*[,/]\s*([\d.]+%?))?\s*\)$/);
|
| 53 |
+
if (m) {
|
| 54 |
+
let a = m[4] === undefined ? 1 : (m[4].endsWith('%') ? parseFloat(m[4]) / 100 : parseFloat(m[4]));
|
| 55 |
+
c = [Math.round(+m[1]), Math.round(+m[2]), Math.round(+m[3]), a];
|
| 56 |
+
} else {
|
| 57 |
+
try { c = canvasParse(s); } catch (e) { c = null; }
|
| 58 |
+
}
|
| 59 |
+
if (c && c[3] <= 0.004) c = null;
|
| 60 |
+
}
|
| 61 |
+
colorCache.set(str, c);
|
| 62 |
+
return c;
|
| 63 |
+
}
|
| 64 |
+
// all colors inside a composite value (box-shadow, gradients)
|
| 65 |
+
const COLOR_RE = /(rgba?\([^)]*\)|hsla?\([^)]*\)|oklch\([^)]*\)|oklab\([^)]*\)|lab\([^)]*\)|lch\([^)]*\)|color\([^)]*\)|#[0-9a-fA-F]{3,8}\b)/g;
|
| 66 |
+
function colorsIn(str) {
|
| 67 |
+
if (!str || str === 'none') return [];
|
| 68 |
+
const out = [];
|
| 69 |
+
for (const m of str.matchAll(COLOR_RE)) { const c = parseColor(m[1]); if (c) out.push(c); }
|
| 70 |
+
return out;
|
| 71 |
+
}
|
| 72 |
+
const hex2 = (n) => n.toString(16).padStart(2, '0');
|
| 73 |
+
const keyOf = (c) => hex2(c[0]) + hex2(c[1]) + hex2(c[2]) + hex2(Math.round(c[3] * 255));
|
| 74 |
+
|
| 75 |
+
// sRGB -> OKLab
|
| 76 |
+
function lin(u) { u /= 255; return u <= 0.04045 ? u / 12.92 : Math.pow((u + 0.055) / 1.055, 2.4); }
|
| 77 |
+
function oklab(c) {
|
| 78 |
+
const r = lin(c[0]), g = lin(c[1]), b = lin(c[2]);
|
| 79 |
+
const l = Math.cbrt(0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b);
|
| 80 |
+
const m = Math.cbrt(0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b);
|
| 81 |
+
const s = Math.cbrt(0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b);
|
| 82 |
+
return [0.2104542553 * l + 0.7936177850 * m - 0.0040720468 * s,
|
| 83 |
+
1.9779984951 * l - 2.4285922050 * m + 0.4505937099 * s,
|
| 84 |
+
0.0259040371 * l + 0.7827717662 * m - 0.8086757660 * s];
|
| 85 |
+
}
|
| 86 |
+
function relLum(c) { return 0.2126 * lin(c[0]) + 0.7152 * lin(c[1]) + 0.0722 * lin(c[2]); }
|
| 87 |
+
function over(top, bottom) { // composite rgba over opaque rgb
|
| 88 |
+
const a = top[3];
|
| 89 |
+
return [top[0] * a + bottom[0] * (1 - a), top[1] * a + bottom[1] * (1 - a), top[2] * a + bottom[2] * (1 - a), 1];
|
| 90 |
+
}
|
| 91 |
+
|
| 92 |
+
// ---------- dark-mode support scan ----------
|
| 93 |
+
const DARK_SEL_RE = /(?:^|[\s,>+~(])((?:html|body|:root)?(?:\.(?:[\w-]*-)?dark(?:-[\w-]+)?|\.(?:[\w-]*-)?night[\w-]*|\[data-[\w-]*(?:theme|mode|scheme|color-mode|appearance)[\w-]*(?:[~|^*$]?=\s*["']?[\w-]*dark[\w-]*["']?)\s*(?:i\s*)?\]|\[(?:theme|mode)\s*=\s*["']?dark["']?\]|\[dark(?:-theme|-mode)?\]))/i;
|
| 94 |
+
function scanSheets(limitRules) {
|
| 95 |
+
const res = {sheets: 0, blocked: 0, rules: 0, darkMediaRules: 0, darkMediaColorDecls: 0, lightMediaRules: 0,
|
| 96 |
+
darkSelectors: {}, ruleColors: [], hoverColors: 0};
|
| 97 |
+
let budget = limitRules;
|
| 98 |
+
const seenRuleColors = new Map();
|
| 99 |
+
const pushRuleColor = (role, c, state) => {
|
| 100 |
+
const k = role + ':' + keyOf(c);
|
| 101 |
+
const e = seenRuleColors.get(k);
|
| 102 |
+
if (e) { e.n++; if (state) e.state++; return; }
|
| 103 |
+
seenRuleColors.set(k, {role, c, n: 1, state: state ? 1 : 0});
|
| 104 |
+
};
|
| 105 |
+
function walk(rules, inDarkMedia, inLightMedia) {
|
| 106 |
+
for (const rule of rules) {
|
| 107 |
+
if (budget-- <= 0) return;
|
| 108 |
+
if (rule.cssRules && rule.media) {
|
| 109 |
+
const mt = rule.media.mediaText || '';
|
| 110 |
+
const dark = /prefers-color-scheme\s*:\s*dark/i.test(mt);
|
| 111 |
+
const light = /prefers-color-scheme\s*:\s*light/i.test(mt);
|
| 112 |
+
if (dark) res.darkMediaRules++;
|
| 113 |
+
if (light) res.lightMediaRules++;
|
| 114 |
+
walk(rule.cssRules, inDarkMedia || dark, inLightMedia || light);
|
| 115 |
+
continue;
|
| 116 |
+
}
|
| 117 |
+
if (rule.cssRules && !rule.style) { walk(rule.cssRules, inDarkMedia, inLightMedia); continue; }
|
| 118 |
+
if (!rule.style) continue;
|
| 119 |
+
res.rules++;
|
| 120 |
+
const st = rule.style;
|
| 121 |
+
let colorDecls = 0;
|
| 122 |
+
const sel = rule.selectorText || '';
|
| 123 |
+
const state = /:(hover|focus|active|focus-visible|focus-within|checked)/.test(sel);
|
| 124 |
+
for (let i = 0; i < st.length; i++) {
|
| 125 |
+
const p = st[i];
|
| 126 |
+
const v = st.getPropertyValue(p);
|
| 127 |
+
if (!v) continue;
|
| 128 |
+
let role = -1;
|
| 129 |
+
if (p === 'background-color' || p === 'background') role = R.bg;
|
| 130 |
+
else if (p === 'color' || p === '-webkit-text-fill-color') role = R.text;
|
| 131 |
+
else if (p.startsWith('border') && p.endsWith('color') || p === 'outline-color') role = R.border;
|
| 132 |
+
else if (p === 'fill') role = R.fill;
|
| 133 |
+
else if (p === 'stroke') role = R.stroke;
|
| 134 |
+
else if (p === 'box-shadow' || p === 'text-shadow') role = R.shadow;
|
| 135 |
+
else if (p === 'background-image') role = R.gradient;
|
| 136 |
+
else if (p.startsWith('--')) { if (COLOR_RE.test(v)) colorDecls++; COLOR_RE.lastIndex = 0; continue; }
|
| 137 |
+
if (role < 0) continue;
|
| 138 |
+
colorDecls++;
|
| 139 |
+
if (inDarkMedia || inLightMedia) continue; // light page colors only
|
| 140 |
+
const cs = (role === R.shadow || role === R.gradient || p === 'background') ? colorsIn(v) : [parseColor(v)].filter(Boolean);
|
| 141 |
+
for (const c of cs) { pushRuleColor(role === R.gradient && p === 'background' ? R.bg : role, c, state); if (state) res.hoverColors++; }
|
| 142 |
+
}
|
| 143 |
+
if (inDarkMedia) res.darkMediaColorDecls += colorDecls;
|
| 144 |
+
if (colorDecls && sel) {
|
| 145 |
+
const m = sel.match(DARK_SEL_RE);
|
| 146 |
+
if (m) {
|
| 147 |
+
const k = m[1].trim();
|
| 148 |
+
res.darkSelectors[k] = (res.darkSelectors[k] || 0) + colorDecls;
|
| 149 |
+
}
|
| 150 |
+
}
|
| 151 |
+
}
|
| 152 |
+
}
|
| 153 |
+
const sheets = [...document.styleSheets, ...(document.adoptedStyleSheets || []), ...(scanSheets.extra || [])];
|
| 154 |
+
for (const sh of sheets) {
|
| 155 |
+
if (scanSheets.filter && !scanSheets.filter(sh)) continue;
|
| 156 |
+
res.sheets++;
|
| 157 |
+
let rules;
|
| 158 |
+
try { rules = sh.cssRules; } catch (e) { res.blocked++; continue; }
|
| 159 |
+
if (!rules) { res.blocked++; continue; }
|
| 160 |
+
walk(rules, sh.media && /prefers-color-scheme\s*:\s*dark/i.test(sh.media.mediaText || ''), false);
|
| 161 |
+
}
|
| 162 |
+
res.ruleColors = [...seenRuleColors.values()];
|
| 163 |
+
return res;
|
| 164 |
+
}
|
| 165 |
+
|
| 166 |
+
function schemeInfo(sheetScan) {
|
| 167 |
+
const meta = document.querySelector('meta[name="color-scheme"]');
|
| 168 |
+
const rootCS = getComputedStyle(document.documentElement).colorScheme || '';
|
| 169 |
+
const darkLinks = [...document.querySelectorAll('link[rel~="stylesheet"][media]')].filter(l => /prefers-color-scheme\s*:\s*dark/i.test(l.media)).length;
|
| 170 |
+
const darkSources = [...document.querySelectorAll('picture source[media]')].filter(s => /prefers-color-scheme\s*:\s*dark/i.test(s.media)).length;
|
| 171 |
+
const html = document.documentElement, body = document.body;
|
| 172 |
+
const attrs = {};
|
| 173 |
+
for (const el of [html, body]) {
|
| 174 |
+
if (!el) continue;
|
| 175 |
+
for (const a of el.attributes) if (/theme|mode|scheme|appearance|dark|light/i.test(a.name + '=' + a.value) && a.name !== 'style') attrs[el.tagName.toLowerCase() + '@' + a.name] = a.value.slice(0, 80);
|
| 176 |
+
}
|
| 177 |
+
const sel = Object.entries(sheetScan.darkSelectors).sort((a, b) => b[1] - a[1]).slice(0, 12);
|
| 178 |
+
return {
|
| 179 |
+
metaColorScheme: meta ? meta.content : null,
|
| 180 |
+
rootColorScheme: rootCS,
|
| 181 |
+
darkMediaRules: sheetScan.darkMediaRules,
|
| 182 |
+
darkMediaColorDecls: sheetScan.darkMediaColorDecls,
|
| 183 |
+
lightMediaRules: sheetScan.lightMediaRules,
|
| 184 |
+
darkLinks, darkSources,
|
| 185 |
+
darkSelectors: sel,
|
| 186 |
+
rootAttrs: attrs,
|
| 187 |
+
rootClass: (html.className && typeof html.className === 'string' ? html.className : '').slice(0, 200),
|
| 188 |
+
bodyClass: (body && typeof body.className === 'string' ? body.className : '').slice(0, 200),
|
| 189 |
+
sheets: sheetScan.sheets, blockedSheets: sheetScan.blocked, rules: sheetScan.rules,
|
| 190 |
+
};
|
| 191 |
+
}
|
| 192 |
+
|
| 193 |
+
// ---------- ads ----------
|
| 194 |
+
// Ad slots are ignored entirely, so a page looks the same to the model with or without an ad blocker.
|
| 195 |
+
const AD_TOKEN = /^(ad|ads|adv|advert|adverts|advertisement|advertising|adsbygoogle|adslot|adunit|adbox|adblock|adbanner|ad[-_](slot|unit|container|wrapper|banner|box|block|placeholder|label|space|zone|frame|holder|area|leaderboard|rail|sidebar|top|bottom|inline|native|mpu|billboard|skyscraper|sticky)|(top|side|bottom|inline|sticky|native|header|footer|right|left)[-_]?ads?|banner[-_]?ads?|gpt[-_]ad.*|dfp[-_]?ad.*|google[-_]ads?.*|taboola.*|trc_.*|outbrain.*|ob[-_]widget.*|sponsored([-_].*)?|promoted[-_]?(content|post|ad)s?|mgid.*|revcontent.*|teads.*|criteo.*|amzn[-_]?ad.*|adthrive.*|ezoic.*|mediavine.*)$/i;
|
| 196 |
+
const AD_ID = /^(google_ads|div-gpt-ad|gpt-ad|ad[-_]|ads[-_]|taboola|outbrain|adthrive|ezoic-pub-ad|mgid|revcontent)/i;
|
| 197 |
+
const AD_SRC = /(doubleclick\.net|googlesyndication\.com|googleadservices\.com|adservice\.google|amazon-adsystem\.com|adnxs\.com|taboola\.com|outbrain\.com|criteo\.|pubmatic\.com|rubiconproject\.com|openx\.net|adsrvr\.org|moatads\.com|mgid\.com|revcontent\.com|teads\.tv|media\.net\/|ezoic|adthrive|mediavine)/i;
|
| 198 |
+
function isAd(el) {
|
| 199 |
+
const tag = el.tagName;
|
| 200 |
+
if (tag === 'INS' && /adsbygoogle/.test(el.className || '')) return true;
|
| 201 |
+
if ((tag === 'IFRAME' || tag === 'IMG' || tag === 'A') && AD_SRC.test(el.getAttribute('src') || el.getAttribute('href') || '')) return true;
|
| 202 |
+
if (el.hasAttribute('data-ad-slot') || el.hasAttribute('data-google-query-id') || el.hasAttribute('data-ad-unit') || el.hasAttribute('data-adunit') || el.hasAttribute('data-ad-client')) return true;
|
| 203 |
+
const id = el.id;
|
| 204 |
+
if (id && typeof id === 'string' && AD_ID.test(id)) return true;
|
| 205 |
+
const cls = typeof el.className === 'string' ? el.className : '';
|
| 206 |
+
if (cls) { for (const c of cls.split(/\s+/)) if (c && AD_TOKEN.test(c)) return true; }
|
| 207 |
+
const al = el.getAttribute('aria-label');
|
| 208 |
+
if (al && /^(advertisement|ad|sponsored|anzeige|publicidad|publicité|pubblicità|реклама|広告)$/i.test(al.trim())) return true;
|
| 209 |
+
return false;
|
| 210 |
+
}
|
| 211 |
+
|
| 212 |
+
// ---------- element walk ----------
|
| 213 |
+
function describe(el) {
|
| 214 |
+
let s = el.tagName.toLowerCase();
|
| 215 |
+
if (el.id && typeof el.id === 'string') s += '#' + el.id.slice(0, 24);
|
| 216 |
+
const cls = typeof el.className === 'string' ? el.className : (el.className && el.className.baseVal) || '';
|
| 217 |
+
if (cls) s += '.' + cls.trim().split(/\s+/).slice(0, 3).join('.').slice(0, 48);
|
| 218 |
+
return s;
|
| 219 |
+
}
|
| 220 |
+
function ownText(el) {
|
| 221 |
+
let t = '';
|
| 222 |
+
for (const n of el.childNodes) if (n.nodeType === 3) t += n.nodeValue;
|
| 223 |
+
t = t.replace(/\s+/g, ' ').trim();
|
| 224 |
+
if (!t && FORM.has(el.tagName)) {
|
| 225 |
+
const v = el.value || el.placeholder || '';
|
| 226 |
+
if (el.type !== 'hidden' && el.type !== 'checkbox' && el.type !== 'radio') t = String(v).slice(0, 60) || ' ';
|
| 227 |
+
}
|
| 228 |
+
return t;
|
| 229 |
+
}
|
| 230 |
+
|
| 231 |
+
async function extract(opts) {
|
| 232 |
+
opts = Object.assign({maxElements: 5000, maxRules: 40000, maxTokens: 384, maxImages: 48, describe: false,
|
| 233 |
+
thumbs: false, thumbSize: 32, bigThumbSize: 0, fetchImage: null, pseudo: true, time: true}, opts || {});
|
| 234 |
+
const t0 = performance.now();
|
| 235 |
+
const vw = innerWidth, vh = innerHeight;
|
| 236 |
+
const sx = scrollX, sy = scrollY;
|
| 237 |
+
const docEl = document.documentElement;
|
| 238 |
+
const docH = Math.max(docEl.scrollHeight, document.body ? document.body.scrollHeight : 0, vh);
|
| 239 |
+
const docW = Math.max(vw, Math.min(docEl.scrollWidth, vw * 1.5));
|
| 240 |
+
const pageArea = docW * Math.min(docH, vh * 8);
|
| 241 |
+
scanSheets.filter = opts.sheetFilter || null;
|
| 242 |
+
scanSheets.extra = opts.extraSheets || [];
|
| 243 |
+
const sheetScan = scanSheets(opts.maxRules);
|
| 244 |
+
|
| 245 |
+
// canvas default background (white unless the page opts into dark color-scheme)
|
| 246 |
+
const rootCS = getComputedStyle(docEl);
|
| 247 |
+
const canvasDark = /dark/.test(rootCS.colorScheme) && !/light/.test(rootCS.colorScheme);
|
| 248 |
+
const canvasColor = canvasDark ? [18, 18, 18, 1] : [255, 255, 255, 1];
|
| 249 |
+
|
| 250 |
+
const tokens = new Map(); // key role:hex -> token
|
| 251 |
+
function tok(role, c) {
|
| 252 |
+
const k = role + ':' + keyOf(c);
|
| 253 |
+
let t = tokens.get(k);
|
| 254 |
+
if (!t) {
|
| 255 |
+
t = {k, role, c, i: -1, area: 0, area1: 0, n: 0, nVis: 0, chars: 0, fsSum: 0, fsMax: 0, bold: 0, root: 0,
|
| 256 |
+
inter: 0, link: 0, head: 0, code: 0, form: 0, header: 0, footer: 0, nav: 0, under: [0, 0, 0, 0], over: [0, 0, 0, 0],
|
| 257 |
+
contrastSum: 0, contrastW: 0, rules: 0, state: 0, depthSum: 0, ex: [], xs: 0, ys: 0, ws: 0, hs: 0, imgUnder: 0, fixed: 0, media: 0};
|
| 258 |
+
tokens.set(k, t);
|
| 259 |
+
}
|
| 260 |
+
return t;
|
| 261 |
+
}
|
| 262 |
+
const rel = new Map(); // "type|a|b" -> weight
|
| 263 |
+
function addRel(type, a, b, w) {
|
| 264 |
+
if (!a || !b || a === b) return;
|
| 265 |
+
const k = type + '|' + a.k + '|' + b.k;
|
| 266 |
+
rel.set(k, (rel.get(k) || 0) + w);
|
| 267 |
+
}
|
| 268 |
+
|
| 269 |
+
// walk DOM (including open shadow roots) in document order
|
| 270 |
+
let adsSkipped = 0;
|
| 271 |
+
const els = [];
|
| 272 |
+
const stack = document.body ? [document.body] : [];
|
| 273 |
+
const deferred = [];
|
| 274 |
+
const parentOf = new Map();
|
| 275 |
+
while ((stack.length || deferred.length) && els.length < opts.maxElements) {
|
| 276 |
+
const el = stack.length ? stack.pop() : deferred.shift();
|
| 277 |
+
els.push(el);
|
| 278 |
+
if (el !== document.body && el.getClientRects().length === 0 && !(el instanceof SVGElement)) {
|
| 279 |
+
// hidden subtree (display:none etc.): visit later, only if budget remains
|
| 280 |
+
const kids = [...el.children].filter(k => !SKIP_TAGS.has(k.tagName) && !isAd(k));
|
| 281 |
+
for (const k of kids) { parentOf.set(k, el); deferred.push(k); }
|
| 282 |
+
continue;
|
| 283 |
+
}
|
| 284 |
+
const kids = [];
|
| 285 |
+
const root = el.shadowRoot || null;
|
| 286 |
+
if (root) for (const k of root.children) kids.push(k);
|
| 287 |
+
for (const k of el.children) kids.push(k);
|
| 288 |
+
for (let i = kids.length - 1; i >= 0; i--) {
|
| 289 |
+
const k = kids[i];
|
| 290 |
+
if (SKIP_TAGS.has(k.tagName)) continue;
|
| 291 |
+
if (isAd(k)) { adsSkipped++; continue; }
|
| 292 |
+
parentOf.set(k, el);
|
| 293 |
+
stack.push(k);
|
| 294 |
+
}
|
| 295 |
+
}
|
| 296 |
+
|
| 297 |
+
// media (video players, embeds, canvases, large images): overlays on these must keep their colours
|
| 298 |
+
const media = [];
|
| 299 |
+
for (const m of document.querySelectorAll('video, iframe, canvas, embed, object, img')) {
|
| 300 |
+
if (m.tagName === 'IFRAME' && AD_SRC.test(m.src || '')) continue;
|
| 301 |
+
const r = m.getBoundingClientRect();
|
| 302 |
+
if (r.width * r.height < (m.tagName === 'IMG' ? 40000 : 6000)) continue;
|
| 303 |
+
media.push({m, x: r.left + sx, y: r.top + sy, w: r.width, h: r.height});
|
| 304 |
+
if (media.length > 64) break;
|
| 305 |
+
}
|
| 306 |
+
const onMedia = (el, x, y, w, h) => {
|
| 307 |
+
if (!media.length || w * h <= 0) return 0;
|
| 308 |
+
for (const md of media) {
|
| 309 |
+
if (md.m === el) continue;
|
| 310 |
+
if (el.contains(md.m) && w * h < md.w * md.h * 2.5) return 1; // player/media container
|
| 311 |
+
const ix = Math.max(0, Math.min(x + w, md.x + md.w) - Math.max(x, md.x));
|
| 312 |
+
const iy = Math.max(0, Math.min(y + h, md.y + md.h) - Math.max(y, md.y));
|
| 313 |
+
if (ix * iy >= 0.6 * w * h) return 1; // drawn on top of media
|
| 314 |
+
}
|
| 315 |
+
return 0;
|
| 316 |
+
};
|
| 317 |
+
// per-element info
|
| 318 |
+
const info = new Map();
|
| 319 |
+
const htmlBg = parseColor(rootCS.backgroundColor);
|
| 320 |
+
const bodyCS = document.body ? getComputedStyle(document.body) : null;
|
| 321 |
+
const bodyBg = bodyCS ? parseColor(bodyCS.backgroundColor) : null;
|
| 322 |
+
// effective root surface: body bg propagates to canvas if html has none
|
| 323 |
+
let rootSurface = htmlBg ? over(htmlBg, canvasColor) : (bodyBg ? over(bodyBg, canvasColor) : canvasColor);
|
| 324 |
+
const rootTok = tok(R.bg, htmlBg || bodyBg || canvasColor);
|
| 325 |
+
rootTok.root = 1;
|
| 326 |
+
rootTok.n++; rootTok.nVis++;
|
| 327 |
+
const rootRec = {surf: rootSurface, surfTok: rootTok, imgBg: false, region: 0, depth: 0, fixed: false};
|
| 328 |
+
const images = [];
|
| 329 |
+
const svgRoots = [];
|
| 330 |
+
let visibleCount = 0;
|
| 331 |
+
|
| 332 |
+
const ec = opts.elemColors ? [] : null;
|
| 333 |
+
for (const el of els) {
|
| 334 |
+
const cs = getComputedStyle(el);
|
| 335 |
+
const ecRow = ec ? [el.tagName, 0, null, null, null, null, null, null, null] : null;
|
| 336 |
+
if (ec) ec.push(ecRow);
|
| 337 |
+
const parent = parentOf.get(el);
|
| 338 |
+
const p = parent ? info.get(parent) : rootRec;
|
| 339 |
+
const pr = p || rootRec;
|
| 340 |
+
if (cs.display === 'none') { info.set(el, Object.assign({}, pr, {hidden: true})); }
|
| 341 |
+
const hidden = cs.display === 'none' || (pr && pr.hidden) || cs.visibility === 'hidden' || cs.opacity === '0';
|
| 342 |
+
const rect = el.getBoundingClientRect();
|
| 343 |
+
const x = rect.left + sx, y = rect.top + sy;
|
| 344 |
+
const w = Math.max(0, Math.min(rect.width, docW - Math.max(0, x)));
|
| 345 |
+
const h = Math.max(0, rect.height);
|
| 346 |
+
const vis = !hidden && w > 1 && h > 1 && x + w > 0 && y + h > 0;
|
| 347 |
+
const area = vis ? Math.min(w * h, docW * docH) : 0;
|
| 348 |
+
const area1 = vis ? Math.max(0, Math.min(x + w, vw) - Math.max(x, 0)) * Math.max(0, Math.min(y + h, vh) - Math.max(y, 0)) : 0;
|
| 349 |
+
if (vis) visibleCount++;
|
| 350 |
+
if (ecRow) ecRow[1] = Math.round(area);
|
| 351 |
+
const tag = el.tagName;
|
| 352 |
+
const isSvg = el instanceof SVGElement;
|
| 353 |
+
// regions
|
| 354 |
+
let region = pr.region;
|
| 355 |
+
if (tag === 'HEADER' || el.getAttribute && el.getAttribute('role') === 'banner') region |= 1;
|
| 356 |
+
if (tag === 'NAV' || el.getAttribute && el.getAttribute('role') === 'navigation') region |= 2;
|
| 357 |
+
if (tag === 'FOOTER' || el.getAttribute && el.getAttribute('role') === 'contentinfo') region |= 4;
|
| 358 |
+
const fixed = pr.fixed || cs.position === 'fixed' || cs.position === 'sticky';
|
| 359 |
+
const om = vis ? onMedia(el, x, y, w, h) : 0;
|
| 360 |
+
const rec = {surf: pr.surf, surfTok: pr.surfTok, imgBg: pr.imgBg, region, depth: pr.depth + 1, hidden: hidden, fixed, om,
|
| 361 |
+
inLink: pr.inLink || tag === 'A', link: tag === 'A' ? el : pr.link};
|
| 362 |
+
info.set(el, rec);
|
| 363 |
+
const flags = (t) => {
|
| 364 |
+
t.n++;
|
| 365 |
+
if (vis) { t.nVis++; t.depthSum += rec.depth; }
|
| 366 |
+
if (INTERACTIVE.has(tag) || el.getAttribute && el.getAttribute('role') === 'button') t.inter++;
|
| 367 |
+
if (rec.inLink) t.link++;
|
| 368 |
+
if (HEADINGS.has(tag)) t.head++;
|
| 369 |
+
if (CODE.has(tag)) t.code++;
|
| 370 |
+
if (FORM.has(tag)) t.form++;
|
| 371 |
+
if (region & 1) t.header++;
|
| 372 |
+
if (region & 2) t.nav++;
|
| 373 |
+
if (region & 4) t.footer++;
|
| 374 |
+
if (fixed) t.fixed++;
|
| 375 |
+
if (om) t.media++;
|
| 376 |
+
if (opts.describe && t.ex.length < 3 && vis) {
|
| 377 |
+
const txt = ownText(el);
|
| 378 |
+
t.ex.push(describe(el) + (txt ? ' "' + txt.slice(0, 40) + '"' : ''));
|
| 379 |
+
}
|
| 380 |
+
};
|
| 381 |
+
|
| 382 |
+
// background
|
| 383 |
+
const bg = parseColor(cs.backgroundColor);
|
| 384 |
+
const bgImg = cs.backgroundImage;
|
| 385 |
+
if (bg) {
|
| 386 |
+
const t = tok(R.bg, bg);
|
| 387 |
+
if (ecRow) ecRow[2] = keyOf(bg);
|
| 388 |
+
flags(t);
|
| 389 |
+
t.area += area; t.area1 += area1;
|
| 390 |
+
if (vis) { t.xs += x * area; t.ys += y * area; t.ws += w * area; t.hs += h * area; }
|
| 391 |
+
const under = pr.surf;
|
| 392 |
+
if (vis) { addRel('bb', t, pr.surfTok, area); }
|
| 393 |
+
const ul = oklab(under);
|
| 394 |
+
t.under[0] += ul[0] * (area + 1); t.under[1] += ul[1] * (area + 1); t.under[2] += ul[2] * (area + 1); t.under[3] += area + 1;
|
| 395 |
+
if (pr.imgBg) t.imgUnder += 1;
|
| 396 |
+
rec.surf = over(bg, under);
|
| 397 |
+
rec.surfTok = t;
|
| 398 |
+
rec.surfEl = el;
|
| 399 |
+
rec.ownArea = area;
|
| 400 |
+
}
|
| 401 |
+
if (bgImg && bgImg !== 'none') {
|
| 402 |
+
if (bgImg.includes('gradient')) {
|
| 403 |
+
if (ecRow) ecRow[7] = colorsIn(bgImg).map(keyOf).join(',');
|
| 404 |
+
for (const c of colorsIn(bgImg)) { const t = tok(R.gradient, c); flags(t); t.area += area / 2; t.area1 += area1 / 2; }
|
| 405 |
+
const gc = colorsIn(bgImg);
|
| 406 |
+
if (gc.length) {
|
| 407 |
+
let avg = [0, 0, 0, 1];
|
| 408 |
+
for (const c of gc) { const o = over(c, pr.surf); avg[0] += o[0] / gc.length; avg[1] += o[1] / gc.length; avg[2] += o[2] / gc.length; }
|
| 409 |
+
rec.surf = avg; rec.imgBg = false;
|
| 410 |
+
}
|
| 411 |
+
}
|
| 412 |
+
const urls = [...bgImg.matchAll(/url\(\s*["']?([^"')]+)["']?\s*\)/g)].map(m => m[1]);
|
| 413 |
+
if (urls.length && vis && area > 64) {
|
| 414 |
+
images.push({el, kind: 'bg', src: urls[0], x, y, w, h, area, area1, rec});
|
| 415 |
+
if (area > 0.02 * vw * vh) rec.imgBg = true;
|
| 416 |
+
}
|
| 417 |
+
}
|
| 418 |
+
if (ecRow) ecRow[8] = keyOf(rec.surf.map((v, i) => i < 3 ? Math.round(v) : 1));
|
| 419 |
+
// text
|
| 420 |
+
const txt = ownText(el);
|
| 421 |
+
if (txt) {
|
| 422 |
+
const col = parseColor(cs.color);
|
| 423 |
+
if (col) {
|
| 424 |
+
const t = tok(R.text, col);
|
| 425 |
+
if (ecRow) ecRow[3] = keyOf(col);
|
| 426 |
+
flags(t);
|
| 427 |
+
const n = txt.length;
|
| 428 |
+
t.chars += vis ? n : 0;
|
| 429 |
+
const fs = parseFloat(cs.fontSize) || 16;
|
| 430 |
+
t.fsSum += fs * n; t.fsMax = Math.max(t.fsMax, vis ? fs : 0);
|
| 431 |
+
if ((parseInt(cs.fontWeight) || 400) >= 600) t.bold += n;
|
| 432 |
+
t.area += vis ? Math.min(area, n * fs * fs * 0.5) : 0;
|
| 433 |
+
t.area1 += area1 ? Math.min(area1, n * fs * fs * 0.5) : 0;
|
| 434 |
+
const under = rec.surf;
|
| 435 |
+
const ul = oklab(under);
|
| 436 |
+
const wgt = vis ? n : 0.1;
|
| 437 |
+
t.under[0] += ul[0] * wgt; t.under[1] += ul[1] * wgt; t.under[2] += ul[2] * wgt; t.under[3] += wgt;
|
| 438 |
+
const fg = over(col, under);
|
| 439 |
+
const l1 = relLum(fg), l2 = relLum(under);
|
| 440 |
+
t.contrastSum += ((Math.max(l1, l2) + 0.05) / (Math.min(l1, l2) + 0.05)) * wgt; t.contrastW += wgt;
|
| 441 |
+
if (rec.imgBg) t.imgUnder += 1;
|
| 442 |
+
if (vis) addRel('tb', t, rec.surfTok, n);
|
| 443 |
+
// reverse: text sitting on this surface
|
| 444 |
+
const st = rec.surfTok;
|
| 445 |
+
const tl = oklab(col);
|
| 446 |
+
st.over[0] += tl[0] * wgt; st.over[1] += tl[1] * wgt; st.over[2] += tl[2] * wgt; st.over[3] += wgt;
|
| 447 |
+
if (bg) addRel('se', t, tok(R.bg, bg), n);
|
| 448 |
+
}
|
| 449 |
+
}
|
| 450 |
+
// borders
|
| 451 |
+
if (!isSvg) {
|
| 452 |
+
const seen = new Set();
|
| 453 |
+
for (const side of ['Top', 'Right', 'Bottom', 'Left']) {
|
| 454 |
+
const bw = parseFloat(cs['border' + side + 'Width']);
|
| 455 |
+
if (!bw || cs['border' + side + 'Style'] === 'none' || cs['border' + side + 'Style'] === 'hidden') continue;
|
| 456 |
+
const bc = parseColor(cs['border' + side + 'Color']);
|
| 457 |
+
if (!bc) continue;
|
| 458 |
+
const k = keyOf(bc);
|
| 459 |
+
const len = (side === 'Top' || side === 'Bottom') ? w : h;
|
| 460 |
+
if (seen.has(k)) { const t = tok(R.border, bc); t.area += vis ? bw * len : 0; continue; }
|
| 461 |
+
seen.add(k);
|
| 462 |
+
if (ecRow && !ecRow[4]) ecRow[4] = k;
|
| 463 |
+
const t = tok(R.border, bc);
|
| 464 |
+
flags(t);
|
| 465 |
+
t.area += vis ? bw * len : 0;
|
| 466 |
+
const ul = oklab(pr.surf);
|
| 467 |
+
t.under[0] += ul[0]; t.under[1] += ul[1]; t.under[2] += ul[2]; t.under[3] += 1;
|
| 468 |
+
if (vis) addRel('rb', t, pr.surfTok, bw * len);
|
| 469 |
+
if (bg) addRel('se', t, tok(R.bg, bg), 1);
|
| 470 |
+
}
|
| 471 |
+
const ow = parseFloat(cs.outlineWidth);
|
| 472 |
+
if (ow && cs.outlineStyle !== 'none') {
|
| 473 |
+
const oc = parseColor(cs.outlineColor);
|
| 474 |
+
if (oc) { const t = tok(R.border, oc); flags(t); t.area += vis ? ow * 2 * (w + h) : 0; }
|
| 475 |
+
}
|
| 476 |
+
// shadows
|
| 477 |
+
for (const sh of [cs.boxShadow, cs.textShadow]) {
|
| 478 |
+
if (!sh || sh === 'none') continue;
|
| 479 |
+
for (const c of colorsIn(sh)) { const t = tok(R.shadow, c); flags(t); t.area += vis ? Math.min(area, 2 * (w + h) * 4) : 0;
|
| 480 |
+
const ul = oklab(pr.surf); t.under[0] += ul[0]; t.under[1] += ul[1]; t.under[2] += ul[2]; t.under[3] += 1; }
|
| 481 |
+
}
|
| 482 |
+
}
|
| 483 |
+
// svg paint
|
| 484 |
+
if (isSvg) {
|
| 485 |
+
const ltag = tag.toLowerCase();
|
| 486 |
+
if (ltag === 'svg' && !(parent instanceof SVGElement) && vis && area >= 400) {
|
| 487 |
+
const meta = ((el.getAttribute('class') || '') + ' ' + (el.id || '') + ' ' + (el.getAttribute('aria-label') || '') + ' ' + ((el.querySelector('title') || {}).textContent || '')).toLowerCase();
|
| 488 |
+
const useExt = [...el.querySelectorAll('use')].some(u => { const h = u.getAttribute('href') || u.getAttribute('xlink:href') || ''; return h && !h.startsWith('#'); });
|
| 489 |
+
if (rec.inLink || (region & 1) || /logo|brand|wordmark/.test(meta) || useExt || area >= 4000)
|
| 490 |
+
images.push({el, kind: 'svg', src: '', x, y, w, h, area, area1, rec, svgMeta: meta, useExt});
|
| 491 |
+
}
|
| 492 |
+
if (SVG_SHAPES.has(ltag) && ltag !== 'svg' && ltag !== 'g') {
|
| 493 |
+
const fc = cs.fill && cs.fill.startsWith('url') ? null : parseColor(cs.fill);
|
| 494 |
+
const sc = cs.stroke && cs.stroke.startsWith('url') ? null : parseColor(cs.stroke);
|
| 495 |
+
const sw = parseFloat(cs.strokeWidth) || 0;
|
| 496 |
+
const fo = parseFloat(cs.fillOpacity); const so = parseFloat(cs.strokeOpacity);
|
| 497 |
+
if (ecRow && fc) ecRow[5] = keyOf(fc);
|
| 498 |
+
if (ecRow && sc) ecRow[6] = keyOf(sc);
|
| 499 |
+
if (fc && fo !== 0) { const t = tok(R.fill, fc); flags(t); t.area += area * 0.5; t.area1 += area1 * 0.5;
|
| 500 |
+
const ul = oklab(rec.surf); t.under[0] += ul[0]; t.under[1] += ul[1]; t.under[2] += ul[2]; t.under[3] += 1; if (vis) addRel('tb', t, rec.surfTok, Math.sqrt(area + 1)); }
|
| 501 |
+
if (sc && sw > 0 && so !== 0) { const t = tok(R.stroke, sc); flags(t); t.area += vis ? sw * (w + h) : 0;
|
| 502 |
+
const ul = oklab(rec.surf); t.under[0] += ul[0]; t.under[1] += ul[1]; t.under[2] += ul[2]; t.under[3] += 1; if (vis) addRel('tb', t, rec.surfTok, Math.sqrt(area + 1)); }
|
| 503 |
+
}
|
| 504 |
+
}
|
| 505 |
+
// images
|
| 506 |
+
if (tag === 'IMG' && vis && area > 16) {
|
| 507 |
+
images.push({el, kind: 'img', src: el.currentSrc || el.src, x, y, w, h, area, area1, rec});
|
| 508 |
+
} else if ((tag === 'IMAGE' || tag === 'image') && vis) {
|
| 509 |
+
const href = el.getAttribute('href') || el.getAttribute('xlink:href');
|
| 510 |
+
if (href) images.push({el, kind: 'svgimage', src: new URL(href, location.href).href, x, y, w, h, area, area1, rec});
|
| 511 |
+
}
|
| 512 |
+
// pseudo elements
|
| 513 |
+
if (opts.pseudo && vis && !isSvg && els.length < 20000) {
|
| 514 |
+
for (const pseudo of ['::before', '::after']) {
|
| 515 |
+
const ps = getComputedStyle(el, pseudo);
|
| 516 |
+
const ct = ps.content;
|
| 517 |
+
if (!ct || ct === 'none' || ct === 'normal') continue;
|
| 518 |
+
const pb = parseColor(ps.backgroundColor);
|
| 519 |
+
const pw = parseFloat(ps.width) || 0, ph = parseFloat(ps.height) || 0;
|
| 520 |
+
const parea = Math.min(area, pw * ph);
|
| 521 |
+
if (pb) { const t = tok(R.bg, pb); flags(t); t.area += parea; addRel('bb', t, rec.surfTok, parea); }
|
| 522 |
+
if (ct !== '""' && ct !== "''") {
|
| 523 |
+
const pc = parseColor(ps.color);
|
| 524 |
+
if (pc) { const t = tok(R.text, pc); flags(t); t.chars += 1; const ul = oklab(rec.surf); t.under[0] += ul[0]; t.under[1] += ul[1]; t.under[2] += ul[2]; t.under[3] += 1; addRel('tb', t, rec.surfTok, 1); }
|
| 525 |
+
}
|
| 526 |
+
const bw = parseFloat(ps.borderTopWidth) || parseFloat(ps.borderBottomWidth) || parseFloat(ps.borderLeftWidth);
|
| 527 |
+
if (bw) { const bc = parseColor(ps.borderTopColor !== 'rgba(0, 0, 0, 0)' ? ps.borderTopColor : ps.borderBottomColor); if (bc) { const t = tok(R.border, bc); flags(t); t.area += bw * Math.max(pw, ph); } }
|
| 528 |
+
const psh = ps.boxShadow;
|
| 529 |
+
if (psh && psh !== 'none') for (const c of colorsIn(psh)) { const t = tok(R.shadow, c); flags(t); }
|
| 530 |
+
if (ps.backgroundImage && ps.backgroundImage.includes('gradient')) for (const c of colorsIn(ps.backgroundImage)) { const t = tok(R.gradient, c); flags(t); t.area += parea / 2; }
|
| 531 |
+
}
|
| 532 |
+
}
|
| 533 |
+
}
|
| 534 |
+
|
| 535 |
+
// own area for surfaces: subtract direct child surfaces
|
| 536 |
+
// (approximation: walk surfaces; area already includes descendants)
|
| 537 |
+
const childSurfArea = new Map();
|
| 538 |
+
for (const el of els) {
|
| 539 |
+
const r = info.get(el);
|
| 540 |
+
if (!r || !r.surfEl || r.surfEl !== el || !r.ownArea) continue;
|
| 541 |
+
// find nearest ancestor surface element
|
| 542 |
+
let p = parentOf.get(el);
|
| 543 |
+
while (p && !(info.get(p) && info.get(p).surfEl === p)) p = parentOf.get(p);
|
| 544 |
+
const key = p || 'root';
|
| 545 |
+
childSurfArea.set(key, (childSurfArea.get(key) || 0) + r.ownArea);
|
| 546 |
+
}
|
| 547 |
+
const ownAreaTok = new Map();
|
| 548 |
+
let rootOwn = Math.max(0, docW * docH - (childSurfArea.get('root') || 0));
|
| 549 |
+
ownAreaTok.set(rootTok, rootOwn);
|
| 550 |
+
for (const el of els) {
|
| 551 |
+
const r = info.get(el);
|
| 552 |
+
if (!r || r.surfEl !== el || !r.ownArea) continue;
|
| 553 |
+
const own = Math.max(0, r.ownArea - (childSurfArea.get(el) || 0));
|
| 554 |
+
ownAreaTok.set(r.surfTok, (ownAreaTok.get(r.surfTok) || 0) + own);
|
| 555 |
+
}
|
| 556 |
+
|
| 557 |
+
// rule-only colors
|
| 558 |
+
for (const rc of sheetScan.ruleColors) {
|
| 559 |
+
const t = tok(rc.role, rc.c);
|
| 560 |
+
t.rules += rc.n; t.state += rc.state;
|
| 561 |
+
}
|
| 562 |
+
// root token area = canvas own area
|
| 563 |
+
rootTok.area += rootOwn;
|
| 564 |
+
|
| 565 |
+
// ---------- darkness measurement ----------
|
| 566 |
+
let bgLsum = 0, bgW = 0, darkW = 0;
|
| 567 |
+
for (const [t, own] of ownAreaTok) {
|
| 568 |
+
if (t.role !== R.bg) continue;
|
| 569 |
+
const eff = over(t.c, t.root ? canvasColor : (t.under[3] ? labToApproxRgb(t.under) : canvasColor));
|
| 570 |
+
const L = oklab(eff)[0];
|
| 571 |
+
bgLsum += L * own; bgW += own; if (L < 0.45) darkW += own;
|
| 572 |
+
}
|
| 573 |
+
let txLsum = 0, txW = 0;
|
| 574 |
+
for (const t of tokens.values()) if (t.role === R.text && t.chars) { txLsum += oklab(t.c)[0] * t.chars; txW += t.chars; }
|
| 575 |
+
const darkness = {bgL: bgW ? bgLsum / bgW : 1, darkFrac: bgW ? darkW / bgW : 0, textL: txW ? txLsum / txW : 0, rootL: oklab(rootSurface)[0]};
|
| 576 |
+
|
| 577 |
+
// ---------- token selection & features ----------
|
| 578 |
+
let list = [...tokens.values()];
|
| 579 |
+
const score = (t) => (t.root ? 1e12 : 0) + Math.log1p(t.area) * 3 + Math.log1p(t.chars) * 2 + Math.log1p(t.nVis) * 2 + Math.log1p(t.n) + Math.log1p(t.rules) * 0.7 + Math.log1p(t.state);
|
| 580 |
+
list.sort((a, b) => score(b) - score(a));
|
| 581 |
+
const dropped = Math.max(0, list.length - opts.maxTokens);
|
| 582 |
+
list = list.slice(0, opts.maxTokens);
|
| 583 |
+
list.forEach((t, i) => t.i = i);
|
| 584 |
+
|
| 585 |
+
// bg lightness rank among visible bg tokens
|
| 586 |
+
const bgVis = list.filter(t => t.role === R.bg && t.area > 0).map(t => oklab(t.c)[0]).sort((a, b) => a - b);
|
| 587 |
+
const rankL = (L) => { if (!bgVis.length) return 0.5; let i = 0; while (i < bgVis.length && bgVis[i] < L) i++; return i / bgVis.length; };
|
| 588 |
+
const totalChars = list.reduce((s, t) => s + (t.role === R.text ? t.chars : 0), 0) || 1;
|
| 589 |
+
|
| 590 |
+
const out = [];
|
| 591 |
+
for (const t of list) {
|
| 592 |
+
const lab = oklab(t.c);
|
| 593 |
+
const C = Math.hypot(lab[1], lab[2]);
|
| 594 |
+
const under = t.under[3] ? [t.under[0] / t.under[3], t.under[1] / t.under[3], t.under[2] / t.under[3]] : oklab(canvasColor);
|
| 595 |
+
const overL = t.over[3] ? [t.over[0] / t.over[3], t.over[1] / t.over[3], t.over[2] / t.over[3]] : [0, 0, 0];
|
| 596 |
+
const n = Math.max(1, t.n);
|
| 597 |
+
const own = ownAreaTok.get(t) || 0;
|
| 598 |
+
const f = [
|
| 599 |
+
lab[0], lab[1] * 4, lab[2] * 4, C * 4, t.c[3], C > 0.02 ? lab[1] / C : 0, C > 0.02 ? lab[2] / C : 0,
|
| 600 |
+
t.c[0] / 255, t.c[1] / 255, t.c[2] / 255,
|
| 601 |
+
Math.log1p(t.area / 1000) / 10, Math.sqrt(Math.min(1, t.area / pageArea)), Math.sqrt(Math.min(1, t.area1 / (vw * vh))), Math.sqrt(Math.min(1, own / pageArea)),
|
| 602 |
+
Math.log1p(t.n) / 8, Math.log1p(t.nVis) / 8, Math.log1p(t.rules) / 6, t.rules ? t.state / (t.rules) : 0, (t.n === 0 && t.rules > 0) ? 1 : 0,
|
| 603 |
+
Math.log1p(t.chars) / 10, t.chars / totalChars, t.chars ? (t.fsSum / t.chars) / 32 : 0, t.fsMax / 48, t.chars ? t.bold / t.chars : 0,
|
| 604 |
+
t.root, t.inter / n, t.link / n, t.head / n, t.code / n, t.form / n, t.header / n, t.nav / n, t.footer / n, t.fixed / n,
|
| 605 |
+
under[0], under[1] * 4, under[2] * 4, t.contrastW ? Math.log(t.contrastSum / t.contrastW) / 3 : 0, Math.min(1, t.imgUnder / n),
|
| 606 |
+
overL[0], overL[1] * 4, overL[2] * 4, t.over[3] ? 1 : 0,
|
| 607 |
+
t.role === R.bg ? rankL(lab[0]) : 0, t.nVis ? Math.min(1, t.depthSum / t.nVis / 30) : 0,
|
| 608 |
+
t.ws && t.area ? Math.min(1, t.ws / t.area / docW) : 0, t.hs && t.area ? Math.min(1, t.hs / t.area / vh) : 0,
|
| 609 |
+
t.ys && t.area ? Math.min(1, t.ys / t.area / docH) : 0,
|
| 610 |
+
t.media / n,
|
| 611 |
+
];
|
| 612 |
+
for (let q = 0; q < f.length; q++) if (!Number.isFinite(f[q])) f[q] = 0; // JSON turns NaN into null; never feed it to the model
|
| 613 |
+
const o = {k: t.k, role: t.role, rgba: [t.c[0], t.c[1], t.c[2], +t.c[3].toFixed(3)], f};
|
| 614 |
+
if (opts.describe) {
|
| 615 |
+
o.d = {area: Math.round(t.area), area1: Math.round(t.area1), own: Math.round(own), n: t.n, nVis: t.nVis, rules: t.rules, state: t.state, chars: t.chars,
|
| 616 |
+
fs: t.chars ? +(t.fsSum / t.chars).toFixed(1) : 0, root: t.root, inter: t.inter, link: t.link, head: t.head, header: t.header, nav: t.nav, footer: t.footer,
|
| 617 |
+
contrast: t.contrastW ? +(t.contrastSum / t.contrastW).toFixed(2) : null, media: t.media, ex: t.ex};
|
| 618 |
+
}
|
| 619 |
+
out.push(o);
|
| 620 |
+
}
|
| 621 |
+
const relOut = [];
|
| 622 |
+
for (const [k, w] of rel) {
|
| 623 |
+
const [type, a, b] = k.split('|');
|
| 624 |
+
const ta = tokens.get(a), tb = tokens.get(b);
|
| 625 |
+
if (!ta || !tb || ta.i < 0 || tb.i < 0) continue;
|
| 626 |
+
relOut.push([['tb', 'bb', 'rb', 'se'].indexOf(type), ta.i, tb.i, +Math.log1p(w).toFixed(3)]);
|
| 627 |
+
}
|
| 628 |
+
|
| 629 |
+
// ---------- images ----------
|
| 630 |
+
const imgSeen = new Map();
|
| 631 |
+
let svgCount = 0;
|
| 632 |
+
// logo-like images first, then by area
|
| 633 |
+
const logoish = (im) => (im.rec.link && im.rec.inLink ? 1 : 0) + ((im.rec.region & 1) ? 1 : 0) + (im.kind === 'svg' && /logo|brand/.test(im.svgMeta || '') ? 2 : 0);
|
| 634 |
+
images.sort((a, b) => (logoish(b) - logoish(a)) || (b.area - a.area));
|
| 635 |
+
const imgOut = [];
|
| 636 |
+
for (const im of images) {
|
| 637 |
+
if (imgOut.length >= opts.maxImages) break;
|
| 638 |
+
if (im.kind !== 'svg' && (!im.src || im.src.startsWith('data:image/gif;base64,R0lGOD'))) continue;
|
| 639 |
+
if (im.kind === 'svg') { im.src = 'svg:' + svgCount++; if (opts.markSvg) im.el.setAttribute('data-nadir-svg', im.src); }
|
| 640 |
+
const key = im.src;
|
| 641 |
+
if (imgSeen.has(key)) { const o = imgOut[imgSeen.get(key)]; o.count++; o.areaSum += im.area; continue; }
|
| 642 |
+
const el = im.el;
|
| 643 |
+
const a = im.rec.link;
|
| 644 |
+
const href = a ? (a.getAttribute('href') || '') : '';
|
| 645 |
+
let homeLink = 0;
|
| 646 |
+
try { if (a) { const u = new URL(href, location.href); homeLink = (u.origin === location.origin && (u.pathname === '/' || u.pathname === '')) ? 1 : 0; } } catch (e) {}
|
| 647 |
+
const meta = im.kind === 'svg' ? (im.svgMeta || '') : ((el.getAttribute && (el.getAttribute('alt') || '')) + ' ' + (typeof el.className === 'string' ? el.className : '') + ' ' + (el.id || '') + ' ' + (im.src.split('?')[0].split('/').pop() || '')).toLowerCase();
|
| 648 |
+
const ext = im.kind === 'svg' ? 'svg' : (im.src.split('?')[0].match(/\.(svg|png|jpe?g|webp|gif|avif|ico)$/i) || [, ''])[1].toLowerCase();
|
| 649 |
+
const under = oklab(im.rec.surf);
|
| 650 |
+
const o = {src: im.src.slice(0, 600), kind: im.kind, x: Math.round(im.x), y: Math.round(im.y), w: Math.round(im.w), h: Math.round(im.h), count: 1, areaSum: im.area,
|
| 651 |
+
nw: el.naturalWidth || 0, nh: el.naturalHeight || 0, homeLink, inLink: im.rec.inLink ? 1 : 0, header: (im.rec.region & 3) ? 1 : 0, footer: (im.rec.region & 4) ? 1 : 0,
|
| 652 |
+
logo: /logo|brand|wordmark/.test(meta) ? 1 : 0, icon: /icon|ico\b|glyph|arrow|chevron|caret|svg/.test(meta) ? 1 : 0, avatar: /avatar|profile|user/.test(meta) ? 1 : 0,
|
| 653 |
+
ext, underTok: im.rec.surfTok && im.rec.surfTok.i >= 0 ? im.rec.surfTok.i : -1, under,
|
| 654 |
+
surf: im.rec.surf.slice(0, 3).map(Math.round), vy: Math.round(im.y - sy), vx: Math.round(im.x - sx), useExt: im.useExt ? 1 : 0,
|
| 655 |
+
alt: opts.describe ? ((el.getAttribute && el.getAttribute('alt')) || '').slice(0, 60) : undefined,
|
| 656 |
+
stats: null, thumb: null};
|
| 657 |
+
if (im.kind === 'img' && el.complete === false) o.loading = 1;
|
| 658 |
+
imgSeen.set(key, imgOut.length);
|
| 659 |
+
imgOut.push(o);
|
| 660 |
+
}
|
| 661 |
+
if (opts.thumbs) await imageStats(imgOut, opts);
|
| 662 |
+
for (const o of imgOut) o.f = imageFeatures(o, vw, vh, docH);
|
| 663 |
+
const res = {
|
| 664 |
+
v: 5, url: location.href, title: document.title.slice(0, 200), vw, vh, docH, docW,
|
| 665 |
+
scheme: schemeInfo(sheetScan), darkness, canvasDark,
|
| 666 |
+
tokens: out, rel: relOut, images: imgOut,
|
| 667 |
+
ec: ec || undefined,
|
| 668 |
+
stats: {ads: adsSkipped, elements: els.length, visible: visibleCount, tokensTotal: tokens.size, tokensDropped: dropped, ms: Math.round(performance.now() - t0)},
|
| 669 |
+
};
|
| 670 |
+
return res;
|
| 671 |
+
}
|
| 672 |
+
|
| 673 |
+
function labToApproxRgb(u) {
|
| 674 |
+
// inverse OKLab for averaged under-colour
|
| 675 |
+
const L = u[0] / u[3], a = u[1] / u[3], b = u[2] / u[3];
|
| 676 |
+
const l_ = L + 0.3963377774 * a + 0.2158037573 * b;
|
| 677 |
+
const m_ = L - 0.1055613458 * a - 0.0638541728 * b;
|
| 678 |
+
const s_ = L - 0.0894841775 * a - 1.2914855480 * b;
|
| 679 |
+
const l = l_ ** 3, m = m_ ** 3, s = s_ ** 3;
|
| 680 |
+
const r = 4.0767416621 * l - 3.3077115913 * m + 0.2309699292 * s;
|
| 681 |
+
const g = -1.2684380046 * l + 2.6097574011 * m - 0.3413193965 * s;
|
| 682 |
+
const bb = -0.0041960863 * l - 0.7034186147 * m + 1.7076147010 * s;
|
| 683 |
+
const enc = (x) => { x = Math.max(0, Math.min(1, x)); return 255 * (x <= 0.0031308 ? 12.92 * x : 1.055 * Math.pow(x, 1 / 2.4) - 0.055); };
|
| 684 |
+
return [enc(r), enc(g), enc(bb), 1];
|
| 685 |
+
}
|
| 686 |
+
|
| 687 |
+
// ---------- image pixels ----------
|
| 688 |
+
function blobToDataURL(blob) {
|
| 689 |
+
return new Promise((res) => { const fr = new FileReader(); fr.onload = () => res(fr.result); fr.onerror = () => res(null); fr.readAsDataURL(blob); });
|
| 690 |
+
}
|
| 691 |
+
function loadViaImg(url) {
|
| 692 |
+
return new Promise((res) => {
|
| 693 |
+
const img = new Image(); img.decoding = 'async';
|
| 694 |
+
const done = (ok) => res(ok && img.naturalWidth ? img : null);
|
| 695 |
+
img.onload = () => done(true); img.onerror = () => done(false); setTimeout(() => done(false), 4000);
|
| 696 |
+
img.src = url;
|
| 697 |
+
});
|
| 698 |
+
}
|
| 699 |
+
// Decode without loading a URL into the page where possible (page CSP often blocks blob: images).
|
| 700 |
+
async function loadImage(src, fetchImage) {
|
| 701 |
+
if (fetchImage && !src.startsWith('data:')) {
|
| 702 |
+
const blob = await fetchImage(src);
|
| 703 |
+
if (!blob || !blob.size) return null;
|
| 704 |
+
const isSvg = /svg/.test(blob.type) || /\.svg(\?|#|$)/i.test(src);
|
| 705 |
+
if (!isSvg) { try { return await createImageBitmap(blob); } catch (e) {} }
|
| 706 |
+
const du = await blobToDataURL(isSvg ? new Blob([blob], {type: 'image/svg+xml'}) : blob);
|
| 707 |
+
return du ? await loadViaImg(du) : null;
|
| 708 |
+
}
|
| 709 |
+
return await loadViaImg(src);
|
| 710 |
+
}
|
| 711 |
+
|
| 712 |
+
async function imageStats(list, opts) {
|
| 713 |
+
const jobs = list.map(async (o) => {
|
| 714 |
+
try {
|
| 715 |
+
const img = await loadImage(o.src, opts.fetchImage);
|
| 716 |
+
if (!img) return;
|
| 717 |
+
const IW = img.naturalWidth || img.width, IH = img.naturalHeight || img.height;
|
| 718 |
+
o.nw = o.nw || IW; o.nh = o.nh || IH;
|
| 719 |
+
statsFromSource(o, img, IW, IH, opts);
|
| 720 |
+
if (img.close) img.close();
|
| 721 |
+
} catch (e) { o.err = String(e).slice(0, 80); }
|
| 722 |
+
});
|
| 723 |
+
await Promise.all(jobs);
|
| 724 |
+
}
|
| 725 |
+
function drawThumbSrc(src, iw, ih, size) {
|
| 726 |
+
const c = document.createElement('canvas'); c.width = size; c.height = size;
|
| 727 |
+
const ctx = c.getContext('2d', {willReadFrequently: true});
|
| 728 |
+
const s = Math.min(size / iw, size / ih);
|
| 729 |
+
const dw = Math.max(1, Math.round(iw * s)), dh = Math.max(1, Math.round(ih * s));
|
| 730 |
+
ctx.drawImage(src, Math.floor((size - dw) / 2), Math.floor((size - dh) / 2), dw, dh);
|
| 731 |
+
return {c, ctx, dw, dh};
|
| 732 |
+
}
|
| 733 |
+
function statsFromSource(o, img, IW, IH, opts) {
|
| 734 |
+
const S = opts.thumbSize;
|
| 735 |
+
const {ctx, dw, dh} = drawThumbSrc(img, IW, IH, S);
|
| 736 |
+
const d = ctx.getImageData(0, 0, S, S).data;
|
| 737 |
+
o.thumb = Array.from(d); // RGBA bytes, letterboxed (transparent pad)
|
| 738 |
+
let n = 0, tr = 0, semi = 0, Ls = 0, L2 = 0, Cs = 0, white = 0, black = 0, ringA = 0, ringL = 0, ringN = 0, total = 0;
|
| 739 |
+
const x0 = Math.floor((S - dw) / 2), y0 = Math.floor((S - dh) / 2);
|
| 740 |
+
const qset = new Set();
|
| 741 |
+
for (let y = y0; y < y0 + dh; y++) for (let x = x0; x < x0 + dw; x++) {
|
| 742 |
+
const i = (y * S + x) * 4; const a = d[i + 3] / 255; total++;
|
| 743 |
+
const ring = (x === x0 || y === y0 || x === x0 + dw - 1 || y === y0 + dh - 1);
|
| 744 |
+
if (ring) { ringA += a; ringN++; }
|
| 745 |
+
if (a < 0.1) { tr++; continue; }
|
| 746 |
+
if (a < 0.9) semi++;
|
| 747 |
+
const lab = oklab([d[i], d[i + 1], d[i + 2]]);
|
| 748 |
+
const C = Math.hypot(lab[1], lab[2]);
|
| 749 |
+
n++; Ls += lab[0]; L2 += lab[0] * lab[0]; Cs += C;
|
| 750 |
+
if (lab[0] > 0.93 && C < 0.03) white++;
|
| 751 |
+
if (lab[0] < 0.3 && C < 0.04) black++;
|
| 752 |
+
if (ring) ringL += lab[0];
|
| 753 |
+
qset.add((d[i] >> 5) << 6 | (d[i + 1] >> 5) << 3 | (d[i + 2] >> 5));
|
| 754 |
+
}
|
| 755 |
+
const mL = n ? Ls / n : 0;
|
| 756 |
+
o.stats = {transp: total ? tr / total : 0, semi: total ? semi / total : 0, meanL: mL, stdL: n ? Math.sqrt(Math.max(0, L2 / n - mL * mL)) : 0,
|
| 757 |
+
meanC: n ? Cs / n : 0, white: n ? white / n : 0, black: n ? black / n : 0, ringA: ringN ? ringA / ringN : 0,
|
| 758 |
+
ringL: ringN ? ringL / Math.max(1, ringN * (ringA / ringN || 1)) : 0, colors: qset.size / 512};
|
| 759 |
+
if (opts.bigThumbSize) o.big = drawThumbSrc(img, IW, IH, opts.bigThumbSize).c.toDataURL('image/png');
|
| 760 |
+
}
|
| 761 |
+
function imageFeatures(o, vw, vh, docH) {
|
| 762 |
+
const s = o.stats || {};
|
| 763 |
+
const ext = o.ext;
|
| 764 |
+
return sanitize([
|
| 765 |
+
Math.log1p(o.w) / 8, Math.log1p(o.h) / 8, Math.log1p(o.areaSum / 100) / 12, Math.min(1, (o.w * o.h) / (vw * vh)), Math.log(Math.max(1, o.w) / Math.max(1, o.h)) / 3,
|
| 766 |
+
Math.log1p(o.nw) / 9, Math.log1p(o.nh) / 9, Math.min(1, o.y / Math.max(1, docH)), o.y < vh ? 1 : 0, Math.log1p(o.count) / 3,
|
| 767 |
+
o.homeLink, o.inLink, o.header, o.footer, o.logo, o.icon, o.avatar,
|
| 768 |
+
ext === 'svg' ? 1 : 0, ext === 'png' ? 1 : 0, (ext === 'jpg' || ext === 'jpeg') ? 1 : 0, ext === 'webp' ? 1 : 0, ext === 'gif' ? 1 : 0, o.kind === 'bg' ? 1 : 0,
|
| 769 |
+
o.under ? o.under[0] : 1, o.under ? o.under[1] * 4 : 0, o.under ? o.under[2] * 4 : 0,
|
| 770 |
+
o.stats ? 1 : 0, s.transp || 0, s.semi || 0, s.meanL || 0, s.stdL || 0, (s.meanC || 0) * 4, s.white || 0, s.black || 0, s.ringA || 0, s.ringL || 0, s.colors || 0,
|
| 771 |
+
o.kind === 'svg' ? 1 : 0, o.useExt || 0,
|
| 772 |
+
]);
|
| 773 |
+
}
|
| 774 |
+
const sanitize = (f) => f.map((v) => (Number.isFinite(+v) ? +v : 0));
|
| 775 |
+
|
| 776 |
+
// Compute pixel stats for a list of image entries outside the page (helper page / extension offscreen document).
|
| 777 |
+
// opts.fetchImage(src) -> Blob; opts.shot: dataURL of the viewport screenshot (for inline-SVG crops).
|
| 778 |
+
async function defaultFetch(src) {
|
| 779 |
+
const ac = new AbortController(); const t = setTimeout(() => ac.abort(), 6000);
|
| 780 |
+
try { const r = await fetch(src, {signal: ac.signal, credentials: 'omit'}); if (!r.ok) return null; const b = await r.blob(); return b.size > 8e6 ? null : b; }
|
| 781 |
+
catch (e) { return null; } finally { clearTimeout(t); }
|
| 782 |
+
}
|
| 783 |
+
async function imageStatsStandalone(list, opts) {
|
| 784 |
+
opts = Object.assign({thumbSize: 32, bigThumbSize: 0, concurrency: 8, fetchImage: defaultFetch}, opts || {});
|
| 785 |
+
let shotImg = null;
|
| 786 |
+
if (opts.shot && list.some(o => o.kind === 'svg')) shotImg = await loadViaImg(opts.shot);
|
| 787 |
+
// inline-SVG crops first (local, fast), then fetched images with limited concurrency
|
| 788 |
+
for (const o of list) {
|
| 789 |
+
if (o.kind !== 'svg' || !shotImg) continue;
|
| 790 |
+
const dpr = opts.dpr || 1;
|
| 791 |
+
if (o.vy < 0 || (o.vy + o.h) * dpr > shotImg.naturalHeight + 1 || o.w < 2 || o.h < 2) continue;
|
| 792 |
+
cropStats(o, shotImg, opts);
|
| 793 |
+
}
|
| 794 |
+
const raster = list.filter(o => o.kind !== 'svg');
|
| 795 |
+
let k = 0;
|
| 796 |
+
const run = async () => { while (k < raster.length) { const o = raster[k++]; await imageStats([o], opts); } };
|
| 797 |
+
await Promise.all(Array.from({length: opts.concurrency}, run));
|
| 798 |
+
return list.map(o => ({stats: o.stats || null, thumb: o.thumb || null, big: o.big || null, nw: o.nw, nh: o.nh, err: o.err}));
|
| 799 |
+
}
|
| 800 |
+
// inline SVGs: crop the rendered element from a screenshot; alpha is keyed off the surface colour beneath it
|
| 801 |
+
function cropStats(o, shot, opts) {
|
| 802 |
+
const S = opts.thumbSize;
|
| 803 |
+
const c = document.createElement('canvas'); c.width = Math.max(1, Math.round(o.w)); c.height = Math.max(1, Math.round(o.h));
|
| 804 |
+
const cx = c.getContext('2d', {willReadFrequently: true});
|
| 805 |
+
const dpr = opts.dpr || 1;
|
| 806 |
+
cx.drawImage(shot, o.vx * dpr, o.vy * dpr, o.w * dpr, o.h * dpr, 0, 0, c.width, c.height);
|
| 807 |
+
const src = cx.getImageData(0, 0, c.width, c.height);
|
| 808 |
+
const sl = oklab(o.surf || [255, 255, 255]);
|
| 809 |
+
const d = src.data;
|
| 810 |
+
for (let i = 0; i < d.length; i += 4) {
|
| 811 |
+
const lab = oklab([d[i], d[i + 1], d[i + 2]]);
|
| 812 |
+
const dist = Math.hypot(lab[0] - sl[0], lab[1] - sl[1], lab[2] - sl[2]);
|
| 813 |
+
d[i + 3] = Math.round(255 * Math.max(0, Math.min(1, (dist - 0.02) / 0.1)));
|
| 814 |
+
}
|
| 815 |
+
cx.putImageData(src, 0, 0);
|
| 816 |
+
o.nw = c.width; o.nh = c.height;
|
| 817 |
+
statsFromSource(o, c, c.width, c.height, opts);
|
| 818 |
+
}
|
| 819 |
+
|
| 820 |
+
globalThis.__nadirImageStats = imageStatsStandalone;
|
| 821 |
+
globalThis.__nadirImageFeatures = imageFeatures;
|
| 822 |
+
globalThis.__nadirExtract = extract;
|
| 823 |
+
globalThis.__nadirOklab = oklab;
|
| 824 |
+
globalThis.__nadirParseColor = parseColor;
|
| 825 |
+
})();
|
runtime/nadir-infer.js
ADDED
|
@@ -0,0 +1,193 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
// Nadir inference glue: page description (from extract.js + image stats) -> model inputs -> dark mapping.
|
| 2 |
+
// Mirrors train.py collate() for a single page. Works in browsers (offscreen document) and Node (parity tests).
|
| 3 |
+
(function (root) {
|
| 4 |
+
'use strict';
|
| 5 |
+
const MAX_TOK = 256, MAX_IMG = 32, FT = 49, FI = 39, K = 11;
|
| 6 |
+
const TREAT = ['keep', 'dim', 'invert', 'light', 'plate'];
|
| 7 |
+
const FG_ROLES = new Set([1, 5, 6]); // text, svg fill, svg stroke
|
| 8 |
+
const lin = (u) => { u /= 255; return u <= 0.04045 ? u / 12.92 : Math.pow((u + 0.055) / 1.055, 2.4); };
|
| 9 |
+
function oklab(r, g, b) {
|
| 10 |
+
r = lin(r); g = lin(g); b = lin(b);
|
| 11 |
+
const l = Math.cbrt(0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b);
|
| 12 |
+
const m = Math.cbrt(0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b);
|
| 13 |
+
const s = Math.cbrt(0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b);
|
| 14 |
+
return [0.2104542553 * l + 0.7936177850 * m - 0.0040720468 * s, 1.9779984951 * l - 2.4285922050 * m + 0.4505937099 * s, 0.0259040371 * l + 0.7827717662 * m - 0.8086757660 * s];
|
| 15 |
+
}
|
| 16 |
+
function labToRgb(L, a, b) {
|
| 17 |
+
const l_ = L + 0.3963377774 * a + 0.2158037573 * b, m_ = L - 0.1055613458 * a - 0.0638541728 * b, s_ = L - 0.0894841775 * a - 1.2914855480 * b;
|
| 18 |
+
const l = l_ ** 3, m = m_ ** 3, s = s_ ** 3;
|
| 19 |
+
const enc = (x) => { x = Math.max(0, Math.min(1, x)); return Math.round(255 * (x <= 0.0031308 ? 12.92 * x : 1.055 * Math.pow(x, 1 / 2.4) - 0.055)); };
|
| 20 |
+
return [enc(4.0767416621 * l - 3.3077115913 * m + 0.2309699292 * s), enc(-1.2684380046 * l + 2.6097574011 * m - 0.3413193965 * s), enc(-0.0041960863 * l - 0.7034186147 * m + 1.7076147010 * s)];
|
| 21 |
+
}
|
| 22 |
+
const relLum = (rgb) => 0.2126 * lin(rgb[0]) + 0.7152 * lin(rgb[1]) + 0.0722 * lin(rgb[2]);
|
| 23 |
+
const contrast = (a, b) => { const x = relLum(a), y = relLum(b); return (Math.max(x, y) + 0.05) / (Math.min(x, y) + 0.05); };
|
| 24 |
+
const hex = (rgb, a) => '#' + rgb.map(v => v.toString(16).padStart(2, '0')).join('') + (a < 0.999 ? Math.round(a * 255).toString(16).padStart(2, '0') : '');
|
| 25 |
+
|
| 26 |
+
function buildInputs(page) {
|
| 27 |
+
const toks = page.tokens.slice(0, MAX_TOK);
|
| 28 |
+
const N = Math.max(1, toks.length);
|
| 29 |
+
const imgs = page.images.slice(0, MAX_IMG);
|
| 30 |
+
const Mreal = imgs.length, M = Math.max(1, Mreal);
|
| 31 |
+
const T = 1 + N + M;
|
| 32 |
+
const tok_f = new Float32Array(N * FT), role = new BigInt64Array(N), lab = new Float32Array(N * 3);
|
| 33 |
+
toks.forEach((t, i) => {
|
| 34 |
+
tok_f.set(t.f.slice(0, FT).map((v) => (Number.isFinite(v) ? v : 0)), i * FT); role[i] = BigInt(t.role);
|
| 35 |
+
const o = oklab(t.rgba[0], t.rgba[1], t.rgba[2]); lab.set(o, i * 3);
|
| 36 |
+
});
|
| 37 |
+
const img_f = new Float32Array(M * FI), thumbs = new Float32Array(M * 4 * 32 * 32);
|
| 38 |
+
imgs.forEach((im, j) => {
|
| 39 |
+
img_f.set(im.f.slice(0, FI).map((v) => (Number.isFinite(v) ? v : 0)), j * FI);
|
| 40 |
+
if (im.thumb && im.thumb.length === 4096) {
|
| 41 |
+
for (let y = 0; y < 32; y++) for (let x = 0; x < 32; x++) for (let c = 0; c < 4; c++)
|
| 42 |
+
thumbs[j * 4096 + c * 1024 + y * 32 + x] = im.thumb[(y * 32 + x) * 4 + c] / 255;
|
| 43 |
+
}
|
| 44 |
+
});
|
| 45 |
+
const st = page.stats || {};
|
| 46 |
+
const page_f = new Float32Array([page.vw / 1280, Math.min(8, page.docH / page.vh) / 8, toks.length / MAX_TOK, Mreal / MAX_IMG,
|
| 47 |
+
page.darkness.bgL, page.darkness.textL, page.canvasDark ? 1 : 0, Math.min(1, (st.tokensDropped || 0) / 200)]);
|
| 48 |
+
// relations
|
| 49 |
+
const rel = new Float32Array(T * T * K);
|
| 50 |
+
const n = toks.length;
|
| 51 |
+
const R = new Float32Array(4 * (n + 1) * (n + 1));
|
| 52 |
+
const P = n + 1;
|
| 53 |
+
const tbBest = new Map();
|
| 54 |
+
for (const r of page.rel) {
|
| 55 |
+
const [typ, a, b, w] = r;
|
| 56 |
+
if (a >= n || b >= n) continue;
|
| 57 |
+
const idx = typ * P * P + (a + 1) * P + (b + 1);
|
| 58 |
+
if (w > R[idx]) R[idx] = w;
|
| 59 |
+
if (typ === 0) { const cur = tbBest.get(a); if (!cur || w > cur[0]) tbBest.set(a, [w, b]); }
|
| 60 |
+
}
|
| 61 |
+
for (let k = 0; k < 4; k++) for (let a = 0; a < P; a++) {
|
| 62 |
+
let mx = 0; for (let b = 0; b < P; b++) mx = Math.max(mx, R[k * P * P + a * P + b]);
|
| 63 |
+
const den = mx + 1e-6;
|
| 64 |
+
for (let b = 0; b < P; b++) {
|
| 65 |
+
const v = R[k * P * P + a * P + b] / den;
|
| 66 |
+
if (!v) continue;
|
| 67 |
+
rel[(a * T + b) * K + k] = v; // a -> b
|
| 68 |
+
rel[(b * T + a) * K + 4 + k] = v; // transpose channel
|
| 69 |
+
}
|
| 70 |
+
}
|
| 71 |
+
imgs.forEach((im, j) => {
|
| 72 |
+
const u = im.underTok;
|
| 73 |
+
if (u != null && u >= 0 && u < n) { rel[((1 + N + j) * T + (1 + u)) * K + 8] = 1; rel[((1 + u) * T + (1 + N + j)) * K + 9] = 1; }
|
| 74 |
+
});
|
| 75 |
+
const keyOf = (t) => t.rgba[0] * 65536 + t.rgba[1] * 256 + t.rgba[2];
|
| 76 |
+
const byKey = new Map();
|
| 77 |
+
toks.forEach((t, i) => { const k = keyOf(t); if (!byKey.has(k)) byKey.set(k, []); byKey.get(k).push(i); });
|
| 78 |
+
for (const list of byKey.values()) if (list.length > 1) for (const i of list) for (const j of list) if (i !== j) rel[((1 + i) * T + (1 + j)) * K + 10] = 1;
|
| 79 |
+
const mask = new Uint8Array(T);
|
| 80 |
+
for (let i = 0; i <= n; i++) mask[i] = 1;
|
| 81 |
+
for (let j = 0; j < Mreal; j++) mask[1 + N + j] = 1;
|
| 82 |
+
return {N, M, T, Mreal, tbBest, feeds: {
|
| 83 |
+
tok_f: [tok_f, [1, N, FT]], tok_role: [role, [1, N]], tok_lab: [lab, [1, N, 3]], img_f: [img_f, [1, M, FI]], img_thumb: [thumbs, [1, M, 4, 32, 32]],
|
| 84 |
+
page_f: [page_f, [1, 8]], src: [new BigInt64Array([0n]), [1]], rel: [rel, [1, T, T, K]], mask: [mask, [1, T]],
|
| 85 |
+
}};
|
| 86 |
+
}
|
| 87 |
+
|
| 88 |
+
// outputs: lab Float32Array(N*3), alpha_logit Float32Array(N), img_logits Float32Array(M*5)
|
| 89 |
+
// brand colours keep their hue: a coloured input whose output drifts more than 45 degrees is turned back to its
|
| 90 |
+
// own hue (keeping the predicted lightness and chroma); neutral inputs get at most a faint tint
|
| 91 |
+
function hueGuard(t, o) {
|
| 92 |
+
const i = oklab(t.rgba[0], t.rgba[1], t.rgba[2]), q = oklab(...o.rgb);
|
| 93 |
+
const cin = Math.hypot(i[1], i[2]), cout = Math.hypot(q[1], q[2]);
|
| 94 |
+
if (cin > 0.03 && cout > 0.03) {
|
| 95 |
+
const d = Math.abs(((Math.atan2(q[2], q[1]) - Math.atan2(i[2], i[1]) + Math.PI) % (2 * Math.PI) + 2 * Math.PI) % (2 * Math.PI) - Math.PI);
|
| 96 |
+
if (d <= Math.PI / 4) return;
|
| 97 |
+
const h = Math.atan2(i[2], i[1]);
|
| 98 |
+
o.rgb = labToRgb(q[0], cout * Math.cos(h), cout * Math.sin(h));
|
| 99 |
+
} else if (cin <= 0.03 && cout > 0.06) {
|
| 100 |
+
o.rgb = labToRgb(q[0], q[1] * 0.06 / cout, q[2] * 0.06 / cout);
|
| 101 |
+
}
|
| 102 |
+
}
|
| 103 |
+
|
| 104 |
+
function decode(page, built, out, opts) {
|
| 105 |
+
opts = opts || {};
|
| 106 |
+
const toks = page.tokens.slice(0, MAX_TOK);
|
| 107 |
+
const n = toks.length;
|
| 108 |
+
const outputs = new Array(n);
|
| 109 |
+
const keep = new Set();
|
| 110 |
+
toks.forEach((t, i) => {
|
| 111 |
+
const L = out.lab[i * 3], a = out.lab[i * 3 + 1], b = out.lab[i * 3 + 2];
|
| 112 |
+
const ain = Math.min(1 - 1e-3, Math.max(1e-3, t.rgba[3]));
|
| 113 |
+
const alpha = 1 / (1 + Math.exp(-(out.alpha_logit[i] + Math.log(ain / (1 - ain)))));
|
| 114 |
+
outputs[i] = {rgb: labToRgb(L, a, b), a: t.rgba[3] >= 0.999 && alpha > 0.97 ? 1 : alpha, L};
|
| 115 |
+
// colours over video/photos that already read as dark mode (dark scrims, saturated player accents, light
|
| 116 |
+
// captions and icons) are shown as-is; anything else over media is redesigned like any other colour
|
| 117 |
+
const f = t.f;
|
| 118 |
+
const bgLike = t.role === 0 || t.role === 3 || t.role === 4;
|
| 119 |
+
if (f[48] >= 0.5 && (bgLike ? (f[24] <= 0.5 && f[11] < 0.25 && (f[0] < 0.5 || f[3] > 0.4)) : f[0] > 0.6)) {
|
| 120 |
+
keep.add(i);
|
| 121 |
+
const rgb = t.rgba.slice(0, 3);
|
| 122 |
+
outputs[i] = {rgb, a: t.rgba[3], L: oklab(...rgb)[0]};
|
| 123 |
+
} else hueGuard(t, outputs[i]);
|
| 124 |
+
});
|
| 125 |
+
// safety: page canvas must be dark
|
| 126 |
+
toks.forEach((t, i) => {
|
| 127 |
+
if (t.f[24] > 0.5 && t.role === 0 && outputs[i].L > 0.35) { const o = oklab(...outputs[i].rgb); outputs[i].rgb = labToRgb(0.2, o[1] * 0.3, o[2] * 0.3); outputs[i].L = 0.2; }
|
| 128 |
+
});
|
| 129 |
+
// things that were invisible against their surface stay invisible: gradient stops with exactly a surface's
|
| 130 |
+
// colour (text fades) take that surface's dark colour, and so do borders that matched their surface
|
| 131 |
+
const keyOf = (t) => t.rgba[0] * 65536 + t.rgba[1] * 256 + t.rgba[2];
|
| 132 |
+
const surfOf = new Map();
|
| 133 |
+
toks.forEach((t, i) => { if (t.role === 0 && t.rgba[3] >= 0.999 && !keep.has(i) && outputs[i].a >= 0.999 && !surfOf.has(keyOf(t))) surfOf.set(keyOf(t), i); });
|
| 134 |
+
toks.forEach((t, i) => {
|
| 135 |
+
if (t.role === 4 && !keep.has(i) && surfOf.has(keyOf(t))) { const j = surfOf.get(keyOf(t)); outputs[i].rgb = outputs[j].rgb.slice(); outputs[i].L = outputs[j].L; }
|
| 136 |
+
});
|
| 137 |
+
const borderOn = new Map();
|
| 138 |
+
for (const [typ, a, b, w] of page.rel) {
|
| 139 |
+
if (typ !== 2 || a >= n || b >= n || toks[a].role !== 2 || toks[b].role !== 0) continue;
|
| 140 |
+
const cur = borderOn.get(a); if (!cur || w > cur[0]) borderOn.set(a, [w, b]);
|
| 141 |
+
}
|
| 142 |
+
for (const [i, [, j]] of borderOn) {
|
| 143 |
+
if (keep.has(i) || keep.has(j)) continue;
|
| 144 |
+
if (contrast(toks[i].rgba, toks[j].rgba) < 1.05 && toks[i].rgba[3] >= 0.5) { outputs[i].rgb = outputs[j].rgb.slice(); outputs[i].L = outputs[j].L; }
|
| 145 |
+
}
|
| 146 |
+
// safety: every foreground colour (text, svg fill/stroke) stays readable on all the surfaces it sits on,
|
| 147 |
+
// never asking for more contrast than the light design had; one lightness per shared token, chosen on a grid
|
| 148 |
+
const surf = new Map();
|
| 149 |
+
for (const [typ, a, b, w] of page.rel) {
|
| 150 |
+
if (typ !== 0 || a >= n || b >= n || toks[b].role !== 0 || outputs[b].a < 0.999) continue;
|
| 151 |
+
if (!surf.has(a)) surf.set(a, []);
|
| 152 |
+
surf.get(a).push([b, w + 0.1]);
|
| 153 |
+
}
|
| 154 |
+
for (const [i, lst] of surf) {
|
| 155 |
+
const t = toks[i];
|
| 156 |
+
if (keep.has(i) || !FG_ROLES.has(t.role) || outputs[i].a < 0.999) continue;
|
| 157 |
+
let need0 = 3;
|
| 158 |
+
if (t.role === 1) { const fs = t.f[21] * 32, bold = t.f[23] > 0.5; need0 = (fs >= 24 || (fs >= 18.5 && bold)) ? 3 : 4.5; }
|
| 159 |
+
if (opts.minContrast) need0 = Math.min(need0, opts.minContrast);
|
| 160 |
+
const reqs = [];
|
| 161 |
+
for (const [j, w] of lst) {
|
| 162 |
+
const need = Math.min(need0, contrast(t.rgba, toks[j].rgba));
|
| 163 |
+
if (need >= 1.5) reqs.push([outputs[j].rgb, need, w]);
|
| 164 |
+
}
|
| 165 |
+
if (!reqs.length) continue;
|
| 166 |
+
const deficit = (rgb) => { let d = 0; for (const [bg, need, w] of reqs) d += w * Math.max(0, Math.log(need) - Math.log(contrast(rgb, bg))); return d; };
|
| 167 |
+
if (deficit(outputs[i].rgb) <= 0) continue;
|
| 168 |
+
const o = oklab(...outputs[i].rgb);
|
| 169 |
+
let W = 0; for (const r of reqs) W += r[2];
|
| 170 |
+
let best = null;
|
| 171 |
+
for (let k = 0; k <= 50; k++) {
|
| 172 |
+
const Lx = k / 50, rgb = labToRgb(Lx, o[1], o[2]);
|
| 173 |
+
const c = deficit(rgb) / W + 0.05 * Math.abs(Lx - o[0]);
|
| 174 |
+
if (!best || c < best[0]) best = [c, rgb, Lx];
|
| 175 |
+
}
|
| 176 |
+
outputs[i].rgb = best[1]; outputs[i].L = best[2];
|
| 177 |
+
}
|
| 178 |
+
const colors = toks.map((t, i) => keep.has(i) ? null : hex(outputs[i].rgb, outputs[i].a));
|
| 179 |
+
const imgs = page.images.slice(0, MAX_IMG);
|
| 180 |
+
const treatments = imgs.map((im, j) => {
|
| 181 |
+
let best = 0;
|
| 182 |
+
for (let c = 1; c < 5; c++) if (out.img_logits[j * 5 + c] > out.img_logits[j * 5 + best]) best = c;
|
| 183 |
+
const t = TREAT[best];
|
| 184 |
+
const u = im.underTok;
|
| 185 |
+
return root.NadirSafeTreatment ? root.NadirSafeTreatment(t, im, u != null && u >= 0 && u < n ? outputs[u].L : null) : t;
|
| 186 |
+
});
|
| 187 |
+
return {colors, keep: [...keep], treatments};
|
| 188 |
+
}
|
| 189 |
+
|
| 190 |
+
const api = {buildInputs, decode, oklab, labToRgb, TREAT, MAX_TOK, MAX_IMG};
|
| 191 |
+
root.NadirInfer = api;
|
| 192 |
+
if (typeof module !== 'undefined') module.exports = api;
|
| 193 |
+
})(typeof globalThis !== 'undefined' ? globalThis : this);
|
runtime/nadir-map.js
ADDED
|
@@ -0,0 +1,244 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
// Nadir colour mapping runtime: turns per-token dark colours (from the model or a teacher) into a lookup
|
| 2 |
+
// for every colour the page might use. Shared by the extension and the evaluation renderer.
|
| 3 |
+
(function () {
|
| 4 |
+
'use strict';
|
| 5 |
+
const ROLES = ['bg', 'text', 'border', 'shadow', 'gradient', 'fill', 'stroke'];
|
| 6 |
+
// which token roles answer a query role, in order of preference
|
| 7 |
+
const CHAIN = {
|
| 8 |
+
bg: ['bg', 'gradient', 'shadow'],
|
| 9 |
+
gradient: ['gradient', 'bg', 'shadow'],
|
| 10 |
+
shadow: ['shadow', 'bg', 'gradient'],
|
| 11 |
+
text: ['text', 'fill', 'stroke'],
|
| 12 |
+
fill: ['fill', 'text', 'stroke'],
|
| 13 |
+
stroke: ['stroke', 'fill', 'text'],
|
| 14 |
+
border: ['border', 'bg', 'stroke'],
|
| 15 |
+
};
|
| 16 |
+
const lin = (u) => { u /= 255; return u <= 0.04045 ? u / 12.92 : Math.pow((u + 0.055) / 1.055, 2.4); };
|
| 17 |
+
function oklab(r, g, b) {
|
| 18 |
+
r = lin(r); g = lin(g); b = lin(b);
|
| 19 |
+
const l = Math.cbrt(0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b);
|
| 20 |
+
const m = Math.cbrt(0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b);
|
| 21 |
+
const s = Math.cbrt(0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b);
|
| 22 |
+
return [0.2104542553 * l + 0.7936177850 * m - 0.0040720468 * s,
|
| 23 |
+
1.9779984951 * l - 2.4285922050 * m + 0.4505937099 * s,
|
| 24 |
+
0.0259040371 * l + 0.7827717662 * m - 0.8086757660 * s];
|
| 25 |
+
}
|
| 26 |
+
function toRgb(L, a, b) {
|
| 27 |
+
const l_ = L + 0.3963377774 * a + 0.2158037573 * b;
|
| 28 |
+
const m_ = L - 0.1055613458 * a - 0.0638541728 * b;
|
| 29 |
+
const s_ = L - 0.0894841775 * a - 1.2914855480 * b;
|
| 30 |
+
const l = l_ ** 3, m = m_ ** 3, s = s_ ** 3;
|
| 31 |
+
const enc = (x) => { x = Math.max(0, Math.min(1, x)); return Math.round(255 * (x <= 0.0031308 ? 12.92 * x : 1.055 * Math.pow(x, 1 / 2.4) - 0.055)); };
|
| 32 |
+
return [enc(4.0767416621 * l - 3.3077115913 * m + 0.2309699292 * s), enc(-1.2684380046 * l + 2.6097574011 * m - 0.3413193965 * s), enc(-0.0041960863 * l - 0.7034186147 * m + 1.7076147010 * s)];
|
| 33 |
+
}
|
| 34 |
+
const hex2 = (n) => n.toString(16).padStart(2, '0');
|
| 35 |
+
const keyOf = (r, g, b, a) => hex2(r) + hex2(g) + hex2(b) + hex2(Math.round(a * 255));
|
| 36 |
+
function css(rgb, a) {
|
| 37 |
+
return a >= 0.999 ? `#${hex2(rgb[0])}${hex2(rgb[1])}${hex2(rgb[2])}` : `rgba(${rgb[0]}, ${rgb[1]}, ${rgb[2]}, ${+a.toFixed(3)})`;
|
| 38 |
+
}
|
| 39 |
+
function parseHex(h) {
|
| 40 |
+
h = h.replace('#', '');
|
| 41 |
+
return [parseInt(h.slice(0, 2), 16), parseInt(h.slice(2, 4), 16), parseInt(h.slice(4, 6), 16), h.length >= 8 ? parseInt(h.slice(6, 8), 16) / 255 : 1];
|
| 42 |
+
}
|
| 43 |
+
|
| 44 |
+
class NadirMap {
|
| 45 |
+
// tokens: [{role, rgba:[r,g,b,a]}]; outputs: same length, each '#rrggbb[aa]' | [r,g,b,a] | null
|
| 46 |
+
// opts.keep: set of token indices that must stay unchanged (media overlays)
|
| 47 |
+
constructor(tokens, outputs, opts) {
|
| 48 |
+
opts = opts || {};
|
| 49 |
+
this.exact = {}; // role -> key -> css
|
| 50 |
+
this.pts = {}; // role -> [{lab, alpha, dl, da, db, dA}]
|
| 51 |
+
for (const r of ROLES) { this.exact[r] = new Map(); this.pts[r] = []; }
|
| 52 |
+
this.fillL = new Map(); // SVG fill colour -> output lightness
|
| 53 |
+
const keep = opts.keep || new Set();
|
| 54 |
+
tokens.forEach((t, i) => {
|
| 55 |
+
let o = outputs[i];
|
| 56 |
+
if (keep.has(i)) o = t.rgba;
|
| 57 |
+
if (o == null) return;
|
| 58 |
+
if (typeof o === 'string') o = parseHex(o);
|
| 59 |
+
const role = typeof t.role === 'number' ? ROLES[t.role] : t.role;
|
| 60 |
+
const [r, g, b, a] = t.rgba;
|
| 61 |
+
this.exact[role].set(keyOf(r, g, b, a), css([o[0], o[1], o[2]], o[3] == null ? a : o[3]));
|
| 62 |
+
const li = oklab(r, g, b), lo = oklab(o[0], o[1], o[2]);
|
| 63 |
+
if (role === 'fill' && !keep.has(i)) this.fillL.set(keyOf(r, g, b, a), lo[0]); // media keeps say nothing about ink
|
| 64 |
+
this.pts[role].push({lab: li, alpha: a, d: [lo[0] - li[0], lo[1] - li[1], lo[2] - li[2]], out: lo, oA: o[3] == null ? a : o[3]});
|
| 65 |
+
});
|
| 66 |
+
this.cache = new Map();
|
| 67 |
+
this.maxDist = opts.maxDist || 0.16;
|
| 68 |
+
}
|
| 69 |
+
// An SVG fill the model saw on this page: drawn as light ink (logo, icon) or as a dark surface (wave, panel)?
|
| 70 |
+
// null when the colour was never seen as a fill: the engine then falls back to its own size heuristic.
|
| 71 |
+
fillSide(r, g, b, a) {
|
| 72 |
+
const L = this.fillL.get(keyOf(r, g, b, a));
|
| 73 |
+
return L == null ? null : (L >= 0.5 ? 'fg' : 'bg');
|
| 74 |
+
}
|
| 75 |
+
lookup(role, r, g, b, a) {
|
| 76 |
+
const ck = role + keyOf(r, g, b, a);
|
| 77 |
+
let v = this.cache.get(ck);
|
| 78 |
+
if (v !== undefined) return v;
|
| 79 |
+
v = this._lookup(role, r, g, b, a);
|
| 80 |
+
this.cache.set(ck, v);
|
| 81 |
+
return v;
|
| 82 |
+
}
|
| 83 |
+
_lookup(role, r, g, b, a) {
|
| 84 |
+
const key = keyOf(r, g, b, a);
|
| 85 |
+
const chain = CHAIN[role] || [role];
|
| 86 |
+
for (const rr of chain) { const e = this.exact[rr].get(key); if (e) return e; }
|
| 87 |
+
// nearest-neighbour transfer within the role chain (first role with points)
|
| 88 |
+
const q = oklab(r, g, b);
|
| 89 |
+
let best = [];
|
| 90 |
+
for (let ci = 0; ci < chain.length; ci++) {
|
| 91 |
+
const pts = this.pts[chain[ci]];
|
| 92 |
+
for (const p of pts) {
|
| 93 |
+
const dl = p.lab[0] - q[0], da = p.lab[1] - q[1], db = p.lab[2] - q[2];
|
| 94 |
+
const d = Math.sqrt(dl * dl * 1.5 + da * da + db * db + (p.alpha - a) * (p.alpha - a) * 0.25) + ci * 0.03;
|
| 95 |
+
best.push([d, p]);
|
| 96 |
+
}
|
| 97 |
+
if (best.length >= 6) break;
|
| 98 |
+
}
|
| 99 |
+
if (!best.length) return null;
|
| 100 |
+
best.sort((x, y) => x[0] - y[0]);
|
| 101 |
+
best = best.slice(0, 4);
|
| 102 |
+
if (best[0][0] > this.maxDist) return null; // too far from anything the model saw: built-in fallback
|
| 103 |
+
let w = 0, d0 = 0, d1 = 0, d2 = 0, dA = 0;
|
| 104 |
+
for (const [d, p] of best) {
|
| 105 |
+
const wi = 1 / (d * d + 1e-4);
|
| 106 |
+
w += wi; d0 += p.d[0] * wi; d1 += p.d[1] * wi; d2 += p.d[2] * wi; dA += (p.oA - p.alpha) * wi;
|
| 107 |
+
}
|
| 108 |
+
const L = Math.max(0, Math.min(1, q[0] + d0 / w));
|
| 109 |
+
// keep chroma direction of the query, scale like the neighbours did
|
| 110 |
+
const rgb = toRgb(L, q[1] + d1 / w, q[2] + d2 / w);
|
| 111 |
+
const alpha = Math.max(0, Math.min(1, a + dA / w));
|
| 112 |
+
return css(rgb, alpha);
|
| 113 |
+
}
|
| 114 |
+
}
|
| 115 |
+
|
| 116 |
+
// image treatments as CSS
|
| 117 |
+
const TREAT_CSS = {
|
| 118 |
+
invert: 'filter: invert(1) hue-rotate(180deg) !important;',
|
| 119 |
+
light: 'filter: brightness(0) invert(0.92) !important;',
|
| 120 |
+
dim: 'filter: brightness(0.88) !important;',
|
| 121 |
+
plate: 'background-color: rgba(255,255,255,0.92) !important; box-shadow: 0 0 0 3px rgba(255,255,255,0.92) !important; border-radius: 4px !important;',
|
| 122 |
+
};
|
| 123 |
+
// Content images must never be colour-altered: only flat graphics (logos, icons, diagrams) may get a treatment.
|
| 124 |
+
// a logo on its own white rectangle: adapt it instead of leaving a white box. Only images the page marks as its
|
| 125 |
+
// logo (logo/brand/wordmark naming, or the link home); product shots, photos and other images are never touched
|
| 126 |
+
// Returns 'adapt:<fallback>': the page tries the adaptation (its pixel check confirms the white background)
|
| 127 |
+
// and uses the fallback treatment when there is nothing to adapt. Without pixel stats yet, logos still try.
|
| 128 |
+
function adaptable(im, s, surfL) {
|
| 129 |
+
if (!(im.logo || (im.homeLink && im.header)) || im.h > 200 || (surfL != null && surfL >= 0.4)) return false;
|
| 130 |
+
if (!s) return im.w * im.h <= 250000;
|
| 131 |
+
return s.transp < 0.05 && (s.ringA || 0) >= 0.95 && (s.ringL || 0) >= 0.85 && (s.white || 0) >= 0.3 && s.colors < 0.2;
|
| 132 |
+
}
|
| 133 |
+
function baseTreatment(t, im, surfL) {
|
| 134 |
+
const s = im && im.stats;
|
| 135 |
+
// a light, flat CSS background image (texture, pale wash, pattern) sits behind content whose text is now light:
|
| 136 |
+
// clear it so the dark surface designed for that element shows through. Photos and detailed images stay.
|
| 137 |
+
// (almost no lightness variation, so never a photo; SVG backgrounds can't be photos, so a smooth gradient qualifies).
|
| 138 |
+
// Logos drawn as backgrounds are never cleared.
|
| 139 |
+
if (im && im.kind === 'bg' && !im.logo && !im.homeLink && s && s.transp < 0.05 && s.meanL > 0.8 && im.w * im.h >= 40000 &&
|
| 140 |
+
s.stdL < (/\.svg$/i.test((im.src || '').split('?')[0]) ? 0.15 : 0.07)) return 'clear';
|
| 141 |
+
// a mostly dark graphic on transparency would vanish on a dark surface: hue-preserving inversion
|
| 142 |
+
if ((!t || t === 'keep') && s && s.transp > 0.15 && s.meanL < 0.35 && im.w * im.h < 200000 && (surfL == null || surfL < 0.4)) return 'invert';
|
| 143 |
+
if (!t || t === 'keep') return 'keep';
|
| 144 |
+
const photoLike = s ? (s.transp < 0.02 && s.colors > 0.18 && s.stdL > 0.12) : (im && im.w * im.h > 60000);
|
| 145 |
+
if (photoLike && t !== 'dim') return 'keep';
|
| 146 |
+
if (t === 'dim' && im && im.w * im.h < 20000) return 'keep';
|
| 147 |
+
// a transparent graphic that is already light reads fine on a dark surface: leave it alone
|
| 148 |
+
if ((t === 'plate' || t === 'light' || t === 'invert') && s && s.transp > 0.15 && s.meanL >= 0.55 && (s.black || 0) < 0.15 &&
|
| 149 |
+
(surfL == null || surfL < 0.4)) return 'keep';
|
| 150 |
+
// no light boxes on a dark page: a dark logo on transparency gets a hue-preserving inversion instead of a plate,
|
| 151 |
+
// an opaque one already carries its own background
|
| 152 |
+
if (t === 'plate') return (!s || s.transp > 0.15) ? 'invert' : 'keep';
|
| 153 |
+
return t;
|
| 154 |
+
}
|
| 155 |
+
function safeTreatment(t, im, surfL) {
|
| 156 |
+
const base = baseTreatment(t, im, surfL);
|
| 157 |
+
return base !== 'clear' && im && adaptable(im, im.stats, surfL) ? 'adapt:' + base : base;
|
| 158 |
+
}
|
| 159 |
+
|
| 160 |
+
// ---------- flat graphics on a white background (logos, icons, diagrams): adapt the image itself ----------
|
| 161 |
+
// The white background becomes transparent so the page's own dark surface shows through, neutral ink flips to
|
| 162 |
+
// light, and brand colours keep their hue and are lifted only as far as needed to read on dark.
|
| 163 |
+
function adaptPixels(d, w, h, info) {
|
| 164 |
+
info = info || {};
|
| 165 |
+
let n = 0, light = 0, sr = 0, sg = 0, sb = 0;
|
| 166 |
+
const ringPx = (p) => {
|
| 167 |
+
const i = p * 4;
|
| 168 |
+
if (d[i + 3] < 250) return;
|
| 169 |
+
n++;
|
| 170 |
+
const o = oklab(d[i], d[i + 1], d[i + 2]);
|
| 171 |
+
if (o[0] > 0.9 && Math.hypot(o[1], o[2]) < 0.04) { light++; sr += d[i]; sg += d[i + 1]; sb += d[i + 2]; }
|
| 172 |
+
};
|
| 173 |
+
for (let x = 0; x < w; x++) { ringPx(x); ringPx((h - 1) * w + x); }
|
| 174 |
+
for (let y = 1; y < h - 1; y++) { ringPx(y * w); ringPx(y * w + w - 1); }
|
| 175 |
+
info.ring = n ? light / n : 0;
|
| 176 |
+
if (!n || light / n < 0.5) { info.why = 'ring'; return false; } // no light background to remove (tight crops touch the edge)
|
| 177 |
+
const bg = oklab(sr / light, sg / light, sb / light);
|
| 178 |
+
const cache = new Map();
|
| 179 |
+
for (let p = 0, P = w * h; p < P; p++) {
|
| 180 |
+
const i = p * 4;
|
| 181 |
+
if (!d[i + 3]) continue;
|
| 182 |
+
const key = (d[i] << 16) | (d[i + 1] << 8) | d[i + 2];
|
| 183 |
+
let r = cache.get(key);
|
| 184 |
+
if (!r) {
|
| 185 |
+
const o = oklab(d[i], d[i + 1], d[i + 2]);
|
| 186 |
+
const dist = Math.hypot(o[0] - bg[0], o[1] - bg[1], o[2] - bg[2]);
|
| 187 |
+
const cov = Math.max(0, Math.min(1, (dist - 0.03) / 0.22)); // 0 = background, 1 = ink, between = anti-aliased edge
|
| 188 |
+
if (cov <= 0) r = [0, 0, 0, 0];
|
| 189 |
+
else {
|
| 190 |
+
// unmix the edge pixel from the background, then restyle the ink for a dark surface
|
| 191 |
+
let L = (o[0] - (1 - cov) * bg[0]) / cov, a = (o[1] - (1 - cov) * bg[1]) / cov, b = (o[2] - (1 - cov) * bg[2]) / cov;
|
| 192 |
+
L = Math.max(0, Math.min(1, L));
|
| 193 |
+
const C = Math.hypot(a, b);
|
| 194 |
+
if (C < 0.05) L = 0.94 - 0.75 * L; // black/gray ink -> light
|
| 195 |
+
else if (L < 0.62) { L = 0.62 + 0.15 * L / 0.62; a *= 0.92; b *= 0.92; } // brand colour: same hue, lifted to read on dark
|
| 196 |
+
const rgb = toRgb(L, a, b);
|
| 197 |
+
r = [rgb[0], rgb[1], rgb[2], cov];
|
| 198 |
+
}
|
| 199 |
+
cache.set(key, r);
|
| 200 |
+
}
|
| 201 |
+
d[i] = r[0]; d[i + 1] = r[1]; d[i + 2] = r[2]; d[i + 3] = Math.round(d[i + 3] * r[3]);
|
| 202 |
+
}
|
| 203 |
+
return true;
|
| 204 |
+
}
|
| 205 |
+
// blob -> data URL of the adapted image at its natural size (so layout never shifts), or null
|
| 206 |
+
async function sniffType(blob) {
|
| 207 |
+
const h = new Uint8Array(await blob.slice(0, 64).arrayBuffer());
|
| 208 |
+
const txt = String.fromCharCode(...h);
|
| 209 |
+
if (h[0] === 0x89 && txt.slice(1, 4) === 'PNG') return 'image/png';
|
| 210 |
+
if (h[0] === 0xff && h[1] === 0xd8) return 'image/jpeg';
|
| 211 |
+
if (txt.startsWith('GIF8')) return 'image/gif';
|
| 212 |
+
if (txt.startsWith('RIFF') && txt.slice(8, 12) === 'WEBP') return 'image/webp';
|
| 213 |
+
if (/<svg|<\?xml/i.test(txt)) return 'image/svg+xml';
|
| 214 |
+
return blob.type;
|
| 215 |
+
}
|
| 216 |
+
async function adaptImage(blob, info) {
|
| 217 |
+
info = info || {};
|
| 218 |
+
const type = await sniffType(blob); // servers mislabel images (e.g. video/webp)
|
| 219 |
+
if (type !== blob.type) blob = new Blob([blob], {type});
|
| 220 |
+
let src = null;
|
| 221 |
+
if (!/svg/.test(blob.type)) { try { src = await createImageBitmap(blob); } catch (e) { src = null; } }
|
| 222 |
+
if (!src) {
|
| 223 |
+
const u = URL.createObjectURL(blob.type ? blob : new Blob([blob], {type: 'image/svg+xml'}));
|
| 224 |
+
src = await new Promise((res) => { const im = new Image(); im.onload = () => res(im); im.onerror = () => res(null); im.src = u; });
|
| 225 |
+
URL.revokeObjectURL(u);
|
| 226 |
+
}
|
| 227 |
+
if (!src) { info.why = 'decode'; return null; }
|
| 228 |
+
const w = src.naturalWidth || src.width, h = src.naturalHeight || src.height;
|
| 229 |
+
if (!w || !h || w * h > 2400 * 2400) { info.why = 'size'; return null; }
|
| 230 |
+
const c = document.createElement('canvas'); c.width = w; c.height = h;
|
| 231 |
+
const ctx = c.getContext('2d', {willReadFrequently: true});
|
| 232 |
+
ctx.drawImage(src, 0, 0, w, h);
|
| 233 |
+
if (src.close) src.close();
|
| 234 |
+
const im = ctx.getImageData(0, 0, w, h);
|
| 235 |
+
if (!adaptPixels(im.data, w, h, info)) return null;
|
| 236 |
+
ctx.putImageData(im, 0, 0);
|
| 237 |
+
return c.toDataURL('image/png');
|
| 238 |
+
}
|
| 239 |
+
|
| 240 |
+
globalThis.NadirMap = NadirMap;
|
| 241 |
+
globalThis.NadirAdaptImage = adaptImage;
|
| 242 |
+
globalThis.NadirImageCSS = TREAT_CSS;
|
| 243 |
+
globalThis.NadirSafeTreatment = safeTreatment;
|
| 244 |
+
})();
|
training/build_ds.py
ADDED
|
@@ -0,0 +1,281 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Build Nadir training data: page features + per-token dark targets (Haiku designs and native dark-mode pairs)."""
|
| 3 |
+
import argparse, collections, difflib, gzip, hashlib, json, math, os, random, sys
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
import numpy as np
|
| 6 |
+
|
| 7 |
+
import safety
|
| 8 |
+
|
| 9 |
+
ROLES = ["bg", "text", "border", "shadow", "gradient", "fill", "stroke"]
|
| 10 |
+
TREAT = ["keep", "dim", "invert", "light", "plate"]
|
| 11 |
+
MAX_TOK, MAX_IMG = 256, 32
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def lin(u):
|
| 15 |
+
u = np.asarray(u, dtype=np.float64) / 255.0
|
| 16 |
+
return np.where(u <= 0.04045, u / 12.92, ((u + 0.055) / 1.055) ** 2.4)
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
def oklab(rgb):
|
| 20 |
+
r, g, b = lin(rgb[..., 0]), lin(rgb[..., 1]), lin(rgb[..., 2])
|
| 21 |
+
l = np.cbrt(0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b)
|
| 22 |
+
m = np.cbrt(0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b)
|
| 23 |
+
s = np.cbrt(0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b)
|
| 24 |
+
return np.stack([0.2104542553 * l + 0.7936177850 * m - 0.0040720468 * s,
|
| 25 |
+
1.9779984951 * l - 2.4285922050 * m + 0.4505937099 * s,
|
| 26 |
+
0.0259040371 * l + 0.7827717662 * m - 0.8086757660 * s], -1)
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def rel_lum(rgb):
|
| 30 |
+
c = lin(rgb)
|
| 31 |
+
return 0.2126 * c[..., 0] + 0.7152 * c[..., 1] + 0.0722 * c[..., 2]
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def hex2rgba(h):
|
| 35 |
+
h = h.lstrip("#")
|
| 36 |
+
return [int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16), int(h[6:8], 16) / 255 if len(h) >= 8 else 1.0]
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def contrast(fg, bg):
|
| 40 |
+
l1, l2 = rel_lum(np.array(fg[:3])), rel_lum(np.array(bg[:3]))
|
| 41 |
+
return (max(l1, l2) + 0.05) / (min(l1, l2) + 0.05)
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def over(top, bottom):
|
| 45 |
+
a = top[3]
|
| 46 |
+
return [top[i] * a + bottom[i] * (1 - a) for i in range(3)] + [1.0]
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def page_arrays(L):
|
| 50 |
+
toks = L["tokens"][:MAX_TOK]
|
| 51 |
+
n = len(toks)
|
| 52 |
+
F = np.array([t["f"] for t in toks], dtype=np.float32)
|
| 53 |
+
role = np.array([t["role"] for t in toks], dtype=np.int64)
|
| 54 |
+
rgba = np.array([t["rgba"] for t in toks], dtype=np.float32)
|
| 55 |
+
rel = np.array([r for r in L["rel"] if r[1] < n and r[2] < n], dtype=np.float32).reshape(-1, 4)
|
| 56 |
+
imgs = L["images"][:MAX_IMG]
|
| 57 |
+
IF = np.array([im["f"] for im in imgs], dtype=np.float32).reshape(len(imgs), 39)
|
| 58 |
+
thumbs = np.zeros((len(imgs), 32, 32, 4), dtype=np.uint8)
|
| 59 |
+
for j, im in enumerate(imgs):
|
| 60 |
+
if im.get("thumb") and len(im["thumb"]) == 4096:
|
| 61 |
+
thumbs[j] = np.array(im["thumb"], dtype=np.uint8).reshape(32, 32, 4)
|
| 62 |
+
img_under = np.array([im.get("underTok", -1) if (im.get("underTok", -1) < n) else -1 for im in imgs], dtype=np.int64)
|
| 63 |
+
w = []
|
| 64 |
+
for t in toks:
|
| 65 |
+
d = t.get("d") or {}
|
| 66 |
+
w.append(1.0 + math.log1p(d.get("area", 0) / 2000) + math.log1p(d.get("chars", 0) / 20) + 0.5 * math.log1p(d.get("nVis", 0)) + 0.3 * math.log1p(d.get("rules", 0)))
|
| 67 |
+
keep = [i for i, t in enumerate(toks) if safety.is_media_keep(t)]
|
| 68 |
+
page = np.array([L["vw"] / 1280, min(8, L["docH"] / L["vh"]) / 8, n / MAX_TOK, len(imgs) / MAX_IMG, L["darkness"]["bgL"], L["darkness"]["textL"],
|
| 69 |
+
1.0 if L.get("canvasDark") else 0.0, min(1, L["stats"].get("tokensDropped", 0) / 200)], dtype=np.float32)
|
| 70 |
+
F, IF, page = (np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) for x in (F, IF, page))
|
| 71 |
+
return dict(F=F, role=role, rgba=rgba, rel=rel, IF=IF, thumbs=thumbs, img_under=img_under, w=np.array(w, dtype=np.float32), page=page, keep=np.array(keep, dtype=np.int64))
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def haiku_targets(L, h, arrs, qc):
|
| 75 |
+
n = len(arrs["role"])
|
| 76 |
+
tgt = np.full((n, 4), np.nan, dtype=np.float32)
|
| 77 |
+
for k, v in h["colors"].items():
|
| 78 |
+
i = int(k)
|
| 79 |
+
if i < n:
|
| 80 |
+
tgt[i] = hex2rgba(v)
|
| 81 |
+
imgt = np.full((len(arrs["IF"]),), -1, dtype=np.int64)
|
| 82 |
+
for k, v in (h.get("images") or {}).items():
|
| 83 |
+
j = int(k)
|
| 84 |
+
if j < len(imgt) and v in TREAT:
|
| 85 |
+
imgt[j] = TREAT.index(v)
|
| 86 |
+
# ---- QC ----
|
| 87 |
+
asked = min(n, h.get("n_tok", n))
|
| 88 |
+
got = np.isfinite(tgt[:asked, 0]).sum()
|
| 89 |
+
if asked == 0 or got / asked < 0.9:
|
| 90 |
+
qc["coverage"] += 1
|
| 91 |
+
return None
|
| 92 |
+
root = [i for i, t in enumerate(L["tokens"][:n]) if (t.get("d") or {}).get("root")]
|
| 93 |
+
if root and np.isfinite(tgt[root[0], 0]):
|
| 94 |
+
rl = oklab(np.array(tgt[root[0], :3]))[0]
|
| 95 |
+
if rl > 0.4:
|
| 96 |
+
qc["light_root"] += 1
|
| 97 |
+
return None
|
| 98 |
+
# text contrast vs mapped dominant surface
|
| 99 |
+
tb = collections.defaultdict(list)
|
| 100 |
+
for typ, a, b, w in arrs["rel"]:
|
| 101 |
+
if int(typ) == 0:
|
| 102 |
+
tb[int(a)].append((w, int(b)))
|
| 103 |
+
bad = tot = 0.0
|
| 104 |
+
for i, t in enumerate(L["tokens"][:n]):
|
| 105 |
+
if t["role"] != 1 or not np.isfinite(tgt[i, 0]):
|
| 106 |
+
continue
|
| 107 |
+
chars = (t.get("d") or {}).get("chars", 0)
|
| 108 |
+
if not chars or not tb.get(i):
|
| 109 |
+
continue
|
| 110 |
+
j = max(tb[i])[1]
|
| 111 |
+
if not np.isfinite(tgt[j, 0]):
|
| 112 |
+
continue
|
| 113 |
+
bgc = tgt[j]
|
| 114 |
+
base = over(list(bgc), [18, 18, 18, 1]) if bgc[3] < 1 else list(bgc)
|
| 115 |
+
fg = over(list(tgt[i]), base)
|
| 116 |
+
c = contrast(fg, base)
|
| 117 |
+
tot += chars
|
| 118 |
+
if c < 3.0:
|
| 119 |
+
bad += chars
|
| 120 |
+
if tot and bad / tot > 0.15:
|
| 121 |
+
qc["contrast"] += 1
|
| 122 |
+
return None
|
| 123 |
+
qc["ok"] += 1
|
| 124 |
+
return tgt, imgt
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
def repair_contrast(arrs, tgt):
|
| 128 |
+
"""Apply the runtime rules (hue guard, invisible-stays-invisible, contrast) to a Haiku design, so the model
|
| 129 |
+
learns what the runtime would show."""
|
| 130 |
+
n = len(arrs["role"])
|
| 131 |
+
toks = [{"role": int(arrs["role"][i]), "rgba": arrs["rgba"][i].tolist(), "f": arrs["F"][i]} for i in range(n)]
|
| 132 |
+
out, keep = [], set(int(k) for k in arrs["keep"])
|
| 133 |
+
for i in range(n):
|
| 134 |
+
if not np.isfinite(tgt[i, 0]):
|
| 135 |
+
keep.add(i)
|
| 136 |
+
out.append({"rgb": [0, 0, 0], "a": 0.0, "L": 0.0})
|
| 137 |
+
continue
|
| 138 |
+
rgb = [int(round(float(x))) for x in tgt[i, :3]]
|
| 139 |
+
out.append({"rgb": rgb, "a": float(tgt[i, 3]), "L": safety.oklab(*rgb)[0]})
|
| 140 |
+
before = [list(o["rgb"]) for o in out]
|
| 141 |
+
for i in range(n):
|
| 142 |
+
if i not in keep:
|
| 143 |
+
safety.hue_guard(toks[i], out[i])
|
| 144 |
+
safety.keep_invisible(toks, out, arrs["rel"].tolist(), keep)
|
| 145 |
+
safety.enforce_contrast(toks, out, arrs["rel"].tolist(), keep)
|
| 146 |
+
changed = 0
|
| 147 |
+
for i in range(n):
|
| 148 |
+
if out[i]["rgb"] != before[i]:
|
| 149 |
+
tgt[i, :3] = out[i]["rgb"]
|
| 150 |
+
changed += 1
|
| 151 |
+
return changed
|
| 152 |
+
|
| 153 |
+
|
| 154 |
+
def native_targets(L, D, arrs, qc):
|
| 155 |
+
le, de = L.get("ec") or [], D.get("ec") or []
|
| 156 |
+
if not le or not de:
|
| 157 |
+
qc["no_ec"] += 1
|
| 158 |
+
return None
|
| 159 |
+
if len(le) == len(de) and all(a[0] == b[0] for a, b in zip(le, de)):
|
| 160 |
+
pairs = list(zip(range(len(le)), range(len(de))))
|
| 161 |
+
else:
|
| 162 |
+
sm = difflib.SequenceMatcher(None, [r[0] for r in le], [r[0] for r in de], autojunk=False)
|
| 163 |
+
pairs = []
|
| 164 |
+
for blk in sm.get_matching_blocks():
|
| 165 |
+
pairs += [(blk.a + k, blk.b + k) for k in range(blk.size)]
|
| 166 |
+
if len(pairs) < 0.8 * len(le):
|
| 167 |
+
qc["align"] += 1
|
| 168 |
+
return None
|
| 169 |
+
votes = collections.defaultdict(collections.Counter)
|
| 170 |
+
for ia, ib in pairs:
|
| 171 |
+
a, b = le[ia], de[ib]
|
| 172 |
+
w = (a[1] or 0) + 1
|
| 173 |
+
if a[2]:
|
| 174 |
+
# bg: effective dark surface at this element (opaque light bg) or raw dark bg (translucent)
|
| 175 |
+
alpha = int(a[2][6:8], 16)
|
| 176 |
+
tgt = (b[8] if alpha >= 250 else b[2]) or b[8]
|
| 177 |
+
votes[(0, a[2])][tgt] += w
|
| 178 |
+
for j, role in ((3, 1), (4, 2), (5, 5), (6, 6)):
|
| 179 |
+
if a[j]:
|
| 180 |
+
votes[(role, a[j])][b[j] or "00000000"] += w
|
| 181 |
+
if a[7] and b[7]:
|
| 182 |
+
la, lb = a[7].split(","), b[7].split(",")
|
| 183 |
+
if len(la) == len(lb):
|
| 184 |
+
for x, y in zip(la, lb):
|
| 185 |
+
votes[(4, x)][y] += w
|
| 186 |
+
n = len(arrs["role"])
|
| 187 |
+
tgt = np.full((n, 4), np.nan, dtype=np.float32)
|
| 188 |
+
hit = 0
|
| 189 |
+
for i, t in enumerate(L["tokens"][:n]):
|
| 190 |
+
key = t["k"].split(":")[1]
|
| 191 |
+
c = votes.get((t["role"], key))
|
| 192 |
+
if not c:
|
| 193 |
+
continue
|
| 194 |
+
best, bw = c.most_common(1)[0]
|
| 195 |
+
if bw / sum(c.values()) < 0.5:
|
| 196 |
+
continue
|
| 197 |
+
tgt[i] = hex2rgba(best)
|
| 198 |
+
hit += 1
|
| 199 |
+
if hit < 4:
|
| 200 |
+
qc["few"] += 1
|
| 201 |
+
return None
|
| 202 |
+
qc["ok"] += 1
|
| 203 |
+
return tgt, np.full((len(arrs["IF"]),), -1, dtype=np.int64)
|
| 204 |
+
|
| 205 |
+
|
| 206 |
+
def main():
|
| 207 |
+
ap = argparse.ArgumentParser()
|
| 208 |
+
ap.add_argument("--crawl", required=True)
|
| 209 |
+
ap.add_argument("--haiku", nargs="*", default=[])
|
| 210 |
+
ap.add_argument("--out", required=True)
|
| 211 |
+
ap.add_argument("--val-frac", type=float, default=0.05)
|
| 212 |
+
args = ap.parse_args()
|
| 213 |
+
crawl = Path(args.crawl)
|
| 214 |
+
index = {}
|
| 215 |
+
for line in (l for ip in sorted(crawl.glob("index*.jsonl")) for l in open(ip)):
|
| 216 |
+
try:
|
| 217 |
+
r = json.loads(line)
|
| 218 |
+
index[r["id"]] = r
|
| 219 |
+
except Exception:
|
| 220 |
+
pass
|
| 221 |
+
samples = []
|
| 222 |
+
qc_h = collections.Counter(); qc_n = collections.Counter()
|
| 223 |
+
# Haiku-designed pages (source 0)
|
| 224 |
+
for hd in args.haiku:
|
| 225 |
+
for fp in Path(hd).glob("*.json"):
|
| 226 |
+
if fp.name in ("ledger.jsonl",):
|
| 227 |
+
continue
|
| 228 |
+
try:
|
| 229 |
+
h = json.loads(fp.read_text())
|
| 230 |
+
except Exception:
|
| 231 |
+
continue
|
| 232 |
+
if h.get("refusal") or "colors" not in h:
|
| 233 |
+
continue
|
| 234 |
+
pid = h["id"]
|
| 235 |
+
try:
|
| 236 |
+
L = json.load(gzip.open(crawl / "pages" / pid / "light.json.gz"))
|
| 237 |
+
except Exception:
|
| 238 |
+
continue
|
| 239 |
+
arrs = page_arrays(L)
|
| 240 |
+
r = haiku_targets(L, h, arrs, qc_h)
|
| 241 |
+
if r is None:
|
| 242 |
+
continue
|
| 243 |
+
arrs["tgt"], arrs["imgt"] = r
|
| 244 |
+
if repair_contrast(arrs, arrs["tgt"]):
|
| 245 |
+
qc_h["repaired"] += 1
|
| 246 |
+
arrs["src"] = 0
|
| 247 |
+
arrs["id"] = pid
|
| 248 |
+
arrs["host"] = index.get(pid, {}).get("url", L["url"]).split("/")[2]
|
| 249 |
+
samples.append(arrs)
|
| 250 |
+
# native dark-mode pairs (source 1)
|
| 251 |
+
for pid, r in index.items():
|
| 252 |
+
if r.get("kind") not in ("native_media", "native_class"):
|
| 253 |
+
continue
|
| 254 |
+
try:
|
| 255 |
+
L = json.load(gzip.open(crawl / "pages" / pid / "light.json.gz"))
|
| 256 |
+
D = json.load(gzip.open(crawl / "pages" / pid / "dark.json.gz"))
|
| 257 |
+
except Exception:
|
| 258 |
+
continue
|
| 259 |
+
arrs = page_arrays(L)
|
| 260 |
+
res = native_targets(L, D, arrs, qc_n)
|
| 261 |
+
if res is None:
|
| 262 |
+
continue
|
| 263 |
+
arrs["tgt"], arrs["imgt"] = res
|
| 264 |
+
arrs["src"] = 1
|
| 265 |
+
arrs["id"] = pid
|
| 266 |
+
arrs["host"] = r["url"].split("/")[2]
|
| 267 |
+
samples.append(arrs)
|
| 268 |
+
print("haiku QC", dict(qc_h), "native QC", dict(qc_n), "samples", len(samples), flush=True)
|
| 269 |
+
hosts = sorted(set(s["host"].replace("www.", "") for s in samples))
|
| 270 |
+
# stable split: a host's side never changes as the crawl grows
|
| 271 |
+
val_hosts = {h for h in hosts if int(hashlib.md5(h.encode()).hexdigest()[:8], 16) % 1000 < args.val_frac * 1000}
|
| 272 |
+
train = [s for s in samples if s["host"].replace("www.", "") not in val_hosts]
|
| 273 |
+
val = [s for s in samples if s["host"].replace("www.", "") in val_hosts]
|
| 274 |
+
import pickle
|
| 275 |
+
with open(args.out, "wb") as f:
|
| 276 |
+
pickle.dump({"train": train, "val": val, "roles": ROLES, "treat": TREAT}, f, protocol=4)
|
| 277 |
+
print("train", len(train), "val", len(val), "haiku", sum(s["src"] == 0 for s in samples), "native", sum(s["src"] == 1 for s in samples))
|
| 278 |
+
|
| 279 |
+
|
| 280 |
+
if __name__ == "__main__":
|
| 281 |
+
main()
|
training/export.py
ADDED
|
@@ -0,0 +1,211 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Export NadirNet to ONNX (fp32/fp16) and build calibrated 8-bit and 4-bit variants; measure each against fp32."""
|
| 3 |
+
import argparse, json, os, pickle, time
|
| 4 |
+
import numpy as np
|
| 5 |
+
import torch
|
| 6 |
+
import torch.nn as nn
|
| 7 |
+
|
| 8 |
+
from model import NadirNet
|
| 9 |
+
from train import collate, prep, evaluate
|
| 10 |
+
|
| 11 |
+
NAMES_IN = ["tok_f", "tok_role", "tok_lab", "img_f", "img_thumb", "page_f", "src", "rel", "mask"]
|
| 12 |
+
NAMES_OUT = ["lab", "alpha_logit", "img_logits"]
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
class Wrapped(nn.Module):
|
| 16 |
+
def __init__(self, m):
|
| 17 |
+
super().__init__()
|
| 18 |
+
self.m = m
|
| 19 |
+
|
| 20 |
+
def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask):
|
| 21 |
+
return self.m(tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask)
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def feeds_for(s, pad=None):
|
| 25 |
+
b = collate([s], "cpu")
|
| 26 |
+
if pad: # fixed shapes for calibration (quantizers stack activations across samples)
|
| 27 |
+
N0, M0 = b["tok_f"].shape[1], b["img_f"].shape[1]
|
| 28 |
+
N, M = pad
|
| 29 |
+
if N0 <= N and M0 <= M:
|
| 30 |
+
T0, T = 1 + N0 + M0, 1 + N + M
|
| 31 |
+
def padto(x, dims, val=0):
|
| 32 |
+
shape = list(x.shape)
|
| 33 |
+
for d, sz in dims.items():
|
| 34 |
+
shape[d] = sz
|
| 35 |
+
y = torch.full(shape, val, dtype=x.dtype)
|
| 36 |
+
y[tuple(slice(0, x.shape[i]) for i in range(x.ndim))] = x
|
| 37 |
+
return y
|
| 38 |
+
idx = list(range(1 + N0)) + list(range(1 + N, 1 + N + M0))
|
| 39 |
+
rel = torch.zeros(1, T, T, b["rel"].shape[-1]); mask = torch.zeros(1, T, dtype=torch.bool)
|
| 40 |
+
ii = torch.tensor(idx)
|
| 41 |
+
rel[0][ii[:, None], ii[None, :]] = b["rel"][0]
|
| 42 |
+
mask[0, ii] = b["mask"][0]
|
| 43 |
+
b = dict(b, tok_f=padto(b["tok_f"], {1: N}), role=padto(b["role"], {1: N}), lab=padto(b["lab"], {1: N}),
|
| 44 |
+
img_f=padto(b["img_f"], {1: M}), thumbs=padto(b["thumbs"], {1: M}), rel=rel, mask=mask)
|
| 45 |
+
M = b["img_f"].shape[1]
|
| 46 |
+
if M == 0: # the runtime always sends at least one (masked) image slot
|
| 47 |
+
T = b["mask"].shape[1]
|
| 48 |
+
b["img_f"] = torch.zeros(1, 1, 39); b["thumbs"] = torch.zeros(1, 1, 4, 32, 32)
|
| 49 |
+
rel = torch.zeros(1, T + 1, T + 1, b["rel"].shape[-1]); rel[:, :T, :T] = b["rel"]; b["rel"] = rel
|
| 50 |
+
b["mask"] = torch.cat([b["mask"], torch.zeros(1, 1, dtype=torch.bool)], 1)
|
| 51 |
+
return {"tok_f": b["tok_f"].numpy(), "tok_role": b["role"].numpy(), "tok_lab": b["lab"].numpy(), "img_f": b["img_f"].numpy(),
|
| 52 |
+
"img_thumb": b["thumbs"].numpy(), "page_f": b["page"].numpy(), "src": b["src"].numpy(), "rel": b["rel"].numpy(), "mask": b["mask"].numpy()}
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def fidelity(sess, ref_sess, samples):
|
| 56 |
+
"""Mean OKLab distance between this model's colours and fp32's, plus image-decision agreement."""
|
| 57 |
+
de, agree, n_img, worst = [], 0, 0, 0.0
|
| 58 |
+
for s in samples:
|
| 59 |
+
f = feeds_for(s)
|
| 60 |
+
a = sess.run(None, f); r = ref_sess.run(None, f)
|
| 61 |
+
n = s["F"].shape[0]
|
| 62 |
+
d = np.linalg.norm(a[0][0, :n] - r[0][0, :n], axis=-1)
|
| 63 |
+
w = s["w"][:n]
|
| 64 |
+
de.append(float((d * w).sum() / w.sum())); worst = max(worst, float(d.max()))
|
| 65 |
+
m = len(s["IF"])
|
| 66 |
+
if m:
|
| 67 |
+
agree += int((a[2][0, :m].argmax(-1) == r[2][0, :m].argmax(-1)).sum()); n_img += m
|
| 68 |
+
return {"dE_vs_fp32": float(np.mean(de)), "dE_max": worst, "img_agree": agree / max(1, n_img)}
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
def bench(sess, samples, reps=3):
|
| 72 |
+
ts = []
|
| 73 |
+
for s in samples[:20]:
|
| 74 |
+
f = feeds_for(s)
|
| 75 |
+
sess.run(None, f)
|
| 76 |
+
t = time.time()
|
| 77 |
+
for _ in range(reps):
|
| 78 |
+
sess.run(None, f)
|
| 79 |
+
ts.append((time.time() - t) / reps * 1000)
|
| 80 |
+
return float(np.median(ts))
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
def main():
|
| 84 |
+
ap = argparse.ArgumentParser()
|
| 85 |
+
ap.add_argument("--ckpt", required=True)
|
| 86 |
+
ap.add_argument("--data", required=True)
|
| 87 |
+
ap.add_argument("--out", required=True)
|
| 88 |
+
ap.add_argument("--calib", type=int, default=256)
|
| 89 |
+
ap.add_argument("--only-float", action="store_true", help="fp32/fp16 only (calibrated 8/4-bit come from quantize.py)")
|
| 90 |
+
args = ap.parse_args()
|
| 91 |
+
os.makedirs(args.out, exist_ok=True)
|
| 92 |
+
import onnx, onnxruntime as ort
|
| 93 |
+
ck = torch.load(args.ckpt, map_location="cpu")
|
| 94 |
+
model = NadirNet(**{k: v for k, v in ck["cfg"].items()}); model.load_state_dict(ck["model"]); model.eval()
|
| 95 |
+
data = pickle.load(open(args.data, "rb"))
|
| 96 |
+
val = [prep(s) for s in data["val"]]
|
| 97 |
+
rng = np.random.default_rng(0)
|
| 98 |
+
calib = [prep(s) for s in rng.choice(np.array(data["train"], dtype=object), size=min(args.calib, len(data["train"])), replace=False)]
|
| 99 |
+
s0 = max(val, key=lambda s: len(s["role"]))
|
| 100 |
+
ex = feeds_for(s0)
|
| 101 |
+
par_samples = val[:8]
|
| 102 |
+
with torch.no_grad():
|
| 103 |
+
ref_out = [model(*[torch.from_numpy(feeds_for(s)[k]) for k in NAMES_IN])[0].numpy() for s in par_samples]
|
| 104 |
+
fp32 = os.path.join(args.out, "nadir-fp32.onnx")
|
| 105 |
+
torch.onnx.export(Wrapped(model), tuple(torch.from_numpy(ex[k]) for k in NAMES_IN), fp32, input_names=NAMES_IN, output_names=NAMES_OUT,
|
| 106 |
+
dynamic_axes={"tok_f": {1: "N"}, "tok_role": {1: "N"}, "tok_lab": {1: "N"}, "img_f": {1: "M"}, "img_thumb": {1: "M"},
|
| 107 |
+
"rel": {1: "T", 2: "T"}, "mask": {1: "T"}, "lab": {1: "N"}, "alpha_logit": {1: "N"}, "img_logits": {1: "M"}},
|
| 108 |
+
opset_version=17, do_constant_folding=True, dynamo=False)
|
| 109 |
+
onnx.checker.check_model(fp32)
|
| 110 |
+
so = ort.SessionOptions(); so.intra_op_num_threads = 4
|
| 111 |
+
ref = ort.InferenceSession(fp32, so, providers=["CPUExecutionProvider"])
|
| 112 |
+
# parity torch vs onnx (reference computed before export)
|
| 113 |
+
parity = max(float(np.abs(r - ref.run(None, feeds_for(s))[0]).max()) for r, s in zip(ref_out, par_samples))
|
| 114 |
+
report = {"params": sum(p.numel() for p in model.parameters()), "parity_max_abs": parity, "variants": {}}
|
| 115 |
+
print("parity", parity, flush=True)
|
| 116 |
+
|
| 117 |
+
# fp16 (weights + compute in fp16, IO kept fp32)
|
| 118 |
+
# fp16: export a half-precision copy with fp32 inputs/outputs (casts inside the graph)
|
| 119 |
+
have16 = True
|
| 120 |
+
try:
|
| 121 |
+
class Half(nn.Module):
|
| 122 |
+
def __init__(self, m):
|
| 123 |
+
super().__init__(); self.m = m.half()
|
| 124 |
+
def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask):
|
| 125 |
+
o = self.m(tok_f.half(), tok_role, tok_lab.half(), img_f.half(), img_thumb.half(), page_f.half(), src, rel.half(), mask)
|
| 126 |
+
return tuple(x.float() for x in o)
|
| 127 |
+
import copy
|
| 128 |
+
torch.onnx.export(Half(copy.deepcopy(model)), tuple(torch.from_numpy(ex[k]) for k in NAMES_IN), os.path.join(args.out, "nadir-fp16.onnx"),
|
| 129 |
+
input_names=NAMES_IN, output_names=NAMES_OUT,
|
| 130 |
+
dynamic_axes={"tok_f": {1: "N"}, "tok_role": {1: "N"}, "tok_lab": {1: "N"}, "img_f": {1: "M"}, "img_thumb": {1: "M"},
|
| 131 |
+
"rel": {1: "T", 2: "T"}, "mask": {1: "T"}, "lab": {1: "N"}, "alpha_logit": {1: "N"}, "img_logits": {1: "M"}},
|
| 132 |
+
opset_version=17, do_constant_folding=True, dynamo=False)
|
| 133 |
+
ort.InferenceSession(os.path.join(args.out, "nadir-fp16.onnx"), providers=["CPUExecutionProvider"])
|
| 134 |
+
except Exception as e:
|
| 135 |
+
print("fp16 failed:", str(e)[:300], flush=True); have16 = False
|
| 136 |
+
|
| 137 |
+
if args.only_float:
|
| 138 |
+
variants = {"fp32": fp32}
|
| 139 |
+
if have16: variants["fp16"] = os.path.join(args.out, "nadir-fp16.onnx")
|
| 140 |
+
for name, path in variants.items():
|
| 141 |
+
sess = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
|
| 142 |
+
r = fidelity(sess, ref, val[:300]); r["size_mb"] = round(os.path.getsize(path) / 1e6, 2); r["cpu_ms_median"] = round(bench(sess, val[:300]), 2)
|
| 143 |
+
report["variants"][name] = r
|
| 144 |
+
print(name, r, flush=True)
|
| 145 |
+
json.dump(report, open(os.path.join(args.out, "export_report.json"), "w"), indent=1)
|
| 146 |
+
return
|
| 147 |
+
from onnxruntime.quantization import quantize_dynamic, quantize_static, QuantType, QuantFormat, CalibrationDataReader, CalibrationMethod
|
| 148 |
+
from onnxruntime.quantization.shape_inference import quant_pre_process
|
| 149 |
+
pre = os.path.join(args.out, "_pre.onnx")
|
| 150 |
+
quant_pre_process(fp32, pre, skip_symbolic_shape=True)
|
| 151 |
+
|
| 152 |
+
class Reader(CalibrationDataReader):
|
| 153 |
+
def __init__(self, ss):
|
| 154 |
+
self.it = iter([feeds_for(s, pad=(256, 32)) for s in ss])
|
| 155 |
+
def get_next(self):
|
| 156 |
+
return next(self.it, None)
|
| 157 |
+
|
| 158 |
+
# 8-bit A: per-channel int8 weights, dynamic activations
|
| 159 |
+
q8d = os.path.join(args.out, "nadir-int8-dynamic.onnx")
|
| 160 |
+
quantize_dynamic(pre, q8d, weight_type=QuantType.QInt8, per_channel=True, op_types_to_quantize=["MatMul", "Gemm"])
|
| 161 |
+
# 8-bit B: static QDQ, activations calibrated on real pages (percentile)
|
| 162 |
+
q8s = os.path.join(args.out, "nadir-int8-static.onnx")
|
| 163 |
+
try:
|
| 164 |
+
quantize_static(pre, q8s, Reader(calib[:128]), quant_format=QuantFormat.QDQ, per_channel=True, activation_type=QuantType.QInt8,
|
| 165 |
+
weight_type=QuantType.QInt8, op_types_to_quantize=["MatMul", "Gemm"], calibrate_method=CalibrationMethod.Percentile,
|
| 166 |
+
extra_options={"CalibPercentile": 99.99, "ActivationSymmetric": True, "WeightSymmetric": True})
|
| 167 |
+
except Exception as e:
|
| 168 |
+
print("static int8 failed:", e, flush=True); q8s = None
|
| 169 |
+
# 4-bit: block-wise weights (MatMulNBits), GPTQ calibrated on real pages, RTN fallback
|
| 170 |
+
q4 = os.path.join(args.out, "nadir-int4.onnx")
|
| 171 |
+
from onnxruntime.quantization import matmul_nbits_quantizer as mnq
|
| 172 |
+
done4 = False
|
| 173 |
+
# GPTQ needs reduction dims that are multiples of the block size; leave tiny matrices (relation bias, 11 wide) in fp32
|
| 174 |
+
pm = onnx.load(pre)
|
| 175 |
+
inits = {i.name: i for i in pm.graph.initializer}
|
| 176 |
+
skip = [n.name for n in pm.graph.node if n.op_type == "MatMul" and len(n.input) > 1 and n.input[1] in inits and inits[n.input[1]].dims[0] % 32 != 0]
|
| 177 |
+
report["int4_fp32_nodes"] = len(skip)
|
| 178 |
+
for algo in ("gptq", "hqq", "rtn"):
|
| 179 |
+
try:
|
| 180 |
+
if algo == "gptq":
|
| 181 |
+
cfg = mnq.GPTQWeightOnlyQuantConfig(calibration_data_reader=Reader(calib), percdamp=0.01, block_size=32, actorder=True, mse=True)
|
| 182 |
+
elif algo == "hqq":
|
| 183 |
+
cfg = mnq.HQQWeightOnlyQuantConfig(block_size=32, bits=4)
|
| 184 |
+
else:
|
| 185 |
+
cfg = mnq.RTNWeightOnlyQuantConfig()
|
| 186 |
+
qz = mnq.MatMulNBitsQuantizer(onnx.load(pre), block_size=32, is_symmetric=False, accuracy_level=4, algo_config=cfg, nodes_to_exclude=skip)
|
| 187 |
+
qz.process()
|
| 188 |
+
qz.model.save_model_to_file(q4, use_external_data_format=False)
|
| 189 |
+
report["int4_algo"] = algo
|
| 190 |
+
done4 = True
|
| 191 |
+
break
|
| 192 |
+
except Exception as e:
|
| 193 |
+
print(f"int4 {algo} failed: {str(e)[:300]}", flush=True)
|
| 194 |
+
variants = {"fp32": fp32, "int8-dynamic": q8d}
|
| 195 |
+
if have16: variants["fp16"] = os.path.join(args.out, "nadir-fp16.onnx")
|
| 196 |
+
if q8s: variants["int8-static"] = q8s
|
| 197 |
+
if done4: variants["int4"] = q4
|
| 198 |
+
evs = val[:300]
|
| 199 |
+
for name, path in variants.items():
|
| 200 |
+
sess = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
|
| 201 |
+
r = fidelity(sess, ref, evs)
|
| 202 |
+
r["size_mb"] = round(os.path.getsize(path) / 1e6, 2)
|
| 203 |
+
r["cpu_ms_median"] = round(bench(sess, evs), 2)
|
| 204 |
+
report["variants"][name] = r
|
| 205 |
+
print(name, r, flush=True)
|
| 206 |
+
os.remove(pre)
|
| 207 |
+
json.dump(report, open(os.path.join(args.out, "export_report.json"), "w"), indent=1)
|
| 208 |
+
|
| 209 |
+
|
| 210 |
+
if __name__ == "__main__":
|
| 211 |
+
main()
|
training/haiku_gen.py
ADDED
|
@@ -0,0 +1,621 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Nadir teacher data: Claude Haiku 5.5 designs a dark theme for each crawled light page.
|
| 3 |
+
|
| 4 |
+
Runs on Vega's host (the API key vault proxy is only reachable there). Writes one JSON per page to
|
| 5 |
+
OUT/<id>.json and keeps a running cost ledger; stops at --budget dollars.
|
| 6 |
+
"""
|
| 7 |
+
import argparse, asyncio, base64, gzip, io, json, math, os, random, re, sys, time
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
|
| 10 |
+
import anthropic
|
| 11 |
+
from PIL import Image, ImageDraw, ImageFont
|
| 12 |
+
|
| 13 |
+
MODEL = "claude-haiku-5-5"
|
| 14 |
+
PRICE_IN, PRICE_OUT = 0.10 / 1e6, 0.50 / 1e6 # prompts <= 100K tokens
|
| 15 |
+
PRICE_CACHE_READ, PRICE_CACHE_WRITE = 0.01 / 1e6, 0.125 / 1e6
|
| 16 |
+
ROLES = ["bg", "text", "border", "shadow", "gradient", "fill", "stroke"]
|
| 17 |
+
TREATMENTS = ["keep", "dim", "invert", "light", "plate"]
|
| 18 |
+
|
| 19 |
+
SYSTEM = """You are a senior product designer. A company has a light-themed website and has asked you to design its dark theme: the version their own design team would ship, not a mechanical inversion. You are given a screenshot of the light page, a contact sheet of its images, and the complete inventory of colors the page uses, each with its role and where it appears. You return the dark-theme replacement for every color and a treatment for every image.
|
| 20 |
+
|
| 21 |
+
How to design it:
|
| 22 |
+
- Surfaces. Choose one base dark surface for the page canvas. Near-black neutrals in the #0d0f12-#1c1e22 range work for most sites; a faint tint of the brand hue is welcome when it suits the brand. Avoid pure #000000 unless the brand is starkly monochrome. Keep the surface hierarchy of the light design but express it in dark-UI terms: anything that sits ON another surface (cards, panels, popovers, inputs, sticky headers, table stripes) is slightly LIGHTER than what is beneath it, in small steps (about 3-7 points of lightness). Alternating light-gray section bands become subtly different dark tones. Never let a nested surface end up darker than its parent unless it is a deliberate inset (code wells, input fields may be a touch darker).
|
| 23 |
+
- Colors that are already dark in the light design (dark heroes, footers, banners, code blocks, dark buttons) stay dark. Harmonize them with the new palette; never flip them to light.
|
| 24 |
+
- Text. Primary text is light but not pure white for body copy (around #e4e6ea-#f0f1f3). Secondary and muted text keep their hierarchy as progressively dimmer grays, and every text color must keep at least 4.5:1 contrast on the surface it actually sits on (3:1 for large or bold headings). Text that sits on a colored or dark-in-light-mode surface keeps working on that surface's new color.
|
| 25 |
+
- Brand and accent colors keep their hue so the brand stays recognizable. When an accent is used as text, a link or an icon on a dark surface, raise its lightness and slightly reduce chroma until it reads comfortably. When an accent is a filled button or badge background, keep it close to the original (deepen slightly if very light) and keep its label readable.
|
| 26 |
+
- Status colors (error, warning, success, info) keep their meaning and become readable on dark.
|
| 27 |
+
- Borders and dividers become subtle lines slightly lighter than the surface they sit on. Focus rings stay clearly visible.
|
| 28 |
+
- Translucent colors (8-digit hex): keep a similar alpha unless it no longer works. Black-tinted overlays used to darken or separate things in light mode usually become white-tinted at a similar or slightly lower alpha (rgba(0,0,0,.05) hover -> #ffffff0f). Shadows and scrims stay dark; shadows may get a little stronger.
|
| 29 |
+
- Hover, focus and active colors (marked as rule-only / state) keep their relationship to the resting state: if hover was darker in light mode, make it visibly lighter (or more emphasized) in dark mode.
|
| 30 |
+
- Overlays drawn on top of videos or images (video player controls, captions on photos, scrims, gradient fades over media) and the black backgrounds of media players are part of the media: return them unchanged.
|
| 31 |
+
- Syntax highlighting and data-visualization colors stay distinguishable from each other and readable on dark.
|
| 32 |
+
- The same input color in the same role always gets the same output. Every listed id must appear in your answer exactly once.
|
| 33 |
+
|
| 34 |
+
Images (choose one treatment each; the contact sheet shows transparency as a checkerboard):
|
| 35 |
+
- keep: photos, video thumbnails, illustrations, screenshots of colorful UI, colorful logos or icons that remain clear on a dark background, avatars, product shots on their own backgrounds.
|
| 36 |
+
- dim: large, very bright images (mostly white backgrounds, bright illustrations) that would glare on a dark page; they are shown slightly dimmed.
|
| 37 |
+
- invert: dark-on-light or dark-on-transparent line art, diagrams, charts, black text images, and mostly-monochrome logos whose few colored parts still look right after a hue-preserving inversion.
|
| 38 |
+
- light: monochrome dark glyphs, icons and wordmarks on transparent backgrounds that should simply become near-white.
|
| 39 |
+
- plate: logos or images with important dark parts on transparency whose colors must not be altered (multicolor brand marks, partner logos); they get a subtle light rounded backing.
|
| 40 |
+
Think about the site logo especially carefully: it must stay recognizable and clearly visible.
|
| 41 |
+
|
| 42 |
+
Answer with JSON only, following the schema. Use "#rrggbb", or "#rrggbbaa" when alpha is below 1."""
|
| 43 |
+
|
| 44 |
+
SCHEMA = {
|
| 45 |
+
"type": "object",
|
| 46 |
+
"properties": {
|
| 47 |
+
"design_notes": {"type": "string"},
|
| 48 |
+
"colors": {"type": "array", "items": {"type": "object", "properties": {"id": {"type": "string"}, "dark": {"type": "string"}},
|
| 49 |
+
"required": ["id", "dark"], "additionalProperties": False}},
|
| 50 |
+
"images": {"type": "array", "items": {"type": "object", "properties": {"id": {"type": "string"}, "treatment": {"type": "string", "enum": TREATMENTS}},
|
| 51 |
+
"required": ["id", "treatment"], "additionalProperties": False}},
|
| 52 |
+
},
|
| 53 |
+
"required": ["design_notes", "colors", "images"],
|
| 54 |
+
"additionalProperties": False,
|
| 55 |
+
}
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def hexa(rgba):
|
| 59 |
+
r, g, b, a = rgba
|
| 60 |
+
h = "#%02x%02x%02x" % (r, g, b)
|
| 61 |
+
if a < 0.999:
|
| 62 |
+
h += "%02x" % round(a * 255)
|
| 63 |
+
return h
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def pct(x):
|
| 67 |
+
if x >= 0.095:
|
| 68 |
+
return f"{x * 100:.0f}%"
|
| 69 |
+
if x >= 0.0005:
|
| 70 |
+
return f"{x * 100:.1f}%"
|
| 71 |
+
return "<0.1%"
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def describe_tokens(page, max_tokens):
|
| 75 |
+
toks = page["tokens"][:max_tokens]
|
| 76 |
+
n = len(toks)
|
| 77 |
+
pa = max(1.0, page["docW"] * min(page["docH"], page["vh"] * 8))
|
| 78 |
+
va = page["vw"] * page["vh"]
|
| 79 |
+
# relations: (type, a, b, w)
|
| 80 |
+
best = {}
|
| 81 |
+
for typ, a, b, w in page["rel"]:
|
| 82 |
+
if a >= n or b >= n:
|
| 83 |
+
continue
|
| 84 |
+
best.setdefault((typ, a), []).append((w, b))
|
| 85 |
+
if typ in (0, 1):
|
| 86 |
+
best.setdefault((10 + typ, b), []).append((w, a)) # reverse: what sits on b
|
| 87 |
+
def top(key, k=1):
|
| 88 |
+
lst = sorted(best.get(key, []), reverse=True)[:k]
|
| 89 |
+
return [f"t{b}" for _, b in lst]
|
| 90 |
+
lines = []
|
| 91 |
+
for i, t in enumerate(toks):
|
| 92 |
+
d = t["d"]
|
| 93 |
+
role = ROLES[t["role"]]
|
| 94 |
+
parts = [f"t{i}", role, hexa(t["rgba"])]
|
| 95 |
+
info = []
|
| 96 |
+
if d["root"]:
|
| 97 |
+
info.append("PAGE CANVAS (html/body background)")
|
| 98 |
+
visible = d["nVis"] > 0
|
| 99 |
+
if role == "bg":
|
| 100 |
+
if d["area"]:
|
| 101 |
+
info.append(f"covers {pct(min(1, d['area'] / pa))} of page (unobstructed {pct(min(1, d['own'] / pa))}), first screen {pct(min(1, d['area1'] / va))}")
|
| 102 |
+
on = top((1, i))
|
| 103 |
+
if on and not d["root"]:
|
| 104 |
+
info.append(f"sits on {on[0]}")
|
| 105 |
+
txt = top((10, i), 3)
|
| 106 |
+
if txt:
|
| 107 |
+
info.append("text on it: " + ",".join(txt))
|
| 108 |
+
elif role in ("text", "fill", "stroke"):
|
| 109 |
+
if role == "text" and d["chars"]:
|
| 110 |
+
info.append(f"{d['chars']} chars, ~{d['fs']:.0f}px")
|
| 111 |
+
elif role != "text":
|
| 112 |
+
info.append("SVG " + ("fill (icons/logos/shapes)" if role == "fill" else "stroke (icons/lines)"))
|
| 113 |
+
on = top((0, i))
|
| 114 |
+
if on:
|
| 115 |
+
info.append(f"on {on[0]}" + (f" (contrast {d['contrast']:.1f}:1)" if d.get("contrast") else ""))
|
| 116 |
+
elif role == "border":
|
| 117 |
+
on = top((2, i))
|
| 118 |
+
if on:
|
| 119 |
+
info.append(f"line on {on[0]}")
|
| 120 |
+
elif role == "shadow":
|
| 121 |
+
info.append("box/text-shadow color")
|
| 122 |
+
elif role == "gradient":
|
| 123 |
+
info.append("gradient stop in a background-image")
|
| 124 |
+
if d["n"]:
|
| 125 |
+
info.append(f"{d['n']} elements" + ("" if visible else ", none visible now (hidden menus etc.)"))
|
| 126 |
+
flags = []
|
| 127 |
+
n_el = max(1, d["n"])
|
| 128 |
+
if d.get("media", 0) / n_el >= 0.5:
|
| 129 |
+
info.append("ON TOP OF / CONTAINS video or images (player UI, caption, scrim)")
|
| 130 |
+
for key, label in (("inter", "interactive"), ("link", "in links"), ("head", "headings"), ("header", "header"), ("nav", "nav"), ("footer", "footer")):
|
| 131 |
+
if d.get(key, 0) / n_el >= 0.3:
|
| 132 |
+
flags.append(label)
|
| 133 |
+
if flags:
|
| 134 |
+
info.append("/".join(flags))
|
| 135 |
+
if d["rules"] and not d["n"]:
|
| 136 |
+
info.append(f"CSS-rule only ({d['rules']} rules" + (f", {d['state']} hover/focus/active" if d["state"] else "") + ")")
|
| 137 |
+
elif d["state"]:
|
| 138 |
+
info.append(f"also in {d['state']} hover/focus/active rules")
|
| 139 |
+
if d.get("ex"):
|
| 140 |
+
info.append("e.g. " + "; ".join(x[:70] for x in d["ex"][:2]))
|
| 141 |
+
lines.append(" ".join(parts) + " | " + " | ".join(info))
|
| 142 |
+
return lines
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
def describe_images(page, max_images):
|
| 146 |
+
out = []
|
| 147 |
+
for j, im in enumerate(page["images"][:max_images]):
|
| 148 |
+
s = im.get("stats") or {}
|
| 149 |
+
bits = [f"i{j}", f"{im['w']}x{im['h']}px", im.get("ext") or im["kind"]]
|
| 150 |
+
if im["kind"] == "bg":
|
| 151 |
+
bits.append("CSS background image")
|
| 152 |
+
where = []
|
| 153 |
+
if im.get("header"): where.append("header/nav")
|
| 154 |
+
if im.get("footer"): where.append("footer")
|
| 155 |
+
if im.get("homeLink"): where.append("links to homepage")
|
| 156 |
+
elif im.get("inLink"): where.append("inside a link")
|
| 157 |
+
if im["y"] < page["vh"]: where.append("first screen")
|
| 158 |
+
if where: bits.append(", ".join(where))
|
| 159 |
+
if im.get("logo"): bits.append("named like a logo")
|
| 160 |
+
if im.get("icon"): bits.append("named like an icon")
|
| 161 |
+
if im.get("alt"): bits.append(f'alt "{im["alt"][:50]}"')
|
| 162 |
+
if s:
|
| 163 |
+
bits.append(f"transparent {pct(s.get('transp', 0))}, opaque mean lightness {s.get('meanL', 0):.2f}, near-white {pct(s.get('white', 0))}, near-black {pct(s.get('black', 0))}, chroma {s.get('meanC', 0):.2f}")
|
| 164 |
+
else:
|
| 165 |
+
bits.append("pixels unavailable")
|
| 166 |
+
if im.get("underTok", -1) >= 0:
|
| 167 |
+
bits.append(f"on surface t{im['underTok']}")
|
| 168 |
+
if im.get("count", 1) > 1:
|
| 169 |
+
bits.append(f"used {im['count']}x")
|
| 170 |
+
out.append(" | ".join(bits))
|
| 171 |
+
return out
|
| 172 |
+
|
| 173 |
+
|
| 174 |
+
def contact_sheet(page, max_images):
|
| 175 |
+
ims = [im for im in page["images"][:max_images]]
|
| 176 |
+
if not ims:
|
| 177 |
+
return None
|
| 178 |
+
cols = 6
|
| 179 |
+
cell, pad, lab = 112, 8, 16
|
| 180 |
+
rows = math.ceil(len(ims) / cols)
|
| 181 |
+
W, H = cols * (cell + pad) + pad, rows * (cell + pad + lab) + pad
|
| 182 |
+
sheet = Image.new("RGB", (W, H), (255, 255, 255))
|
| 183 |
+
dr = ImageDraw.Draw(sheet)
|
| 184 |
+
try:
|
| 185 |
+
font = ImageFont.load_default(size=13)
|
| 186 |
+
except Exception:
|
| 187 |
+
font = ImageFont.load_default()
|
| 188 |
+
for j, im in enumerate(ims):
|
| 189 |
+
x = pad + (j % cols) * (cell + pad)
|
| 190 |
+
y = pad + (j // cols) * (cell + pad + lab)
|
| 191 |
+
# checkerboard
|
| 192 |
+
for cy in range(0, cell, 8):
|
| 193 |
+
for cx in range(0, cell, 8):
|
| 194 |
+
c = (204, 204, 204) if ((cx + cy) // 8) % 2 else (245, 245, 245)
|
| 195 |
+
dr.rectangle([x + cx, y + lab + cy, x + cx + 7, y + lab + cy + 7], fill=c)
|
| 196 |
+
dr.text((x, y), f"i{j}", fill=(200, 0, 0), font=font)
|
| 197 |
+
if im.get("big"):
|
| 198 |
+
try:
|
| 199 |
+
b = base64.b64decode(im["big"].split(",", 1)[1])
|
| 200 |
+
th = Image.open(io.BytesIO(b)).convert("RGBA").resize((cell, cell), Image.LANCZOS)
|
| 201 |
+
sheet.paste(th, (x, y + lab), th)
|
| 202 |
+
except Exception:
|
| 203 |
+
dr.text((x + 4, y + lab + 40), "n/a", fill=(0, 0, 0), font=font)
|
| 204 |
+
else:
|
| 205 |
+
dr.text((x + 4, y + lab + 40), "n/a", fill=(0, 0, 0), font=font)
|
| 206 |
+
buf = io.BytesIO()
|
| 207 |
+
sheet.save(buf, "PNG", optimize=True)
|
| 208 |
+
return buf.getvalue()
|
| 209 |
+
|
| 210 |
+
|
| 211 |
+
def screenshot(pdir, max_w=900, max_h=1400):
|
| 212 |
+
im = Image.open(pdir / "shot.jpg").convert("RGB")
|
| 213 |
+
s = max_w / im.width
|
| 214 |
+
im = im.resize((max_w, int(im.height * s)), Image.LANCZOS)
|
| 215 |
+
if im.height > max_h:
|
| 216 |
+
im = im.crop((0, 0, max_w, max_h))
|
| 217 |
+
buf = io.BytesIO()
|
| 218 |
+
im.save(buf, "JPEG", quality=82)
|
| 219 |
+
return buf.getvalue()
|
| 220 |
+
|
| 221 |
+
|
| 222 |
+
def build_request(page, pdir, args):
|
| 223 |
+
lines = describe_tokens(page, args.max_tokens)
|
| 224 |
+
imgs = describe_images(page, args.max_images)
|
| 225 |
+
sheet = contact_sheet(page, args.max_images) if imgs else None
|
| 226 |
+
host = re.sub(r"^https?://", "", page["url"]).split("/")[0]
|
| 227 |
+
text = [f"Website: {host} — \"{page['title'][:100]}\"", f"Viewport 1280x{page['vh']}, page height {page['docH']}px.",
|
| 228 |
+
"", f"COLOR INVENTORY ({len(lines)} colors; id role color | usage):"] + lines
|
| 229 |
+
if imgs:
|
| 230 |
+
text += ["", f"IMAGES ({len(imgs)}; thumbnails in the contact sheet):"] + imgs
|
| 231 |
+
text += ["", f"Return the dark-theme color for every id t0-t{len(lines) - 1}" + (f" and a treatment for every image i0-i{len(imgs) - 1}." if imgs else " (images: return an empty list).")]
|
| 232 |
+
content = [{"type": "text", "text": "Screenshot of the light page (top section):"},
|
| 233 |
+
{"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": base64.b64encode(screenshot(pdir)).decode()}}]
|
| 234 |
+
if sheet:
|
| 235 |
+
content += [{"type": "text", "text": "Contact sheet of the page's images:"},
|
| 236 |
+
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": base64.b64encode(sheet).decode()}}]
|
| 237 |
+
content.append({"type": "text", "text": "\n".join(text)})
|
| 238 |
+
return content, len(lines), len(imgs)
|
| 239 |
+
|
| 240 |
+
|
| 241 |
+
HEX_RE = re.compile(r"^#?([0-9a-fA-F]{6})([0-9a-fA-F]{2})?$")
|
| 242 |
+
|
| 243 |
+
|
| 244 |
+
def parse_answer(msg, n_tok, n_img):
|
| 245 |
+
txt = "".join(b.text for b in msg.content if b.type == "text")
|
| 246 |
+
data = json.loads(txt)
|
| 247 |
+
colors, images = {}, {}
|
| 248 |
+
for c in data.get("colors", []):
|
| 249 |
+
m = HEX_RE.match(c.get("dark", "").strip())
|
| 250 |
+
mid = re.match(r"^t(\d+)$", c.get("id", ""))
|
| 251 |
+
if m and mid and int(mid.group(1)) < n_tok:
|
| 252 |
+
colors[int(mid.group(1))] = "#" + m.group(1).lower() + (m.group(2) or "").lower()
|
| 253 |
+
for c in data.get("images", []):
|
| 254 |
+
mid = re.match(r"^i(\d+)$", c.get("id", ""))
|
| 255 |
+
if mid and int(mid.group(1)) < n_img and c.get("treatment") in TREATMENTS:
|
| 256 |
+
images[int(mid.group(1))] = c["treatment"]
|
| 257 |
+
return data.get("design_notes", ""), colors, images
|
| 258 |
+
|
| 259 |
+
|
| 260 |
+
class Gate:
|
| 261 |
+
"""Adaptive concurrency: grow slowly while healthy, back off on the org concurrency cap and stay below it."""
|
| 262 |
+
def __init__(self, start, cap):
|
| 263 |
+
self.limit, self.cap, self.active = float(start), cap, 0
|
| 264 |
+
self.cond = asyncio.Condition()
|
| 265 |
+
self.last429 = 0.0
|
| 266 |
+
|
| 267 |
+
async def acquire(self):
|
| 268 |
+
async with self.cond:
|
| 269 |
+
while self.active >= int(self.limit):
|
| 270 |
+
await self.cond.wait()
|
| 271 |
+
self.active += 1
|
| 272 |
+
|
| 273 |
+
async def release(self, ok=True, limited=False):
|
| 274 |
+
async with self.cond:
|
| 275 |
+
self.active -= 1
|
| 276 |
+
if limited:
|
| 277 |
+
seen = self.active + 1
|
| 278 |
+
self.cap = max(2, min(self.cap, int(seen * 0.75)))
|
| 279 |
+
self.limit = max(2.0, min(self.limit * 0.7, self.cap))
|
| 280 |
+
self.last429 = time.time()
|
| 281 |
+
elif ok and time.time() - self.last429 > 120:
|
| 282 |
+
self.limit = min(self.cap, self.limit + 0.2)
|
| 283 |
+
self.cond.notify_all()
|
| 284 |
+
|
| 285 |
+
|
| 286 |
+
class Ledger:
|
| 287 |
+
def __init__(self, path, budget):
|
| 288 |
+
self.path, self.budget = Path(path), budget
|
| 289 |
+
self.cost = 0.0
|
| 290 |
+
self.n = 0
|
| 291 |
+
if self.path.exists():
|
| 292 |
+
for line in self.path.open():
|
| 293 |
+
try:
|
| 294 |
+
r = json.loads(line)
|
| 295 |
+
self.cost += r["cost"]; self.n += 1
|
| 296 |
+
except Exception:
|
| 297 |
+
pass
|
| 298 |
+
self.f = self.path.open("a")
|
| 299 |
+
|
| 300 |
+
def add(self, rec):
|
| 301 |
+
self.cost += rec["cost"]; self.n += 1
|
| 302 |
+
self.f.write(json.dumps(rec) + "\n"); self.f.flush()
|
| 303 |
+
|
| 304 |
+
|
| 305 |
+
def usage_cost(u):
|
| 306 |
+
cr = getattr(u, "cache_read_input_tokens", 0) or 0
|
| 307 |
+
cw = getattr(u, "cache_creation_input_tokens", 0) or 0
|
| 308 |
+
return u.input_tokens * PRICE_IN + u.output_tokens * PRICE_OUT + cr * PRICE_CACHE_READ + cw * PRICE_CACHE_WRITE
|
| 309 |
+
|
| 310 |
+
|
| 311 |
+
ERRLOG = {}
|
| 312 |
+
|
| 313 |
+
|
| 314 |
+
def logerr(kind, e):
|
| 315 |
+
ERRLOG[kind] = ERRLOG.get(kind, 0) + 1
|
| 316 |
+
if ERRLOG[kind] <= 3:
|
| 317 |
+
print(f"[{kind}] {type(e).__name__}: {str(e)[:300]}", flush=True)
|
| 318 |
+
|
| 319 |
+
|
| 320 |
+
async def one(client, sem, pid, crawl_dir, out_dir, ledger, args, stats):
|
| 321 |
+
out = out_dir / f"{pid}.json"
|
| 322 |
+
if out.exists():
|
| 323 |
+
return
|
| 324 |
+
pdir = crawl_dir / "pages" / pid
|
| 325 |
+
try:
|
| 326 |
+
page = json.load(gzip.open(pdir / "light.json.gz"))
|
| 327 |
+
content, n_tok, n_img = build_request(page, pdir, args)
|
| 328 |
+
except Exception as e:
|
| 329 |
+
stats["prep_err"] = stats.get("prep_err", 0) + 1
|
| 330 |
+
return
|
| 331 |
+
for attempt in range(8):
|
| 332 |
+
if ledger.cost >= args.budget:
|
| 333 |
+
return
|
| 334 |
+
await sem.acquire()
|
| 335 |
+
ok, limited = False, False
|
| 336 |
+
t0 = time.time()
|
| 337 |
+
try:
|
| 338 |
+
async with client.messages.stream(
|
| 339 |
+
model=MODEL, max_tokens=args.max_out, system=SYSTEM,
|
| 340 |
+
thinking={"type": "adaptive"}, output_config={"effort": args.effort, "format": {"type": "json_schema", "schema": SCHEMA}},
|
| 341 |
+
messages=[{"role": "user", "content": content}],
|
| 342 |
+
) as stream:
|
| 343 |
+
msg = await stream.get_final_message()
|
| 344 |
+
ok = True
|
| 345 |
+
except anthropic.RateLimitError as e:
|
| 346 |
+
stats["429"] = stats.get("429", 0) + 1
|
| 347 |
+
limited = "concurrent" in str(e).lower()
|
| 348 |
+
logerr("429", e)
|
| 349 |
+
except anthropic.APIStatusError as e:
|
| 350 |
+
stats[f"http{e.status_code}"] = stats.get(f"http{e.status_code}", 0) + 1
|
| 351 |
+
logerr(f"http{e.status_code}", e)
|
| 352 |
+
if e.status_code in (400, 401, 403, 404, 413):
|
| 353 |
+
(out_dir / "errors").mkdir(exist_ok=True)
|
| 354 |
+
(out_dir / "errors" / f"{pid}.txt").write_text(str(e)[:2000])
|
| 355 |
+
await sem.release(False)
|
| 356 |
+
return
|
| 357 |
+
except (anthropic.APIConnectionError, anthropic.APITimeoutError) as e:
|
| 358 |
+
stats["conn"] = stats.get("conn", 0) + 1
|
| 359 |
+
logerr("conn", e)
|
| 360 |
+
except Exception as e:
|
| 361 |
+
stats["other"] = stats.get("other", 0) + 1
|
| 362 |
+
logerr("other", e)
|
| 363 |
+
await sem.release(ok, limited)
|
| 364 |
+
if not ok:
|
| 365 |
+
await asyncio.sleep(min(60, 4 * (attempt + 1)) + random.random() * 4)
|
| 366 |
+
continue
|
| 367 |
+
cost = usage_cost(msg.usage)
|
| 368 |
+
rec = {"id": pid, "cost": cost, "in": msg.usage.input_tokens, "out": msg.usage.output_tokens, "stop": msg.stop_reason,
|
| 369 |
+
"secs": round(time.time() - t0, 1), "t": time.time()}
|
| 370 |
+
ledger.add(rec)
|
| 371 |
+
if msg.stop_reason == "refusal":
|
| 372 |
+
stats["refusal"] = stats.get("refusal", 0) + 1
|
| 373 |
+
out.write_text(json.dumps({"id": pid, "refusal": True, "usage": rec}))
|
| 374 |
+
return
|
| 375 |
+
try:
|
| 376 |
+
notes, colors, images = parse_answer(msg, n_tok, n_img)
|
| 377 |
+
except Exception as e:
|
| 378 |
+
stats["parse"] = stats.get("parse", 0) + 1
|
| 379 |
+
(out_dir / "errors").mkdir(exist_ok=True)
|
| 380 |
+
(out_dir / "errors" / f"{pid}.txt").write_text("parse: " + str(e)[:500])
|
| 381 |
+
return
|
| 382 |
+
out.write_text(json.dumps({"id": pid, "url": page["url"], "n_tok": n_tok, "n_img": n_img, "notes": notes,
|
| 383 |
+
"colors": colors, "images": images, "usage": rec, "effort": args.effort, "model": MODEL}))
|
| 384 |
+
stats["ok"] = stats.get("ok", 0) + 1
|
| 385 |
+
return
|
| 386 |
+
|
| 387 |
+
|
| 388 |
+
# ---------------- batch mode ----------------
|
| 389 |
+
import urllib.request
|
| 390 |
+
|
| 391 |
+
PROXY = os.environ.get("ANTHROPIC_BASE_URL", "https://api.anthropic.com")
|
| 392 |
+
HDRS = {"x-api-key": os.environ.get("ANTHROPIC_API_KEY", ""), "anthropic-version": "2023-06-01", "content-type": "application/json"}
|
| 393 |
+
|
| 394 |
+
|
| 395 |
+
def http(method, path, body=None, timeout=300):
|
| 396 |
+
req = urllib.request.Request(PROXY + path, data=json.dumps(body).encode() if body is not None else None, headers=HDRS, method=method)
|
| 397 |
+
with urllib.request.urlopen(req, timeout=timeout) as r:
|
| 398 |
+
return r.read()
|
| 399 |
+
|
| 400 |
+
|
| 401 |
+
def params_for(content, args):
|
| 402 |
+
return {"model": MODEL, "max_tokens": args.max_out, "system": SYSTEM, "thinking": {"type": "adaptive"},
|
| 403 |
+
"output_config": {"effort": args.effort, "format": {"type": "json_schema", "schema": SCHEMA}},
|
| 404 |
+
"messages": [{"role": "user", "content": content}]}
|
| 405 |
+
|
| 406 |
+
|
| 407 |
+
class Msg: # minimal adapter so parse_answer works on batch result dicts
|
| 408 |
+
def __init__(self, d):
|
| 409 |
+
self.content = [type("B", (), {"type": b.get("type"), "text": b.get("text", "")}) for b in d.get("content", [])]
|
| 410 |
+
|
| 411 |
+
|
| 412 |
+
def collect_batch(bid, meta, crawl_dir, out_dir, ledger, stats):
|
| 413 |
+
raw = http("GET", f"/v1/messages/batches/{bid}/results", timeout=600).decode()
|
| 414 |
+
seen = set()
|
| 415 |
+
for line in raw.splitlines():
|
| 416 |
+
if not line.strip():
|
| 417 |
+
continue
|
| 418 |
+
r = json.loads(line)
|
| 419 |
+
pid = r["custom_id"]
|
| 420 |
+
seen.add(pid)
|
| 421 |
+
res = r["result"]
|
| 422 |
+
if res["type"] != "succeeded":
|
| 423 |
+
stats["b_" + res["type"]] = stats.get("b_" + res["type"], 0) + 1
|
| 424 |
+
continue
|
| 425 |
+
m = res["message"]
|
| 426 |
+
u = m.get("usage", {})
|
| 427 |
+
cost = 0.5 * (u.get("input_tokens", 0) * PRICE_IN + u.get("output_tokens", 0) * PRICE_OUT
|
| 428 |
+
+ (u.get("cache_read_input_tokens") or 0) * PRICE_CACHE_READ + (u.get("cache_creation_input_tokens") or 0) * PRICE_CACHE_WRITE)
|
| 429 |
+
rec = {"id": pid, "cost": cost, "in": u.get("input_tokens", 0), "out": u.get("output_tokens", 0), "stop": m.get("stop_reason"), "batch": bid, "t": time.time()}
|
| 430 |
+
ledger.add(rec)
|
| 431 |
+
out = out_dir / f"{pid}.json"
|
| 432 |
+
if m.get("stop_reason") == "refusal":
|
| 433 |
+
stats["refusal"] = stats.get("refusal", 0) + 1
|
| 434 |
+
out.write_text(json.dumps({"id": pid, "refusal": True, "usage": rec}))
|
| 435 |
+
continue
|
| 436 |
+
n_tok, n_img = meta[pid]
|
| 437 |
+
try:
|
| 438 |
+
notes, colors, images = parse_answer(Msg(m), n_tok, n_img)
|
| 439 |
+
except Exception as e:
|
| 440 |
+
stats["parse"] = stats.get("parse", 0) + 1
|
| 441 |
+
continue
|
| 442 |
+
url = ""
|
| 443 |
+
out.write_text(json.dumps({"id": pid, "n_tok": n_tok, "n_img": n_img, "notes": notes, "colors": colors, "images": images,
|
| 444 |
+
"usage": rec, "effort": args_global.effort, "model": MODEL}))
|
| 445 |
+
stats["ok"] = stats.get("ok", 0) + 1
|
| 446 |
+
return seen
|
| 447 |
+
|
| 448 |
+
|
| 449 |
+
args_global = None
|
| 450 |
+
|
| 451 |
+
|
| 452 |
+
def batch_main(args):
|
| 453 |
+
global args_global
|
| 454 |
+
args_global = args
|
| 455 |
+
crawl_dir, out_dir = Path(args.crawl), Path(args.out)
|
| 456 |
+
out_dir.mkdir(parents=True, exist_ok=True)
|
| 457 |
+
ledger = Ledger(out_dir / "ledger.jsonl", args.budget)
|
| 458 |
+
state_p = out_dir / "batches.jsonl"
|
| 459 |
+
batches = {}
|
| 460 |
+
if state_p.exists():
|
| 461 |
+
for line in state_p.open():
|
| 462 |
+
b = json.loads(line)
|
| 463 |
+
batches.setdefault(b["id"], {}).update(b)
|
| 464 |
+
submitted = set(pid for b in batches.values() for pid in b.get("ids", []))
|
| 465 |
+
stats = {}
|
| 466 |
+
est_per = 0.0035
|
| 467 |
+
last_submit = 0
|
| 468 |
+
while True:
|
| 469 |
+
# collect finished batches
|
| 470 |
+
pending_est = 0.0
|
| 471 |
+
for bid, b in list(batches.items()):
|
| 472 |
+
if b.get("collected"):
|
| 473 |
+
continue
|
| 474 |
+
try:
|
| 475 |
+
info = json.loads(http("GET", f"/v1/messages/batches/{bid}"))
|
| 476 |
+
except Exception as e:
|
| 477 |
+
print("poll error", bid, str(e)[:200], flush=True)
|
| 478 |
+
pending_est += len(b["ids"]) * est_per
|
| 479 |
+
continue
|
| 480 |
+
if info["processing_status"] == "ended":
|
| 481 |
+
try:
|
| 482 |
+
collect_batch(bid, b["meta"], crawl_dir, out_dir, ledger, stats)
|
| 483 |
+
b["collected"] = True
|
| 484 |
+
with state_p.open("a") as f:
|
| 485 |
+
f.write(json.dumps({"id": bid, "collected": True}) + "\n")
|
| 486 |
+
print(f"[{time.strftime('%H:%M:%S')}] collected {bid} {info['request_counts']}", flush=True)
|
| 487 |
+
except Exception as e:
|
| 488 |
+
print("collect error", bid, str(e)[:300], flush=True)
|
| 489 |
+
pending_est += len(b["ids"]) * est_per
|
| 490 |
+
else:
|
| 491 |
+
pending_est += len(b["ids"]) * est_per
|
| 492 |
+
if ledger.n > 50:
|
| 493 |
+
est_per = max(0.002, ledger.cost / ledger.n)
|
| 494 |
+
# candidates
|
| 495 |
+
cands = []
|
| 496 |
+
if args.ids:
|
| 497 |
+
cands = [l.strip() for l in open(args.ids) if l.strip()]
|
| 498 |
+
else:
|
| 499 |
+
for line in (l for ip in sorted(crawl_dir.glob("index*.jsonl")) for l in open(ip)):
|
| 500 |
+
try:
|
| 501 |
+
r = json.loads(line)
|
| 502 |
+
except Exception:
|
| 503 |
+
continue
|
| 504 |
+
if r.get("kind") in args.kinds.split(",") and (r.get("n_tokens") or 0) >= 8 and (r.get("status") or 200) < 400:
|
| 505 |
+
cands.append(r["id"])
|
| 506 |
+
new = [c for c in cands if c not in submitted and not (out_dir / f"{c}.json").exists()]
|
| 507 |
+
room = args.budget - ledger.cost - pending_est
|
| 508 |
+
can = int(max(0, room) / est_per)
|
| 509 |
+
ready = len(new) >= args.batch_size or (new and time.time() - last_submit > args.max_wait)
|
| 510 |
+
if ready and can > 0:
|
| 511 |
+
chunk = new[: min(args.batch_size, can)]
|
| 512 |
+
reqs, meta = [], {}
|
| 513 |
+
for pid in chunk:
|
| 514 |
+
pdir = crawl_dir / "pages" / pid
|
| 515 |
+
try:
|
| 516 |
+
page = json.load(gzip.open(pdir / "light.json.gz"))
|
| 517 |
+
content, n_tok, n_img = build_request(page, pdir, args)
|
| 518 |
+
except Exception as e:
|
| 519 |
+
submitted.add(pid)
|
| 520 |
+
continue
|
| 521 |
+
reqs.append({"custom_id": pid, "params": params_for(content, args)})
|
| 522 |
+
meta[pid] = (n_tok, n_img)
|
| 523 |
+
if reqs:
|
| 524 |
+
try:
|
| 525 |
+
info = json.loads(http("POST", "/v1/messages/batches", {"requests": reqs}, timeout=900))
|
| 526 |
+
b = {"id": info["id"], "ids": list(meta), "meta": meta, "t": time.time()}
|
| 527 |
+
batches[info["id"]] = b
|
| 528 |
+
with state_p.open("a") as f:
|
| 529 |
+
f.write(json.dumps(b) + "\n")
|
| 530 |
+
submitted.update(meta)
|
| 531 |
+
last_submit = time.time()
|
| 532 |
+
print(f"[{time.strftime('%H:%M:%S')}] submitted {info['id']} with {len(reqs)} (spent ${ledger.cost:.2f}, pending ~${pending_est:.2f}, est/page ${est_per:.4f})", flush=True)
|
| 533 |
+
except Exception as e:
|
| 534 |
+
print("submit error", str(e)[:300], flush=True)
|
| 535 |
+
time.sleep(60)
|
| 536 |
+
print(f"[{time.strftime('%H:%M:%S')}] spent ${ledger.cost:.3f} on {ledger.n} | open batches {sum(1 for b in batches.values() if not b.get('collected'))} | waiting pages {len(new)} | {stats}", flush=True)
|
| 537 |
+
open_batches = [b for b in batches.values() if not b.get("collected")]
|
| 538 |
+
if not args.follow and not open_batches and (not new or can <= 0):
|
| 539 |
+
break
|
| 540 |
+
if room <= 0 and not open_batches:
|
| 541 |
+
print("budget reached", flush=True)
|
| 542 |
+
break
|
| 543 |
+
time.sleep(args.poll)
|
| 544 |
+
print(f"DONE ${ledger.cost:.3f} {ledger.n} {stats}", flush=True)
|
| 545 |
+
|
| 546 |
+
|
| 547 |
+
async def main():
|
| 548 |
+
ap = argparse.ArgumentParser()
|
| 549 |
+
ap.add_argument("--crawl", required=True)
|
| 550 |
+
ap.add_argument("--out", required=True)
|
| 551 |
+
ap.add_argument("--ids", help="file with page ids (default: all light pages in index)")
|
| 552 |
+
ap.add_argument("--budget", type=float, default=95.0)
|
| 553 |
+
ap.add_argument("--concurrency", type=int, default=64)
|
| 554 |
+
ap.add_argument("--effort", default="high")
|
| 555 |
+
ap.add_argument("--max-tokens", type=int, default=140)
|
| 556 |
+
ap.add_argument("--max-images", type=int, default=24)
|
| 557 |
+
ap.add_argument("--max-out", type=int, default=48000)
|
| 558 |
+
ap.add_argument("--limit", type=int, default=0)
|
| 559 |
+
ap.add_argument("--follow", action="store_true", help="keep polling the crawl index for new pages")
|
| 560 |
+
ap.add_argument("--mode", default="batch", choices=["batch", "realtime"])
|
| 561 |
+
ap.add_argument("--batch-size", type=int, default=250)
|
| 562 |
+
ap.add_argument("--max-wait", type=int, default=900)
|
| 563 |
+
ap.add_argument("--poll", type=int, default=60)
|
| 564 |
+
ap.add_argument("--kinds", default="light")
|
| 565 |
+
args = ap.parse_args()
|
| 566 |
+
if args.mode == "batch":
|
| 567 |
+
return batch_main(args)
|
| 568 |
+
crawl_dir, out_dir = Path(args.crawl), Path(args.out)
|
| 569 |
+
out_dir.mkdir(parents=True, exist_ok=True)
|
| 570 |
+
ledger = Ledger(out_dir / "ledger.jsonl", args.budget)
|
| 571 |
+
client = anthropic.AsyncAnthropic(base_url=PROXY, api_key=os.environ.get("ANTHROPIC_API_KEY", ""),
|
| 572 |
+
max_retries=0, timeout=600)
|
| 573 |
+
sem = Gate(min(8, args.concurrency), args.concurrency)
|
| 574 |
+
stats = {}
|
| 575 |
+
started = set()
|
| 576 |
+
tasks = []
|
| 577 |
+
t_start = time.time()
|
| 578 |
+
|
| 579 |
+
def candidates():
|
| 580 |
+
if args.ids:
|
| 581 |
+
return [l.strip() for l in open(args.ids) if l.strip()]
|
| 582 |
+
ids = []
|
| 583 |
+
for line in (l for ip in sorted(crawl_dir.glob("index*.jsonl")) for l in open(ip)):
|
| 584 |
+
try:
|
| 585 |
+
r = json.loads(line)
|
| 586 |
+
except Exception:
|
| 587 |
+
continue
|
| 588 |
+
if r.get("kind") == "light" and (r.get("n_tokens") or 0) >= 8 and (r.get("status") or 200) < 400:
|
| 589 |
+
ids.append(r["id"])
|
| 590 |
+
return ids
|
| 591 |
+
|
| 592 |
+
async def reporter():
|
| 593 |
+
while True:
|
| 594 |
+
await asyncio.sleep(60)
|
| 595 |
+
el = time.time() - t_start
|
| 596 |
+
print(f"[{time.strftime('%H:%M:%S')}] ${ledger.cost:.3f} spent, {ledger.n} calls, {stats}, {ledger.n / max(1, el) * 3600:.0f}/h, conc {sem.active}/{sem.limit:.1f} cap {sem.cap}", flush=True)
|
| 597 |
+
rep = asyncio.create_task(reporter())
|
| 598 |
+
while True:
|
| 599 |
+
ids = candidates()
|
| 600 |
+
new = [i for i in ids if i not in started]
|
| 601 |
+
if args.limit:
|
| 602 |
+
new = new[: max(0, args.limit - len(started))]
|
| 603 |
+
for pid in new:
|
| 604 |
+
started.add(pid)
|
| 605 |
+
tasks.append(asyncio.create_task(one(client, sem, pid, crawl_dir, out_dir, ledger, args, stats)))
|
| 606 |
+
tasks = [t for t in tasks if not t.done()]
|
| 607 |
+
if ledger.cost >= args.budget:
|
| 608 |
+
print("budget reached", flush=True)
|
| 609 |
+
break
|
| 610 |
+
if not args.follow and not tasks:
|
| 611 |
+
break
|
| 612 |
+
if args.limit and len(started) >= args.limit and not tasks:
|
| 613 |
+
break
|
| 614 |
+
await asyncio.sleep(30 if args.follow else 2)
|
| 615 |
+
await asyncio.gather(*tasks, return_exceptions=True)
|
| 616 |
+
rep.cancel()
|
| 617 |
+
print(f"DONE ${ledger.cost:.3f} {ledger.n} calls {stats}", flush=True)
|
| 618 |
+
|
| 619 |
+
|
| 620 |
+
if __name__ == "__main__":
|
| 621 |
+
asyncio.run(main())
|
training/judge.py
ADDED
|
@@ -0,0 +1,110 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Blind pairwise judging of dark-mode renders with Claude Haiku 5.5 (Batch API). Runs on Vega's host."""
|
| 3 |
+
import argparse, base64, io, json, random, time
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
from PIL import Image
|
| 6 |
+
|
| 7 |
+
import haiku_gen as hg
|
| 8 |
+
|
| 9 |
+
JUDGE_SYSTEM = """You are a meticulous product-design reviewer. You compare two dark-mode versions (A and B) of the same web page, given the original light page for reference. Judge which dark mode a demanding user would rather browse with, considering in this order: (1) readability: every piece of text clearly legible against its background, nothing lost; (2) nothing broken: no large light patches left glaring, logos and icons still visible and recognizable, photos and videos not color-distorted; (3) design quality: surfaces, hierarchy and elevation preserved, brand colors recognizable, pleasant and coherent palette, no muddy or harsh colors. Small differences in exact shades matter less than legibility and breakage. Answer with JSON only."""
|
| 10 |
+
|
| 11 |
+
SCHEMA = {"type": "object", "properties": {
|
| 12 |
+
"winner": {"type": "string", "enum": ["A", "B", "tie"]},
|
| 13 |
+
"score_A": {"type": "integer"}, "score_B": {"type": "integer"},
|
| 14 |
+
"reason": {"type": "string"}},
|
| 15 |
+
"required": ["winner", "score_A", "score_B", "reason"], "additionalProperties": False}
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
def img_block(path, w=800, h=1000):
|
| 19 |
+
im = Image.open(path).convert("RGB")
|
| 20 |
+
s = w / im.width
|
| 21 |
+
im = im.resize((w, int(im.height * s)), Image.LANCZOS).crop((0, 0, w, min(h, int(im.height * s))))
|
| 22 |
+
buf = io.BytesIO(); im.save(buf, "JPEG", quality=80)
|
| 23 |
+
return {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": base64.b64encode(buf.getvalue()).decode()}}
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def main():
|
| 27 |
+
ap = argparse.ArgumentParser()
|
| 28 |
+
ap.add_argument("--renders", required=True, help="render output dir with one subdir per tag")
|
| 29 |
+
ap.add_argument("--crawl", required=True)
|
| 30 |
+
ap.add_argument("--pairs", default="nadir:dr,nadir:haiku")
|
| 31 |
+
ap.add_argument("--out", required=True)
|
| 32 |
+
ap.add_argument("--effort", default="medium")
|
| 33 |
+
ap.add_argument("--budget", type=float, default=3.0)
|
| 34 |
+
ap.add_argument("--chunk", type=int, default=400, help="requests per batch (3 images each, ~150 KB)")
|
| 35 |
+
args = ap.parse_args()
|
| 36 |
+
R, crawl, out = Path(args.renders), Path(args.crawl), Path(args.out)
|
| 37 |
+
out.mkdir(parents=True, exist_ok=True)
|
| 38 |
+
rng = random.Random(0)
|
| 39 |
+
reqs, meta = [], {}
|
| 40 |
+
for pair in args.pairs.split(","):
|
| 41 |
+
x, y = pair.split(":")
|
| 42 |
+
for rx in sorted((R / x).glob("*.json")):
|
| 43 |
+
pid = rx.stem
|
| 44 |
+
ry = R / y / f"{pid}.json"
|
| 45 |
+
if not ry.exists():
|
| 46 |
+
continue
|
| 47 |
+
jx, jy = json.loads(rx.read_text()), json.loads(ry.read_text())
|
| 48 |
+
if not (jx.get("ok") and jy.get("ok")) or not (R / x / f"{pid}.jpg").exists() or not (R / y / f"{pid}.jpg").exists():
|
| 49 |
+
continue
|
| 50 |
+
flip = rng.random() < 0.5
|
| 51 |
+
a, b = (y, x) if flip else (x, y)
|
| 52 |
+
cid = f"{pid}-{x}-{y}"
|
| 53 |
+
content = [{"type": "text", "text": "Original light page:"}, img_block(crawl / "pages" / pid / "shot.jpg"),
|
| 54 |
+
{"type": "text", "text": "Dark mode A:"}, img_block(R / a / f"{pid}.jpg"),
|
| 55 |
+
{"type": "text", "text": "Dark mode B:"}, img_block(R / b / f"{pid}.jpg"),
|
| 56 |
+
{"type": "text", "text": "Which dark mode is better? Scores are 1-10."}]
|
| 57 |
+
reqs.append({"custom_id": cid, "params": {"model": hg.MODEL, "max_tokens": 16000, "system": JUDGE_SYSTEM, "thinking": {"type": "adaptive"},
|
| 58 |
+
"output_config": {"effort": args.effort, "format": {"type": "json_schema", "schema": SCHEMA}},
|
| 59 |
+
"messages": [{"role": "user", "content": content}]}})
|
| 60 |
+
meta[cid] = {"pid": pid, "x": x, "y": y, "A": a, "B": b}
|
| 61 |
+
print(len(reqs), "comparisons", flush=True)
|
| 62 |
+
est = len(reqs) * 0.0015
|
| 63 |
+
if est > args.budget:
|
| 64 |
+
reqs = reqs[: int(args.budget / 0.0015)]
|
| 65 |
+
results, cost = [], 0.0
|
| 66 |
+
bids = []
|
| 67 |
+
for k in range(0, len(reqs), args.chunk): # submit every chunk first, then wait for all of them together
|
| 68 |
+
info = json.loads(hg.http("POST", "/v1/messages/batches", {"requests": reqs[k:k + args.chunk]}, timeout=900))
|
| 69 |
+
bids.append(info["id"])
|
| 70 |
+
print("submitted", info["id"], len(reqs[k:k + args.chunk]), flush=True)
|
| 71 |
+
pending = set(bids)
|
| 72 |
+
while pending:
|
| 73 |
+
time.sleep(30)
|
| 74 |
+
for bid in list(pending):
|
| 75 |
+
if json.loads(hg.http("GET", f"/v1/messages/batches/{bid}"))["processing_status"] == "ended":
|
| 76 |
+
pending.discard(bid)
|
| 77 |
+
for bid in bids:
|
| 78 |
+
raw = hg.http("GET", f"/v1/messages/batches/{bid}/results", timeout=600).decode()
|
| 79 |
+
for line in raw.splitlines():
|
| 80 |
+
if not line.strip():
|
| 81 |
+
continue
|
| 82 |
+
r = json.loads(line)
|
| 83 |
+
if r["result"]["type"] != "succeeded":
|
| 84 |
+
continue
|
| 85 |
+
m = r["result"]["message"]; u = m["usage"]
|
| 86 |
+
cost += 0.5 * (u["input_tokens"] * hg.PRICE_IN + u["output_tokens"] * hg.PRICE_OUT)
|
| 87 |
+
txt = "".join(b.get("text", "") for b in m["content"] if b.get("type") == "text")
|
| 88 |
+
try:
|
| 89 |
+
v = json.loads(txt)
|
| 90 |
+
except Exception:
|
| 91 |
+
continue
|
| 92 |
+
mt = meta[r["custom_id"]]
|
| 93 |
+
win = {"A": mt["A"], "B": mt["B"], "tie": "tie"}[v["winner"]]
|
| 94 |
+
sc = {mt["A"]: v["score_A"], mt["B"]: v["score_B"]}
|
| 95 |
+
results.append({**mt, "winner": win, "score_x": sc[mt["x"]], "score_y": sc[mt["y"]], "reason": v["reason"]})
|
| 96 |
+
summary = {}
|
| 97 |
+
for pair in args.pairs.split(","):
|
| 98 |
+
x, y = pair.split(":")
|
| 99 |
+
rs = [r for r in results if r["x"] == x and r["y"] == y]
|
| 100 |
+
if not rs:
|
| 101 |
+
continue
|
| 102 |
+
wx = sum(r["winner"] == x for r in rs); wy = sum(r["winner"] == y for r in rs); t = sum(r["winner"] == "tie" for r in rs)
|
| 103 |
+
summary[pair] = {"n": len(rs), f"{x}_wins": wx, f"{y}_wins": wy, "ties": t, f"{x}_winrate_excl_ties": round(wx / max(1, wx + wy), 3),
|
| 104 |
+
f"mean_score_{x}": round(sum(r["score_x"] for r in rs) / len(rs), 2), f"mean_score_{y}": round(sum(r["score_y"] for r in rs) / len(rs), 2)}
|
| 105 |
+
json.dump({"summary": summary, "cost": cost, "results": results}, open(out / "judge.json", "w"), indent=1)
|
| 106 |
+
print(json.dumps(summary, indent=1), f"cost ${cost:.3f}", flush=True)
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
if __name__ == "__main__":
|
| 110 |
+
main()
|
training/model.py
ADDED
|
@@ -0,0 +1,141 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""NadirNet: a relation-biased set transformer over a page's colour palette and images.
|
| 2 |
+
|
| 3 |
+
Tokens: [PAGE] + one token per (colour, role) the page uses + one token per image.
|
| 4 |
+
Attention logits get a learned per-head bias from the page's structure graph (text-on-surface,
|
| 5 |
+
surface-on-surface, border-on-surface, same-element, image-on-surface, same-colour), so the model
|
| 6 |
+
reasons about hierarchy and contrast the way a designer would. Outputs a dark colour per token
|
| 7 |
+
(OKLab, parameterised relative to the input so hue is kept unless the model chooses otherwise),
|
| 8 |
+
an alpha, and a treatment class per image.
|
| 9 |
+
"""
|
| 10 |
+
import math
|
| 11 |
+
import torch
|
| 12 |
+
import torch.nn as nn
|
| 13 |
+
import torch.nn.functional as F
|
| 14 |
+
|
| 15 |
+
FT, FI, FP = 49, 39, 8
|
| 16 |
+
N_ROLES, N_TREAT = 7, 5
|
| 17 |
+
K_REL = 11
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def fourier_lab(lab, n=6):
|
| 21 |
+
x = torch.stack([lab[..., 0], lab[..., 1] * 4, lab[..., 2] * 4], -1) # (..., 3)
|
| 22 |
+
freqs = (2.0 ** torch.arange(n, device=lab.device, dtype=lab.dtype)) * math.pi
|
| 23 |
+
ang = x[..., None] * freqs # (...,3,n)
|
| 24 |
+
return torch.cat([torch.sin(ang), torch.cos(ang)], -1).flatten(-2) # (..., 6n)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
class MLP(nn.Module):
|
| 28 |
+
def __init__(self, i, h, o, drop=0.0):
|
| 29 |
+
super().__init__()
|
| 30 |
+
self.net = nn.Sequential(nn.Linear(i, h), nn.GELU(), nn.Dropout(drop), nn.Linear(h, o))
|
| 31 |
+
|
| 32 |
+
def forward(self, x):
|
| 33 |
+
return self.net(x)
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
class ThumbCNN(nn.Module):
|
| 37 |
+
"""32x32 RGBA thumbnail -> embedding. Alpha matters: transparent logos behave differently from photos."""
|
| 38 |
+
def __init__(self, o=128):
|
| 39 |
+
super().__init__()
|
| 40 |
+
c = [5, 32, 64, 96, 128]
|
| 41 |
+
layers = []
|
| 42 |
+
for a, b in zip(c[:-1], c[1:]):
|
| 43 |
+
layers += [nn.Conv2d(a, b, 3, padding=1), nn.GELU(), nn.Conv2d(b, b, 3, padding=1, stride=2), nn.GELU()]
|
| 44 |
+
self.net = nn.Sequential(*layers)
|
| 45 |
+
self.proj = nn.Linear(128 * 2 * 2, o)
|
| 46 |
+
|
| 47 |
+
def forward(self, x): # (B, M, 4, 32, 32) in [0,1]
|
| 48 |
+
B, M = x.shape[:2]
|
| 49 |
+
x = x.reshape(B * M, 4, 32, 32)
|
| 50 |
+
a = x[:, 3:4]
|
| 51 |
+
lum = (0.2126 * x[:, 0:1] + 0.7152 * x[:, 1:2] + 0.0722 * x[:, 2:3]) * a # premultiplied luminance
|
| 52 |
+
x = torch.cat([x[:, :3] * a, a, lum], 1)
|
| 53 |
+
x = self.net(x) # (BM,128,2,2)
|
| 54 |
+
return self.proj(x.flatten(1)).reshape(B, M, -1)
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
class Block(nn.Module):
|
| 58 |
+
def __init__(self, d, h, ff, drop):
|
| 59 |
+
super().__init__()
|
| 60 |
+
self.h = h
|
| 61 |
+
self.ln1 = nn.LayerNorm(d)
|
| 62 |
+
self.qkv = nn.Linear(d, 3 * d)
|
| 63 |
+
self.proj = nn.Linear(d, d)
|
| 64 |
+
self.rel = nn.Linear(K_REL, h, bias=False)
|
| 65 |
+
nn.init.normal_(self.rel.weight, std=0.5)
|
| 66 |
+
self.ln2 = nn.LayerNorm(d)
|
| 67 |
+
self.ff = nn.Sequential(nn.Linear(d, ff * d), nn.GELU(), nn.Dropout(drop), nn.Linear(ff * d, d))
|
| 68 |
+
self.drop = nn.Dropout(drop)
|
| 69 |
+
|
| 70 |
+
def forward(self, x, rel, pad_bias):
|
| 71 |
+
B, T, D = x.shape
|
| 72 |
+
q, k, v = self.qkv(self.ln1(x)).reshape(B, T, 3, self.h, D // self.h).permute(2, 0, 3, 1, 4)
|
| 73 |
+
bias = self.rel(rel).permute(0, 3, 1, 2) + pad_bias # (B,H,T,T)
|
| 74 |
+
a = F.scaled_dot_product_attention(q, k, v, attn_mask=bias.to(q.dtype))
|
| 75 |
+
x = x + self.drop(self.proj(a.transpose(1, 2).reshape(B, T, D)))
|
| 76 |
+
return x + self.drop(self.ff(self.ln2(x)))
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
class NadirNet(nn.Module):
|
| 80 |
+
def __init__(self, d=384, layers=8, heads=6, ff=4, drop=0.1):
|
| 81 |
+
super().__init__()
|
| 82 |
+
self.cfg = dict(d=d, layers=layers, heads=heads, ff=ff, drop=drop)
|
| 83 |
+
self.tok_in = MLP(FT + N_ROLES + 36, 2 * d, d)
|
| 84 |
+
self.cnn = ThumbCNN(128)
|
| 85 |
+
self.img_in = MLP(FI + 128, 2 * d, d)
|
| 86 |
+
self.page_in = MLP(FP, d, d)
|
| 87 |
+
self.src = nn.Embedding(2, d)
|
| 88 |
+
self.typ = nn.Embedding(3, d)
|
| 89 |
+
self.blocks = nn.ModuleList([Block(d, heads, ff, drop) for _ in range(layers)])
|
| 90 |
+
self.ln = nn.LayerNorm(d)
|
| 91 |
+
self.color_head = MLP(d + 3, d, 6)
|
| 92 |
+
self.img_head = MLP(d, d // 2, N_TREAT)
|
| 93 |
+
|
| 94 |
+
def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask):
|
| 95 |
+
"""tok_f (B,N,FT) tok_role (B,N) long, tok_lab (B,N,3), img_f (B,M,FI), img_thumb (B,M,4,32,32) float 0-1,
|
| 96 |
+
page_f (B,FP), src (B,) long, rel (B,T,T,K) with T=1+N+M, mask (B,T) bool (True=real)."""
|
| 97 |
+
B, N = tok_role.shape
|
| 98 |
+
M = img_f.shape[1]
|
| 99 |
+
tok = self.tok_in(torch.cat([tok_f, F.one_hot(tok_role, N_ROLES).to(tok_f.dtype), fourier_lab(tok_lab)], -1)) + self.typ.weight[1]
|
| 100 |
+
if M > 0:
|
| 101 |
+
img = self.img_in(torch.cat([img_f, self.cnn(img_thumb)], -1)) + self.typ.weight[2]
|
| 102 |
+
else:
|
| 103 |
+
img = tok.new_zeros(B, 0, tok.shape[-1])
|
| 104 |
+
cls = self.page_in(page_f)[:, None] + self.typ.weight[0]
|
| 105 |
+
x = torch.cat([cls, tok, img], 1) + self.src(src)[:, None]
|
| 106 |
+
pad_bias = torch.zeros(mask.shape, dtype=x.dtype, device=x.device).masked_fill(~mask, -1e4)[:, None, None, :]
|
| 107 |
+
for blk in self.blocks:
|
| 108 |
+
x = blk(x, rel, pad_bias)
|
| 109 |
+
x = self.ln(x)
|
| 110 |
+
ht = x[:, 1:1 + N]
|
| 111 |
+
z = self.color_head(torch.cat([ht, tok_lab], -1))
|
| 112 |
+
L = torch.sigmoid(z[..., 0])
|
| 113 |
+
s = 1.6 * torch.sigmoid(z[..., 1])
|
| 114 |
+
r = 0.6 * torch.tanh(z[..., 2])
|
| 115 |
+
a_in, b_in = tok_lab[..., 1], tok_lab[..., 2]
|
| 116 |
+
a = s * a_in - r * b_in + 0.12 * torch.tanh(z[..., 3])
|
| 117 |
+
b = r * a_in + s * b_in + 0.12 * torch.tanh(z[..., 4])
|
| 118 |
+
lab = torch.stack([L, a, b], -1)
|
| 119 |
+
alpha_logit = z[..., 5]
|
| 120 |
+
img_logits = self.img_head(x[:, 1 + N:]) if M > 0 else x.new_zeros(B, 0, N_TREAT)
|
| 121 |
+
return lab, alpha_logit, img_logits
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
# ---------- colour maths (differentiable) ----------
|
| 125 |
+
OKLAB_TO_LMS = torch.tensor([[1.0, 0.3963377774, 0.2158037573], [1.0, -0.1055613458, -0.0638541728], [1.0, -0.0894841775, -1.2914855480]])
|
| 126 |
+
LMS_TO_RGB = torch.tensor([[4.0767416621, -3.3077115913, 0.2309699292], [-1.2684380046, 2.6097574011, -0.3413193965], [-0.0041960863, -0.7034186147, 1.7076147010]])
|
| 127 |
+
|
| 128 |
+
|
| 129 |
+
def oklab_to_linrgb(lab):
|
| 130 |
+
lms = lab @ OKLAB_TO_LMS.to(lab).T
|
| 131 |
+
return (lms ** 3) @ LMS_TO_RGB.to(lab).T
|
| 132 |
+
|
| 133 |
+
|
| 134 |
+
def luminance(lab):
|
| 135 |
+
rgb = oklab_to_linrgb(lab).clamp(0, 1)
|
| 136 |
+
return rgb @ torch.tensor([0.2126, 0.7152, 0.0722]).to(lab)
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
def contrast_ratio(lab1, lab2):
|
| 140 |
+
l1, l2 = luminance(lab1), luminance(lab2)
|
| 141 |
+
return (torch.maximum(l1, l2) + 0.05) / (torch.minimum(l1, l2) + 0.05)
|
training/predict.py
ADDED
|
@@ -0,0 +1,93 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Run a trained NadirNet over held-out pages and write render jobs (Python port of nadir-infer.js decode)."""
|
| 3 |
+
import argparse, gzip, json, math, pickle
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
import numpy as np
|
| 6 |
+
import torch
|
| 7 |
+
|
| 8 |
+
from model import NadirNet
|
| 9 |
+
from train import collate, prep
|
| 10 |
+
from mkjobs import job_from
|
| 11 |
+
|
| 12 |
+
TREAT = ["keep", "dim", "invert", "light", "plate"]
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
from safety import oklab, lab_to_rgb, contrast, safe_treatment, enforce_contrast, is_media_keep, hue_guard, keep_invisible
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
def decode(L, lab, alpha_logit, img_logits, rels):
|
| 19 |
+
toks = L["tokens"][:256]
|
| 20 |
+
out, keep = [], set()
|
| 21 |
+
for i, t in enumerate(toks):
|
| 22 |
+
ain = min(1 - 1e-3, max(1e-3, t["rgba"][3]))
|
| 23 |
+
alpha = 1 / (1 + math.exp(-(alpha_logit[i] + math.log(ain / (1 - ain)))))
|
| 24 |
+
a = 1.0 if (t["rgba"][3] >= 0.999 and alpha > 0.97) else alpha
|
| 25 |
+
out.append({"rgb": lab_to_rgb(*lab[i]), "a": a, "L": float(lab[i][0])})
|
| 26 |
+
if is_media_keep(t): # shown as-is: later safety passes see its real colour
|
| 27 |
+
keep.add(i)
|
| 28 |
+
rgb = [int(x) for x in t["rgba"][:3]]
|
| 29 |
+
out[i] = {"rgb": rgb, "a": float(t["rgba"][3]), "L": oklab(*rgb)[0]}
|
| 30 |
+
else:
|
| 31 |
+
hue_guard(t, out[i])
|
| 32 |
+
for i, t in enumerate(toks):
|
| 33 |
+
if t["f"][24] > 0.5 and t["role"] == 0 and out[i]["L"] > 0.35:
|
| 34 |
+
o = oklab(*out[i]["rgb"]); out[i]["rgb"] = lab_to_rgb(0.2, o[1] * 0.3, o[2] * 0.3); out[i]["L"] = 0.2
|
| 35 |
+
keep_invisible(toks, out, rels, keep)
|
| 36 |
+
enforce_contrast(toks, out, rels, keep)
|
| 37 |
+
def hx(o):
|
| 38 |
+
h = "#%02x%02x%02x" % tuple(o["rgb"])
|
| 39 |
+
return h + ("%02x" % round(o["a"] * 255) if o["a"] < 0.999 else "")
|
| 40 |
+
colors = [None if i in keep else hx(o) for i, o in enumerate(out)]
|
| 41 |
+
treat = {}
|
| 42 |
+
for j, im in enumerate(L["images"][:32]):
|
| 43 |
+
t = TREAT[int(np.argmax(img_logits[j]))] if j < len(img_logits) else "keep"
|
| 44 |
+
u = im.get("underTok", -1)
|
| 45 |
+
treat[str(j)] = safe_treatment(t, im, out[u]["L"] if u is not None and 0 <= u < len(out) else None)
|
| 46 |
+
return colors, sorted(keep), treat
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def main():
|
| 50 |
+
ap = argparse.ArgumentParser()
|
| 51 |
+
ap.add_argument("--ckpt", required=True)
|
| 52 |
+
ap.add_argument("--data", required=True)
|
| 53 |
+
ap.add_argument("--crawl", required=True)
|
| 54 |
+
ap.add_argument("--out", required=True)
|
| 55 |
+
ap.add_argument("--n", type=int, default=150)
|
| 56 |
+
ap.add_argument("--src", type=int, default=0)
|
| 57 |
+
ap.add_argument("--tag", default="nadir")
|
| 58 |
+
ap.add_argument("--also", default="haiku,dr", help="extra render tags to emit for the same pages")
|
| 59 |
+
ap.add_argument("--haiku", default="")
|
| 60 |
+
args = ap.parse_args()
|
| 61 |
+
ck = torch.load(args.ckpt, map_location="cpu")
|
| 62 |
+
model = NadirNet(**ck["cfg"]); model.load_state_dict(ck["model"]); model.eval()
|
| 63 |
+
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
| 64 |
+
model.to(dev)
|
| 65 |
+
data = pickle.load(open(args.data, "rb"))
|
| 66 |
+
val = [s for s in data["val"] if s["src"] == args.src][: args.n]
|
| 67 |
+
with open(args.out, "w") as f:
|
| 68 |
+
for s in val:
|
| 69 |
+
s = prep(s)
|
| 70 |
+
L = json.load(gzip.open(Path(args.crawl) / "pages" / s["id"] / "light.json.gz"))
|
| 71 |
+
b = collate([s], dev)
|
| 72 |
+
with torch.no_grad():
|
| 73 |
+
lab, al, il = model(b["tok_f"], b["role"], b["lab"], b["img_f"], b["thumbs"], b["page"], torch.zeros_like(b["src"]), b["rel"], b["mask"])
|
| 74 |
+
n = len(s["role"])
|
| 75 |
+
colors, keep, treat = decode(L, lab[0, :n].float().cpu().numpy().tolist(), al[0, :n].float().cpu().numpy().tolist(),
|
| 76 |
+
il[0].float().cpu().numpy().tolist(), s["rel"])
|
| 77 |
+
colors = colors + [None] * (len(L["tokens"]) - len(colors))
|
| 78 |
+
j = job_from(L, s["id"], args.tag, colors, treat)
|
| 79 |
+
j["keep"] = sorted(set(j["keep"]) | set(keep))
|
| 80 |
+
f.write(json.dumps(j) + "\n")
|
| 81 |
+
if "dr" in args.also:
|
| 82 |
+
f.write(json.dumps(job_from(L, s["id"], "dr", [None] * len(L["tokens"]), {}, baseline=True)) + "\n")
|
| 83 |
+
if "haiku" in args.also and args.haiku:
|
| 84 |
+
hp = Path(args.haiku) / f"{s['id']}.json"
|
| 85 |
+
if hp.exists():
|
| 86 |
+
h = json.loads(hp.read_text())
|
| 87 |
+
outs = [h["colors"].get(str(i)) for i in range(len(L["tokens"]))]
|
| 88 |
+
f.write(json.dumps(job_from(L, s["id"], "haiku", outs, h["images"])) + "\n")
|
| 89 |
+
print("wrote", len(val), "pages")
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
if __name__ == "__main__":
|
| 93 |
+
main()
|
training/quantize.py
ADDED
|
@@ -0,0 +1,330 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Calibrated weight-only quantization of NadirNet into ONNX MatMulNBits (8-bit and 4-bit).
|
| 3 |
+
|
| 4 |
+
Layer-sequential GPTQ: each layer's Hessian comes from real crawled pages flowing through the already
|
| 5 |
+
quantized earlier layers. Each group of 32 weights gets asymmetric uint codes whose clipping range is
|
| 6 |
+
chosen by an MSE search. Codes are packed in ORT's MatMulNBits layout, and every variant is checked
|
| 7 |
+
three ways: against a PyTorch fake-quant copy (packing parity), against fp32 (fidelity) and against
|
| 8 |
+
the Haiku targets (quality)."""
|
| 9 |
+
import argparse, copy, json, os, pickle, time
|
| 10 |
+
import numpy as np
|
| 11 |
+
import torch
|
| 12 |
+
import torch.nn as nn
|
| 13 |
+
|
| 14 |
+
from model import NadirNet
|
| 15 |
+
from train import prep
|
| 16 |
+
from export import feeds_for, fidelity, bench, NAMES_IN
|
| 17 |
+
|
| 18 |
+
GS = 32
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
class Stop(Exception):
|
| 22 |
+
pass
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def group_params(w, bits, maxshrink, grid=100):
|
| 26 |
+
"""Per-row scale/zero-point for one group (rows, GS), minimizing |q(w)-w|^2.4 over clipping ranges."""
|
| 27 |
+
maxq = 2 ** bits - 1
|
| 28 |
+
xmin = w.min(1).values.clamp(max=0)
|
| 29 |
+
xmax = w.max(1).values.clamp(min=0)
|
| 30 |
+
best = torch.full_like(xmin, float("inf"))
|
| 31 |
+
bs, bz = torch.ones_like(xmin), torch.zeros_like(xmin)
|
| 32 |
+
for i in range(int(maxshrink * grid) + 1):
|
| 33 |
+
p = 1 - i / grid
|
| 34 |
+
lo, hi = p * xmin, p * xmax
|
| 35 |
+
s = ((hi - lo) / maxq).clamp(min=1e-10)
|
| 36 |
+
z = torch.round(-lo / s).clamp(0, maxq)
|
| 37 |
+
q = (torch.clamp(torch.round(w / s[:, None]) + z[:, None], 0, maxq) - z[:, None]) * s[:, None]
|
| 38 |
+
err = (q - w).abs().pow(2.4).sum(1)
|
| 39 |
+
m = err < best
|
| 40 |
+
best[m], bs[m], bz[m] = err[m], s[m], z[m]
|
| 41 |
+
return bs, bz
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def gptq(W, H, bits, percdamp=0.01, blocksize=128, gs=GS):
|
| 45 |
+
"""W (out, Kpad) with Kpad % gs == 0, H (Kpad, Kpad); gs=None quantizes per output channel (one group per row).
|
| 46 |
+
Returns dequantized W, codes, scales, zero points."""
|
| 47 |
+
maxq = 2 ** bits - 1
|
| 48 |
+
maxshrink = 0.5 if bits == 4 else 0.2
|
| 49 |
+
W = W.clone().double()
|
| 50 |
+
H = H.clone().double()
|
| 51 |
+
rows, cols = W.shape
|
| 52 |
+
dead = torch.diag(H) == 0
|
| 53 |
+
H[dead, dead] = 1
|
| 54 |
+
W[:, dead] = 0
|
| 55 |
+
H += percdamp * torch.mean(torch.diag(H)) * torch.eye(cols, dtype=H.dtype)
|
| 56 |
+
Hinv = torch.linalg.cholesky(torch.cholesky_inverse(torch.linalg.cholesky(H)), upper=True)
|
| 57 |
+
Q = torch.zeros_like(W)
|
| 58 |
+
codes = torch.zeros(rows, cols, dtype=torch.uint8)
|
| 59 |
+
if gs is None: # per channel: one range per output row, chosen on the whole row up front
|
| 60 |
+
s, z = group_params(W, bits, maxshrink)
|
| 61 |
+
S, Z = s[:, None].clone(), z[:, None].clone()
|
| 62 |
+
else:
|
| 63 |
+
S = torch.zeros(rows, cols // gs, dtype=torch.float64)
|
| 64 |
+
Z = torch.zeros(rows, cols // gs, dtype=torch.float64)
|
| 65 |
+
for i1 in range(0, cols, blocksize):
|
| 66 |
+
i2 = min(i1 + blocksize, cols)
|
| 67 |
+
W1 = W[:, i1:i2].clone()
|
| 68 |
+
E1 = torch.zeros_like(W1)
|
| 69 |
+
H1 = Hinv[i1:i2, i1:i2]
|
| 70 |
+
for i in range(i2 - i1):
|
| 71 |
+
c = i1 + i
|
| 72 |
+
if gs is not None and c % gs == 0: # groups never straddle blocks (blocksize % gs == 0)
|
| 73 |
+
s, z = group_params(W1[:, i:i + gs], bits, maxshrink)
|
| 74 |
+
S[:, c // gs], Z[:, c // gs] = s, z
|
| 75 |
+
w = W1[:, i]
|
| 76 |
+
qc = torch.clamp(torch.round(w / s) + z, 0, maxq)
|
| 77 |
+
q = (qc - z) * s
|
| 78 |
+
codes[:, c] = qc.to(torch.uint8)
|
| 79 |
+
Q[:, c] = q
|
| 80 |
+
err = (w - q) / H1[i, i]
|
| 81 |
+
W1[:, i:] -= err[:, None] * H1[i, i:][None, :]
|
| 82 |
+
E1[:, i] = err
|
| 83 |
+
W[:, i2:] -= E1 @ Hinv[i1:i2, i2:]
|
| 84 |
+
return Q.float(), codes, S.float(), Z.to(torch.uint8)
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
def pack(codes, S, Z, bits):
|
| 88 |
+
"""MatMulNBits layout: B (N, k_blocks, blob), scales (N, k_blocks), zero points (N, ceil(k_blocks*bits/8))."""
|
| 89 |
+
c = codes.numpy()
|
| 90 |
+
N, Kp = c.shape
|
| 91 |
+
kb = Kp // GS
|
| 92 |
+
z = Z.numpy().astype(np.uint8)
|
| 93 |
+
if bits == 8:
|
| 94 |
+
B = c.reshape(N, kb, GS).astype(np.uint8)
|
| 95 |
+
zp = z
|
| 96 |
+
else:
|
| 97 |
+
c = c.reshape(N, kb, GS)
|
| 98 |
+
B = (c[..., 0::2] | (c[..., 1::2] << 4)).astype(np.uint8)
|
| 99 |
+
if kb % 2:
|
| 100 |
+
z = np.concatenate([z, np.full((N, 1), 8, np.uint8)], 1)
|
| 101 |
+
zp = (z[:, 0::2] | (z[:, 1::2] << 4)).astype(np.uint8)
|
| 102 |
+
return B, S.numpy().astype(np.float32), zp
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
def linear_order(model, sample):
|
| 106 |
+
"""Linear modules in the order the forward pass calls them."""
|
| 107 |
+
order, hooks = [], []
|
| 108 |
+
for name, mod in model.named_modules():
|
| 109 |
+
if isinstance(mod, nn.Linear):
|
| 110 |
+
hooks.append(mod.register_forward_pre_hook(lambda m, a, n=name: order.append(n) if n not in order else None))
|
| 111 |
+
with torch.no_grad():
|
| 112 |
+
model(*[torch.from_numpy(sample[k]) for k in NAMES_IN])
|
| 113 |
+
for h in hooks:
|
| 114 |
+
h.remove()
|
| 115 |
+
return order
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
def onnx_matmuls(model, onnx_model):
|
| 119 |
+
"""Map each Linear to the MatMul node whose initializer is its transposed weight (Gemm exports stay fp32)."""
|
| 120 |
+
from onnx import numpy_helper
|
| 121 |
+
inits = {i.name: i for i in onnx_model.graph.initializer}
|
| 122 |
+
by_shape = {}
|
| 123 |
+
for node in onnx_model.graph.node:
|
| 124 |
+
if node.op_type == "MatMul" and node.input[1] in inits:
|
| 125 |
+
arr = numpy_helper.to_array(inits[node.input[1]])
|
| 126 |
+
by_shape.setdefault(arr.shape, []).append((node, arr))
|
| 127 |
+
out = {}
|
| 128 |
+
for name, mod in model.named_modules():
|
| 129 |
+
if isinstance(mod, nn.Linear):
|
| 130 |
+
wt = mod.weight.detach().numpy().T
|
| 131 |
+
for node, arr in by_shape.get(wt.shape, []):
|
| 132 |
+
if np.array_equal(arr, wt):
|
| 133 |
+
out[name] = node
|
| 134 |
+
break
|
| 135 |
+
return out
|
| 136 |
+
|
| 137 |
+
|
| 138 |
+
def hessian(model, name, calib):
|
| 139 |
+
"""2/n * X^T X of the inputs to one Linear, over all real token rows of the calibration pages."""
|
| 140 |
+
mod = dict(model.named_modules())[name]
|
| 141 |
+
K = mod.in_features
|
| 142 |
+
H = torch.zeros(K, K, dtype=torch.float64)
|
| 143 |
+
n = [0]
|
| 144 |
+
|
| 145 |
+
def hook(m, a):
|
| 146 |
+
x = a[0].detach().reshape(-1, K).double()
|
| 147 |
+
H.add_(x.T @ x)
|
| 148 |
+
n[0] += x.shape[0]
|
| 149 |
+
raise Stop
|
| 150 |
+
|
| 151 |
+
h = mod.register_forward_pre_hook(hook)
|
| 152 |
+
with torch.no_grad():
|
| 153 |
+
for f in calib:
|
| 154 |
+
try:
|
| 155 |
+
model(*[torch.from_numpy(f[k]) for k in NAMES_IN])
|
| 156 |
+
except Stop:
|
| 157 |
+
pass
|
| 158 |
+
h.remove()
|
| 159 |
+
return H * (2.0 / max(1, n[0]))
|
| 160 |
+
|
| 161 |
+
|
| 162 |
+
def quantize_model(model, calib, bits_for, order):
|
| 163 |
+
"""Sequential GPTQ over `order`; returns the fake-quant model and packed tensors per layer."""
|
| 164 |
+
qm = copy.deepcopy(model).eval()
|
| 165 |
+
mods = dict(qm.named_modules())
|
| 166 |
+
packed = {}
|
| 167 |
+
for name in order:
|
| 168 |
+
bits = bits_for.get(name)
|
| 169 |
+
if not bits:
|
| 170 |
+
continue
|
| 171 |
+
mod = mods[name]
|
| 172 |
+
K = mod.in_features
|
| 173 |
+
H = hessian(qm, name, calib)
|
| 174 |
+
t = time.time()
|
| 175 |
+
# blocks of 32 along K (zero-padded to a multiple of 32)
|
| 176 |
+
Kp = -(-K // GS) * GS
|
| 177 |
+
Hp = torch.zeros(Kp, Kp, dtype=torch.float64); Hp[:K, :K] = H
|
| 178 |
+
Wp = torch.zeros(mod.out_features, Kp); Wp[:, :K] = mod.weight.detach()
|
| 179 |
+
Q, codes, S, Z = gptq(Wp, Hp, bits)
|
| 180 |
+
if bits == 8:
|
| 181 |
+
# 8-bit: uint8 codes as a (K/32, 32, N) tensor with per-block scale and zero point, dequantized in-graph
|
| 182 |
+
N, kb = mod.out_features, Kp // GS
|
| 183 |
+
packed[name] = ("dq8", K, codes.numpy().T.reshape(kb, GS, N).copy(), S.numpy().T[:, None, :].astype(np.float32).copy(),
|
| 184 |
+
Z.numpy().T[:, None, :].astype(np.uint8).copy())
|
| 185 |
+
else:
|
| 186 |
+
# 4-bit: packed as MatMulNBits
|
| 187 |
+
packed[name] = ("nbits", K) + pack(codes, S, Z, bits)
|
| 188 |
+
with torch.no_grad():
|
| 189 |
+
mod.weight.copy_(Q[:, :K])
|
| 190 |
+
print(f" {name} {tuple(mod.weight.shape)} {bits}-bit {time.time() - t:.1f}s", flush=True)
|
| 191 |
+
return qm, packed
|
| 192 |
+
|
| 193 |
+
|
| 194 |
+
def build_onnx(fp32_path, mapping, packed, out_path, accuracy_level=0):
|
| 195 |
+
import onnx
|
| 196 |
+
from onnx import helper, numpy_helper
|
| 197 |
+
m = onnx.load(fp32_path)
|
| 198 |
+
g = m.graph
|
| 199 |
+
drop_inits = set()
|
| 200 |
+
for name, (kind, K, B, S, zp) in packed.items():
|
| 201 |
+
node = mapping[name]
|
| 202 |
+
wname = node.input[1]
|
| 203 |
+
idx = list(g.node).index(node)
|
| 204 |
+
if kind == "dq8":
|
| 205 |
+
# uint8 codes dequantized with plain Cast/Sub/Mul: every onnxruntime constant-folds this into exact fp32
|
| 206 |
+
# weights at load. (DequantizeLinear->MatMul gets rewritten by onnxruntime into an activation-quantized
|
| 207 |
+
# kernel, which costs accuracy.)
|
| 208 |
+
kb, gs, N = B.shape
|
| 209 |
+
inits = [numpy_helper.from_array(B, wname + "_q8"), numpy_helper.from_array(zp, wname + "_zero_point"),
|
| 210 |
+
numpy_helper.from_array(S, wname + "_scale"), numpy_helper.from_array(np.array([kb * gs, N], np.int64), wname + "_kn")]
|
| 211 |
+
g.initializer.extend(inits)
|
| 212 |
+
base = (node.name or wname) + "_deq8"
|
| 213 |
+
F = onnx.TensorProto.FLOAT
|
| 214 |
+
deq = [helper.make_node("Cast", [inits[0].name], [wname + "_f"], name=base + "_cast", to=F),
|
| 215 |
+
helper.make_node("Cast", [inits[1].name], [wname + "_zf"], name=base + "_castz", to=F),
|
| 216 |
+
helper.make_node("Sub", [wname + "_f", wname + "_zf"], [wname + "_c"], name=base + "_sub"),
|
| 217 |
+
helper.make_node("Mul", [wname + "_c", inits[2].name], [wname + "_d3"], name=base + "_mul"),
|
| 218 |
+
helper.make_node("Reshape", [wname + "_d3", inits[3].name], [wname + ("_dq" if kb * gs == K else "_dp")], name=base + "_reshape")]
|
| 219 |
+
if kb * gs != K: # drop the zero padding rows
|
| 220 |
+
lims = [numpy_helper.from_array(np.array(v, np.int64), wname + n) for v, n in (([0], "_s0"), ([K], "_sK"), ([0], "_ax"))]
|
| 221 |
+
g.initializer.extend(lims)
|
| 222 |
+
deq.append(helper.make_node("Slice", [wname + "_dp"] + [i.name for i in lims], [wname + "_dq"], name=base + "_slice"))
|
| 223 |
+
g.node[idx].input[1] = wname + "_dq" # `node` belongs to another copy of the graph: edit this one
|
| 224 |
+
for k, dn in enumerate(deq):
|
| 225 |
+
g.node.insert(idx + k, dn)
|
| 226 |
+
else:
|
| 227 |
+
N = B.shape[0]
|
| 228 |
+
inits = [numpy_helper.from_array(B, wname + "_Q4"), numpy_helper.from_array(S, wname + "_scales"),
|
| 229 |
+
numpy_helper.from_array(zp, wname + "_zero_points")]
|
| 230 |
+
g.initializer.extend(inits)
|
| 231 |
+
kw = dict(K=K, N=N, bits=4, block_size=GS)
|
| 232 |
+
if accuracy_level:
|
| 233 |
+
kw["accuracy_level"] = accuracy_level
|
| 234 |
+
new = helper.make_node("MatMulNBits", [node.input[0]] + [i.name for i in inits], list(node.output),
|
| 235 |
+
name=(node.name or wname) + "_Q4", domain="com.microsoft", **kw)
|
| 236 |
+
g.node.remove(node)
|
| 237 |
+
g.node.insert(idx, new)
|
| 238 |
+
drop_inits.add(wname)
|
| 239 |
+
keep = [i for i in g.initializer if i.name not in drop_inits]
|
| 240 |
+
del g.initializer[:]
|
| 241 |
+
g.initializer.extend(keep)
|
| 242 |
+
if any(n.domain == "com.microsoft" for n in g.node) and not any(o.domain == "com.microsoft" for o in m.opset_import):
|
| 243 |
+
m.opset_import.append(helper.make_opsetid("com.microsoft", 1))
|
| 244 |
+
onnx.checker.check_model(m)
|
| 245 |
+
onnx.save(m, out_path)
|
| 246 |
+
|
| 247 |
+
|
| 248 |
+
def quality(sess, samples):
|
| 249 |
+
"""dE vs the training targets (Haiku for src 0, native dark for src 1) and image-treatment accuracy."""
|
| 250 |
+
agg = {0: [0.0, 0.0], 1: [0.0, 0.0]}
|
| 251 |
+
ok = n = 0
|
| 252 |
+
for s in samples:
|
| 253 |
+
f = feeds_for(s)
|
| 254 |
+
lab, _, imgl = sess.run(None, f)
|
| 255 |
+
k = len(s["role"])
|
| 256 |
+
m = s["tmask"][:k]
|
| 257 |
+
if m.any():
|
| 258 |
+
d = np.linalg.norm(lab[0, :k] - s["lab_t"][:k], axis=-1)
|
| 259 |
+
w = s["w"][:k] * m
|
| 260 |
+
agg[int(s["src"])][0] += float((d * w).sum() / max(w.sum(), 1e-9)); agg[int(s["src"])][1] += 1
|
| 261 |
+
mi = len(s["IF"])
|
| 262 |
+
if mi:
|
| 263 |
+
t = np.asarray(s["imgt"][:mi])
|
| 264 |
+
v = t >= 0
|
| 265 |
+
ok += int((imgl[0, :mi].argmax(-1)[v] == t[v]).sum()); n += int(v.sum())
|
| 266 |
+
return {"haiku_dE": agg[0][0] / max(1, agg[0][1]), "native_dE": agg[1][0] / max(1, agg[1][1]), "img_acc": ok / max(1, n)}
|
| 267 |
+
|
| 268 |
+
|
| 269 |
+
def main():
|
| 270 |
+
ap = argparse.ArgumentParser()
|
| 271 |
+
ap.add_argument("--ckpt", required=True)
|
| 272 |
+
ap.add_argument("--data", required=True)
|
| 273 |
+
ap.add_argument("--fp32", required=True, help="fp32 ONNX from export.py")
|
| 274 |
+
ap.add_argument("--out", required=True)
|
| 275 |
+
ap.add_argument("--calib", type=int, default=256)
|
| 276 |
+
ap.add_argument("--eval", type=int, default=400)
|
| 277 |
+
ap.add_argument("--threads", type=int, default=8)
|
| 278 |
+
args = ap.parse_args()
|
| 279 |
+
torch.set_num_threads(args.threads)
|
| 280 |
+
os.makedirs(args.out, exist_ok=True)
|
| 281 |
+
import onnx, onnxruntime as ort
|
| 282 |
+
ck = torch.load(args.ckpt, map_location="cpu")
|
| 283 |
+
model = NadirNet(**ck["cfg"]); model.load_state_dict(ck["model"]); model.eval()
|
| 284 |
+
data = pickle.load(open(args.data, "rb"))
|
| 285 |
+
hv = [s for s in data["val"] if s["src"] == 0]; nv = [s for s in data["val"] if s["src"] == 1]
|
| 286 |
+
val = [prep(s) for s in hv[: args.eval * 3 // 4] + nv[: args.eval - min(len(hv), args.eval * 3 // 4)]] # both sources
|
| 287 |
+
rng = np.random.default_rng(0)
|
| 288 |
+
pick = rng.choice(len(data["train"]), size=min(args.calib, len(data["train"])), replace=False)
|
| 289 |
+
calib = [feeds_for(prep(data["train"][i])) for i in pick]
|
| 290 |
+
order = linear_order(model, calib[0])
|
| 291 |
+
mapping = onnx_matmuls(model, onnx.load(args.fp32))
|
| 292 |
+
quantizable = [n for n in order if n in mapping and not n.endswith(".rel")]
|
| 293 |
+
print(f"{len(quantizable)}/{len(order)} linears map to MatMul; fp32 kept: {[n for n in order if n not in quantizable]}", flush=True)
|
| 294 |
+
edges = [n for n in quantizable if n.startswith(("tok_in.", "img_in.", "color_head.", "img_head."))]
|
| 295 |
+
plans = {"int8": {n: 8 for n in quantizable},
|
| 296 |
+
"int4": {n: (8 if n in edges else 4) for n in quantizable}}
|
| 297 |
+
so = ort.SessionOptions(); so.intra_op_num_threads = 4
|
| 298 |
+
ref = ort.InferenceSession(args.fp32, so, providers=["CPUExecutionProvider"])
|
| 299 |
+
report = {"format": {"8-bit layers": f"uint8 blocks of {GS} with fp32 scale + uint8 zero point, dequantized in-graph (Cast/Sub/Mul, folded at load)", "4-bit layers": f"uint4 blocks of {GS}, MatMulNBits"},
|
| 300 |
+
"group_size": GS, "calib_pages": len(calib), "eval_pages": len(val), "fp32_layers": [n for n in order if n not in quantizable],
|
| 301 |
+
"int4_8bit_edge_layers": edges, "variants": {}}
|
| 302 |
+
r = quality(ref, val); r["size_mb"] = round(os.path.getsize(args.fp32) / 1e6, 2); r["cpu_ms_median"] = round(bench(ref, val), 2)
|
| 303 |
+
report["variants"]["fp32"] = r
|
| 304 |
+
print("fp32", r, flush=True)
|
| 305 |
+
for vname, bits_for in plans.items():
|
| 306 |
+
print(f"== {vname}", flush=True)
|
| 307 |
+
qm, packed = quantize_model(model, calib, bits_for, quantizable)
|
| 308 |
+
path = os.path.join(args.out, f"nadir-{vname}.onnx")
|
| 309 |
+
build_onnx(args.fp32, mapping, packed, path)
|
| 310 |
+
sess = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
|
| 311 |
+
# packing parity: ORT MatMulNBits graph vs the PyTorch fake-quant model
|
| 312 |
+
par = 0.0
|
| 313 |
+
with torch.no_grad():
|
| 314 |
+
for s in val[:8]:
|
| 315 |
+
f = feeds_for(s)
|
| 316 |
+
tq = qm(*[torch.from_numpy(f[k]) for k in NAMES_IN])[0].numpy()
|
| 317 |
+
par = max(par, float(np.abs(tq - sess.run(None, f)[0]).max()))
|
| 318 |
+
r = fidelity(sess, ref, val)
|
| 319 |
+
r.update(quality(sess, val))
|
| 320 |
+
r["pack_parity_max_abs"] = par
|
| 321 |
+
r["size_mb"] = round(os.path.getsize(path) / 1e6, 2)
|
| 322 |
+
r["cpu_ms_median"] = round(bench(sess, val), 2)
|
| 323 |
+
report["variants"][vname] = r
|
| 324 |
+
print(vname, r, flush=True)
|
| 325 |
+
torch.save({"cfg": ck["cfg"], "model": qm.state_dict()}, os.path.join(args.out, f"fakequant-{vname}.pt"))
|
| 326 |
+
json.dump(report, open(os.path.join(args.out, "quant_report.json"), "w"), indent=1)
|
| 327 |
+
|
| 328 |
+
|
| 329 |
+
if __name__ == "__main__":
|
| 330 |
+
main()
|
training/safety.py
ADDED
|
@@ -0,0 +1,188 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Deterministic safety rules shared by the dataset builder (target repair) and the decoder (runtime).
|
| 2 |
+
Mirrors nadir-infer.js / nadir-map.js."""
|
| 3 |
+
import math
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
def lin(u):
|
| 7 |
+
u = u / 255.0
|
| 8 |
+
return u / 12.92 if u <= 0.04045 else ((u + 0.055) / 1.055) ** 2.4
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def oklab(r, g, b):
|
| 12 |
+
r, g, b = lin(r), lin(g), lin(b)
|
| 13 |
+
l = (0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b) ** (1 / 3)
|
| 14 |
+
m = (0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b) ** (1 / 3)
|
| 15 |
+
s = (0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b) ** (1 / 3)
|
| 16 |
+
return [0.2104542553 * l + 0.7936177850 * m - 0.0040720468 * s, 1.9779984951 * l - 2.4285922050 * m + 0.4505937099 * s,
|
| 17 |
+
0.0259040371 * l + 0.7827717662 * m - 0.8086757660 * s]
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def lab_to_rgb(L, a, b):
|
| 21 |
+
l_, m_, s_ = L + 0.3963377774 * a + 0.2158037573 * b, L - 0.1055613458 * a - 0.0638541728 * b, L - 0.0894841775 * a - 1.2914855480 * b
|
| 22 |
+
l, m, s = l_ ** 3, m_ ** 3, s_ ** 3
|
| 23 |
+
def enc(x):
|
| 24 |
+
x = max(0.0, min(1.0, x))
|
| 25 |
+
return round(255 * (12.92 * x if x <= 0.0031308 else 1.055 * x ** (1 / 2.4) - 0.055))
|
| 26 |
+
return [enc(4.0767416621 * l - 3.3077115913 * m + 0.2309699292 * s), enc(-1.2684380046 * l + 2.6097574011 * m - 0.3413193965 * s),
|
| 27 |
+
enc(-0.0041960863 * l - 0.7034186147 * m + 1.7076147010 * s)]
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def rel_lum(c):
|
| 31 |
+
return 0.2126 * lin(c[0]) + 0.7152 * lin(c[1]) + 0.0722 * lin(c[2])
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def contrast(a, b):
|
| 35 |
+
x, y = rel_lum(a), rel_lum(b)
|
| 36 |
+
return (max(x, y) + 0.05) / (min(x, y) + 0.05)
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def adaptable(im, s, surf_L=None):
|
| 40 |
+
"""A logo on its own white rectangle: adapt it instead of leaving a white box. Only images the page marks as its
|
| 41 |
+
logo (logo/brand/wordmark naming, or the link home); product shots, photos and other images are never touched.
|
| 42 |
+
Without pixel stats a logo still tries (the adaptation itself checks for the white background)."""
|
| 43 |
+
if not (im.get("logo") or (im.get("homeLink") and im.get("header"))) or im["h"] > 200 or (surf_L is not None and surf_L >= 0.4):
|
| 44 |
+
return False
|
| 45 |
+
if not s:
|
| 46 |
+
return im["w"] * im["h"] <= 250000
|
| 47 |
+
return s["transp"] < 0.05 and s.get("ringA", 0) >= 0.95 and s.get("ringL", 0) >= 0.85 and s.get("white", 0) >= 0.3 and s["colors"] < 0.2
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def is_light_backdrop(im, s):
|
| 51 |
+
"""Opaque, light, flat (almost no lightness variation, so never a photo; SVG backgrounds can't be photos, so a
|
| 52 |
+
smooth gradient qualifies too) and big enough to sit behind content. Logos drawn as backgrounds are never cleared."""
|
| 53 |
+
if im.get("kind") != "bg" or im.get("logo") or im.get("homeLink") or not s or s["transp"] >= 0.05 or s["meanL"] <= 0.8 or im["w"] * im["h"] < 40000:
|
| 54 |
+
return False
|
| 55 |
+
svg = im.get("src", "").split("?")[0].lower().endswith(".svg")
|
| 56 |
+
return s["stdL"] < (0.15 if svg else 0.07)
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def safe_treatment(t, im, surf_L=None):
|
| 60 |
+
"""'adapt:<fallback>' for logos on a white rectangle (fallback used when there is nothing to adapt)."""
|
| 61 |
+
base = _base_treatment(t, im, surf_L)
|
| 62 |
+
return "adapt:" + base if base != "clear" and adaptable(im, im.get("stats"), surf_L) else base
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def _base_treatment(t, im, surf_L=None):
|
| 66 |
+
s = im.get("stats")
|
| 67 |
+
# a light, flat CSS background image (texture, pale wash, pattern) sits behind content whose text is now light:
|
| 68 |
+
# clear it so the dark surface designed for that element shows through. Photos and detailed images stay.
|
| 69 |
+
if is_light_backdrop(im, s):
|
| 70 |
+
return "clear"
|
| 71 |
+
# a mostly dark graphic on transparency would vanish on a dark surface: hue-preserving inversion
|
| 72 |
+
if (not t or t == "keep") and s and s["transp"] > 0.15 and s["meanL"] < 0.35 and im["w"] * im["h"] < 200000 and (surf_L is None or surf_L < 0.4):
|
| 73 |
+
return "invert"
|
| 74 |
+
if not t or t == "keep":
|
| 75 |
+
return "keep"
|
| 76 |
+
photo = (s["transp"] < 0.02 and s["colors"] > 0.18 and s["stdL"] > 0.12) if s else (im["w"] * im["h"] > 60000)
|
| 77 |
+
if photo and t != "dim":
|
| 78 |
+
return "keep"
|
| 79 |
+
if t == "dim" and im["w"] * im["h"] < 20000:
|
| 80 |
+
return "keep"
|
| 81 |
+
# a transparent graphic that is already light reads fine on a dark surface: leave it alone
|
| 82 |
+
if t in ("plate", "light", "invert") and s and s["transp"] > 0.15 and s["meanL"] >= 0.55 and s.get("black", 0) < 0.15 and (surf_L is None or surf_L < 0.4):
|
| 83 |
+
return "keep"
|
| 84 |
+
# no light boxes on a dark page: a dark logo on transparency gets a hue-preserving inversion instead of a plate,
|
| 85 |
+
# an opaque one already carries its own background
|
| 86 |
+
if t == "plate":
|
| 87 |
+
return "invert" if (s is None or s["transp"] > 0.15) else "keep"
|
| 88 |
+
return t
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
FG_ROLES = (1, 5, 6) # text, svg fill, svg stroke
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
def hue_guard(t, o):
|
| 95 |
+
"""Brand colours keep their hue: a coloured input whose output drifts more than 45 degrees is turned back to
|
| 96 |
+
its own hue (keeping the predicted lightness and chroma); neutral inputs get at most a faint tint."""
|
| 97 |
+
rgb = [int(x) for x in t["rgba"][:3]]
|
| 98 |
+
i = oklab(*rgb)
|
| 99 |
+
q = oklab(*o["rgb"])
|
| 100 |
+
cin, cout = math.hypot(i[1], i[2]), math.hypot(q[1], q[2])
|
| 101 |
+
if cin > 0.03 and cout > 0.03:
|
| 102 |
+
d = abs((math.atan2(q[2], q[1]) - math.atan2(i[2], i[1]) + math.pi) % (2 * math.pi) - math.pi)
|
| 103 |
+
if d <= math.radians(45):
|
| 104 |
+
return
|
| 105 |
+
h = math.atan2(i[2], i[1])
|
| 106 |
+
o["rgb"] = lab_to_rgb(q[0], cout * math.cos(h), cout * math.sin(h))
|
| 107 |
+
elif cin <= 0.03 and cout > 0.06:
|
| 108 |
+
o["rgb"] = lab_to_rgb(q[0], q[1] * 0.06 / cout, q[2] * 0.06 / cout)
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
def keep_invisible(toks, out, rels, keep):
|
| 112 |
+
"""Things that were invisible against their surface in the light design stay invisible: a gradient stop
|
| 113 |
+
with exactly a surface's colour (text fades, edge fades) takes that surface's dark colour, and a border that
|
| 114 |
+
matched the surface it sits on takes the surface's dark colour too (otherwise: dark slabs, white edge bars)."""
|
| 115 |
+
n = len(toks)
|
| 116 |
+
key = lambda t: tuple(int(x) for x in t["rgba"][:3])
|
| 117 |
+
surf = {}
|
| 118 |
+
for i, t in enumerate(toks): # first (most prominent) opaque surface of each colour
|
| 119 |
+
if t["role"] == 0 and t["rgba"][3] >= 0.999 and i not in keep and out[i]["a"] >= 0.999:
|
| 120 |
+
surf.setdefault(key(t), i)
|
| 121 |
+
for i, t in enumerate(toks):
|
| 122 |
+
if t["role"] == 4 and i not in keep and key(t) in surf:
|
| 123 |
+
j = surf[key(t)]
|
| 124 |
+
out[i]["rgb"], out[i]["L"] = list(out[j]["rgb"]), out[j]["L"]
|
| 125 |
+
best = {}
|
| 126 |
+
for typ, a, b, w in rels:
|
| 127 |
+
a, b = int(a), int(b)
|
| 128 |
+
if int(typ) == 2 and a < n and b < n and toks[a]["role"] == 2 and toks[b]["role"] == 0 and (a not in best or w > best[a][0]):
|
| 129 |
+
best[a] = (w, b)
|
| 130 |
+
for i, (_, j) in best.items():
|
| 131 |
+
if i in keep or j in keep:
|
| 132 |
+
continue
|
| 133 |
+
if contrast(toks[i]["rgba"][:3], toks[j]["rgba"][:3]) < 1.05 and toks[i]["rgba"][3] >= 0.5:
|
| 134 |
+
out[i]["rgb"], out[i]["L"] = list(out[j]["rgb"]), out[j]["L"]
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
def is_media_keep(t):
|
| 138 |
+
"""Colours drawn over video/photos that already read as dark mode stay exactly as they are: dark scrims and
|
| 139 |
+
player bars, light captions and icons. Everything else over media (a white page or a frosted white button
|
| 140 |
+
laid over a decorative canvas) is redesigned like any other colour."""
|
| 141 |
+
f = t["f"]
|
| 142 |
+
if f[48] < 0.5:
|
| 143 |
+
return False
|
| 144 |
+
if t["role"] in (0, 3, 4): # bg, shadow, gradient: dark scrims, or saturated accents such as a progress bar
|
| 145 |
+
return f[24] <= 0.5 and f[11] < 0.25 and (f[0] < 0.5 or f[3] > 0.4)
|
| 146 |
+
return f[0] > 0.6
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
def enforce_contrast(toks, out, rels, keep):
|
| 150 |
+
"""Keep every foreground colour readable on all the surfaces it sits on, never asking for more contrast
|
| 151 |
+
than the light design had. One lightness per token (tokens are shared), chosen on a grid to minimise the
|
| 152 |
+
weighted log-contrast deficit, with a small pull towards the model's own lightness."""
|
| 153 |
+
n = len(toks)
|
| 154 |
+
surf = {}
|
| 155 |
+
for typ, a, b, w in rels:
|
| 156 |
+
a, b = int(a), int(b)
|
| 157 |
+
if int(typ) == 0 and a < n and b < n and toks[b]["role"] == 0 and out[b]["a"] >= 0.999:
|
| 158 |
+
surf.setdefault(a, []).append((b, float(w) + 0.1)) # log-area weights: small buttons still count
|
| 159 |
+
for i, lst in surf.items():
|
| 160 |
+
t = toks[i]
|
| 161 |
+
if i in keep or t["role"] not in FG_ROLES or out[i]["a"] < 0.999:
|
| 162 |
+
continue
|
| 163 |
+
if t["role"] == 1:
|
| 164 |
+
fs, bold = t["f"][21] * 32, t["f"][23] > 0.5
|
| 165 |
+
need0 = 3.0 if (fs >= 24 or (fs >= 18.5 and bold)) else 4.5
|
| 166 |
+
else:
|
| 167 |
+
need0 = 3.0
|
| 168 |
+
reqs = []
|
| 169 |
+
for j, w in lst:
|
| 170 |
+
need = min(need0, contrast(t["rgba"][:3], toks[j]["rgba"][:3]))
|
| 171 |
+
if need >= 1.5:
|
| 172 |
+
reqs.append((out[j]["rgb"], need, w))
|
| 173 |
+
if not reqs:
|
| 174 |
+
continue
|
| 175 |
+
def deficit(rgb):
|
| 176 |
+
return sum(w * max(0.0, math.log(need) - math.log(contrast(rgb, bg))) for bg, need, w in reqs)
|
| 177 |
+
if deficit(out[i]["rgb"]) <= 0:
|
| 178 |
+
continue
|
| 179 |
+
o = oklab(*out[i]["rgb"])
|
| 180 |
+
W = sum(w for _, _, w in reqs)
|
| 181 |
+
best = None
|
| 182 |
+
for k in range(51):
|
| 183 |
+
Lx = k / 50
|
| 184 |
+
rgb = lab_to_rgb(Lx, o[1], o[2])
|
| 185 |
+
c = deficit(rgb) / W + 0.05 * abs(Lx - o[0])
|
| 186 |
+
if best is None or c < best[0]:
|
| 187 |
+
best = (c, rgb, Lx)
|
| 188 |
+
out[i]["rgb"], out[i]["L"] = best[1], best[2]
|
training/train.py
ADDED
|
@@ -0,0 +1,292 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Train NadirNet on Haiku-designed and native dark-mode targets."""
|
| 3 |
+
import argparse, json, math, os, pickle, random, time
|
| 4 |
+
from collections import defaultdict
|
| 5 |
+
import numpy as np
|
| 6 |
+
import torch
|
| 7 |
+
import torch.nn.functional as F
|
| 8 |
+
|
| 9 |
+
from model import NadirNet, K_REL, contrast_ratio
|
| 10 |
+
from build_ds import oklab
|
| 11 |
+
|
| 12 |
+
COLOR_IDX = dict(lab=(0, 1, 2), C=3, alpha=4, cs=(5, 6), rgb=(7, 8, 9), under=(34, 35, 36), over=(39, 40, 41))
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
def lab_to_rgb_np(lab):
|
| 16 |
+
M1 = np.array([[1.0, 0.3963377774, 0.2158037573], [1.0, -0.1055613458, -0.0638541728], [1.0, -0.0894841775, -1.2914855480]])
|
| 17 |
+
M2 = np.array([[4.0767416621, -3.3077115913, 0.2309699292], [-1.2684380046, 2.6097574011, -0.3413193965], [-0.0041960863, -0.7034186147, 1.7076147010]])
|
| 18 |
+
lin = np.clip((lab @ M1.T) ** 3 @ M2.T, 0, 1)
|
| 19 |
+
return np.where(lin <= 0.0031308, 12.92 * lin, 1.055 * np.power(lin, 1 / 2.4) - 0.055)
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def prep(s):
|
| 23 |
+
"""Precompute OKLab of inputs/targets once per sample."""
|
| 24 |
+
if "lab_in" in s:
|
| 25 |
+
return s
|
| 26 |
+
s["lab_in"] = oklab(s["rgba"][:, :3]).astype(np.float32)
|
| 27 |
+
t = s["tgt"]
|
| 28 |
+
ok = np.isfinite(t[:, 0])
|
| 29 |
+
lab_t = np.zeros((len(t), 3), np.float32)
|
| 30 |
+
if ok.any():
|
| 31 |
+
lab_t[ok] = oklab(t[ok, :3]).astype(np.float32)
|
| 32 |
+
s["lab_t"], s["tmask"] = lab_t, ok
|
| 33 |
+
s["alpha_t"] = np.where(ok, t[:, 3], 1).astype(np.float32)
|
| 34 |
+
km = np.zeros(len(t), bool)
|
| 35 |
+
if len(s["keep"]):
|
| 36 |
+
km[s["keep"]] = True
|
| 37 |
+
s["tmask"] = s["tmask"] & ~km
|
| 38 |
+
return s
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def augment(s, rng):
|
| 42 |
+
s = dict(s)
|
| 43 |
+
n = len(s["role"])
|
| 44 |
+
# token dropout (simulates ad blockers, A/B variants, partial loads)
|
| 45 |
+
if rng.random() < 0.6 and n > 8:
|
| 46 |
+
p = rng.uniform(0, 0.25)
|
| 47 |
+
root = s["F"][:, 24] > 0.5
|
| 48 |
+
keepm = (rng.random(n) > p) | root
|
| 49 |
+
idx = np.nonzero(keepm)[0]
|
| 50 |
+
remap = -np.ones(n, np.int64); remap[idx] = np.arange(len(idx))
|
| 51 |
+
for k in ("F", "role", "rgba", "w", "lab_in", "lab_t", "tmask", "alpha_t"):
|
| 52 |
+
s[k] = s[k][idx]
|
| 53 |
+
rel = s["rel"]
|
| 54 |
+
if len(rel):
|
| 55 |
+
a, b = rel[:, 1].astype(int), rel[:, 2].astype(int)
|
| 56 |
+
good = (remap[a] >= 0) & (remap[b] >= 0)
|
| 57 |
+
rel = rel[good].copy(); rel[:, 1] = remap[a[good]]; rel[:, 2] = remap[b[good]]
|
| 58 |
+
s["rel"] = rel
|
| 59 |
+
iu = s["img_under"].copy(); m = iu >= 0; iu[m] = remap[iu[m]]; s["img_under"] = iu
|
| 60 |
+
n = len(idx)
|
| 61 |
+
# hue rotation of the whole design (inputs and targets together)
|
| 62 |
+
if rng.random() < 0.3:
|
| 63 |
+
th = rng.uniform(0, 2 * math.pi)
|
| 64 |
+
c, si = math.cos(th), math.sin(th)
|
| 65 |
+
rot = lambda a, b: (c * a - si * b, si * a + c * b)
|
| 66 |
+
F_ = s["F"].copy()
|
| 67 |
+
for (ia, ib) in ((1, 2), (35, 36), (40, 41), (5, 6)):
|
| 68 |
+
F_[:, ia], F_[:, ib] = rot(F_[:, ia], F_[:, ib])
|
| 69 |
+
lab = s["lab_in"].copy(); lab[:, 1], lab[:, 2] = rot(lab[:, 1], lab[:, 2])
|
| 70 |
+
F_[:, 7:10] = lab_to_rgb_np(lab.astype(np.float64)).astype(np.float32)
|
| 71 |
+
s["F"], s["lab_in"] = F_, lab.astype(np.float32)
|
| 72 |
+
lt = s["lab_t"].copy(); lt[:, 1], lt[:, 2] = rot(lt[:, 1], lt[:, 2]); s["lab_t"] = lt.astype(np.float32)
|
| 73 |
+
return s
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def collate(batch, device):
|
| 77 |
+
B = len(batch)
|
| 78 |
+
N = max(1, max(len(s["role"]) for s in batch))
|
| 79 |
+
M = max(len(s["IF"]) for s in batch)
|
| 80 |
+
T = 1 + N + M
|
| 81 |
+
tok_f = np.zeros((B, N, 49), np.float32); role = np.zeros((B, N), np.int64); lab = np.zeros((B, N, 3), np.float32)
|
| 82 |
+
lab_t = np.zeros((B, N, 3), np.float32); tmask = np.zeros((B, N), bool); w = np.zeros((B, N), np.float32)
|
| 83 |
+
alpha_in = np.ones((B, N), np.float32); alpha_t = np.ones((B, N), np.float32)
|
| 84 |
+
img_f = np.zeros((B, M, 39), np.float32); thumbs = np.zeros((B, M, 4, 32, 32), np.float32); imgt = -np.ones((B, M), np.int64)
|
| 85 |
+
page = np.zeros((B, 8), np.float32); src = np.zeros(B, np.int64)
|
| 86 |
+
rel = np.zeros((B, T, T, K_REL), np.float32); mask = np.zeros((B, T), bool)
|
| 87 |
+
tb_pair = -np.ones((B, N), np.int64)
|
| 88 |
+
for i, s in enumerate(batch):
|
| 89 |
+
n, m = len(s["role"]), len(s["IF"])
|
| 90 |
+
tok_f[i, :n] = s["F"]; role[i, :n] = s["role"]; lab[i, :n] = s["lab_in"]; lab_t[i, :n] = s["lab_t"]; tmask[i, :n] = s["tmask"]
|
| 91 |
+
w[i, :n] = s["w"]; alpha_in[i, :n] = s["rgba"][:, 3]; alpha_t[i, :n] = s["alpha_t"]
|
| 92 |
+
if m:
|
| 93 |
+
img_f[i, :m] = s["IF"]; thumbs[i, :m] = s["thumbs"].transpose(0, 3, 1, 2) / 255.0; imgt[i, :m] = s["imgt"]
|
| 94 |
+
page[i] = s["page"]; src[i] = s["src"]
|
| 95 |
+
mask[i, :1 + n] = True; mask[i, 1 + N:1 + N + m] = True
|
| 96 |
+
# relations
|
| 97 |
+
r = s["rel"]
|
| 98 |
+
if len(r):
|
| 99 |
+
typ, a, b, wt = r[:, 0].astype(int), r[:, 1].astype(int) + 1, r[:, 2].astype(int) + 1, r[:, 3]
|
| 100 |
+
R = np.zeros((4, n + 1, n + 1), np.float32)
|
| 101 |
+
np.maximum.at(R, (typ, a, b), wt)
|
| 102 |
+
R /= R.max(axis=2, keepdims=True) + 1e-6
|
| 103 |
+
rel[i, :n + 1, :n + 1, 0:4] = R.transpose(1, 2, 0)
|
| 104 |
+
rel[i, :n + 1, :n + 1, 4:8] = R.transpose(2, 1, 0)
|
| 105 |
+
# dominant surface for each text token (contrast loss)
|
| 106 |
+
tbm = typ == 0
|
| 107 |
+
best = {}
|
| 108 |
+
for aa, bb, ww in zip(a[tbm] - 1, b[tbm] - 1, wt[tbm]):
|
| 109 |
+
if ww > best.get(aa, (-1, -1))[0]:
|
| 110 |
+
best[aa] = (ww, bb)
|
| 111 |
+
for aa, (ww, bb) in best.items():
|
| 112 |
+
tb_pair[i, aa] = bb
|
| 113 |
+
for j in range(m):
|
| 114 |
+
u = s["img_under"][j]
|
| 115 |
+
if u >= 0:
|
| 116 |
+
rel[i, 1 + N + j, 1 + u, 8] = 1.0
|
| 117 |
+
rel[i, 1 + u, 1 + N + j, 9] = 1.0
|
| 118 |
+
# same colour, different role
|
| 119 |
+
key = s["rgba"][:, :3] @ np.array([65536, 256, 1], np.float32)
|
| 120 |
+
same = (key[:, None] == key[None, :]) & ~np.eye(n, dtype=bool)
|
| 121 |
+
rel[i, 1:1 + n, 1:1 + n, 10] = same
|
| 122 |
+
if device is None:
|
| 123 |
+
return dict(tok_f=tok_f, role=role, lab=lab, lab_t=lab_t, tmask=tmask, w=w, alpha_in=alpha_in, alpha_t=alpha_t,
|
| 124 |
+
img_f=img_f, thumbs=thumbs, imgt=imgt, page=page, src=src, rel=rel, mask=mask, tb_pair=tb_pair)
|
| 125 |
+
t = lambda x: torch.from_numpy(x).to(device, non_blocking=True)
|
| 126 |
+
return dict(tok_f=t(tok_f), role=t(role), lab=t(lab), lab_t=t(lab_t), tmask=t(tmask), w=t(w), alpha_in=t(alpha_in), alpha_t=t(alpha_t),
|
| 127 |
+
img_f=t(img_f), thumbs=t(thumbs), imgt=t(imgt), page=t(page), src=t(src), rel=t(rel), mask=t(mask), tb_pair=t(tb_pair))
|
| 128 |
+
|
| 129 |
+
|
| 130 |
+
class Batches(torch.utils.data.Dataset):
|
| 131 |
+
"""Each item is a whole (augmented) batch, assembled in a worker process."""
|
| 132 |
+
def __init__(self, data, order, bs, aug, seed):
|
| 133 |
+
self.data, self.order, self.bs, self.aug, self.seed = data, order, bs, aug, seed
|
| 134 |
+
|
| 135 |
+
def __len__(self):
|
| 136 |
+
return math.ceil(len(self.order) / self.bs)
|
| 137 |
+
|
| 138 |
+
def __getitem__(self, k):
|
| 139 |
+
rng = np.random.default_rng((self.seed, k))
|
| 140 |
+
items = [self.data[i] for i in self.order[k * self.bs:(k + 1) * self.bs]]
|
| 141 |
+
if self.aug:
|
| 142 |
+
items = [augment(s, rng) for s in items]
|
| 143 |
+
return collate(items, None)
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
def to_dev(b, device):
|
| 147 |
+
return {k: (v if torch.is_tensor(v) else torch.from_numpy(v)).to(device, non_blocking=True) for k, v in b.items()}
|
| 148 |
+
|
| 149 |
+
|
| 150 |
+
def huber(x, d):
|
| 151 |
+
ax = x.abs()
|
| 152 |
+
return torch.where(ax < d, 0.5 * x * x / d, ax - 0.5 * d)
|
| 153 |
+
|
| 154 |
+
|
| 155 |
+
def losses(model, b, img_w=None):
|
| 156 |
+
lab, alpha_logit, img_logits = model(b["tok_f"], b["role"], b["lab"], b["img_f"], b["thumbs"], b["page"], b["src"], b["rel"], b["mask"])
|
| 157 |
+
lab = lab.float(); alpha_logit = alpha_logit.float(); img_logits = img_logits.float()
|
| 158 |
+
m = b["tmask"].float() * b["w"]
|
| 159 |
+
d = lab - b["lab_t"]
|
| 160 |
+
per = 2.0 * huber(d[..., 0], 0.02) + huber(d[..., 1], 0.01) + huber(d[..., 2], 0.01)
|
| 161 |
+
lc = (per * m).sum() / m.sum().clamp_min(1)
|
| 162 |
+
# alpha (only where translucency is involved)
|
| 163 |
+
am = m * ((b["alpha_in"] < 0.999) | (b["alpha_t"] < 0.999)).float()
|
| 164 |
+
a_in = b["alpha_in"].clamp(1e-3, 1 - 1e-3)
|
| 165 |
+
alpha = torch.sigmoid(alpha_logit + torch.log(a_in / (1 - a_in)))
|
| 166 |
+
la = (huber(alpha - b["alpha_t"], 0.05) * am).sum() / am.sum().clamp_min(1)
|
| 167 |
+
# readability: text must not lose contrast against its own surface
|
| 168 |
+
pair = b["tb_pair"]
|
| 169 |
+
pm = (pair >= 0) & b["tmask"] & (b["role"] == 1) & (b["alpha_in"] > 0.99)
|
| 170 |
+
if pm.any():
|
| 171 |
+
idx = pair.clamp_min(0)
|
| 172 |
+
bg_pred = torch.gather(lab, 1, idx[..., None].expand(-1, -1, 3))
|
| 173 |
+
bg_t = torch.gather(b["lab_t"], 1, idx[..., None].expand(-1, -1, 3))
|
| 174 |
+
bgm = torch.gather(b["tmask"], 1, idx)
|
| 175 |
+
pm = pm & bgm
|
| 176 |
+
c_pred = contrast_ratio(lab, bg_pred)
|
| 177 |
+
c_t = contrast_ratio(b["lab_t"], bg_t).clamp(max=4.5)
|
| 178 |
+
lcon = (F.relu(c_t - c_pred) * pm.float()).sum() / pm.float().sum().clamp_min(1)
|
| 179 |
+
else:
|
| 180 |
+
lcon = lab.new_zeros(())
|
| 181 |
+
# images
|
| 182 |
+
im = b["imgt"] >= 0
|
| 183 |
+
if im.any():
|
| 184 |
+
li = F.cross_entropy(img_logits[im], b["imgt"][im], weight=img_w, label_smoothing=0.05)
|
| 185 |
+
acc = (img_logits[im].argmax(-1) == b["imgt"][im]).float().mean()
|
| 186 |
+
else:
|
| 187 |
+
li = lab.new_zeros(()); acc = lab.new_zeros(())
|
| 188 |
+
with torch.no_grad():
|
| 189 |
+
de = (d.norm(dim=-1) * m).sum() / m.sum().clamp_min(1)
|
| 190 |
+
dl = (d[..., 0].abs() * m).sum() / m.sum().clamp_min(1)
|
| 191 |
+
return lc, la, lcon, li, dict(dE=de.item(), dL=dl.item(), img_acc=acc.item(), n_img=int(im.sum()))
|
| 192 |
+
|
| 193 |
+
|
| 194 |
+
def evaluate(model, val, device, bs, img_w):
|
| 195 |
+
model.eval()
|
| 196 |
+
agg = defaultdict(float); cnt = defaultdict(float)
|
| 197 |
+
with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
|
| 198 |
+
for src in (0, 1):
|
| 199 |
+
vs = [s for s in val if s["src"] == src]
|
| 200 |
+
for k in range(0, len(vs), bs):
|
| 201 |
+
b = collate(vs[k:k + bs], device)
|
| 202 |
+
lc, la, lcon, li, st = losses(model, b, img_w)
|
| 203 |
+
n = len(vs[k:k + bs])
|
| 204 |
+
agg[f"s{src}_dE"] += st["dE"] * n; agg[f"s{src}_dL"] += st["dL"] * n; agg[f"s{src}_con"] += lcon.item() * n
|
| 205 |
+
cnt[f"s{src}"] += n
|
| 206 |
+
if st["n_img"]:
|
| 207 |
+
agg["img_acc"] += st["img_acc"] * st["n_img"]; cnt["img"] += st["n_img"]
|
| 208 |
+
out = {}
|
| 209 |
+
for k, v in agg.items():
|
| 210 |
+
if k == "img_acc":
|
| 211 |
+
out[k] = v / max(1, cnt["img"])
|
| 212 |
+
else:
|
| 213 |
+
out[k] = v / max(1, cnt[k.split("_")[0]])
|
| 214 |
+
model.train()
|
| 215 |
+
return out
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
def main():
|
| 219 |
+
ap = argparse.ArgumentParser()
|
| 220 |
+
ap.add_argument("--data", required=True)
|
| 221 |
+
ap.add_argument("--out", required=True)
|
| 222 |
+
ap.add_argument("--d", type=int, default=384)
|
| 223 |
+
ap.add_argument("--layers", type=int, default=8)
|
| 224 |
+
ap.add_argument("--heads", type=int, default=6)
|
| 225 |
+
ap.add_argument("--epochs", type=int, default=40)
|
| 226 |
+
ap.add_argument("--bs", type=int, default=48)
|
| 227 |
+
ap.add_argument("--lr", type=float, default=4e-4)
|
| 228 |
+
ap.add_argument("--wd", type=float, default=0.05)
|
| 229 |
+
ap.add_argument("--drop", type=float, default=0.1)
|
| 230 |
+
ap.add_argument("--native-weight", type=float, default=0.5)
|
| 231 |
+
ap.add_argument("--seed", type=int, default=0)
|
| 232 |
+
ap.add_argument("--workers", type=int, default=8)
|
| 233 |
+
args = ap.parse_args()
|
| 234 |
+
torch.manual_seed(args.seed); rng = np.random.default_rng(args.seed)
|
| 235 |
+
os.makedirs(args.out, exist_ok=True)
|
| 236 |
+
data = pickle.load(open(args.data, "rb"))
|
| 237 |
+
train, val = [prep(s) for s in data["train"]], [prep(s) for s in data["val"]]
|
| 238 |
+
print(f"train {len(train)} val {len(val)} (haiku {sum(s['src'] == 0 for s in train)}, native {sum(s['src'] == 1 for s in train)})", flush=True)
|
| 239 |
+
device = "cuda"
|
| 240 |
+
model = NadirNet(d=args.d, layers=args.layers, heads=args.heads, drop=args.drop).to(device)
|
| 241 |
+
nparam = sum(p.numel() for p in model.parameters())
|
| 242 |
+
print(f"params {nparam / 1e6:.2f}M", flush=True)
|
| 243 |
+
# image class weights (inverse sqrt frequency)
|
| 244 |
+
cnt = np.ones(5)
|
| 245 |
+
for s in train:
|
| 246 |
+
for v in s["imgt"]:
|
| 247 |
+
if v >= 0:
|
| 248 |
+
cnt[v] += 1
|
| 249 |
+
img_w = torch.tensor((cnt.sum() / cnt) ** 0.5, dtype=torch.float32, device=device)
|
| 250 |
+
img_w = img_w / img_w.mean()
|
| 251 |
+
print("image label counts", cnt.astype(int).tolist(), flush=True)
|
| 252 |
+
opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.wd, betas=(0.9, 0.98))
|
| 253 |
+
steps_per = math.ceil(len(train) / args.bs)
|
| 254 |
+
total = steps_per * args.epochs
|
| 255 |
+
warm = min(800, total // 20)
|
| 256 |
+
sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda st: min(1, (st + 1) / warm) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1, st / total)))))
|
| 257 |
+
# sample weights: native pages count less (different designers, noisier pairing)
|
| 258 |
+
sw = np.array([1.0 if s["src"] == 0 else args.native_weight for s in train]); sw /= sw.sum()
|
| 259 |
+
best = 1e9; step = 0; log = []
|
| 260 |
+
t0 = time.time()
|
| 261 |
+
for ep in range(args.epochs):
|
| 262 |
+
order = rng.choice(len(train), size=len(train), p=sw)
|
| 263 |
+
run = defaultdict(float); nb = 0
|
| 264 |
+
loader = torch.utils.data.DataLoader(Batches(train, order, args.bs, True, args.seed * 1000 + ep), batch_size=None, shuffle=False,
|
| 265 |
+
num_workers=args.workers, persistent_workers=False, prefetch_factor=4 if args.workers else None)
|
| 266 |
+
for bnp in loader:
|
| 267 |
+
b = to_dev(bnp, device)
|
| 268 |
+
with torch.autocast("cuda", dtype=torch.bfloat16):
|
| 269 |
+
lc, la, lcon, li, st = losses(model, b, img_w)
|
| 270 |
+
loss = lc + 0.5 * la + 0.02 * lcon + 0.05 * li
|
| 271 |
+
opt.zero_grad(set_to_none=True)
|
| 272 |
+
if not torch.isfinite(loss): # one bad page must never poison the weights
|
| 273 |
+
run["skipped"] += 1
|
| 274 |
+
continue
|
| 275 |
+
loss.backward()
|
| 276 |
+
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
| 277 |
+
opt.step(); sched.step(); step += 1
|
| 278 |
+
run["loss"] += loss.item(); run["dE"] += st["dE"]; run["con"] += lcon.item(); run["img"] += li.item(); nb += 1
|
| 279 |
+
ev = evaluate(model, val, device, 64, img_w)
|
| 280 |
+
rec = {"epoch": ep, "step": step, "secs": round(time.time() - t0), **{k: round(v / nb, 5) for k, v in run.items()}, **{k: round(v, 5) for k, v in ev.items()}}
|
| 281 |
+
log.append(rec)
|
| 282 |
+
print(json.dumps(rec), flush=True)
|
| 283 |
+
score = ev.get("s0_dE", 9) + 0.3 * ev.get("s1_dE", 0) + 0.02 * ev.get("s0_con", 0) + 0.02 * (1 - ev.get("img_acc", 0))
|
| 284 |
+
if score < best:
|
| 285 |
+
best = score
|
| 286 |
+
torch.save({"model": model.state_dict(), "cfg": model.cfg, "epoch": ep, "eval": ev}, os.path.join(args.out, "best.pt"))
|
| 287 |
+
torch.save({"model": model.state_dict(), "cfg": model.cfg, "epoch": args.epochs - 1}, os.path.join(args.out, "last.pt"))
|
| 288 |
+
json.dump({"args": vars(args), "params": nparam, "log": log, "best_score": best}, open(os.path.join(args.out, "train_log.json"), "w"), indent=1)
|
| 289 |
+
|
| 290 |
+
|
| 291 |
+
if __name__ == "__main__":
|
| 292 |
+
main()
|