open-standards/aac/aac-rs/tools/ae/eval/run.mjs

// Measure first letters on test replies, on this app's own worker.
//
//   node tools/ae/eval/run.mjs ORT TOK PACKS ITEMS.json OUT.json [SUMMARY.json]
//
// ORT: a site's vendor/onnxruntime-web-<version>-<rev>/ (third_party/
// onnxruntime-web/vendor.sh writes it); TOK: its vendor/hf-tokenizers-
// <version>-<rev>/ (third_party/hf-tokenizers/vendor.sh); PACKS: what
// tools/ae/build_ae.py wrote (aac-ae-packs/); ITEMS: test items, each {set,
// context, abbrev, target, typoCost} (the model library's
// onnx/ae-en/check_full.json and check.json: tools/ae/README.md).
//
// Each item goes through the app's own worker (ae-worker.js), in installed
// Chrome on its WebGPU, exactly as the app runs it: the worker reads the
// runtime, the tokenizer and the model from the model store, and is asked as
// the app asks it (the letters, what the partner said, and whether slips to
// a neighbouring key are forgiven: the items with typing slips are asked
// with it on, as a person who makes them would set it). Each "clean" set is
// asked again with slips forgiven, to show what that costs someone who
// types accurately. A reply counts when one of the five sentences reads as it
// does (case, spacing and final punctuation aside: the model library's
// measure); "words" counts it when the words are the same, punctuation aside
// ("Yes, she does" for "Yes she does").
//
// OUT gets every item (the test sentences: keep it outside git, as Turk
// Dialogues is not ours to publish); SUMMARY, if given, only the numbers.
//
// Needs Chrome with WebGPU on a machine with a graphics chip, and node with
// Playwright (aac-aurasense-web's node_modules: run from there).

import { createServer } from 'node:http';
import { readFileSync, existsSync, statSync, writeFileSync } from 'node:fs';
import { extname, join, resolve } from 'node:path';
import { createRequire } from 'node:module';

const [ortDir, tokDir, packs, itemsFile, out, summaryFile] = process.argv.slice(2);
if (!out) {
  console.error('usage: node run.mjs ORT TOK PACKS ITEMS.json OUT.json [SUMMARY.json]');
  process.exit(2);
}
const require = createRequire(join(process.cwd(), 'package.json'));
const { chromium } = require('playwright');
const items = JSON.parse(readFileSync(itemsFile, 'utf8'));

// The page: the model store filled, the worker started from it, and the
// items asked through its own messages.
const PAGE = `<!doctype html><meta charset="utf-8"><script type="module">
const CACHE = 'aac-models-v1';
const abs = (u) => new URL(u, location.href).href;
let worker, nextId = 1; const calls = new Map();
const call = (msg) => new Promise((resolve, reject) => { const id = nextId++; calls.set(id, { resolve, reject }); worker.postMessage({ ...msg, id }); });
window.load = async () => {
  const cache = await caches.open(CACHE);
  const idx = await (await fetch('/ae/index.json')).json();
  const cat = await (await fetch('/ae/' + idx.current)).json();
  const [id, m] = Object.entries(cat.models)[0];
  const parts = [m.graph, ...m.data.flatMap((d) => d.parts), m.tokenizer, m.tokenizer_config, m.config, m.vocab, m.vocab_table];
  const files = ['/tok/ae-worker.js', '/tok/tokenizers.min.mjs', '/ort/ort.jspi.bundle.min.mjs', '/ort/ort-wasm-simd-threaded.jspi.wasm',
    ...parts.map((p) => '/ae/' + p.url)];
  for (const f of files) if (!(await cache.match(abs(f)))) await cache.put(abs(f), await fetch(f));
  const src = await (await cache.match(abs('/tok/ae-worker.js'))).blob();
  worker = new Worker(URL.createObjectURL(src), { type: 'module' });
  worker.onmessage = (e) => {
    const msg = e.data;
    const c = calls.get(msg.id); calls.delete(msg.id);
    msg.ok ? c.resolve(msg) : c.reject(new Error(msg.error));
  };
  await call({ type: 'init', bundle: abs('/ort/ort.jspi.bundle.min.mjs'), wasm: abs('/ort/ort-wasm-simd-threaded.jspi.wasm'), tokenizers: abs('/tok/tokenizers.min.mjs') });
  const t0 = performance.now();
  const u = (p) => abs('/ae/' + p.url);
  await call({ type: 'load', key: id, model: {
    graph: u(m.graph), data: m.data.map((d) => ({ name: d.name, parts: d.parts.map((p) => ({ url: u(p), bytes: p.bytes })) })),
    tokenizer: u(m.tokenizer), tokenizer_config: u(m.tokenizer_config), config: u(m.config), vocab: u(m.vocab), vocab_table: u(m.vocab_table),
  } });
  const load_ms = Math.round(performance.now() - t0);
  await call({ type: 'check' });
  const adapter = await navigator.gpu.requestAdapter();
  return { model: id, load_ms, adapter: adapter && adapter.info ? (adapter.info.vendor + ' ' + adapter.info.architecture).trim() : null, ua: navigator.userAgent };
};
window.ask = async (letters, partner, near_keys) => {
  const t0 = performance.now();
  const r = await call({ type: 'expand', letters, partner, near_keys });
  return { sentences: r.sentences, ms: Math.round(performance.now() - t0) };
};
window.ready = true;
</script>`;

