mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-23 20:22:37 +00:00
fix(node): respect transformers tokenizer option (#4457)
## Summary `tokenizer` is a documented option, but initialization always passed the model name to `AutoTokenizer.from_pretrained`. Preserve the resolved tokenizer option, fall back to the model name when it is omitted, and identify the selected tokenizer in loading errors. Add five initialization regressions covering default model/tokenizer selection, an explicit tokenizer, registry variables, and tokenizer loading errors. They run in native Node child processes so the real ESM import works without changing Jest/CI flags. Model loaders are mocked and local/remote model access is disabled; these tests do not download models or test inference. ## Validation - Before: three regressions fail and two default-selection controls pass. After: all five pass. - New tests plus the existing registry and embedding-registry suites: 48 passed. - TypeScript compilation, changed-file Biome checks, and `pnpm run docs` pass; generated tracked docs are unchanged. - Checked with Node 24.14.1 / pnpm 11.1.1 on Linux, both with default Jest flags and the Linux CI ESM flag. The native binding and declarations came from the official `0.40.0-beta.13` packages, with registry integrity verified. Rebuilding Rust from source timed out during Cargo metadata resolution, before compilation; no Rust code is changed.
This commit is contained in:
@@ -0,0 +1,66 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
const assert = require("node:assert/strict");
|
||||
const {
|
||||
TransformersEmbeddingFunction,
|
||||
} = require("../../dist/embedding/transformers");
|
||||
const { getRegistry } = require("../../dist/embedding/registry");
|
||||
|
||||
async function main() {
|
||||
const transformers = await import("@huggingface/transformers");
|
||||
// Never fetch a real model, including if a loader mock stops intercepting.
|
||||
transformers.env.allowRemoteModels = false;
|
||||
transformers.env.allowLocalModels = false;
|
||||
const modelCalls = [];
|
||||
const tokenizerCalls = [];
|
||||
const testCase = Number(process.argv[2]);
|
||||
const defaults = "Xenova/all-MiniLM-L6-v2";
|
||||
const cases = [
|
||||
[{}, defaults, defaults],
|
||||
[{ model: "test/model" }, "test/model", "test/model"],
|
||||
[
|
||||
{ model: "test/model", tokenizer: "test/tokenizer" },
|
||||
"test/model",
|
||||
"test/tokenizer",
|
||||
],
|
||||
[
|
||||
{ model: "test/model", tokenizer: "$var:tokenizer" },
|
||||
"test/model",
|
||||
"test/tokenizer",
|
||||
],
|
||||
[
|
||||
{ model: "test/model", tokenizer: "test/tokenizer" },
|
||||
"test/model",
|
||||
"test/tokenizer",
|
||||
],
|
||||
];
|
||||
assert.ok(
|
||||
Number.isInteger(testCase) && testCase >= 0 && testCase < cases.length,
|
||||
);
|
||||
transformers.AutoModel.from_pretrained = async (...args) => {
|
||||
modelCalls.push(args);
|
||||
return {};
|
||||
};
|
||||
transformers.AutoTokenizer.from_pretrained = async (...args) => {
|
||||
tokenizerCalls.push(args);
|
||||
if (testCase === 4) throw new Error("missing tokenizer");
|
||||
return {};
|
||||
};
|
||||
getRegistry().setVar("tokenizer", "test/tokenizer");
|
||||
const [options, model, tokenizer] = cases[testCase];
|
||||
const fn = new TransformersEmbeddingFunction(options);
|
||||
if (testCase === 4) {
|
||||
await assert.rejects(fn.init(), /error loading tokenizer test\/tokenizer/);
|
||||
} else {
|
||||
await fn.init();
|
||||
}
|
||||
assert.deepEqual(modelCalls, [[model, { dtype: "fp32" }]]);
|
||||
assert.deepEqual(tokenizerCalls, [[tokenizer]]);
|
||||
assert.deepEqual(fn.toJSON(), options);
|
||||
}
|
||||
|
||||
main().catch((error) => {
|
||||
console.error(error);
|
||||
process.exitCode = 1;
|
||||
});
|
||||
@@ -0,0 +1,26 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import { execFileSync } from "node:child_process";
|
||||
import { resolve } from "node:path";
|
||||
|
||||
describe("Transformers embedding tokenizer selection", () => {
|
||||
// Run init's native ESM import outside Jest's CommonJS sandbox, like a user.
|
||||
// The fixture mocks model loading and forbids local/remote model access.
|
||||
test.each([
|
||||
["default model", 0],
|
||||
["default tokenizer", 1],
|
||||
["custom tokenizer", 2],
|
||||
["tokenizer registry variable", 3],
|
||||
["tokenizer loading error", 4],
|
||||
])("handles %s", (_name, testCase) => {
|
||||
execFileSync(
|
||||
process.execPath,
|
||||
[
|
||||
resolve(__dirname, "fixtures", "transformers_tokenizer.cjs"),
|
||||
String(testCase),
|
||||
],
|
||||
{ stdio: "pipe", timeout: 15_000 },
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -38,6 +38,7 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction<
|
||||
#model?: import("@huggingface/transformers").PreTrainedModel;
|
||||
#tokenizer?: import("@huggingface/transformers").PreTrainedTokenizer;
|
||||
#modelName: XenovaTransformerOptions["model"];
|
||||
#tokenizerName: string;
|
||||
#initialized = false;
|
||||
#tokenizerOptions: XenovaTransformerOptions["tokenizerOptions"];
|
||||
#ndims?: number;
|
||||
@@ -58,6 +59,7 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction<
|
||||
|
||||
this.#ndims = options.ndims;
|
||||
this.#modelName = modelName;
|
||||
this.#tokenizerName = options.tokenizer ?? modelName;
|
||||
}
|
||||
|
||||
async init() {
|
||||
@@ -86,11 +88,11 @@ export class TransformersEmbeddingFunction extends EmbeddingFunction<
|
||||
}
|
||||
try {
|
||||
this.#tokenizer = await transformers.AutoTokenizer.from_pretrained(
|
||||
this.#modelName,
|
||||
this.#tokenizerName,
|
||||
);
|
||||
} catch (e) {
|
||||
throw new Error(
|
||||
`error loading tokenizer for ${this.#modelName}. Make sure you are using a wasm compatible model:\nReason: ${e}`,
|
||||
`error loading tokenizer ${this.#tokenizerName}. Make sure you are using a wasm compatible model:\nReason: ${e}`,
|
||||
);
|
||||
}
|
||||
this.#initialized = true;
|
||||
|
||||
Reference in New Issue
Block a user