diff --git a/package-lock.json b/package-lock.json index 7cc211a..472d4a6 100644 --- a/package-lock.json +++ b/package-lock.json @@ -8,15 +8,21 @@ "name": "interscript-ts", "version": "0.1.0", "license": "BSD-2-Clause", + "dependencies": { + "fflate": "^0.8.3", + "js-yaml": "^5.3.0" + }, "bin": { "interscript-ts": "dist/cli.js" }, "devDependencies": { "@eslint/js": "^10.0.1", + "@types/js-yaml": "^4.0.9", "@types/node": "^22.0.0", "@vitest/coverage-v8": "^4.1.10", "eslint": "^10.8.0", "eslint-config-prettier": "^10.1.0", + "onnxruntime-node": "^1.27.0", "playwright": "^1.62.0", "prettier": "^3.9.0", "typescript": "^5.6.0", @@ -726,6 +732,13 @@ "dev": true, "license": "MIT" }, + "node_modules/@types/js-yaml": { + "version": "4.0.9", + "resolved": "https://registry.npmjs.org/@types/js-yaml/-/js-yaml-4.0.9.tgz", + "integrity": "sha512-k4MGaQl5TGo/iipqb2UDG2UwjXziSWkh0uysQelTlJpX1qGlpUZYm8PnO4DxG1qBomtJUdYJ6qR6xdIah10JLg==", + "dev": true, + "license": "MIT" + }, "node_modules/@types/json-schema": { "version": "7.0.15", "resolved": "https://registry.npmjs.org/@types/json-schema/-/json-schema-7.0.15.tgz", @@ -960,6 +973,16 @@ "acorn": "^6.0.0 || ^7.0.0 || ^8.0.0" } }, + "node_modules/adm-zip": { + "version": "0.5.18", + "resolved": "https://registry.npmjs.org/adm-zip/-/adm-zip-0.5.18.tgz", + "integrity": "sha512-ufJnssQGbxzLNS1Ho9bCtX4rQKCCvoVuDLHoJyc3F9dOGDB4BkWs2Ci0kv53lqocAEQ/Cbi+I2XCsNYGqVYqng==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=12.0" + } + }, "node_modules/ajv": { "version": "6.15.0", "resolved": "https://registry.npmjs.org/ajv/-/ajv-6.15.0.tgz", @@ -977,6 +1000,12 @@ "url": "https://github.com/sponsors/epoberezkin" } }, + "node_modules/argparse": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/argparse/-/argparse-2.0.1.tgz", + "integrity": "sha512-8+9WqebbFzpX9OR+Wa6O29asIogeRMzcGtAINdpMHHyAg10f05aSFVBbcEqGf/PXw1EjAZ+q2/bEBg3DvurK3Q==", + "license": "Python-2.0" + }, "node_modules/assertion-error": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/assertion-error/-/assertion-error-2.0.1.tgz", @@ -1079,6 +1108,42 @@ "dev": true, "license": "MIT" }, + "node_modules/define-data-property": { + "version": "1.1.4", + "resolved": "https://registry.npmjs.org/define-data-property/-/define-data-property-1.1.4.tgz", + "integrity": "sha512-rBMvIzlpA8v6E+SJZoo++HAYqsLrkg7MSfIinMPFhmkorw7X+dOXVJQs+QT69zGkzMyfDnIMN2Wid1+NbL3T+A==", + "dev": true, + "license": "MIT", + "dependencies": { + "es-define-property": "^1.0.0", + "es-errors": "^1.3.0", + "gopd": "^1.0.1" + }, + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/define-properties": { + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/define-properties/-/define-properties-1.2.1.tgz", + "integrity": "sha512-8QmQKqEASLd5nx0U1B1okLElbUuuttJ/AnYmRXbbbGDWh6uS208EjD4Xqq/I9wK7u0v6O08XhTWnt5XtEbR6Dg==", + "dev": true, + "license": "MIT", + "dependencies": { + "define-data-property": "^1.0.1", + "has-property-descriptors": "^1.0.0", + "object-keys": "^1.1.1" + }, + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, "node_modules/detect-libc": { "version": "2.1.2", "resolved": "https://registry.npmjs.org/detect-libc/-/detect-libc-2.1.2.tgz", @@ -1089,6 +1154,26 @@ "node": ">=8" } }, + "node_modules/es-define-property": { + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/es-define-property/-/es-define-property-1.0.1.tgz", + "integrity": "sha512-e3nRfgfUZ4rNGL232gUgX06QNyyez04KdjFrF+LTRoOXmrOgFKDg4BCdsjW8EnT69eqdYGmRpJwiPVYNrCaW3g==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.4" + } + }, + "node_modules/es-errors": { + "version": "1.3.0", + "resolved": "https://registry.npmjs.org/es-errors/-/es-errors-1.3.0.tgz", + "integrity": "sha512-Zf5H2Kxt2xjTvbJvP2ZWLEICxA6j+hAmMzIlypy4xcBg1vKVnx89Wy0GbS+kf5cwCVFFzdCFh2XSCFNULS6csw==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.4" + } + }, "node_modules/es-module-lexer": { "version": "2.3.1", "resolved": "https://registry.npmjs.org/es-module-lexer/-/es-module-lexer-2.3.1.tgz", @@ -1339,6 +1424,12 @@ } } }, + "node_modules/fflate": { + "version": "0.8.3", + "resolved": "https://registry.npmjs.org/fflate/-/fflate-0.8.3.tgz", + "integrity": "sha512-tbZNuJrLwGUp3zshBtdy4W+ORxZuIh8a5ilyIEQDC5rY1f3U20JMry0Ll3WBzU58EZKsEuJFXhb5gwv8CsPvgA==", + "license": "MIT" + }, "node_modules/file-entry-cache": { "version": "8.0.0", "resolved": "https://registry.npmjs.org/file-entry-cache/-/file-entry-cache-8.0.0.tgz", @@ -1418,6 +1509,52 @@ "node": ">=10.13.0" } }, + "node_modules/global-agent": { + "version": "4.1.3", + "resolved": "https://registry.npmjs.org/global-agent/-/global-agent-4.1.3.tgz", + "integrity": "sha512-KUJEViiuFT3I97t+GYMikLPJS2Lfo/S2F+DQuBWzuzaMPnvt5yyZePzArx36fBzpGTxZjIpDbXLeySLgh+k76g==", + "dev": true, + "license": "BSD-3-Clause", + "dependencies": { + "globalthis": "^1.0.2", + "matcher": "^4.0.0", + "semver": "^7.3.5", + "serialize-error": "^8.1.0" + }, + "engines": { + "node": ">=10.0" + } + }, + "node_modules/globalthis": { + "version": "1.0.4", + "resolved": "https://registry.npmjs.org/globalthis/-/globalthis-1.0.4.tgz", + "integrity": "sha512-DpLKbNU4WylpxJykQujfCcwYWiV/Jhm50Goo0wrVILAv5jOr9d+H+UR3PhSCD2rCCEIg0uc+G+muBTwD54JhDQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "define-properties": "^1.2.1", + "gopd": "^1.0.1" + }, + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, + "node_modules/gopd": { + "version": "1.2.0", + "resolved": "https://registry.npmjs.org/gopd/-/gopd-1.2.0.tgz", + "integrity": "sha512-ZUKRh6/kUFoAiTAtTYPZJ3hw9wNxx+BIBOijnlG9PnrJsCcSjs1wyyD6vJpaYtgnzDrKYRSqf3OO6Rfa93xsRg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.4" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, "node_modules/has-flag": { "version": "4.0.0", "resolved": "https://registry.npmjs.org/has-flag/-/has-flag-4.0.0.tgz", @@ -1428,6 +1565,19 @@ "node": ">=8" } }, + "node_modules/has-property-descriptors": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/has-property-descriptors/-/has-property-descriptors-1.0.2.tgz", + "integrity": "sha512-55JNKuIW+vq4Ke1BjOTjM2YctQIvCT7GFzHwmfZPGo5wnrgkid0YQtnAleFSqumZm4az3n2BS+erby5ipJdgrg==", + "dev": true, + "license": "MIT", + "dependencies": { + "es-define-property": "^1.0.0" + }, + "funding": { + "url": "https://github.com/sponsors/ljharb" + } + }, "node_modules/html-escaper": { "version": "2.0.2", "resolved": "https://registry.npmjs.org/html-escaper/-/html-escaper-2.0.2.tgz", @@ -1531,6 +1681,28 @@ "dev": true, "license": "MIT" }, + "node_modules/js-yaml": { + "version": "5.3.0", + "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-5.3.0.tgz", + "integrity": "sha512-muutsYr+e2+d3rTgUGslq5rxbBlUy3cJ61IsHag2QNDQV+7zXWjkUpmALIajhrlLlrgRUiymj6U3zUr/TMK84Q==", + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/puzrin" + }, + { + "type": "github", + "url": "https://github.com/sponsors/nodeca" + } + ], + "license": "MIT", + "dependencies": { + "argparse": "^2.0.1" + }, + "bin": { + "js-yaml": "bin/js-yaml.mjs" + } + }, "node_modules/json-buffer": { "version": "3.0.1", "resolved": "https://registry.npmjs.org/json-buffer/-/json-buffer-3.0.1.tgz", @@ -1891,6 +2063,22 @@ "url": "https://github.com/sponsors/sindresorhus" } }, + "node_modules/matcher": { + "version": "4.0.0", + "resolved": "https://registry.npmjs.org/matcher/-/matcher-4.0.0.tgz", + "integrity": "sha512-S6x5wmcDmsDRRU/c2dkccDwQPXoFczc5+HpQ2lON8pnvHlnvHAHj5WlLVvw6n6vNyHuVugYrFohYxbS+pvFpKQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "escape-string-regexp": "^4.0.0" + }, + "engines": { + "node": ">=10" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" + } + }, "node_modules/minimatch": { "version": "10.2.6", "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-10.2.6.tgz", @@ -1940,6 +2128,16 @@ "dev": true, "license": "MIT" }, + "node_modules/object-keys": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/object-keys/-/object-keys-1.1.1.tgz", + "integrity": "sha512-NuAESUOUMrlIXOfHKzD6bpPu3tYt3xvjNdRIQ+FeT0lNb4K8WR70CaDxhuNguS2XG+GjkyMwOzsN5ZktImfhLA==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 0.4" + } + }, "node_modules/obug": { "version": "2.1.4", "resolved": "https://registry.npmjs.org/obug/-/obug-2.1.4.tgz", @@ -1954,6 +2152,31 @@ "node": ">=12.20.0" } }, + "node_modules/onnxruntime-common": { + "version": "1.27.0", + "resolved": "https://registry.npmjs.org/onnxruntime-common/-/onnxruntime-common-1.27.0.tgz", + "integrity": "sha512-3KxL5wIVqa8Ex08jxSzncm9CMgw8CjOFyOQ7SxvG9o0cVLlhTNKXyIQuTbtX4tGPJEf73OER2xrjt4HJSBL4ow==", + "dev": true, + "license": "MIT" + }, + "node_modules/onnxruntime-node": { + "version": "1.27.0", + "resolved": "https://registry.npmjs.org/onnxruntime-node/-/onnxruntime-node-1.27.0.tgz", + "integrity": "sha512-QEzGwrvNBgv4uPVdnbHsOGG4G6T96mdlcFI8aAKPjMU8wOPpVocPXb6k3QGkaZagVTv2G9Bnnbo6Z3JdXr1fQw==", + "dev": true, + "hasInstallScript": true, + "license": "MIT", + "os": [ + "win32", + "darwin", + "linux" + ], + "dependencies": { + "adm-zip": "^0.5.16", + "global-agent": "^4.1.3", + "onnxruntime-common": "1.27.0" + } + }, "node_modules/optionator": { "version": "0.9.4", "resolved": "https://registry.npmjs.org/optionator/-/optionator-0.9.4.tgz", @@ -2210,6 +2433,22 @@ "node": ">=10" } }, + "node_modules/serialize-error": { + "version": "8.1.0", + "resolved": "https://registry.npmjs.org/serialize-error/-/serialize-error-8.1.0.tgz", + "integrity": "sha512-3NnuWfM6vBYoy5gZFvHiYsVbafvI9vZv/+jlIigFn4oP4zjNPK3LhcY0xSCgeb1a5L8jO71Mit9LlNoi2UfDDQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "type-fest": "^0.20.2" + }, + "engines": { + "node": ">=10" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" + } + }, "node_modules/shebang-command": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/shebang-command/-/shebang-command-2.0.0.tgz", @@ -2355,6 +2594,19 @@ "node": ">= 0.8.0" } }, + "node_modules/type-fest": { + "version": "0.20.2", + "resolved": "https://registry.npmjs.org/type-fest/-/type-fest-0.20.2.tgz", + "integrity": "sha512-Ne+eE4r0/iWnpAxD852z3A+N0Bt5RN//NjJwRd2VFHEmrywxf5vsZlh4R6lixl6B+wz/8d+maTSAkN1FIkI3LQ==", + "dev": true, + "license": "(MIT OR CC0-1.0)", + "engines": { + "node": ">=10" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" + } + }, "node_modules/typescript": { "version": "5.9.3", "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.9.3.tgz", diff --git a/package.json b/package.json index 95089c1..3971fd1 100644 --- a/package.json +++ b/package.json @@ -66,14 +66,20 @@ "homepage": "https://github.com/interscript/interscript-ts#readme", "devDependencies": { "@eslint/js": "^10.0.1", + "@types/js-yaml": "^4.0.9", "@types/node": "^22.0.0", "@vitest/coverage-v8": "^4.1.10", "eslint": "^10.8.0", "eslint-config-prettier": "^10.1.0", + "onnxruntime-node": "^1.27.0", "playwright": "^1.62.0", "prettier": "^3.9.0", "typescript": "^5.6.0", "typescript-eslint": "^8.65.0", "vitest": "^4.1.0" + }, + "dependencies": { + "fflate": "^0.8.3", + "js-yaml": "^5.3.0" } } diff --git a/src/ml/imf/index.ts b/src/ml/imf/index.ts new file mode 100644 index 0000000..68c8d55 --- /dev/null +++ b/src/ml/imf/index.ts @@ -0,0 +1,6 @@ +/** @interscript/ml — the IMF v1 runtime for TypeScript. */ + +export { IMFModel } from "./model.js" +export { IMFError, parseManifest, verifyAndRead, type IMFManifest } from "./loader.js" +export { resolve, DEFAULT_INDEX_URL, RegistryError, type IndexEntry } from "./registry.js" +export { encode, decode, BYTE_OFFSET, EOS_ID, PAD_ID, UNK_ID } from "./tokens.js" diff --git a/src/ml/imf/loader.ts b/src/ml/imf/loader.ts new file mode 100644 index 0000000..1549035 --- /dev/null +++ b/src/ml/imf/loader.ts @@ -0,0 +1,68 @@ +/** + * IMF v1 zip loading: manifest parse, per-graph sha256 verification. + * Corrupt downloads fail loudly, before any session is created. + */ + +import { unzipSync } from "fflate" +import { load as loadYaml } from "js-yaml" + +export interface IMFManifest { + format: string + id: string + task: string + decoder: string + precision: string + opset: number + sha256: Record +} + +export class IMFError extends Error {} + +export function parseManifest(zipBytes: Uint8Array): IMFManifest { + const files = unzipSync(zipBytes) + const metaBytes = files["metadata.yaml"] + if (!metaBytes) throw new IMFError("missing metadata.yaml") + const raw = loadYaml(new TextDecoder().decode(metaBytes)) as Record + if (raw["format"] !== "imf-v1") throw new IMFError(`unsupported format: ${String(raw["format"])}`) + if (raw["tokenizer"] !== "bytes") { + throw new IMFError(`tokenizer ${String(raw["tokenizer"])}: this runtime is byte-level only`) + } + if (!files["encoder.onnx"] || !files["decoder.onnx"]) { + throw new IMFError("missing encoder.onnx / decoder.onnx") + } + return { + format: "imf-v1", + id: String(raw["id"]), + task: String(raw["task"]), + decoder: raw["decoder"] === undefined ? "plain" : String(raw["decoder"]), + precision: raw["precision"] === undefined ? "fp32" : String(raw["precision"]), + opset: raw["opset"] === undefined ? 14 : Number(raw["opset"]), + sha256: (raw["sha256"] as Record) ?? {}, + } +} + +async function sha256(data: Uint8Array): Promise { + const buffer = new Uint8Array(data).buffer as ArrayBuffer + const digest = await crypto.subtle.digest("SHA-256", buffer) + return Array.from(new Uint8Array(digest)) + .map((b) => b.toString(16).padStart(2, "0")) + .join("") +} + +/** Unzip + sha256-verify every .onnx member; returns name -> bytes. */ +export async function verifyAndRead(zipBytes: Uint8Array): Promise> { + const manifest = parseManifest(zipBytes) + const files = unzipSync(zipBytes) + const graphs = new Map() + for (const [name, bytes] of Object.entries(files)) { + if (!name.endsWith(".onnx")) continue + const recorded = manifest.sha256[name] + if (!recorded) throw new IMFError(`${name} is not covered by metadata sha256`) + const actual = await sha256(bytes) + if (actual !== recorded) { + throw new IMFError(`${name} sha256 mismatch: zip has ${actual}, metadata says ${recorded}`) + } + graphs.set(name, bytes) + } + return graphs +} diff --git a/src/ml/imf/model.ts b/src/ml/imf/model.ts new file mode 100644 index 0000000..53f1457 --- /dev/null +++ b/src/ml/imf/model.ts @@ -0,0 +1,202 @@ +/** + * IMFModel — a loaded, checksum-verified IMF v1 model. Greedy KV-cache + * decode when the zip ships decoder-kv.onnx (default), plain + * full-recompute fallback otherwise. The decode loop is the shared + * cross-runtime contract; outputs are byte-identical with the Python + * reference on the golden sets. + */ + +import { createSession, type InferenceSession } from "../session/index.js" +import type { Tensor } from "../types.js" +import { verifyAndRead, parseManifest, type IMFManifest } from "./loader.js" +import { resolve } from "./registry.js" +import { EOS_ID, PAD_ID, decode, encode } from "./tokens.js" + +interface InputMeta { + readonly name: string + readonly type: string + readonly shape: ReadonlyArray +} + +interface MetadataSession extends InferenceSession { + readonly inputMetadata?: readonly InputMeta[] +} + +export class IMFModel { + readonly id: string + private readonly manifest: IMFManifest + private readonly encoder: InferenceSession + private readonly decoder: MetadataSession + private readonly kv: boolean + private readonly pasts: ReadonlyArray<{ name: string; dims: readonly number[] }> + + private constructor( + manifest: IMFManifest, + encoder: InferenceSession, + decoder: MetadataSession, + ) { + this.manifest = manifest + this.id = manifest.id + this.encoder = encoder + this.decoder = decoder + this.kv = manifest.decoder === "kv" && decoder.inputNames.some((n) => n.startsWith("past_")) + this.pasts = this.kv ? this.zeroPastSpecs() : [] + } + + static async fromZipBytes(zipBytes: Uint8Array): Promise { + const manifest = parseManifest(zipBytes) + const graphs = await verifyAndRead(zipBytes) + const encoder = await createSession(graphs.get("encoder.onnx")!) + const decoderName = manifest.decoder === "kv" && graphs.has("decoder-kv.onnx") + ? "decoder-kv.onnx" + : "decoder.onnx" + const decoder = (await createSession(graphs.get(decoderName)!)) as MetadataSession + return new IMFModel(manifest, encoder, decoder) + } + + /** Accepts a zip path (Node), raw zip bytes, or a models.yaml model id. */ + static async load(source: string | Uint8Array, indexUrl?: string): Promise { + if (source instanceof Uint8Array) return IMFModel.fromZipBytes(source) + if (source.endsWith(".zip")) { + const fs = (await import("node:fs")) as { readFileSync(path: string): Uint8Array } + return IMFModel.fromZipBytes(fs.readFileSync(source)) + } + const resolved = await resolve(source, indexUrl) + return IMFModel.fromZipBytes(resolved.bytes) + } + + async translate(text: string, maxLen = 256): Promise { + const ids = encode(text) + if (ids.length === 1) return "" + const hidden = await this.runEncoder(ids) + const tokens = this.kv ? await this.greedyKv(hidden, maxLen) : await this.greedyPlain(hidden, maxLen) + return decode(tokens) + } + + async dispose(): Promise { + await this.encoder.dispose() + await this.decoder.dispose() + } + + private async runEncoder(ids: readonly number[]): Promise { + const outputs = await this.encoder.run({ + input_ids: { + name: "input_ids", + type: "int64", + data: new BigInt64Array(ids.map((n) => BigInt(n))), + dims: [1, ids.length], + }, + }) + return outputs["last_hidden_state"]! + } + + private zeroPastSpecs(): ReadonlyArray<{ name: string; dims: readonly number[] }> { + const specs: { name: string; dims: number[] }[] = [] + const metadata = this.decoder.inputMetadata + for (const name of this.decoder.inputNames) { + if (!name.startsWith("past_")) continue + // [batch, heads, past_seq, d_kv]: heads and d_kv are static + let heads = 6 + let dKv = 64 + const meta = metadata?.find((m) => m.name === name) + if (meta && typeof meta.shape[1] === "number") heads = meta.shape[1] as number + if (meta && typeof meta.shape[3] === "number") dKv = meta.shape[3] as number + specs.push({ name, dims: [1, heads, 0, dKv] }) + } + return specs + } + + private pastTensors( + present: ReadonlyMap | undefined, + ): Record { + const feeds: Record = {} + for (const spec of this.pasts) { + const value = present?.get(spec.name) + if (value) { + feeds[spec.name] = { name: spec.name, type: value.type, data: value.data, dims: value.dims } + } else { + feeds[spec.name] = { + name: spec.name, + type: "float32", + data: new Float32Array(0), + dims: [...spec.dims], + } + } + } + return feeds + } + + private argmaxLastStep(logits: Tensor): number { + const dims = logits.dims + const classes = dims[dims.length - 1]! + const data = logits.data as Float32Array | BigInt64Array + const base = (dims[dims.length - 2]! - 1) * classes + let best = 0 + let bestVal = -Infinity + for (let c = 0; c < classes; c++) { + const v = typeof data[base + c] === "bigint" ? Number(data[base + c]) : (data[base + c] as number) + if (v > bestVal) { + bestVal = v + best = c + } + } + return best + } + + private async greedyKv(hidden: Tensor, maxLen: number): Promise { + const generated: number[] = [] + let current = [PAD_ID] + let present: ReadonlyMap | undefined + for (let step = 0; step < maxLen; step++) { + const outputs = await this.decoder.run({ + input_ids: { + name: "input_ids", + type: "int64", + data: new BigInt64Array(current.map((n) => BigInt(n))), + dims: [1, current.length], + }, + encoder_hidden_states: { + name: "encoder_hidden_states", + type: hidden.type, + data: hidden.data, + dims: hidden.dims, + }, + ...this.pastTensors(present), + }) + const token = this.argmaxLastStep(outputs["logits"]!) + if (token === EOS_ID) break + generated.push(token) + present = new Map( + this.pasts.map((spec) => [spec.name, outputs[spec.name.replace("past_", "present_")]!]), + ) + current = [token] + } + return generated + } + + private async greedyPlain(hidden: Tensor, maxLen: number): Promise { + const generated: number[] = [] + const decoderIds: number[] = [PAD_ID] + for (let step = 0; step < maxLen; step++) { + const outputs = await this.decoder.run({ + input_ids: { + name: "input_ids", + type: "int64", + data: new BigInt64Array(decoderIds.map((n) => BigInt(n))), + dims: [1, decoderIds.length], + }, + encoder_hidden_states: { + name: "encoder_hidden_states", + type: hidden.type, + data: hidden.data, + dims: hidden.dims, + }, + }) + const token = this.argmaxLastStep(outputs["logits"]!) + if (token === EOS_ID) break + generated.push(token) + decoderIds.push(token) + } + return generated + } +} diff --git a/src/ml/imf/registry.ts b/src/ml/imf/registry.ts new file mode 100644 index 0000000..a6eb3a7 --- /dev/null +++ b/src/ml/imf/registry.ts @@ -0,0 +1,188 @@ +/** + * models.yaml resolution (the dynamic-fetch contract shared with the + * Python and Ruby runtimes): resolve id -> channel URL, verify a cached + * copy against the index sha256, or download -> verify -> install. + * + * Node persists to ~/.cache/interscript/models// (fs, atomic + * rename); browsers keep the verified bytes in memory (the Cache API + * integration is future work). Overrides: INTERSCRIPT_ML_INDEX, + * INTERSCRIPT_ML_CACHE. + */ + +import { load as loadYaml } from "js-yaml" + +export const DEFAULT_INDEX_URL = + "https://raw.githubusercontent.com/interscript/ml-models/main/models.yaml" + +export interface Part { + url: string + sha256: string + size: number +} + +export interface IndexEntry { + filename: string + url: string + sha256: string + parts?: Part[] +} + +export class RegistryError extends Error {} + +interface NodeFs { + readFileSync(path: string): Uint8Array + writeFileSync(path: string, data: Uint8Array): void + appendFileSync(path: string, data: Uint8Array): void + mkdirSync(path: string, opts: { recursive: boolean }): void + renameSync(from: string, to: string): void + existsSync(path: string): boolean +} + +async function nodeFs(): Promise { + const g = globalThis as { process?: { versions?: { node?: string } } } + if (!g.process?.versions?.node) return undefined + return (await import("node:fs")) as unknown as NodeFs +} + +function cacheDir(): string { + const home = process.env["HOME"] ?? process.env["USERPROFILE"] ?? "." + return process.env["INTERSCRIPT_ML_CACHE"] ?? `${home}/.cache/interscript` +} + +async function fetchIndex(source: string): Promise> { + const text = source.startsWith("http://") || source.startsWith("https://") + ? await (await fetch(source)).text() + : new TextDecoder().decode((await nodeFs())!.readFileSync(source)) + const raw = loadYaml(text) as { + version?: number + models?: Record> + } + if (raw.version !== 1) throw new RegistryError("index must have version: 1") + const entries: Record = {} + for (const [id, spec] of Object.entries(raw.models ?? {})) { + entries[id] = { + filename: spec["filename"] as string, + url: (spec["url"] as string) ?? "", + sha256: spec["sha256"] as string, + parts: spec["parts"] as Part[] | undefined, + } + } + return entries +} + +async function sha256Hex(data: Uint8Array): Promise { + const digest = await crypto.subtle.digest("SHA-256", new Uint8Array(data).buffer as ArrayBuffer) + return Array.from(new Uint8Array(digest)).map((b) => b.toString(16).padStart(2, "0")).join("") +} + +async function partBytes(fs: NodeFs | undefined, url: string): Promise { + if (url.startsWith("file://")) { + if (!fs) throw new RegistryError("file:// parts require a Node host") + return fs.readFileSync(url.replace(/^file:\/\//, "")) + } + return new Uint8Array(await (await fetch(url)).arrayBuffer()) +} + +async function fetchParts( + entry: IndexEntry, + fs: NodeFs | undefined, + onPart: (bytes: Uint8Array) => void, +): Promise { + const parts = entry.parts! + for (let i = 0; i < parts.length; i++) { + const bytes = await partBytes(fs, parts[i]!.url) + const actual = await sha256Hex(bytes) + if (actual !== parts[i]!.sha256) { + throw new RegistryError( + `part ${i} of ${entry.filename} sha256 mismatch: got ${actual}, index says ${parts[i]!.sha256}`, + ) + } + onPart(bytes) + } +} + +export interface ResolvedZip { + bytes: Uint8Array + path?: string +} + +export async function resolve( + modelId: string, + indexUrl?: string, +): Promise { + const source = indexUrl ?? process.env["INTERSCRIPT_ML_INDEX"] ?? DEFAULT_INDEX_URL + const entries = await fetchIndex(source) + const entry = entries[modelId] + if (!entry) { + throw new RegistryError(`unknown model id '${modelId}' (known: ${Object.keys(entries).sort().join(", ")})`) + } + + const fs = await nodeFs() + const target = `${cacheDir()}/models/${modelId}/${entry.filename}` + if (fs?.existsSync(target)) { + const cached = fs.readFileSync(target) + if ((await sha256Hex(cached)) === entry.sha256) return { bytes: cached, path: target } + } + + if (entry.parts?.length) { + if (fs) { + const { createHash } = await import("node:crypto") + const whole = createHash("sha256") + const dir = target.substring(0, target.lastIndexOf("/")) + fs.mkdirSync(dir, { recursive: true }) + const tmp = `${target}.part.${process.pid}` + fs.writeFileSync(tmp, new Uint8Array(0)) + await fetchParts(entry, fs, (bytes) => { + whole.update(bytes) + fs.appendFileSync(tmp, bytes) + }) + const actual = whole.digest("hex") + if (actual !== entry.sha256) { + throw new RegistryError( + `assembled ${entry.filename} sha256 mismatch: got ${actual}, index says ${entry.sha256}`, + ) + } + fs.renameSync(tmp, target) + return { bytes: fs.readFileSync(target), path: target } + } + // browser: parts are all small (client-tier models); concat in memory + const chunks: Uint8Array[] = [] + let total = 0 + await fetchParts(entry, fs, (bytes) => { + chunks.push(bytes) + total += bytes.length + }) + const bytes = new Uint8Array(total) + let offset = 0 + for (const chunk of chunks) { + bytes.set(chunk, offset) + offset += chunk.length + } + const actual = await sha256Hex(bytes) + if (actual !== entry.sha256) { + throw new RegistryError( + `assembled ${entry.filename} sha256 mismatch: got ${actual}, index says ${entry.sha256}`, + ) + } + return { bytes } + } + + const bytes = entry.url.startsWith("file://") + ? fs!.readFileSync(entry.url.replace(/^file:\/\//, "")) + : new Uint8Array(await (await fetch(entry.url)).arrayBuffer()) + const actual = await sha256Hex(bytes) + if (actual !== entry.sha256) { + throw new RegistryError( + `downloaded ${entry.filename} sha256 mismatch: got ${actual}, index says ${entry.sha256}`, + ) + } + if (fs) { + const dir = target.substring(0, target.lastIndexOf("/")) + fs.mkdirSync(dir, { recursive: true }) + const tmp = `${target}.part.${process.pid}` + fs.writeFileSync(tmp, bytes) + fs.renameSync(tmp, target) + return { bytes, path: target } + } + return { bytes } +} diff --git a/src/ml/imf/tokens.ts b/src/ml/imf/tokens.ts new file mode 100644 index 0000000..7705583 --- /dev/null +++ b/src/ml/imf/tokens.ts @@ -0,0 +1,32 @@ +/** + * The canonical ByT5 byte table (shared contract with the Python and + * Ruby runtimes): byte b -> token id b+3, trailing EOS; pad=0, unk=2. + * Token ids are NOT raw byte values — feeding TextEncoder output + * directly as ids silently produces garbage on real checkpoints. + */ + +export const BYTE_OFFSET = 3 +export const PAD_ID = 0 +export const EOS_ID = 1 +export const UNK_ID = 2 + +const encoder = new TextEncoder() +const decoder = new TextDecoder("utf-8", { fatal: false }) + +export function encode(text: string): number[] { + const bytes = encoder.encode(text) + const ids: number[] = new Array(bytes.length + 1) + for (let i = 0; i < bytes.length; i++) ids[i] = bytes[i]! + BYTE_OFFSET + ids[bytes.length] = EOS_ID + return ids +} + +export function decode(tokenIds: readonly number[]): string { + const out: number[] = [] + for (const token of tokenIds) { + if (token === EOS_ID) break + if (token === PAD_ID || token === UNK_ID) continue + out.push((token - BYTE_OFFSET) % 256) + } + return decoder.decode(new Uint8Array(out)) +} diff --git a/src/ml/session/onnx-node.ts b/src/ml/session/onnx-node.ts index 3a539e1..8f055e7 100644 --- a/src/ml/session/onnx-node.ts +++ b/src/ml/session/onnx-node.ts @@ -46,12 +46,17 @@ class NodeInferenceSession implements InferenceSession { private readonly ort: OnnxNodeModule readonly inputNames: readonly string[] readonly outputNames: readonly string[] + readonly inputMetadata?: ReadonlyArray<{ name: string; type: string; shape: ReadonlyArray }> | undefined private constructor(session: OnnxNodeSession, ort: OnnxNodeModule) { this.session = session this.ort = ort this.inputNames = session.inputNames this.outputNames = session.outputNames + const raw = (session as { inputMetadata?: unknown }).inputMetadata + if (Array.isArray(raw)) { + this.inputMetadata = raw as NodeInferenceSession["inputMetadata"] + } } static async create(modelData: ArrayBuffer | Uint8Array): Promise { diff --git a/src/ml/types.ts b/src/ml/types.ts index c65f42b..c6e118f 100644 --- a/src/ml/types.ts +++ b/src/ml/types.ts @@ -47,11 +47,20 @@ export interface InferenceOutputs { * An inference session wraps an ONNX model. Both Node and browser * backends implement this interface; callers don't know which. */ +export interface SessionInputMetadata { + readonly name: string + readonly type: string + readonly shape: ReadonlyArray +} + export interface InferenceSession { /** Run inference with named inputs. Returns named outputs. */ run(inputs: InferenceInputs): Promise /** Input names the model accepts. */ inputNames: readonly string[] + /** Input shapes/types when the backend exposes them (used for + * zero-length KV pasts whose static dims must match the graph). */ + readonly inputMetadata?: readonly SessionInputMetadata[] | undefined /** Output names the model produces. */ outputNames: readonly string[] /** Free native resources held by the session. */ diff --git a/test/fixtures/tiny-imf.zip b/test/fixtures/tiny-imf.zip new file mode 100644 index 0000000..c179973 Binary files /dev/null and b/test/fixtures/tiny-imf.zip differ diff --git a/test/imf.test.ts b/test/imf.test.ts new file mode 100644 index 0000000..c135b01 --- /dev/null +++ b/test/imf.test.ts @@ -0,0 +1,186 @@ +/** + * IMF v1 runtime tests — the TypeScript side of the cross-runtime + * contract. Tiny-graph zips for CI (no download); the golden e2e runs + * when INTERSCRIPT_TS_E2E_ZIP points at a real zip. + */ + +import { readFileSync } from "node:fs" +import { describe, expect, it } from "vitest" +import { IMFModel, IMFError, RegistryError, decode, encode, resolve } from "../src/ml/imf/index.js" +import { parseManifest, verifyAndRead } from "../src/ml/imf/loader.js" + +const fixtureZip = new Uint8Array(readFileSync("test/fixtures/tiny-imf.zip")) + +describe("canonical ByT5 byte table", () => { + it("encodes byte+3 with trailing EOS", () => { + expect(encode("rok")).toEqual([117, 114, 110, 1]) + }) + it("decodes with EOS stop", () => { + expect(decode([117, 114, 110])).toBe("rok") + expect(decode([117, 1, 114])).toBe("r") + }) +}) + +describe("IMF zip loading", () => { + it("parses and verifies the manifest", async () => { + const manifest = parseManifest(fixtureZip) + expect(manifest.format).toBe("imf-v1") + expect(manifest.id).toBe("tiny-1.0") + const graphs = await verifyAndRead(fixtureZip) + expect([...graphs.keys()].sort()).toEqual(["decoder.onnx", "encoder.onnx"]) + }) + + it("rejects tampered bytes loudly", async () => { + const tampered = new Uint8Array(fixtureZip) + tampered[tampered.length - 5] ^= 0xff + await expect(verifyAndRead(tampered)).rejects.toThrow(IMFError) + }) + + it("loads sessions from verified bytes", async () => { + const model = await IMFModel.fromZipBytes(fixtureZip) + expect(model.id).toBe("tiny-1.0") + await model.dispose() + }, + 600_000, + ) +}) + +describe("registry", () => { + it("resolves, verifies, and caches from a local index", async () => { + const { mkdtempSync, writeFileSync, mkdirSync, existsSync, rmSync } = await import("node:fs") + const { tmpdir } = await import("node:os") + const { join } = await import("node:path") + const dir = mkdtempSync(join(tmpdir(), "imf-reg-")) + try { + mkdirSync(join(dir, "channel")) + writeFileSync(join(dir, "channel", "tiny.zip"), fixtureZip) + const { createHash } = await import("node:crypto") + const sha = createHash("sha256").update(fixtureZip).digest("hex") + writeFileSync( + join(dir, "models.yaml"), + `version: 1\nmodels:\n tiny-1.0:\n filename: tiny.zip\n url: file://${dir}/channel/tiny.zip\n sha256: ${sha}\n`, + ) + const cache = join(dir, "cache") + process.env["INTERSCRIPT_ML_CACHE"] = cache + try { + const resolved = await resolve("tiny-1.0", join(dir, "models.yaml")) + expect(resolved.path).toBe(join(cache, "models", "tiny-1.0", "tiny.zip")) + expect(existsSync(resolved.path!)).toBe(true) + rmSync(join(dir, "channel", "tiny.zip")) + const again = await resolve("tiny-1.0", join(dir, "models.yaml")) + expect(again.path).toBe(resolved.path) + } finally { + delete process.env["INTERSCRIPT_ML_CACHE"] + } + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) + + it("raises for unknown ids", async () => { + const { mkdtempSync, writeFileSync, rmSync } = await import("node:fs") + const { tmpdir } = await import("node:os") + const { join } = await import("node:path") + const dir = mkdtempSync(join(tmpdir(), "imf-unk-")) + const index = join(dir, "models.yaml") + writeFileSync(index, "version: 1\nmodels: {}\n") + try { + await expect(resolve("nope-1.0", index)).rejects.toThrow(RegistryError) + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) + + it("assembles split parts and verifies per-part + whole-file sha256", async () => { + const { mkdtempSync, writeFileSync, mkdirSync, rmSync, readFileSync } = await import("node:fs") + const { tmpdir } = await import("node:os") + const { join } = await import("node:path") + const dir = mkdtempSync(join(tmpdir(), "imf-parts-")) + try { + mkdirSync(join(dir, "channel")) + const blob = fixtureZip + const split = Math.floor(blob.length / 2) + 3 + const partA = blob.subarray(0, split) + const partB = blob.subarray(split) + writeFileSync(join(dir, "channel", "tiny.zip.part-00"), partA) + writeFileSync(join(dir, "channel", "tiny.zip.part-01"), partB) + const { createHash } = await import("node:crypto") + const sha = (b: Uint8Array) => createHash("sha256").update(b).digest("hex") + writeFileSync( + join(dir, "models.yaml"), + `version: 1\nmodels:\n tiny-1.0:\n filename: tiny.zip\n sha256: ${sha(blob)}\n parts:\n` + + ` - url: file://${dir}/channel/tiny.zip.part-00\n sha256: ${sha(partA)}\n size: ${partA.length}\n` + + ` - url: file://${dir}/channel/tiny.zip.part-01\n sha256: ${sha(partB)}\n size: ${partB.length}\n`, + ) + const cache = join(dir, "cache") + process.env["INTERSCRIPT_ML_CACHE"] = cache + try { + const resolved = await resolve("tiny-1.0", join(dir, "models.yaml")) + expect(resolved.path).toBe(join(cache, "models", "tiny-1.0", "tiny.zip")) + expect(new Uint8Array(readFileSync(resolved.path!))).toEqual(blob) + rmSync(join(dir, "channel", "tiny.zip.part-00")) + const again = await resolve("tiny-1.0", join(dir, "models.yaml")) + expect(again.path).toBe(resolved.path) + } finally { + delete process.env["INTERSCRIPT_ML_CACHE"] + } + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) + + it("rejects a corrupt part by index", async () => { + const { mkdtempSync, writeFileSync, mkdirSync, rmSync } = await import("node:fs") + const { tmpdir } = await import("node:os") + const { join } = await import("node:path") + const dir = mkdtempSync(join(tmpdir(), "imf-badpart-")) + try { + mkdirSync(join(dir, "channel")) + const partA = fixtureZip.subarray(0, 7) + const partB = fixtureZip.subarray(7) + writeFileSync(join(dir, "channel", "tiny.zip.part-00"), partA) + writeFileSync(join(dir, "channel", "tiny.zip.part-01"), partB) + const { createHash } = await import("node:crypto") + const sha = (b: Uint8Array) => createHash("sha256").update(b).digest("hex") + writeFileSync( + join(dir, "models.yaml"), + `version: 1\nmodels:\n tiny-1.0:\n filename: tiny.zip\n sha256: ${sha(fixtureZip)}\n parts:\n` + + ` - url: file://${dir}/channel/tiny.zip.part-00\n sha256: ${"0".repeat(64)}\n size: ${partA.length}\n` + + ` - url: file://${dir}/channel/tiny.zip.part-01\n sha256: ${sha(partB)}\n size: ${partB.length}\n`, + ) + process.env["INTERSCRIPT_ML_CACHE"] = join(dir, "cache") + try { + await expect(resolve("tiny-1.0", join(dir, "models.yaml"))).rejects.toThrow( + /part 0 .* sha256 mismatch/, + ) + } finally { + delete process.env["INTERSCRIPT_ML_CACHE"] + } + } finally { + rmSync(dir, { recursive: true, force: true }) + } + }) +}) + +const e2eZip = process.env["INTERSCRIPT_TS_E2E_ZIP"] + +describe("golden set e2e", () => { + it.skipIf(!e2eZip)( + "matches the Python reference byte-for-byte", + async () => { + const zip = e2eZip! + const model = await IMFModel.load(zip) + const goldenPath = process.env["INTERSCRIPT_TS_GOLDEN"] + ?? "/Users/mulgogi/src/interscript/ml-models/golden/khm-latn-100.jsonl" + const rows = readFileSync(goldenPath, "utf-8") + .split("\n") + .filter((l) => l.trim()) + .map((l) => JSON.parse(l) as { input: string; output: string }) + for (const row of rows) { + expect(await model.translate(row.input, 128), row.input).toBe(row.output) + } + await model.dispose() + }, + 600_000, + ) +})