canopy-reader-68m / node /reader.ts
gpriday's picture
Canopy reader: ONNX (INT8, FP32), readers, model card
6fc6edb verified
Raw History Blame Contribute Delete
12.2 kB
// The Canopy reader on one screen, in Node: the exported ONNX model (scripts/export_reader.py) plus its
// tokenizer, encoded exactly as canopy/classifier/model.py:encode() and decoded as canopy/reader/evaluate.py.
// A reference for Daintree's integration; src/parity.ts checks it against the Python outputs.
import { readFileSync } from "node:fs";
import { basename, dirname, join, resolve } from "node:path";
import * as ort from "onnxruntime-node";
import { AutoTokenizer, env } from "@huggingface/transformers";
export interface ReaderMeta {
max_len: number;
header_tokens: number;
line_tokens: number;
question_candidate_rows: number;
status_candidate_rows: number;
inputs: string[];
outputs: string[];
categories: string[];
classes: Record<string, string[]>;
thresholds: { summary_first: number; summary_more: number; option: number; gate: number; gate_categories: string[] };
}
export interface Reading {
category: string;
category_probs: number[];
needs_attention: number; // gated and rescaled: >= 0.5 means the card needs the user
needs_attention_raw: number;
blocked_on_user: number;
notable: number;
question_line: number; // line indices into the prepared screen, -1 for none
status_line: number; // the headline
action_line: number;
task_line: number;
summary_lines: number[]; // most important first
option_lines: number[]; // screen order
key_lines: number[];
tests: string;
changes: string;
risk: string;
instruction: string;
progress: string;
failure_repeats: string;
score: string;
}
const cp = (s: string) => [...s].length; // Python len(): code points, not UTF-16 units
// Python's str.strip() whitespace (str.isspace), which differs from JS \s: it includes \x1c-\x1f and \x85 and
// excludes \ufeff. Training used Python, so blank-line skipping and candidate rows must use the same set.
const PY_WS = "\t\n\v\f\r\x1c\x1d\x1e\x1f \x85\xa0\u1680\u2000-\u200a\u2028\u2029\u202f\u205f\u3000";
const STRIP = new RegExp(`^[${PY_WS}]+|[${PY_WS}]+$`, "gu");
const strip = (s: string) => s.replace(STRIP, "");
const sigmoid = (x: number) => 1 / (1 + Math.exp(-x));
function softmax(xs: number[]): number[] {
const m = Math.max(...xs);
const e = xs.map((x) => Math.exp(x - m));
const s = e.reduce((a, b) => a + b, 0);
return e.map((x) => x / s);
}
export function lineCandidates(lines: string[], rows: number): number[] {
const start = Math.max(0, lines.length - rows);
const out: number[] = [];
for (let i = start; i < lines.length; i++) if (cp(strip(lines[i])) > 3) out.push(i);
return out;
}
export class CanopyReader {
meta: ReaderMeta;
session: ort.InferenceSession;
tok: any;
cls: number;
sep: number;
pad: number;
private constructor(meta: ReaderMeta, session: ort.InferenceSession, tok: any) {
this.meta = meta;
this.session = session;
this.tok = tok;
this.cls = tok.model.tokens_to_ids.get(tok.cls_token ?? "[CLS]");
this.sep = tok.model.tokens_to_ids.get(tok.sep_token ?? "[SEP]");
this.pad = tok.model.tokens_to_ids.get(tok.pad_token ?? "[PAD]");
}
/** `dir` holds reader.json, the tokenizer files and model.<variant>.onnx. */
static async load(dir: string, variant = "int8", threads = 0): Promise<CanopyReader> {
const meta = JSON.parse(readFileSync(join(dir, "reader.json"), "utf8")) as ReaderMeta;
env.allowRemoteModels = false;
env.localModelPath = dirname(resolve(dir)) + "/";
const tok = await AutoTokenizer.from_pretrained(basename(resolve(dir)), { local_files_only: true } as any);
const opts: ort.InferenceSession.SessionOptions = { graphOptimizationLevel: "all" };
if (threads) opts.intraOpNumThreads = threads;
const session = await ort.InferenceSession.create(join(dir, `model.${variant}.onnx`), opts);
return new CanopyReader(meta, session, tok);
}
private ids(text: string): number[] {
return this.tok.encode(text, { add_special_tokens: false });
}
/** As encode(): [CLS] state-json [SEP] then "[i] line\n" pieces filled from the bottom up, [SEP]. */
/** `state` must hold integer numbers only (as Daintree sends): JSON.stringify writes 1.0 as 1, where the
* Python training encoder wrote 1.0, so a float in the state would tokenize differently. */
encode(state: object, lines: string[]) {
const m = this.meta;
const header = this.ids(JSON.stringify(state)).slice(0, m.header_tokens);
let budget = m.max_len - header.length - 3;
const pieces: [number, number[]][] = [];
for (let i = lines.length - 1; i >= 0; i--) {
if (!strip(lines[i])) continue;
const ids = this.ids(`[${i}] ${lines[i]}\n`).slice(0, m.line_tokens);
if (ids.length > budget) break;
budget -= ids.length;
pieces.push([i, ids]);
}
pieces.reverse();
const ids = [this.cls, ...header, this.sep];
const lineIndex: number[] = [];
const linePos: number[] = [];
for (const [i, p] of pieces) {
lineIndex.push(i);
linePos.push(ids.length);
ids.push(...p);
}
ids.push(this.sep);
const q = new Set(lineCandidates(lines, m.question_candidate_rows));
const s = new Set(lineCandidates(lines, m.status_candidate_rows));
return { ids, lineIndex, linePos, qMask: lineIndex.map((i) => q.has(i)), sMask: lineIndex.map((i) => s.has(i)) };
}
async read(state: object, lines: string[]): Promise<Reading> {
const e = this.encode(state, lines);
const n = e.ids.length;
const k = Math.max(1, e.lineIndex.length);
const pad = <T>(xs: T[], fill: T) => [...xs, ...Array(k - xs.length).fill(fill)];
const feeds: Record<string, ort.Tensor> = {
input_ids: new ort.Tensor("int64", BigInt64Array.from(e.ids.map(BigInt)), [1, n]),
attention_mask: new ort.Tensor("int64", BigInt64Array.from(e.ids.map(() => 1n)), [1, n]),
line_pos: new ort.Tensor("int64", BigInt64Array.from(pad(e.linePos, 0).map(BigInt)), [1, k]),
line_mask: new ort.Tensor("bool", Uint8Array.from(pad(e.lineIndex.map(() => 1), 0)), [1, k]),
q_mask: new ort.Tensor("bool", Uint8Array.from(pad(e.qMask.map(Number), 0)), [1, k]),
s_mask: new ort.Tensor("bool", Uint8Array.from(pad(e.sMask.map(Number), 0)), [1, k]),
};
const out = await this.session.run(feeds);
const arr = (name: string) => Array.from(out[name].data as Float32Array, Number);
return this.decode(arr, e.lineIndex, e.qMask);
}
decode(arr: (name: string) => number[], idx: number[], qMask: boolean[]): Reading {
const m = idx.length;
const t = this.meta.thresholds;
const pick = (name: string) => {
const scores = arr(name);
const cand = [...scores.slice(0, m), scores[scores.length - 1]];
let j = 0;
for (let x = 1; x < cand.length; x++) if (cand[x] > cand[j]) j = x;
return j === m ? -1 : idx[j];
};
const logit = arr("summary_lines").slice(0, m);
const pSum = logit.map(sigmoid);
const ranked = [...logit.keys()].sort((a, b) => logit[b] - logit[a] || a - b); // by logit: no saturation ties
const summary: number[] = [];
for (const j of ranked.slice(0, 3)) {
if (pSum[j] >= (summary.length ? t.summary_more : t.summary_first)) summary.push(idx[j]);
else break;
}
const pOpt = arr("option_lines").slice(0, m).map(sigmoid);
const key = arr("key_lines").slice(0, m);
const probs = softmax(arr("category"));
const raw = sigmoid(arr("needs_attention")[0]);
const needs = t.gate_categories.reduce((a, c) => a + probs[this.meta.categories.indexOf(c)], 0);
const g = Math.max(raw, needs);
const argmax = (xs: number[]) => xs.reduce((b, x, i) => (x > xs[b] ? i : b), 0);
const cls = (name: string) => this.meta.classes[name][argmax(arr(name))];
return {
category: this.meta.categories[argmax(probs)],
category_probs: probs,
needs_attention: g < t.gate ? (0.5 * g) / t.gate : 0.5 + (0.5 * (g - t.gate)) / (1 - t.gate),
needs_attention_raw: raw,
blocked_on_user: sigmoid(arr("blocked_on_user")[0]),
notable: sigmoid(arr("notable")[0]),
question_line: pick("question_line"),
status_line: pick("status_line"),
action_line: pick("action_line"),
task_line: pick("task_line"),
summary_lines: summary,
option_lines: idx.filter((_, j) => qMask[j] && pOpt[j] >= t.option),
key_lines: idx.filter((_, j) => key[j] > 0),
tests: cls("tests"),
changes: cls("changes"),
risk: cls("risk"),
instruction: cls("instruction"),
progress: cls("progress"),
failure_repeats: cls("failure_repeats"),
score: cls("score"),
};
}
}
// --- The card as Daintree shows it: line picks turned into clean text ---------------------------------
// Bullets, tree and box glyphs, spinners, prompt markers: stripped only when whitespace (or the end) follows,
// so `**bold**` and `>file` survive.
const LEAD = /^(?:[⏺•●✦⎿└├│┃║■✻✶✢✳✽✺·*❯›>▸▶◆◇](?:\s+|$))+/u;
const BOX_TAIL = /\s*[│┃║]\s*$/u;
const TIMER = /\s*\((?:\d[^)]*)?(?:esc to \w+|thinking|thought for|ctrl\+\w to \w+)[^)]*\)\s*$/iu;
const SENTENCE_END = /[.!?:;)"'`]$/u;
/** One screen line as card text: glyphs, borders and timers off, words kept. null when nothing is left. */
export function displayText(line: string | undefined): string | null {
if (line == null) return null;
const text = strip(line).replace(LEAD, "").replace(BOX_TAIL, "").replace(TIMER, "").trim();
return cp(text) >= 2 ? text : null;
}
export interface Card {
category: string;
needsAttention: boolean;
score: number;
headline: string | null;
summary: string[]; // up to 3, most important first; a sentence wrapped over lines is one entry
question: string | null;
options: string[];
action: string | null;
task: string | null;
tests: string;
changes: string;
progress: number | null;
}
/** Does screen line `j` carry on the sentence of the line above it? Yes when it was picked with it, or
* when it has no bullet of its own and starts in lowercase or indented (an agent's own line break). */
function continues(lines: string[], j: number, picks: number[]): boolean {
const next = lines[j];
if (next == null || !strip(next)) return false;
if (LEAD.test(strip(next))) return false; // a new bullet is a new item, even when it was picked too
if (picks.includes(j)) return true;
return /^\s{2,}\S/u.test(next) || /^\p{Ll}/u.test(strip(next));
}
/** Line `i` as text, with up to two following lines joined when the sentence carries on over them. */
function sentenceAt(lines: string[], i: number, picks: number[], used: Set<number>): string | null {
let text = displayText(lines[i]);
used.add(i);
for (let j = i + 1; text && j <= i + 2 && !SENTENCE_END.test(text) && !used.has(j) && continues(lines, j, picks); j++) {
used.add(j);
const more = displayText(lines[j]);
if (more) text += " " + more;
}
return text;
}
/** Summary picks as text, most important first, each a whole sentence where it wraps. */
function summaryText(lines: string[], picks: number[], used: Set<number>): string[] {
const out: string[] = [];
for (const i of picks) {
if (used.has(i)) continue;
const text = sentenceAt(lines, i, picks, used);
if (text) out.push(text);
}
return out;
}
export function card(r: Reading, lines: string[]): Card {
const asks = r.category === "approval" || r.category === "question";
const line = (i: number) => (i >= 0 ? displayText(lines[i]) : null);
const used = new Set<number>();
const headline = r.status_line >= 0 ? sentenceAt(lines, r.status_line, r.summary_lines, used) : null;
return {
category: r.category,
needsAttention: r.needs_attention >= 0.5,
score: Number(r.score),
headline,
summary: summaryText(lines, r.summary_lines, used),
question: asks ? line(r.question_line) : null,
options: asks ? r.option_lines.map((i) => displayText(lines[i])).filter((x): x is string => !!x) : [],
action: r.category === "approval" ? line(r.action_line) : null,
task: line(r.task_line),
tests: r.tests,
changes: r.changes,
progress: r.progress === "none" ? null : Number(r.progress),
};
}