feat(cli): search the catalog by meaning, on this machine (#3089)

* feat(cli): search the catalog by meaning, in three named tiers

Browsing the registry means matching names and tags, which fails whenever the
author's wording differs from yours. "make the pace feel faster" finds nothing
when the move is described as "velocity-driven blur". This ranks by meaning
instead.

Three tiers, and the command always says which one answered:

  words       shared vocabulary, free, offline, no account
  on-device   bge-small, free, offline, one opt-in download
  hosted      Gemini, free for signed-in HeyGen users

The tier is stated because a quietly worse answer looks exactly like a good
one. --json carries it as a token alongside dropped, shown, total and
top_score, so an agent reads provenance as data rather than matching English
that is written to be reworded.

Two consents, asked once each, and never conflated. Sending a query is a
privacy question, so the prompt says the query is sent. Downloading a model is
a disk and bandwidth question, so that prompt talks about size. Neither fires
without a terminal: an unattended run sends nothing and downloads nothing
unless a flag records that a person agreed.

The catalog is derived from registry-item.json rather than from a separate
document, so the set that is ranked and the set that can be installed are the
same object by construction. Only the on-device vectors are committed; the
hosted vectors are nine megabytes and belong on the server.

top_score is reported and never acted on. A "nothing matched" threshold looked
clean on long briefs and collapsed on the short queries people type: "a logo
appears" scores 0.6181 and keyboard mash scores 0.6417, so any cut that catches
the noise rejects the real query. The measurement is in the evals directory
rather than in this branch.

Not covered here. The published recall figures were measured against a separate
hand-written document, not against registry text, so they should not be quoted
for this catalog until re-measured. The offline tier needs a normal install: a
single-file build cannot load the native ONNX runtime, which the command now
reports instead of silently degrading. And the drop-detection path has never
been observed firing outside its author's tests.

* fix(cli): make this branch pass the repo's own gates

Three things `bun run lint` and `fallow audit --base origin/main` rejected.
CI runs both, so none of this branch would have gone green. Found by running
them, not by reading the diff.

process.exit in catalog.ts, twice: an invalid --type and a cancelled picker.
check:cli-process-ownership reserves that for cli.ts, and the rule is not
cosmetic — process.exit tears the process down where it stands, so anything
cli.ts has queued to run on the way out is dropped. finishCommand throws a
CliResultSignal that cli.ts turns into the exit code, which is what init.ts
already does for a cancelled prompt.

Three exports with no consumers. normalize keeps its body and loses its export;
localEmbedder is the only caller. modelsDirectory goes entirely, having no
caller inside its file or out. The WordPieceConfig re-export goes, and with it
the import it existed to forward: the type is exported from wordpiece.ts, where
its consumers already take it from.

Complexity. prepareOnDeviceTier is lifted out of run(), which took run from 64
cyclomatic and CRAP 948 to 54 and 684. That block is one decision — can the
offline tier run, and if not, why not — and its only product is a list of
warnings, so it reads and tests as a unit, which it could not do inline.

The rest is suppressed rather than refactored, each with its reason on the line
above. Finishing run() means extracting its three output paths, and that is a
refactor of a command this branch already changes for other reasons: a separate
initiative, not something to absorb here. Every suppression says what shape the
function has and why; a bare marker on a function nobody can justify is how a
threshold stops meaning anything.

Verified: `bun run lint` exits 0, fallow reports no issues across 27 changed
files, and 2540 CLI tests pass.

* feat(cli): ship the local search tiers only, drop the hosted one

Search now has two tiers, both local: shared-vocabulary word matching, and the
opt-in on-device model. The hosted tier, which sent the query to a HeyGen
endpoint and ranked it with a hosted model, is removed.

This is a scope decision, not a defect. The endpoint works and its own change is
reviewed and green; it is simply not what we want to ship first. Landing local
only means the feature has no backend dependency, no auth requirement, and
nothing leaves the machine unless someone opts into downloading a model.

Gone: registry/smartSearch.ts and its test, the --smart and --no-smart flags,
the outcome plumbing through the command, the remote branch of applySearch, the
remote tier, and the hosted-only JSON fields (ranking, catalog_version,
top_score). Also the smartSearchEnabled consent field in telemetry config, which
was the persisted storage behind the hosted consent and would otherwise have
been left as dead configuration surface.

Kept exactly as they were: both local tiers, the --on-device and --yes flags,
the download consent prompt, and the runtime check that happens before the
download rather than after it. The --json envelope still reports query, tier,
tier_detail, shown, total, dropped, warnings and results, so an agent can still
tell which tier answered and why. tierToken now distinguishes on-device from
words.

Verified: lint exits 0, fallow reports no issues, 2522 CLI tests pass, and the
command was exercised directly. A query answers on the on-device tier where the
model is installed and falls back to word matching where it is not, reporting
that fallback in warnings rather than silently. An unknown --type still exits 1
with a readable message, and --smart is now rejected as an unknown flag.

* fix(cli): count only moves this registry cannot install as dropped

The dropped count was computed against the list left after the user's own
--type and --tag filters, so every move the user excluded was reported as one
the registry is missing. Filtering made the number go up: the same query
reported 277 unfiltered and 302 with --type block.

The count exists so a caller can tell "nothing matched your words" apart from
"the ranker suggested things this project cannot install". Conflating it with
user filtering destroys exactly that signal, and worse, genuine index skew and a
self-inflicted filter printed a byte-identical line with opposite remedies --
one means refresh the shelf, the other means drop a flag, and refreshing does
nothing.

Now counted against the registry rather than the filtered view. The manifest is
already fetched whole and narrowed in memory, so keeping the unnarrowed name set
costs no extra request, and item loading still runs only on the filtered subset.

Verified against ground truth rather than by eye: the vector artifact holds 411
names, the registry holds 168 installable items, and 134 of those names exist in
both, so 277 are genuinely uninstallable. The count now reads 277 unfiltered,
277 under --type block, 277 under --type component and 277 under --tag, and the
skew it reports is real -- the artifact predates dropping the UI primitives and
still ranks moves that are no longer on the shelf.

Reported by Vance Ingalls, who also noted this closes an item the status doc
listed as unverified. Two earlier sweeps could not make the count fire because
neither combined a filter with a query.

Tests pin the three cases: a genuinely absent name counts, a filter-excluded
name does not, and a fully installable ranking reports zero.

* fix(cli): tell the user when meaning search cannot see the catalog

The on-device index was fetched once and never revalidated: the only
freshness check was two existsSync calls. A move added after that fetch was
invisible to meaning search permanently, not down-ranked but absent from the
candidate set. The registry manifest on the same command carries a 24h TTL,
so the two halves of one feature disagreed about staleness.

The dropped count reported over-coverage only, names the index has that the
registry lacks. Under-coverage was never computed, so the harmless direction
was instrumented and the costly one was silent. Reproduced with an index
truncated to 120 of 168 moves: dropped read 0, perfect health, while 48
moves were unreachable.

Counts under-coverage from the name list the artifact already carries, so no
extra request. Warns only when non-zero, and names the remedy.

The remedy had to be made true: --on-device could not refresh a stale index
because hasLocalVectors short-circuited the fetch. That flag now refetches
when the index is absent or no longer covering.

Two defects the reproduction surfaced. A failed refresh reported the tier
unavailable while the old vectors were still on disk and still ranking. And
the fetch wrote its two files one at a time, so failing between them paired
a new name list with an old matrix, a hard load error rather than stale
data. It now writes both or neither, which matters more once refresh runs on
staleness.

top_score returns, scoped to the on-device tier and set to the score of the
best result actually shown rather than the ranking head, which can describe
a row the caller never received.

Also: scripts/ is now typechecked. It never was, which is how a build script
that crashes after the paid embedding call, and two scripts whose imports do
not resolve at all, went unnoticed. 43 errors fixed, no suppressions.

And the docs stop describing a --smart hosted tier that was deleted, an
item that does not exist, and a registry refresh that cannot fix a stale
vector index.

* ci: fail when the search index stops covering the registry

The catalog vector artifact is regenerated by hand. Nothing in CI, in
package.json or in a hook rebuilds it, because embedding needs the 32 MB
model. So adding a registry item silently makes it invisible to meaning
search until someone remembers to regenerate.

The failure is asymmetric, which is what makes it easy to miss. Removing an
item is self-healing: the ranker still scores the dead vector, then filters
the name before display, so a user is never offered something they cannot
install. Adding one is not: the item is absent from the candidate set
entirely, not ranked low.

Comparing the two name lists needs neither the model nor a network call, so
the gate runs in seconds. CI checks rather than fixes, for the same reason
it cannot regenerate.

Scoped to blocks and components. Examples are starter projects a user
scaffolds, never something catalog ranks, and the artifact carries no vector
for them, so demanding one would keep this gate permanently red and it would
be ignored within a week.

Verified in both directions rather than assumed: adding an unindexed item
exits 1 and names it, restoring the registry exits 0.

* fix(catalog): rebuild the search index from the registry

build-local-vectors.ts read registry/catalog-artifact/catalog.json, a file no script in this repo writes and which is not committed, so the documented regeneration command failed on a missing path. That is why the index could drift from the registry with nothing to run to fix it.

It now reads registry/blocks/* and registry/components/* through catalogFromRegistry, the existing helper that already produced the right shape but had no caller. Rebuilding reproduces the shipped 168 rows byte for byte.

A lefthook catalog-index command regenerates and re-stages both artifact files whenever a staged registry-item.json changes, mirroring the skills-manifest pattern, so adding or removing an item keeps the index in sync without anyone remembering to. Verified end to end: staging a new item took the artifact 168 to 169 rows and staged it in 0.80s.

* fix(cli): refuse a half-downloaded vector cache

The two artifact files have to agree on how many rows there are, and until now nothing checked that before writing them. A truncated or wrong-model response landed in the cache and only failed at load, on every later search, until someone cleared it by hand. The pair is now checked first and refused as a unit, and the cache is created 0o700 with 0o600 files rather than inheriting the umask of a directory the caller may have pointed anywhere.

Also lifts the capture setup the two preview generators had drifted into sharing into scripts/preview-capture.ts, and splits the vector builders batching and packing out of main. Both were findings the audit attributed to this branch.

* fix(cli): keep the catalog vitest run with the tests it runs

Restacking took the base package.json wholesale, which dropped the vitest dependency and the scripts/catalog run this PR adds. Both belong here rather than under it.

* fix(cli): stop the declined model download from happening anyway

Answering no to the on-device download offer recorded no and warned, then carried on. The guard below it is localModelConsent() !== false, which the decline had just made false, so it was skipped rather than taken: control reached recordLocalModelConsent(true), overwrote the answer with yes, and fetched the 32 MB model the user had refused. Next run it never asked again.

No test could catch it. The stub pinned localModelStatus to ready, so the prompt never fired, and recordLocalModelConsent was a no-op that recorded nothing.

Two tests now cover the offer, and they need three things the old stubs did not model: the run has to look like a terminal, because off one the command treats --on-device as the consent and never asks; the ONNX probe has to answer true, or an accepted offer returns at the runtime guard before it can download; and the status has to follow the recorded answer, or the second offer later in the run fires as well. Removing the return makes the decline test fail.

* fix(catalog): let someone without the model still add a component

The pre-commit hook rebuilds the search index, and rebuilding needs the 32 MB embedding model. An outside contributor adding a registry item does not have it, so their commit died inside the ONNX loader on an ENOENT naming a path they never set, and the CI gate then told them to run the command that had just crashed.

The model is an opt-in for search, not a build dependency, so nobody is charged for it to contribute. The builder checks first and explains itself, exiting 3 for cannot as distinct from 1 for failed. The hook treats 3 as skip and lets the commit through. The gate now names both paths: regenerate if you have the model, leave it if you do not and a maintainer will.

Verified both ways: with no model the builder explains and the hook exits 0; with the model it still regenerates byte-identically.

* docs: say that anyone can add a registry item, and stop hand-editing a generated file

Two defects, one of them the reason 64 stale entries survived in registry.json.

The checklist told contributors to add their item to registry/registry.json. That file is generated from the item directories, so an entry added by hand survives until the next regeneration and then vanishes, and one left behind for a directory that no longer exists is worse: hyperframes add resolves the name and then fails on missing files. Both CONTRIBUTING.md and the agent-facing skill reference now run the generator instead.

Nothing said contribution was maintainer-only, but nothing said it was not either, and two steps do need assets an outside contributor has no reason to install. Those are now named in a table with what happens if you do not have them, matching how the preview image was already handled. The search index is the new one: the model behind it is a 32 MB opt-in for search, not a build dependency.

* fix(cli): harden on-device catalog search

* fix(cli): refresh stale catalog vectors

* test: create catalog vector temp dirs securely
This commit is contained in:
Miguel Ángel
2026-08-09 22:59:15 -07:00
committed by GitHub
parent f28bc80a1d
commit 68205dbbc1
35 changed files with 8069 additions and 28 deletions
File diff suppressed because it is too large Load Diff
+122
View File
@@ -0,0 +1,122 @@
/**
* On-device embeddings for local semantic search.
*
* Runs bge-small-en-v1.5 through the ONNX runtime the CLI already depends on,
* with the WordPiece tokenizer in `wordpiece.ts`. Nothing is sent anywhere and
* nothing costs money once the model is cached.
*
* Two details are specific to bge and easy to miss. Pooling takes the CLS token
* rather than a mean over the sequence, and short queries carry an instruction
* prefix that passages do not. Getting either wrong degrades retrieval quietly
* rather than failing, so both are asserted in tests.
*/
import { readFileSync } from "node:fs";
import {
LOCAL_MODEL_DIMENSIONS,
QUERY_INSTRUCTION,
localModelPath,
localTokenizerPath,
} from "./localModel.js";
import { configFromTokenizerJson, encode } from "./wordpiece.js";
export interface LocalEmbedder {
embed(texts: string[], options?: { isQuery?: boolean }): Promise<number[][]>;
}
/** Long inputs are truncated rather than rejected; the model has a 512 token limit. */
const MAX_TOKENS = 512;
/**
* Is the native ONNX runtime reachable in this install?
*
* It cannot be bundled: a single-file build leaves `onnxruntime-node` as an
* unresolved import, so on-device search is unavailable there no matter what
* else is on disk. Checking costs one dynamic import and saves a pointless
* 32 MB download.
*/
export async function localRuntimeAvailable(): Promise<boolean> {
try {
await import("onnxruntime-node");
return true;
} catch {
return false;
}
}
export async function loadLocalEmbedder(): Promise<LocalEmbedder> {
// Imported lazily so the CLI does not pay for the ONNX runtime on every run.
const ort = await import("onnxruntime-node");
const config = configFromTokenizerJson(readFileSync(localTokenizerPath(), "utf-8"));
const session = await ort.InferenceSession.create(localModelPath());
return {
async embed(texts, options) {
if (texts.length === 0) return [];
const prefix = options?.isQuery ? QUERY_INSTRUCTION : "";
const encodings = texts.map((text) => truncate(encode(prefix + text, config)));
const width = Math.max(...encodings.map((e) => e.ids.length));
const batch = encodings.length;
const ids = new BigInt64Array(batch * width);
const mask = new BigInt64Array(batch * width);
const types = new BigInt64Array(batch * width);
encodings.forEach((encoding, row) => {
encoding.ids.forEach((id, column) => {
const at = row * width + column;
ids[at] = BigInt(id);
mask[at] = 1n;
types[at] = 0n;
});
// Remaining positions stay zero: padded ids with a zero attention mask,
// which the model must not attend to.
});
const dims = [batch, width];
const output = await session.run({
input_ids: new ort.Tensor("int64", ids, dims),
attention_mask: new ort.Tensor("int64", mask, dims),
token_type_ids: new ort.Tensor("int64", types, dims),
});
const hidden = output["last_hidden_state"];
if (!hidden) throw new Error("model returned no last_hidden_state");
const data = hidden.data as Float32Array;
const hiddenSize = hidden.dims[2] as number;
if (hiddenSize !== LOCAL_MODEL_DIMENSIONS) {
throw new Error(
`model produced ${hiddenSize} dimensions, expected ${LOCAL_MODEL_DIMENSIONS}`,
);
}
return encodings.map((_, row) => {
// CLS pooling: bge trains the first token as the sequence representation.
// Mean pooling here would produce vectors that look fine and rank worse.
const start = row * width * hiddenSize;
return normalize(Array.from(data.subarray(start, start + hiddenSize)));
});
},
};
}
function truncate(encoding: ReturnType<typeof encode>): ReturnType<typeof encode> {
if (encoding.ids.length <= MAX_TOKENS) return encoding;
// Keep the closing separator so the sequence still ends the way the model expects.
const ids = [
...encoding.ids.slice(0, MAX_TOKENS - 1),
encoding.ids[encoding.ids.length - 1] as number,
];
return { ids, attentionMask: ids.map(() => 1), tokenTypeIds: ids.map(() => 0) };
}
function normalize(vector: number[]): number[] {
const norm = Math.sqrt(vector.reduce((sum, value) => sum + value * value, 0));
return norm === 0 ? vector : vector.map((value) => value / norm);
}
export function cosine(a: number[], b: number[]): number {
let sum = 0;
for (let i = 0; i < a.length; i += 1) sum += (a[i] as number) * (b[i] as number);
return sum;
}
@@ -0,0 +1,206 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
const readConfig = vi.fn();
const writeConfig = vi.fn();
const existsSync = vi.fn();
const readFileSync = vi.fn();
const unlinkSync = vi.fn();
const downloadFile = vi.fn();
const digest = vi.fn();
vi.mock("../telemetry/config.js", () => ({
readConfig: () => readConfig(),
writeConfig: (c: unknown) => writeConfig(c),
}));
vi.mock("node:fs", async (importOriginal) => ({
...(await importOriginal<typeof import("node:fs")>()),
existsSync: (p: string) => existsSync(p),
mkdirSync: vi.fn(),
readFileSync: (p: string) => readFileSync(p),
unlinkSync: (p: string) => unlinkSync(p),
}));
vi.mock("node:crypto", () => ({
createHash: () => {
const hash = {
update: () => hash,
digest: () => digest(),
};
return hash;
},
}));
vi.mock("../utils/download.js", () => ({
downloadFile: (...args: unknown[]) => downloadFile(...args),
}));
const {
LOCAL_MODEL_DIMENSIONS,
LOCAL_MODEL_ARTIFACTS,
LOCAL_MODEL_SIZE_MB,
QUERY_INSTRUCTION,
downloadOfferMessage,
ensureLocalModel,
isLocalModelReady,
localModelConsent,
localModelStatus,
localTokenizerPath,
localModelPath,
recordLocalModelConsent,
} = await import("./localModel.js");
function returnMatchingDigests(rounds = 1): void {
for (let round = 0; round < rounds; round += 1) {
for (const artifact of LOCAL_MODEL_ARTIFACTS) digest.mockReturnValueOnce(artifact.sha256);
}
}
function makeDownloadsAppearOnDisk(): void {
const present = new Set<string>();
existsSync.mockImplementation((path: string) => present.has(path));
downloadFile.mockImplementation(async (_url: string, path: string) => {
present.add(path);
});
}
beforeEach(() => {
readConfig.mockReturnValue({});
existsSync.mockReturnValue(false);
writeConfig.mockReset();
readFileSync.mockReturnValue(Buffer.from("artifact"));
unlinkSync.mockReset();
downloadFile.mockReset();
digest.mockReset();
});
afterEach(() => vi.clearAllMocks());
describe("consent state", () => {
it("reports never-asked as undefined", () => {
expect(localModelConsent()).toBeUndefined();
});
it("persists a yes", () => {
readConfig.mockReturnValue({ telemetryEnabled: true });
recordLocalModelConsent(true);
expect(writeConfig).toHaveBeenCalledWith(
expect.objectContaining({ localEmbeddingEnabled: true }),
);
});
it("persists a no without discarding other settings", () => {
readConfig.mockReturnValue({ telemetryEnabled: false, anonymousId: "abc" });
recordLocalModelConsent(false);
expect(writeConfig).toHaveBeenCalledWith(
expect.objectContaining({ localEmbeddingEnabled: false, anonymousId: "abc" }),
);
});
});
describe("readiness", () => {
it("needs both the model and its tokenizer", () => {
// A model without its tokenizer cannot turn text into tensors, so half a
// download is not a usable state.
existsSync.mockImplementation((p: string) => p === localModelPath());
expect(isLocalModelReady()).toBe(false);
existsSync.mockImplementation((p: string) => p === localTokenizerPath());
expect(isLocalModelReady()).toBe(false);
existsSync.mockReturnValue(true);
returnMatchingDigests();
expect(isLocalModelReady()).toBe(true);
});
it("rejects present files whose digest does not match the pinned revision", () => {
existsSync.mockReturnValue(true);
digest.mockReturnValue("wrong");
expect(isLocalModelReady()).toBe(false);
});
it("bounds each download and rejects a digest mismatch before reporting ready", async () => {
makeDownloadsAppearOnDisk();
digest.mockReturnValue("wrong");
expect(await ensureLocalModel()).toBe(false);
expect(downloadFile).toHaveBeenCalledWith(
expect.stringContaining("model_quantized.onnx"),
localModelPath(),
{ maxBytes: 34_014_426 },
);
expect(unlinkSync).toHaveBeenCalledWith(localModelPath());
});
it("accepts only two bounded downloads that match both pinned digests", async () => {
makeDownloadsAppearOnDisk();
returnMatchingDigests(2);
expect(await ensureLocalModel()).toBe(true);
expect(downloadFile).toHaveBeenNthCalledWith(
2,
expect.stringContaining("tokenizer.json"),
localTokenizerPath(),
{ maxBytes: 711_396 },
);
});
});
describe("status", () => {
it("is declined when the user said no, even if files somehow exist", () => {
readConfig.mockReturnValue({ localEmbeddingEnabled: false });
existsSync.mockReturnValue(true);
expect(localModelStatus()).toEqual({ status: "declined" });
});
it("is ready when consented and downloaded", () => {
readConfig.mockReturnValue({ localEmbeddingEnabled: true });
existsSync.mockReturnValue(true);
returnMatchingDigests();
expect(localModelStatus()).toEqual({ status: "ready" });
});
it("is not-asked when never answered and nothing downloaded", () => {
expect(localModelStatus()).toEqual({ status: "not-asked" });
});
it("is unavailable when consented but not yet downloaded", () => {
readConfig.mockReturnValue({ localEmbeddingEnabled: true });
expect(localModelStatus().status).toBe("unavailable");
});
it("never downloads or prompts by itself", () => {
// The caller owns the prompt, because only it knows whether word matching
// already answered well enough to make the offer pointless.
expect(() => localModelStatus()).not.toThrow();
expect(writeConfig).not.toHaveBeenCalled();
});
});
describe("the offer text", () => {
it("names the cost rather than implying a privacy tradeoff", () => {
const message = downloadOfferMessage(0);
expect(message).toContain(`${LOCAL_MODEL_SIZE_MB} MB`);
expect(message).toMatch(/stays on your machine/);
expect(message).toMatch(/nothing is sent/);
});
it("leads with what the user just saw", () => {
expect(downloadOfferMessage(0)).toMatch(/^No matches/);
expect(downloadOfferMessage(1)).toMatch(/^Only 1 match\b/);
expect(downloadOfferMessage(3)).toMatch(/^Only 3 matches/);
});
it("does not claim word search already ran for an explicit on-device request", () => {
expect(downloadOfferMessage()).toMatch(/^Use on-device meaning search/);
});
});
describe("model contract", () => {
it("uses the dimension the bundled vectors are built at", () => {
expect(LOCAL_MODEL_DIMENSIONS).toBe(384);
});
it("carries the query instruction bge retrieval expects", () => {
// Omitting this measurably degrades short-query retrieval for bge models.
expect(QUERY_INSTRUCTION).toMatch(/searching relevant passages/);
});
});
+184
View File
@@ -0,0 +1,184 @@
/**
* The on-device embedding model behind local semantic search.
*
* Downloading roughly 33 MB to someone's disk is worth asking about. This is
* not a privacy question, because nothing leaves the machine, so the prompt
* says what it actually costs: disk and bandwidth. Borrowing privacy language
* here would misdescribe it.
*
* The CLI already downloads a much larger model for background removal without
* asking, and that is defensible there: running that command is itself a
* request for a model. Searching a catalog is not, so a download arriving
* mid-search would be a surprise. That difference is why this asks and that
* does not.
*/
import { createHash } from "node:crypto";
import { existsSync, mkdirSync, readFileSync, unlinkSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";
import { downloadFile } from "../utils/download.js";
import { readConfig, writeConfig } from "../telemetry/config.js";
/**
* bge-small-en-v1.5: 62.17 MTEB average against 62.3 for the hosted model this
* substitutes for, at 384 dimensions instead of 1536. Chosen for being close to
* parity while small enough to ship, not for being a weaker tier.
*/
export const LOCAL_MODEL_ID = "bge-small-en-v1.5";
export const LOCAL_MODEL_DIMENSIONS = 384;
/** Measured: 32 MB quantized model plus a 0.7 MB tokenizer. */
export const LOCAL_MODEL_SIZE_MB = 33;
/** bge retrieval expects short queries to carry this prefix; passages do not. */
export const QUERY_INSTRUCTION = "Represent this sentence for searching relevant passages: ";
const MODELS_DIR = join(homedir(), ".hyperframes", "models");
/**
* Where the two halves come from.
*
* Pinned to a revision rather than to `main`: a model that silently changes
* under a cached vector set would return confident nonsense, because the
* catalog vectors were produced by a specific set of weights.
*/
const MODEL_REPO = "Xenova/bge-small-en-v1.5";
export const LOCAL_MODEL_REVISION = "ea104dacec62c0de699686887e3f920caeb4f3e3";
/**
* The quantized export, 32 MB, not the 126 MB full-precision one.
*
* The size this feature quotes has always described the quantized build. The
* catalog vectors must be produced by this same file: embedding the catalog
* with one precision and the query with another degrades ranking silently,
* which is the failure mode this whole feature keeps having to defend against.
*/
interface ModelFile {
readonly url: string;
readonly dest: () => string;
readonly bytes: number;
readonly sha256: string;
}
export const LOCAL_MODEL_ARTIFACTS: ReadonlyArray<ModelFile> = [
{
url: `https://huggingface.co/${MODEL_REPO}/resolve/${LOCAL_MODEL_REVISION}/onnx/model_quantized.onnx`,
dest: () => localModelPath(),
bytes: 34_014_426,
sha256: "6c9c6101a956d62dfb5e7190c538226c0c5bb9cb27b651234b6df063ee7dbfe4",
},
{
url: `https://huggingface.co/${MODEL_REPO}/resolve/${LOCAL_MODEL_REVISION}/tokenizer.json`,
dest: () => localTokenizerPath(),
bytes: 711_396,
sha256: "d241a60d5e8f04cc1b2b3e9ef7a4921b27bf526d9f6050ab90f9267a1f9e5c66",
},
];
export type LocalModelDecision = boolean | undefined;
export function localModelConsent(): LocalModelDecision {
return readConfig().localEmbeddingEnabled;
}
export function recordLocalModelConsent(enabled: boolean): void {
writeConfig({ ...readConfig(), localEmbeddingEnabled: enabled });
}
export function localModelPath(): string {
return join(MODELS_DIR, `${LOCAL_MODEL_ID}.onnx`);
}
export function localTokenizerPath(): string {
return join(MODELS_DIR, `${LOCAL_MODEL_ID}.tokenizer.json`);
}
function modelFileIsValid(file: ModelFile): boolean {
const path = file.dest();
if (!existsSync(path)) return false;
try {
return createHash("sha256").update(readFileSync(path)).digest("hex") === file.sha256;
} catch {
return false;
}
}
/** Both halves must match the pinned revision. Presence alone cannot establish readiness. */
export function isLocalModelReady(): boolean {
return LOCAL_MODEL_ARTIFACTS.every(modelFileIsValid);
}
export type LocalModelStatus =
| { status: "ready" }
| { status: "declined" }
| { status: "not-asked" }
| { status: "unavailable"; reason: string };
/**
* Resolve whether local semantic search can run right now.
*
* Deliberately does not prompt. The caller owns that, because only the caller
* knows whether word matching already answered well enough to make the offer
* pointless. Consent has one owner, which is the lesson the smart-search flag
* taught when it was checked in two places and the override became a no-op.
*/
export function localModelStatus(): LocalModelStatus {
const consent = localModelConsent();
if (consent === false) return { status: "declined" };
if (isLocalModelReady()) return { status: "ready" };
if (consent === undefined) return { status: "not-asked" };
return { status: "unavailable", reason: "model not downloaded yet" };
}
/**
* Put the model on disk, once.
*
* Returns false rather than throwing: a failed download costs the offline tier,
* never the command, and the caller says so. Both halves are fetched before
* either is reported ready, because a model without its tokenizer embeds
* nothing and would fail later at a less obvious place.
*/
export async function ensureLocalModel(): Promise<boolean> {
if (isLocalModelReady()) return true;
try {
mkdirSync(MODELS_DIR, { recursive: true });
for (const file of LOCAL_MODEL_ARTIFACTS) {
if (modelFileIsValid(file)) continue;
const dest = file.dest();
await downloadFile(file.url, dest, { maxBytes: file.bytes });
if (!modelFileIsValid(file)) {
unlinkSync(dest);
return false;
}
}
return isLocalModelReady();
} catch {
return false;
}
}
/**
* What to print when nobody can be asked.
*
* An agent or CI run has no one to prompt, and it must not decide a download
* onto someone's disk by itself. So it is told what to ask and which flag
* records the answer. Passing the flag is the human's yes, not the agent's.
*/
export function nonInteractiveConsentMessage(): string {
return (
`On-device search needs a one-time ${LOCAL_MODEL_SIZE_MB} MB download that stays on this machine. ` +
"Nothing is sent anywhere. Ask the person you are working for, and pass --on-device --yes once they agree."
);
}
/** What the prompt should say. Costs, not privacy: nothing is sent anywhere. */
export function downloadOfferMessage(matchCount?: number): string {
const found =
matchCount === undefined
? "Use on-device meaning search."
: matchCount === 0
? "No matches from word search."
: `Only ${matchCount} match${matchCount === 1 ? "" : "es"} from word search.`;
return `${found} Download a ${LOCAL_MODEL_SIZE_MB} MB search model for better offline results? It stays on your machine and nothing is sent.`;
}
@@ -0,0 +1,97 @@
import { describe, expect, it } from "vitest";
import { rankByWords, searchByWords, tokenize } from "./localSearch.js";
interface Item {
name: string;
text: string;
}
const ITEMS: Item[] = [
{ name: "whip-pan", text: "A fast camera whip blurs between two shots at speed." },
{ name: "number-wheel", text: "A rolling digit counter settles on its final value." },
{ name: "text-reveal", text: "Type reveals line by line beneath a mask." },
];
const textOf = (item: Item) => item.text;
describe("tokenize", () => {
it("drops stop words", () => {
expect(tokenize("the camera and the shot")).toEqual(["camera", "shot"]);
});
it("drops tokens of three characters or fewer", () => {
expect(tokenize("ab abc abcd")).toEqual(["abc", "abcd"]);
});
it("lowercases and strips punctuation and digits", () => {
expect(tokenize("Camera, pushes 42 times!")).toEqual(["camera", "pushes", "times"]);
});
it("returns nothing for a query made only of stop words", () => {
expect(tokenize("the and or of")).toEqual([]);
});
});
describe("rankByWords", () => {
it("ranks the entry sharing the most vocabulary first", () => {
expect(rankByWords("rolling counter digit", ITEMS, textOf)[0]?.item.name).toBe("number-wheel");
});
it("returns every item, not only matches", () => {
expect(rankByWords("camera", ITEMS, textOf)).toHaveLength(ITEMS.length);
});
it("sorts non-matching items last with a zero score", () => {
const ranked = rankByWords("rolling counter", ITEMS, textOf);
expect(ranked.at(-1)?.score).toBe(0);
});
it("does not let the wordiest entry win on length alone", () => {
// Without the sqrt divisor this padded entry outranks the real match.
const padded: Item[] = [
...ITEMS,
{
name: "padded",
text: `camera ${Array.from({ length: 300 }, (_, i) => `filler${i}`).join(" ")}`,
},
];
expect(rankByWords("fast camera whip speed", padded, textOf)[0]?.item.name).toBe("whip-pan");
});
it("breaks ties on descending name, matching the evaluation", () => {
const tied: Item[] = [
{ name: "alpha", text: "nothing shared here" },
{ name: "zulu", text: "nothing shared here" },
];
expect(rankByWords("unrelated", tied, textOf).map((s) => s.item.name)).toEqual([
"zulu",
"alpha",
]);
});
it("treats a stop-word-only query as matching nothing rather than everything", () => {
expect(rankByWords("the and of", ITEMS, textOf).every((s) => s.score === 0)).toBe(true);
});
});
describe("searchByWords", () => {
it("keeps only entries sharing at least one token", () => {
expect(searchByWords("rolling counter", ITEMS, textOf).map((i) => i.name)).toEqual([
"number-wheel",
]);
});
it("beats substring matching on a natural phrasing", () => {
// The motivating case. No description contains this phrase, so a substring
// test returns nothing; shared vocabulary still finds the speed entry.
const phrase = "make the shot feel fast";
const substring = ITEMS.filter((i) => i.text.toLowerCase().includes(phrase.toLowerCase()));
expect(substring).toHaveLength(0);
expect(searchByWords(phrase, ITEMS, textOf)[0]?.name).toBe("whip-pan");
});
it("returns nothing when no vocabulary is shared", () => {
expect(searchByWords("quantum entanglement", ITEMS, textOf)).toEqual([]);
});
});
+75
View File
@@ -0,0 +1,75 @@
/**
* Local catalog search that costs nothing and needs no account.
*
* Replaces a substring test. `description.includes(query)` cannot answer "make
* the pace suddenly feel faster" because no description contains that phrase,
* so the honest result was zero matches. Scoring shared vocabulary answers it
* partially, offline, with no model and no network.
*
* The scorer is the one the retrieval evaluation used, reproduced so the local
* arm and the remote fallback rank identically rather than merely similarly:
* lowercase alphabetic tokens, stop words dropped, anything three characters or
* shorter dropped, and the shared-token count divided by the square root of the
* entry's token count. That divisor is load-bearing. Without it the wordiest
* entry wins every query on sheer surface area.
*
* Ties break on descending name, matching the evaluation's sort.
*/
/** Dropped from both sides before scoring. Without this every entry matches on "the". */
const STOP = new Set(
(
"the a an and or of to in on at is are be it its for with that this as by from into " +
"one two must not no all over under across while when where which who whom whose they " +
"them their we our you your he she his her but if then than so such can may might will " +
"would should each other another same both few more most some any every"
).split(" "),
);
export function tokenize(text: string): string[] {
const words = text.toLowerCase().match(/[a-z]+/g) ?? [];
return words.filter((word) => word.length > 2 && !STOP.has(word));
}
export interface Scored<T> {
item: T;
score: number;
}
/**
* Rank every item by shared vocabulary, best first.
*
* Returns all items rather than only matches, so a caller can decide where to
* cut. Items scoring zero sort last.
*/
export function rankByWords<T>(
query: string,
items: T[],
textOf: (item: T) => string,
): Scored<T>[] {
const want = new Set(tokenize(query));
if (want.size === 0) return items.map((item) => ({ item, score: 0 }));
return items
.map((item) => {
const have = new Set(tokenize(textOf(item)));
let shared = 0;
for (const token of want) if (have.has(token)) shared += 1;
// sqrt normalization: a longer entry has more chances to overlap, and
// without this the wordiest blurb ranks first for every query.
return { item, score: shared / (Math.sqrt(have.size) || 1) };
})
.sort((a, b) => b.score - a.score || nameOf(b.item).localeCompare(nameOf(a.item)));
}
/** Only items sharing at least one token, best first. */
export function searchByWords<T>(query: string, items: T[], textOf: (item: T) => string): T[] {
return rankByWords(query, items, textOf)
.filter((scored) => scored.score > 0)
.map((scored) => scored.item);
}
function nameOf(item: unknown): string {
const named = item as { name?: unknown };
return typeof named?.name === "string" ? named.name : "";
}
@@ -0,0 +1,110 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { cachedLocalVectorRevision, fetchLocalVectors } from "./localSemantic.js";
import { LOCAL_MODEL_DIMENSIONS } from "./localModel.js";
describe("fetchLocalVectors", () => {
let dir: string;
beforeEach(() => {
dir = mkdtempSync(join(tmpdir(), "hf-vec-"));
});
afterEach(() => {
rmSync(dir, { recursive: true, force: true });
vi.unstubAllGlobals();
});
/** A metadata/matrix pair that agrees: one name, one row of the real width. */
const servePair = (names: string[], dimensions: number, floats: number, revision?: string) =>
vi.stubGlobal(
"fetch",
vi.fn(async (url: string) => ({
ok: true,
arrayBuffer: async () =>
url.endsWith(".json")
? new TextEncoder().encode(JSON.stringify({ names, dimensions, revision })).buffer
: new Float32Array(floats).buffer,
})),
);
it("writes both files into the cache directory", async () => {
servePair(["whip-pan"], LOCAL_MODEL_DIMENSIONS, LOCAL_MODEL_DIMENSIONS);
expect(await fetchLocalVectors("http://registry.test/", { directory: dir })).toBe(true);
expect(fetch).toHaveBeenCalledWith(expect.any(String), {
signal: expect.any(AbortSignal),
});
expect(existsSync(join(dir, "local-vectors.bin"))).toBe(true);
expect(existsSync(join(dir, "local-vectors.json"))).toBe(true);
});
it("caches nothing when the matrix is short of the names it claims", async () => {
// Half a download is the case worth refusing: written, it loads as an
// error on every later search until someone clears the cache by hand.
servePair(["whip-pan", "rack-focus"], LOCAL_MODEL_DIMENSIONS, LOCAL_MODEL_DIMENSIONS);
expect(await fetchLocalVectors("http://registry.test/", { directory: dir })).toBe(false);
expect(existsSync(join(dir, "local-vectors.bin"))).toBe(false);
expect(existsSync(join(dir, "local-vectors.json"))).toBe(false);
});
it("caches nothing when the vectors came from a different model", async () => {
servePair(["whip-pan"], 1536, 1536);
expect(await fetchLocalVectors("http://registry.test/", { directory: dir })).toBe(false);
expect(existsSync(join(dir, "local-vectors.json"))).toBe(false);
});
it("reports failure instead of throwing, so the command survives", async () => {
// A tier the user switched on that silently never runs is the failure
// being guarded: the caller needs a false to be able to say so.
vi.stubGlobal(
"fetch",
vi.fn(async () => ({ ok: false, arrayBuffer: async () => new ArrayBuffer(0) })),
);
expect(await fetchLocalVectors("http://registry.test", { directory: dir })).toBe(false);
});
it("reports failure when the network throws", async () => {
vi.stubGlobal(
"fetch",
vi.fn(async () => {
throw new Error("offline");
}),
);
expect(await fetchLocalVectors("http://registry.test", { directory: dir })).toBe(false);
});
it("refuses an unexpected revision without replacing the previous pair", async () => {
mkdirSync(dir, { recursive: true });
writeFileSync(
join(dir, "local-vectors.json"),
JSON.stringify({ names: ["old"], dimensions: LOCAL_MODEL_DIMENSIONS, revision: "old" }),
);
writeFileSync(join(dir, "local-vectors.bin"), new Float32Array(LOCAL_MODEL_DIMENSIONS));
servePair(["new"], LOCAL_MODEL_DIMENSIONS, LOCAL_MODEL_DIMENSIONS, "old");
expect(
await fetchLocalVectors("http://registry.test", {
directory: dir,
expectedRevision: "new",
}),
).toBe(false);
expect(JSON.parse(readFileSync(join(dir, "local-vectors.json"), "utf-8"))).toEqual({
names: ["old"],
dimensions: LOCAL_MODEL_DIMENSIONS,
revision: "old",
});
});
it("reads the revision only from a complete cached pair", () => {
mkdirSync(dir, { recursive: true });
writeFileSync(
join(dir, "local-vectors.json"),
JSON.stringify({ names: ["whip-pan"], dimensions: LOCAL_MODEL_DIMENSIONS, revision: "r1" }),
);
expect(cachedLocalVectorRevision(dir)).toBeUndefined();
writeFileSync(join(dir, "local-vectors.bin"), new Float32Array(LOCAL_MODEL_DIMENSIONS));
expect(cachedLocalVectorRevision(dir)).toBe("r1");
});
});
+229
View File
@@ -0,0 +1,229 @@
/**
* Rank the catalog by meaning, on the user's machine, for free.
*
* Sits between word matching and the hosted endpoint. Word matching cannot
* connect "make the pace feel faster" to a whip pan because they share no
* words; the hosted endpoint can, but needs an account and a network. This
* closes that gap with a 33 MB model the user opted into.
*
* The vectors here are 384-dimension and were produced by a different model
* than the hosted 1536-dimension set. The two are not comparable and are never
* mixed: a query is embedded by whichever model produced the vectors it is
* being compared against.
*/
import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";
import { cosine, loadLocalEmbedder } from "./localEmbedder.js";
import { LOCAL_MODEL_DIMENSIONS, isLocalModelReady } from "./localModel.js";
interface LocalVectorSet {
names: string[];
dimensions: number;
vectors: Float32Array;
}
interface LocalVectorMetadata {
names: string[];
dimensions: number;
revision?: string;
}
export interface FetchLocalVectorOptions {
directory?: string;
expectedRevision?: string;
}
const CATALOG_ARTIFACT_TIMEOUT_MS = 30_000;
/** Where the bundled vector set lives, overridable for development. */
function localVectorDirectory(): string {
return (
process.env["HYPERFRAMES_CATALOG_ARTIFACT_DIR"] ?? join(homedir(), ".hyperframes", "catalog")
);
}
/**
* Put the catalog vectors in the user's cache, once.
*
* They are fetched rather than bundled: every install would otherwise carry a
* copy of a file only people who opt into offline ranking will ever read, and
* the vectors have to track the registry anyway, so shipping them inside the
* package would freeze them to the release instead.
*
* Returns false rather than throwing. A download that fails costs the offline
* tier, never the command, and the caller says so.
*/
/**
* Do a freshly fetched metadata/matrix pair describe the same index?
*
* `names.length * dimensions` floats, at four bytes each, is the whole
* contract. A pair that fails it is a truncated download or a different
* model, never something worth caching.
*/
function vectorPairAgrees(fetched: Array<[string, Buffer]>): boolean {
const meta = fetched.find(([file]) => file === "local-vectors.json")?.[1];
const bin = fetched.find(([file]) => file === "local-vectors.bin")?.[1];
if (!meta || !bin) return false;
try {
const parsed = JSON.parse(meta.toString("utf-8")) as {
names?: string[];
dimensions?: number;
};
if (parsed.dimensions !== LOCAL_MODEL_DIMENSIONS) return false;
return bin.byteLength === (parsed.names?.length ?? -1) * LOCAL_MODEL_DIMENSIONS * 4;
} catch {
return false;
}
}
/** The registry manifest and fetched pair must describe the same generation. */
function vectorRevisionAgrees(
fetched: Array<[string, Buffer]>,
expectedRevision?: string,
): boolean {
if (expectedRevision === undefined) return true;
const meta = fetched.find(([file]) => file === "local-vectors.json")?.[1];
if (!meta) return false;
try {
const parsed = JSON.parse(meta.toString("utf-8")) as { revision?: unknown };
return parsed.revision === expectedRevision;
} catch {
return false;
}
}
export async function fetchLocalVectors(
registryBaseUrl: string,
options: FetchLocalVectorOptions = {},
): Promise<boolean> {
const directory = options.directory ?? localVectorDirectory();
const base = registryBaseUrl.replace(/\/+$/, "");
try {
mkdirSync(directory, { recursive: true, mode: 0o700 });
// Downloaded in full before anything is written. The two files have to
// agree on how many rows there are, so a fetch that fails halfway through
// must leave the previous pair intact rather than pairing a new name list
// with an old matrix, which loads as an error instead of as stale data.
const fetched: Array<[string, Buffer]> = [];
for (const file of ["local-vectors.json", "local-vectors.bin"] as const) {
const response = await fetch(`${base}/catalog-artifact/${file}`, {
signal: AbortSignal.timeout(CATALOG_ARTIFACT_TIMEOUT_MS),
});
if (!response.ok) return false;
fetched.push([file, Buffer.from(await response.arrayBuffer())]);
}
// Check the pair agrees BEFORE either file lands. Writing first and
// discovering the mismatch at load time leaves a cache that fails every
// subsequent search until someone deletes it by hand, and it is the only
// point where a truncated or wrong-model response can still be refused.
if (!vectorPairAgrees(fetched) || !vectorRevisionAgrees(fetched, options.expectedRevision)) {
return false;
}
// 0o600: the cache is this user's, and the directory may be world-writable
// when the caller overrides it.
for (const [file, bytes] of fetched) {
writeFileSync(join(directory, file), bytes, { mode: 0o600 });
}
return hasLocalVectors(directory);
} catch {
return false;
}
}
export function hasLocalVectors(directory = localVectorDirectory()): boolean {
return (
existsSync(join(directory, "local-vectors.bin")) &&
existsSync(join(directory, "local-vectors.json"))
);
}
function loadLocalVectors(directory = localVectorDirectory()): LocalVectorSet {
const meta = JSON.parse(
readFileSync(join(directory, "local-vectors.json"), "utf-8"),
) as LocalVectorMetadata;
const buffer = readFileSync(join(directory, "local-vectors.bin"));
const vectors = new Float32Array(buffer.buffer, buffer.byteOffset, buffer.byteLength / 4);
const expected = meta.names.length * meta.dimensions;
if (vectors.length !== expected) {
throw new Error(`local vectors hold ${vectors.length} floats, expected ${expected}`);
}
if (meta.dimensions !== LOCAL_MODEL_DIMENSIONS) {
// A dimension mismatch means the vectors and the model disagree, which
// produces confident nonsense rather than an error.
throw new Error(
`local vectors are ${meta.dimensions}-dimension, model produces ${LOCAL_MODEL_DIMENSIONS}`,
);
}
return { names: meta.names, dimensions: meta.dimensions, vectors };
}
/**
* The names the on-device index holds, without loading the vector matrix.
*
* Callers compare this against the live registry to see what meaning search
* cannot reach. Reads the metadata file only, so it costs a small JSON parse
* rather than the whole matrix, and answers empty rather than throwing: an
* absent or unreadable artifact covers nothing, which is a coverage answer
* rather than a reason to fail the search that asked.
*/
export function localVectorNames(directory = localVectorDirectory()): string[] {
if (!hasLocalVectors(directory)) return [];
try {
const meta = JSON.parse(readFileSync(join(directory, "local-vectors.json"), "utf-8")) as {
names?: string[];
};
return meta.names ?? [];
} catch {
return [];
}
}
/** Revision of the last complete pair, absent for pre-revision and unreadable caches. */
export function cachedLocalVectorRevision(directory = localVectorDirectory()): string | undefined {
if (!hasLocalVectors(directory)) return undefined;
try {
const meta = JSON.parse(readFileSync(join(directory, "local-vectors.json"), "utf-8")) as {
revision?: unknown;
};
return typeof meta.revision === "string" ? meta.revision : undefined;
} catch {
return undefined;
}
}
/**
* Names ranked best first, each with the similarity that placed it there.
* Returns null when local semantic search is not available.
*
* The score rides along because this ranker has no relevance threshold: every
* query returns the whole set in some order, so a caller holding names alone
* cannot tell a real match from the least-bad row of a query that matched
* nothing at all.
*/
export async function localSemanticRanking(
query: string,
directory = localVectorDirectory(),
): Promise<Array<{ name: string; score: number }> | null> {
if (!isLocalModelReady() || !hasLocalVectors(directory)) return null;
const set = loadLocalVectors(directory);
const embedder = await loadLocalEmbedder();
const [queryVector] = await embedder.embed([query], { isQuery: true });
if (!queryVector) return null;
const scored = set.names.map((name, row) => {
const start = row * set.dimensions;
return {
name,
score: cosine(queryVector, Array.from(set.vectors.subarray(start, start + set.dimensions))),
};
});
// Ties break on descending name, matching every other ranking in this system.
scored.sort((a, b) => b.score - a.score || b.name.localeCompare(a.name));
return scored;
}
+152
View File
@@ -0,0 +1,152 @@
import { describe, expect, it } from "vitest";
import { existsSync, readFileSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";
import { configFromTokenizerJson, encode, normalize, preTokenize, wordPiece } from "./wordpiece.js";
const REFERENCE_PATH = join(import.meta.dirname, "__fixtures__", "wordpiece-reference.json");
const TOKENIZER_PATH = join(
homedir(),
".hyperframes",
"models",
"bge-small-en-v1.5.tokenizer.json",
);
/** A tiny vocab keeps the unit tests readable and independent of the 30k file. */
const CONFIG = configFromTokenizerJson(
JSON.stringify({
model: {
vocab: {
"[UNK]": 0,
"[CLS]": 1,
"[SEP]": 2,
fast: 3,
camera: 4,
"##era": 5,
cam: 6,
".": 7,
whip: 8,
pan: 9,
},
unk_token: "[UNK]",
continuing_subword_prefix: "##",
max_input_chars_per_word: 100,
},
}),
);
describe("normalize", () => {
it("lowercases", () => {
expect(normalize("Whip PAN")).toBe("whip pan");
});
it("strips accents", () => {
expect(normalize("café naïve")).toBe("cafe naive");
});
it("collapses control characters and whitespace variants to spaces", () => {
expect(normalize("a\tb\nc")).toBe("a b c");
});
it("isolates CJK characters so each becomes its own token", () => {
expect(normalize("a中b")).toBe("a 中 b");
});
});
describe("preTokenize", () => {
it("splits on whitespace", () => {
expect(preTokenize("fast camera pan")).toEqual(["fast", "camera", "pan"]);
});
it("isolates punctuation into its own token", () => {
expect(preTokenize("group: transitions.")).toEqual(["group", ":", "transitions", "."]);
});
it("drops empty chunks from repeated spaces", () => {
expect(preTokenize("a b")).toEqual(["a", "b"]);
});
});
describe("wordPiece", () => {
it("returns a whole word when the vocab has it", () => {
expect(wordPiece("camera", CONFIG)).toEqual(["camera"]);
});
it("splits into continuation pieces, longest match first", () => {
// "camera" is present whole, so force the split path with a word that is not.
expect(wordPiece("camera", { ...CONFIG, vocab: { cam: 6, "##era": 5 } })).toEqual([
"cam",
"##era",
]);
});
it("marks the whole word unknown when any piece is unmatchable", () => {
// Emitting partial pieces would silently change the embedding rather than
// failing, so a single bad piece invalidates the word.
expect(wordPiece("zzzz", CONFIG)).toEqual(["[UNK]"]);
});
it("marks a word longer than the limit unknown without scanning it", () => {
expect(wordPiece("x".repeat(101), CONFIG)).toEqual(["[UNK]"]);
});
});
describe("encode", () => {
it("wraps the sequence in CLS and SEP", () => {
const { ids } = encode("fast", CONFIG);
expect(ids[0]).toBe(CONFIG.vocab["[CLS]"]);
expect(ids.at(-1)).toBe(CONFIG.vocab["[SEP]"]);
});
it("returns masks matching the id length", () => {
const encoding = encode("fast camera", CONFIG);
expect(encoding.attentionMask).toHaveLength(encoding.ids.length);
expect(encoding.tokenTypeIds).toHaveLength(encoding.ids.length);
expect(new Set(encoding.attentionMask)).toEqual(new Set([1]));
expect(new Set(encoding.tokenTypeIds)).toEqual(new Set([0]));
});
it("encodes empty input as just the special tokens", () => {
expect(encode("", CONFIG).ids).toEqual([1, 2]);
});
});
/**
* The test that actually matters.
*
* Wrong token ids produce plausible embeddings and plausible rankings, so
* nothing downstream fails visibly. Comparing against a real tokenizer over the
* whole catalog is the only way to know this implementation is correct rather
* than merely reasonable.
*/
describe("parity with the reference tokenizer", () => {
const hasFixture = existsSync(REFERENCE_PATH);
const hasVocab = existsSync(TOKENIZER_PATH);
it.runIf(hasFixture && hasVocab)("matches reference ids across the whole corpus", () => {
const config = configFromTokenizerJson(readFileSync(TOKENIZER_PATH, "utf-8"));
const reference = JSON.parse(readFileSync(REFERENCE_PATH, "utf-8")) as {
text: string;
ids: number[];
}[];
expect(reference.length).toBeGreaterThan(500);
const mismatches = reference
.map(({ text, ids }) => ({ text, expected: ids, actual: encode(text, config).ids }))
.filter(({ expected, actual }) => expected.join(",") !== actual.join(","));
expect(
mismatches
.slice(0, 3)
.map(
(m) => `${JSON.stringify(m.text.slice(0, 60))}\n want ${m.expected}\n got ${m.actual}`,
),
).toEqual([]);
expect(mismatches).toHaveLength(0);
});
it.skipIf(hasFixture && hasVocab)("is skipped without the fixture and vocab", () => {
expect(true).toBe(true);
});
});
+194
View File
@@ -0,0 +1,194 @@
/**
* BERT WordPiece tokenization, implemented directly rather than pulled in.
*
* The obvious dependencies both failed: the standalone tokenizer package ships
* builder components with no way to load a `tokenizer.json`, and the full
* wrapper costs 359 MB installed because it pins an exact ONNX runtime and ends
* up with a second native copy plus an unused web runtime. This is ~120 lines
* with no dependencies at all.
*
* Hand-writing a tokenizer is normally a bad idea, because wrong token ids
* still produce plausible-looking embeddings and plausible-looking rankings.
* Nothing fails loudly. That is why `wordpiece.parity.test.ts` compares this
* implementation's ids against a reference implementation over a corpus rather
* than asserting a handful of cases by eye.
*
* The pipeline reproduces the config in the model's own tokenizer.json:
* BertNormalizer (clean text, isolate CJK, strip accents, lowercase), then
* BertPreTokenizer (whitespace split, punctuation isolated), then greedy
* longest-match-first WordPiece, then [CLS] ... [SEP].
*/
export interface WordPieceConfig {
vocab: Record<string, number>;
unkToken: string;
continuingSubwordPrefix: string;
maxInputCharsPerWord: number;
clsToken: string;
sepToken: string;
}
export interface Encoding {
ids: number[];
attentionMask: number[];
tokenTypeIds: number[];
}
/** Read the pieces this tokenizer needs out of a HuggingFace tokenizer.json. */
export function configFromTokenizerJson(raw: string): WordPieceConfig {
const parsed = JSON.parse(raw) as {
model: {
vocab: Record<string, number>;
unk_token?: string;
continuing_subword_prefix?: string;
max_input_chars_per_word?: number;
};
};
const model = parsed.model;
if (!model?.vocab) throw new Error("tokenizer.json has no WordPiece vocab");
return {
vocab: model.vocab,
unkToken: model.unk_token ?? "[UNK]",
continuingSubwordPrefix: model.continuing_subword_prefix ?? "##",
maxInputCharsPerWord: model.max_input_chars_per_word ?? 100,
clsToken: "[CLS]",
sepToken: "[SEP]",
};
}
/**
* BertNormalizer.
*
* Control characters are dropped and every whitespace variant collapses to a
* plain space. CJK characters are isolated so each becomes its own token, which
* is what the reference does even for an English-only model.
*/
export function normalize(text: string): string {
let out = "";
for (const char of text) {
const code = char.codePointAt(0) as number;
if (code === 0 || code === 0xfffd) continue;
if (isControl(char, code)) continue;
if (isWhitespace(char, code)) {
out += " ";
continue;
}
out += isChinese(code) ? ` ${char} ` : char;
}
// Strip accents after lowercasing: NFD splits a letter from its mark, and the
// marks are then dropped. Mirrors strip_accents defaulting to lowercase.
return out
.toLowerCase()
.normalize("NFD")
.replace(/\p{Mn}/gu, "");
}
/** BertPreTokenizer: split on whitespace, then break punctuation into its own token. */
export function preTokenize(normalized: string): string[] {
const words: string[] = [];
for (const chunk of normalized.split(/\s+/)) {
if (!chunk) continue;
let current = "";
for (const char of chunk) {
if (isPunctuation(char)) {
if (current) words.push(current);
words.push(char);
current = "";
} else {
current += char;
}
}
if (current) words.push(current);
}
return words;
}
/** Greedy longest-match-first, the reference WordPiece algorithm. */
export function wordPiece(word: string, config: WordPieceConfig): string[] {
if ([...word].length > config.maxInputCharsPerWord) return [config.unkToken];
const pieces: string[] = [];
let start = 0;
while (start < word.length) {
let end = word.length;
let match: string | null = null;
while (start < end) {
const candidate =
start === 0
? word.slice(start, end)
: config.continuingSubwordPrefix + word.slice(start, end);
if (candidate in config.vocab) {
match = candidate;
break;
}
end -= 1;
}
// A single unmatchable piece makes the whole word unknown, not just that
// piece. Emitting partial pieces here would silently change the embedding.
if (match === null) return [config.unkToken];
pieces.push(match);
start = end;
}
return pieces;
}
export function encode(text: string, config: WordPieceConfig): Encoding {
const tokens = [config.clsToken];
for (const word of preTokenize(normalize(text))) tokens.push(...wordPiece(word, config));
tokens.push(config.sepToken);
const unk = config.vocab[config.unkToken] as number;
const ids = tokens.map((token) => config.vocab[token] ?? unk);
return {
ids,
attentionMask: ids.map(() => 1),
tokenTypeIds: ids.map(() => 0),
};
}
function isControl(char: string, code: number): boolean {
if (char === "\t" || char === "\n" || char === "\r") return false;
return code < 32 || code === 127 || /\p{Cc}|\p{Cf}/u.test(char);
}
function isWhitespace(char: string, code: number): boolean {
return (
char === " " ||
char === "\t" ||
char === "\n" ||
char === "\r" ||
code === 0x0b ||
code === 0x0c ||
/\p{Zs}/u.test(char)
);
}
function isPunctuation(char: string): boolean {
const code = char.codePointAt(0) as number;
// ASCII non-alphanumerics, plus Unicode punctuation categories only.
// Symbol categories are deliberately excluded: the reference leaves a symbol
// attached to the preceding token, so "360°" becomes 360 + ##° rather than
// two separate words. Including \p{S} here silently changed the embedding of
// every description containing one, which the parity corpus caught.
const asciiSymbol =
(code >= 33 && code <= 47) ||
(code >= 58 && code <= 64) ||
(code >= 91 && code <= 96) ||
(code >= 123 && code <= 126);
return asciiSymbol || /\p{P}/u.test(char);
}
// a Unicode block test; the ranges are data, and 16 cyclomatic is what checking them costs
// fallow-ignore-next-line complexity
function isChinese(code: number): boolean {
return (
(code >= 0x4e00 && code <= 0x9fff) ||
(code >= 0x3400 && code <= 0x4dbf) ||
(code >= 0x20000 && code <= 0x2a6df) ||
(code >= 0x2a700 && code <= 0x2b73f) ||
(code >= 0x2b740 && code <= 0x2b81f) ||
(code >= 0x2b820 && code <= 0x2ceaf) ||
(code >= 0xf900 && code <= 0xfaff) ||
(code >= 0x2f800 && code <= 0x2fa1f)
);
}