const roots = { '/ort/': resolve(ortDir), '/tok/': resolve(tokDir), '/ae/': resolve(packs) };
const server = createServer((req, res) => {
  const path = decodeURIComponent(new URL(req.url, 'http://x').pathname);
  if (path === '/' || path === '/eval.html') {
    res.writeHead(200, { 'Content-Type': 'text/html' });
    return res.end(PAGE);
  }
  for (const [prefix, root] of Object.entries(roots)) {
    if (!path.startsWith(prefix)) continue;
    const file = join(root, path.slice(prefix.length));
    if (!file.startsWith(root) || !existsSync(file) || !statSync(file).isFile()) break;
    const type = { '.js': 'text/javascript', '.mjs': 'text/javascript', '.json': 'application/json', '.wasm': 'application/wasm' }[extname(file)] || 'application/octet-stream';
    res.writeHead(200, { 'Content-Type': type });
    return res.end(readFileSync(file));
  }
  res.writeHead(404);
  res.end();
});
await new Promise((r) => server.listen(0, '127.0.0.1', r));
const base = `http://127.0.0.1:${server.address().port}`;

// As the app's worker compares: case, spacing and final punctuation aside.
const isAlnum = (ch) => /[\p{L}\p{N}]/u.test(ch);
const norm = (s) => {
  s = s.normalize('NFKC').replace(/[’‘ʼ]/g, "'").replace(/[“”]/g, '"').replace(/—/g, ' - ').replace(/–/g, '-');
  s = s.split(/\s+/).filter(Boolean).join(' ').trimEnd();
  let i = s.length;
  while (i > 0 && !isAlnum(s[i - 1])) i--;
  return s.slice(0, i).toLowerCase();
};
const words = (s) => norm(s).replace(/[^\p{L}\p{N}' ]+/gu, ' ').split(/\s+/).filter(Boolean).join(' ');

const browser = await chromium.launch({ channel: 'chrome', args: ['--enable-unsafe-webgpu'] });
const page = await browser.newPage();
page.on('pageerror', (e) => console.error('page error:', e.message));
await page.goto(`${base}/eval.html`);
await page.waitForFunction(() => window.ready);
const result = { date: new Date().toLocaleDateString('en-CA'), runtime: await page.evaluate(() => window.load()), rows: [] };
console.log(`loaded ${result.runtime.model} in ${result.runtime.load_ms} ms on ${result.runtime.adapter}`);
// Each item as the app asks it; each clean item again with slips forgiven.
const runs = [];
for (const it of items) {
  const noisy = it.typoCost !== null && it.typoCost !== undefined;
  runs.push({ it, set: it.set, near_keys: noisy });
  if (!noisy) runs.push({ it, set: `${it.set}+near-keys`, near_keys: true });
}
for (const [n, r] of runs.entries()) {
  const got = await page.evaluate(([l, p, k]) => window.ask(l, p, k), [r.it.abbrev, r.it.context || null, r.near_keys]);
  const gold = norm(r.it.target);
  const at = got.sentences.findIndex((s) => norm(s) === gold);
  const atWords = got.sentences.findIndex((s) => words(s) === words(r.it.target));
  result.rows.push({ set: r.set, context: r.it.context, abbrev: r.it.abbrev, target: r.it.target, near_keys: r.near_keys, sentences: got.sentences, at, atWords, ms: got.ms });
  if ((n + 1) % 100 === 0) console.log(`${n + 1}/${runs.length}`);
}
await browser.close();
server.close();

const sets = {};
for (const r of result.rows) (sets[r.set] ||= []).push(r);
const pct = (k, n) => Math.round((1000 * k) / n) / 10;
const summary = { date: result.date, runtime: result.runtime, items: itemsFile.split('/').slice(-2).join('/'), sets: {} };
for (const [name, rows] of Object.entries(sets)) {
  const ms = rows.map((r) => r.ms).sort((a, b) => a - b);
  summary.sets[name] = {
    n: rows.length,
    top1: pct(rows.filter((r) => r.at === 0).length, rows.length),
    top5: pct(rows.filter((r) => r.at >= 0).length, rows.length),
    top5_words: pct(rows.filter((r) => r.atWords >= 0).length, rows.length),
    ms_median: ms[Math.floor(ms.length / 2)],
    ms_p90: ms[Math.floor(ms.length * 0.9)],
  };
}
result.summary = summary;
writeFileSync(out, JSON.stringify(result, null, 1) + '\n');
if (summaryFile) writeFileSync(summaryFile, JSON.stringify(summary, null, 1) + '\n');
console.log(JSON.stringify(summary.sets, null, 1));