From b98a961d416522fdf2b23c3d8fe55bec07f262e1 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Tue, 28 Jul 2026 13:44:49 +0200 Subject: [PATCH 1/5] extracted onnx/ort into its own backend --- .github/workflows/publish.yml | 10 +- CONTRIBUTING.md | 64 ++- packages/transformers-onnx/README.md | 9 + packages/transformers-onnx/jest.config.mjs | 5 + packages/transformers-onnx/package.json | 61 +++ packages/transformers-onnx/scripts/build.mjs | 58 +++ packages/transformers-onnx/src/empty.ts | 1 + packages/transformers-onnx/src/host.ts | 87 ++++ packages/transformers-onnx/src/index.ts | 4 + packages/transformers-onnx/src/provider.ts | 449 ++++++++++++++++++ .../src/runtime.ts} | 19 +- packages/transformers-onnx/src/tensor-ops.ts | 218 +++++++++ packages/transformers-onnx/src/testing.ts | 44 ++ packages/transformers-onnx/src/wasm-cache.ts | 50 ++ .../transformers-onnx/tests/provider.test.js | 11 + packages/transformers-onnx/tsconfig.json | 18 + packages/transformers/package.json | 3 +- .../transformers/scripts/build/buildAll.mjs | 5 - .../scripts/build/buildAllWithWatch.mjs | 6 +- .../transformers/scripts/build/constants.mjs | 9 +- .../build/plugins/ignoreModulesPlugin.mjs | 1 - .../scripts/build/plugins/postBuildPlugin.mjs | 43 -- .../transformers/scripts/build/targets.mjs | 4 - .../transformers/src/backends/artifacts.js | 26 + packages/transformers/src/backends/default.js | 24 + .../transformers/src/backends/inference.js | 135 ++++++ .../src/backends/utils/cacheWasm.js | 101 ---- packages/transformers/src/env.js | 12 +- .../transformers/src/generation/controller.js | 357 ++++++++++++++ .../transformers/src/generation/runtime.js | 261 ++++++++++ .../src/models/auto/modeling_auto.js | 29 +- .../transformers/src/models/modeling_utils.js | 186 +++----- packages/transformers/src/models/session.js | 283 +---------- .../modeling_voxtral_realtime.js | 2 +- packages/transformers/src/ops/registry.js | 177 +------ packages/transformers/src/pipelines.js | 48 +- packages/transformers/src/transformers.js | 11 + packages/transformers/src/utils/dtypes.js | 124 +---- packages/transformers/src/utils/hub.js | 9 +- .../transformers/src/utils/model-loader.js | 111 ----- .../model_registry/get_available_dtypes.js | 35 +- .../utils/model_registry/get_file_metadata.js | 2 + .../src/utils/model_registry/get_files.js | 4 +- .../utils/model_registry/get_model_files.js | 57 +-- .../model_registry/get_pipeline_files.js | 5 +- packages/transformers/src/utils/tensor.js | 79 ++- .../tests/generation_controller.test.js | 214 +++++++++ .../tests/inference_backends.test.js | 126 +++++ packages/transformers/tests/init.js | 51 +- pnpm-lock.yaml | 37 +- types/webgpu-kernels.local.demo.d.ts | 14 + types/webgpu-kernels.local.demo.d.ts.map | 1 + webgpu-compat.local.md | 362 ++++++++++++++ webgpu-kernels.local.md | 310 ++++++++++++ 54 files changed, 3226 insertions(+), 1146 deletions(-) create mode 100644 packages/transformers-onnx/README.md create mode 100644 packages/transformers-onnx/jest.config.mjs create mode 100644 packages/transformers-onnx/package.json create mode 100644 packages/transformers-onnx/scripts/build.mjs create mode 100644 packages/transformers-onnx/src/empty.ts create mode 100644 packages/transformers-onnx/src/host.ts create mode 100644 packages/transformers-onnx/src/index.ts create mode 100644 packages/transformers-onnx/src/provider.ts rename packages/{transformers/src/backends/onnx.js => transformers-onnx/src/runtime.ts} (97%) create mode 100644 packages/transformers-onnx/src/tensor-ops.ts create mode 100644 packages/transformers-onnx/src/testing.ts create mode 100644 packages/transformers-onnx/src/wasm-cache.ts create mode 100644 packages/transformers-onnx/tests/provider.test.js create mode 100644 packages/transformers-onnx/tsconfig.json delete mode 100644 packages/transformers/scripts/build/plugins/postBuildPlugin.mjs create mode 100644 packages/transformers/src/backends/artifacts.js create mode 100644 packages/transformers/src/backends/default.js create mode 100644 packages/transformers/src/backends/inference.js delete mode 100644 packages/transformers/src/backends/utils/cacheWasm.js create mode 100644 packages/transformers/src/generation/controller.js create mode 100644 packages/transformers/src/generation/runtime.js delete mode 100644 packages/transformers/src/utils/model-loader.js create mode 100644 packages/transformers/tests/generation_controller.test.js create mode 100644 packages/transformers/tests/inference_backends.test.js create mode 100644 types/webgpu-kernels.local.demo.d.ts create mode 100644 types/webgpu-kernels.local.demo.d.ts.map create mode 100644 webgpu-compat.local.md create mode 100644 webgpu-kernels.local.md diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index e4eb0411e..5b31683c7 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -11,17 +11,21 @@ jobs: build: runs-on: ubuntu-latest steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - name: Install pnpm - uses: pnpm/action-setup@b906affcce14559ad1aafd4ab0e942779e9f58b1 # v4 + uses: pnpm/action-setup@b906affcce14559ad1aafd4ab0e942779e9f58b1 # v4 # Setup .npmrc file to publish to npm - - uses: actions/setup-node@53b83947a5a98c8d113130e565377fae1a50d02f # v6 + - uses: actions/setup-node@53b83947a5a98c8d113130e565377fae1a50d02f # v6 with: node-version: "24.10.0" registry-url: "https://registry.npmjs.org" cache: "pnpm" - run: pnpm install --frozen-lockfile - run: pnpm build + - name: Publish ONNX provider to npm + run: pnpm --filter @huggingface/transformers-onnx publish --access public --no-git-checks + env: + NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }} - name: Publish to npm run: pnpm --filter @huggingface/transformers publish --access public --no-git-checks env: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c98dc60eb..3e6de6f38 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -27,10 +27,10 @@ helped you, or simply ⭐️ the repository to say thank you. There are several ways you can contribute to 🤗 Transformers.js: -* Fix outstanding issues with the existing code. -* Submit issues related to bugs or desired new features. -* Implement new models. -* Contribute to the examples or to the documentation. +- Fix outstanding issues with the existing code. +- Submit issues related to bugs or desired new features. +- Implement new models. +- Contribute to the examples or to the documentation. ## Fixing outstanding issues @@ -55,9 +55,9 @@ To create a new issue, please [use one of the templates](https://github.com/hugg If there is a new feature you'd like to see in 🤗 Transformers.js, please open an issue and describe: -1. What is the *motivation* behind this feature? Is it related to a problem or frustration with the library? Is it a feature related to something you need for a project? Is it something you worked on and think it could benefit the community? Whatever it is, we'd love to hear about it! +1. What is the _motivation_ behind this feature? Is it related to a problem or frustration with the library? Is it a feature related to something you need for a project? Is it something you worked on and think it could benefit the community? Whatever it is, we'd love to hear about it! 2. Describe your requested feature in as much detail as possible. The more you can tell us about it, the better we'll be able to help you. -3. Provide a *code snippet* that demonstrates the feature's usage. +3. Provide a _code snippet_ that demonstrates the feature's usage. 4. If the feature is related to a paper, please include a link. If your issue is well written we're already 80% of the way there by the time you create it. @@ -98,7 +98,7 @@ Every model file exports a base class and one or more task heads. For the vast m **Decoder-only LLM:** ```js -import { PreTrainedModel } from '../modeling_utils.js'; +import { PreTrainedModel } from "../modeling_utils.js"; export class MyModelPreTrainedModel extends PreTrainedModel {} export class MyModelModel extends MyModelPreTrainedModel {} @@ -108,22 +108,25 @@ export class MyModelForCausalLM extends MyModelPreTrainedModel {} **Encoder-only model:** ```js -import { PreTrainedModel } from '../modeling_utils.js'; -import { MaskedLMOutput, SequenceClassifierOutput } from '../modeling_outputs.js'; +import { PreTrainedModel } from "../modeling_utils.js"; +import { + MaskedLMOutput, + SequenceClassifierOutput, +} from "../modeling_outputs.js"; export class MyModelPreTrainedModel extends PreTrainedModel {} export class MyModelModel extends MyModelPreTrainedModel {} export class MyModelForMaskedLM extends MyModelPreTrainedModel { - async _call(model_inputs) { - return new MaskedLMOutput(await super._call(model_inputs)); - } + async _call(model_inputs) { + return new MaskedLMOutput(await super._call(model_inputs)); + } } export class MyModelForSequenceClassification extends MyModelPreTrainedModel { - async _call(model_inputs) { - return new SequenceClassifierOutput(await super._call(model_inputs)); - } + async _call(model_inputs) { + return new SequenceClassifierOutput(await super._call(model_inputs)); + } } ``` @@ -133,11 +136,11 @@ Only add the task heads the model actually supports. The available output classe Most models reuse an existing tokenizer (e.g. all Llama-family models use `LlamaTokenizer`). Only create a new one if the model genuinely needs custom tokenization or preprocessing logic. -| What | File | Barrel to update | -| --- | --- | --- | -| Custom tokenizer | `src/models//tokenization_.js` | `src/models/tokenizers.js` | -| Custom image processor | `src/models//image_processing_.js` | `src/models/image_processors.js` | -| Custom multimodal processor | `src/models//processing_.js` | `src/models/processors.js` | +| What | File | Barrel to update | +| ------------------------------ | ------------------------------------------------ | ---------------------------------- | +| Custom tokenizer | `src/models//tokenization_.js` | `src/models/tokenizers.js` | +| Custom image processor | `src/models//image_processing_.js` | `src/models/image_processors.js` | +| Custom multimodal processor | `src/models//processing_.js` | `src/models/processors.js` | | Custom audio/feature extractor | `src/models//feature_extraction_.js` | `src/models/feature_extractors.js` | The class name must match the `tokenizer_class` or `processor_class` field in the model's `tokenizer_config.json` / `preprocessor_config.json` on the Hub. @@ -210,7 +213,6 @@ pnpm test pnpm --filter @huggingface/transformers test -t "MyModelForCausalLM" ``` - ## Create a Pull Request Before writing any code, we strongly advise you to search through the existing PRs or @@ -230,6 +232,7 @@ You'll need the following tools installed to contribute to 🤗 Transformers.js: - **[pnpm](https://pnpm.io/)** - Fast, disk space efficient package manager To install pnpm: + ```bash npm install -g pnpm ``` @@ -267,6 +270,7 @@ Follow the steps below to start contributing: the pull request. ### Pull request checklist + ☐ The pull request title should summarize your contribution. ☐ If your pull request addresses an issue, please mention the issue number in the pull request description to make sure they are linked (and people viewing the issue know you @@ -280,19 +284,23 @@ useful to avoid duplicated work, and to differentiate it from PRs ready to be me ☐ If your changes affect user-facing functionality, update the relevant documentation. ### Tests + We are using [Jest](https://jestjs.io/) to execute unit-tests. All tests can be found in `packages/transformers/tests` and have to end with `.test.js` Execute all tests + ```bash pnpm test ``` Execute tests for a specific package + ```bash pnpm --filter @huggingface/transformers test ``` Execute a specific test file + ```bash cd packages/transformers pnpm test -- ./tests/models.test.js @@ -301,14 +309,17 @@ pnpm test -- ./tests/models.test.js ### Style guide #### Code formatting + We use [Prettier](https://prettier.io/) to maintain consistent code formatting across the project. Please ensure your code is formatted before submitting a pull request. **Format all files:** + ```bash pnpm format ``` **Check formatting without making changes:** + ```bash pnpm format:check ``` @@ -318,6 +329,7 @@ pnpm format:check We recommend setting up Prettier in your IDE to format on save: **Visual Studio Code:** + 1. Install the [Prettier extension](https://marketplace.visualstudio.com/items?itemName=esbenp.prettier-vscode) 2. Open Settings (Ctrl+, or Cmd+,) 3. Search for "format on save" @@ -325,6 +337,7 @@ We recommend setting up Prettier in your IDE to format on save: 5. Set Prettier as your default formatter: search for "default formatter" and select "Prettier - Code formatter" **IntelliJ IDEA / WebStorm:** + 1. Go to `Settings` → `Languages & Frameworks` → `JavaScript` → `Prettier` 2. Set the Prettier package path (usually `node_modules/prettier`) 3. Check "On save" under "Run for files" @@ -333,9 +346,10 @@ We recommend setting up Prettier in your IDE to format on save: ## Project Structure -This project uses **pnpm workspaces** to manage multiple packages in a monorepo. Currently, there is one workspace: +This project uses **pnpm workspaces** to manage multiple packages in a monorepo: - `packages/transformers` - The main Transformers.js library +- `packages/transformers-onnx` - The TypeScript ONNX Runtime inference provider This structure allows for better organization and makes it easier to add framework-specific integrations in the future. @@ -346,22 +360,26 @@ This structure allows for better organization and makes it easier to add framewo The recommended way to develop and test changes is to use the watch mode build and install from the local package: 1. Start the build in watch mode: + ```bash pnpm dev ``` + This will automatically rebuild the library whenever you make changes to the source code. 2. Create a separate test project and install transformers.js from your local development directory: + ```bash mkdir my-test-project cd my-test-project npm init -y npm install file:/path/to/transformers.js/packages/transformers ``` + Replace `/path/to/transformers.js` with the actual path to your cloned repository. 3. Make your changes to the transformers.js source code in the main repository. The watch mode will automatically rebuild the library. 4. Test your changes in your test project. The changes will be automatically reflected since the package is linked via the `file:` protocol. -This workflow allows for rapid iteration and testing during development. \ No newline at end of file +This workflow allows for rapid iteration and testing during development. diff --git a/packages/transformers-onnx/README.md b/packages/transformers-onnx/README.md new file mode 100644 index 000000000..f0705bb07 --- /dev/null +++ b/packages/transformers-onnx/README.md @@ -0,0 +1,9 @@ +# @huggingface/transformers-onnx + +ONNX Runtime inference provider for Transformers.js. + +```js +import { OnnxInferenceProvider } from '@huggingface/transformers-onnx'; + +const provider = OnnxInferenceProvider.from_modelId('onnx-community/model-ONNX'); +``` diff --git a/packages/transformers-onnx/jest.config.mjs b/packages/transformers-onnx/jest.config.mjs new file mode 100644 index 000000000..45493f1cb --- /dev/null +++ b/packages/transformers-onnx/jest.config.mjs @@ -0,0 +1,5 @@ +export default { + testEnvironment: "node", + roots: ["./tests/"], + transform: {}, +}; diff --git a/packages/transformers-onnx/package.json b/packages/transformers-onnx/package.json new file mode 100644 index 000000000..67c94f9f0 --- /dev/null +++ b/packages/transformers-onnx/package.json @@ -0,0 +1,61 @@ +{ + "name": "@huggingface/transformers-onnx", + "version": "0.1.0", + "description": "ONNX Runtime inference provider for Transformers.js", + "type": "module", + "main": "./dist/transformers-onnx.node.cjs", + "types": "./types/index.d.ts", + "exports": { + ".": { + "node": { + "import": { + "types": "./types/index.d.ts", + "default": "./dist/transformers-onnx.node.mjs" + }, + "require": { + "types": "./types/index.d.ts", + "default": "./dist/transformers-onnx.node.cjs" + } + }, + "default": { + "types": "./types/index.d.ts", + "default": "./dist/transformers-onnx.web.js" + } + }, + "./testing": { + "types": "./types/testing.d.ts", + "import": "./dist/testing.mjs", + "require": "./dist/testing.cjs" + } + }, + "scripts": { + "format": "prettier --write . --ignore-path ../../.prettierignore", + "format:check": "prettier --check . --ignore-path ../../.prettierignore", + "typegen": "tsc --build", + "build": "node scripts/build.mjs && pnpm typegen", + "test": "node --experimental-vm-modules node_modules/jest/bin/jest.js --runInBand" + }, + "dependencies": { + "onnxruntime-common": "1.24.3", + "onnxruntime-node": "1.24.3", + "onnxruntime-web": "1.26.0-dev.20260416-b7804b056c" + }, + "devDependencies": { + "@types/node": "^24.1.0", + "@webgpu/types": "^0.1.69", + "esbuild": "^0.27.2", + "jest": "^30.2.0", + "typescript": "5.9.3" + }, + "files": [ + "src", + "dist", + "types", + "README.md", + "!**/*.tsbuildinfo" + ], + "publishConfig": { + "access": "public" + }, + "license": "Apache-2.0" +} diff --git a/packages/transformers-onnx/scripts/build.mjs b/packages/transformers-onnx/scripts/build.mjs new file mode 100644 index 000000000..b5fe605cb --- /dev/null +++ b/packages/transformers-onnx/scripts/build.mjs @@ -0,0 +1,58 @@ +import { build } from "esbuild"; +import { mkdir } from "node:fs/promises"; + +await mkdir(new URL("../dist/", import.meta.url), { recursive: true }); + +const shared = { + entryPoints: [new URL("../src/index.ts", import.meta.url).pathname], + bundle: true, + sourcemap: false, + logLevel: "warning", +}; + +await Promise.all([ + build({ + ...shared, + outfile: new URL("../dist/transformers-onnx.node.mjs", import.meta.url).pathname, + platform: "node", + format: "esm", + external: ["onnxruntime-common", "onnxruntime-node"], + alias: { "onnxruntime-web/webgpu": "./src/empty.ts" }, + }), + build({ + ...shared, + outfile: new URL("../dist/transformers-onnx.node.cjs", import.meta.url).pathname, + platform: "node", + format: "cjs", + external: ["onnxruntime-common", "onnxruntime-node"], + alias: { "onnxruntime-web/webgpu": "./src/empty.ts" }, + }), + build({ + ...shared, + outfile: new URL("../dist/transformers-onnx.web.js", import.meta.url).pathname, + platform: "browser", + format: "esm", + external: ["onnxruntime-common", "onnxruntime-web"], + alias: { "onnxruntime-node": "./src/empty.ts" }, + }), + build({ + entryPoints: [new URL("../src/testing.ts", import.meta.url).pathname], + bundle: true, + sourcemap: false, + logLevel: "warning", + outfile: new URL("../dist/testing.mjs", import.meta.url).pathname, + platform: "node", + format: "esm", + external: ["onnxruntime-common", "onnxruntime-node"], + }), + build({ + entryPoints: [new URL("../src/testing.ts", import.meta.url).pathname], + bundle: true, + sourcemap: false, + logLevel: "warning", + outfile: new URL("../dist/testing.cjs", import.meta.url).pathname, + platform: "node", + format: "cjs", + external: ["onnxruntime-common", "onnxruntime-node"], + }), +]); diff --git a/packages/transformers-onnx/src/empty.ts b/packages/transformers-onnx/src/empty.ts new file mode 100644 index 000000000..ff8b4c563 --- /dev/null +++ b/packages/transformers-onnx/src/empty.ts @@ -0,0 +1 @@ +export default {}; diff --git a/packages/transformers-onnx/src/host.ts b/packages/transformers-onnx/src/host.ts new file mode 100644 index 000000000..c8edc0ce5 --- /dev/null +++ b/packages/transformers-onnx/src/host.ts @@ -0,0 +1,87 @@ +export interface BackendTensorStorage { + readonly backend: string; + readonly handle: unknown; + readonly type: string; + dims: number[]; + readonly data: unknown; + readonly size: number; + readonly location: string; + dispose(): void; +} + +export interface OnnxProviderHost { + readonly env: any; + readonly apis: any; + readonly logger: any; + getModelFile(modelId: string, file: string, fatal: boolean, options: any, returnPath?: boolean): Promise; + getCacheNames(config: any, options: any): Set; + createBackendTensor(storage: BackendTensorStorage): any; + getBackendTensorStorage(tensor: any): BackendTensorStorage | null; + getCache?(): Promise; +} + +let configuredHost: OnnxProviderHost | null = null; + +const fallbackEnvironment: any = { + backends: { onnx: {} }, + logLevel: 30, + useWasmCache: typeof caches !== 'undefined', + fetch: (...args: any[]) => (globalThis.fetch as any)(...args), +}; + +const environment = new Proxy(fallbackEnvironment, { + get(target, property) { + return (configuredHost?.env ?? target)[property]; + }, + set(target, property, value) { + (configuredHost?.env ?? target)[property] = value; + return true; + }, +}); + +const apis = { + IS_NODE_ENV: typeof process !== 'undefined' && process?.release?.name === 'node', + IS_WEB_ENV: typeof window !== 'undefined' || typeof self !== 'undefined', + IS_WEBGPU_AVAILABLE: typeof navigator !== 'undefined' && !!navigator.gpu, + IS_WEBNN_AVAILABLE: typeof navigator !== 'undefined' && 'ml' in navigator, + IS_DENO_WEB_RUNTIME: 'Deno' in globalThis && typeof window !== 'undefined', + IS_SAFARI_BELOW_26: false, + IS_SERVICE_WORKER_ENV: + 'ServiceWorkerGlobalScope' in globalThis && globalThis instanceof (globalThis as any).ServiceWorkerGlobalScope, + IS_CHROME_AVAILABLE: 'chrome' in globalThis, +}; + +const logger = new Proxy(console, { + get(target, property) { + return (configuredHost?.logger ?? target)[property]; + }, +}); + +const fallbackHost: OnnxProviderHost = { + env: environment, + apis, + logger, + async getModelFile() { + throw new Error('OnnxInferenceProvider host does not provide model file loading.'); + }, + getCacheNames() { + return new Set(); + }, + createBackendTensor() { + throw new Error('OnnxInferenceProvider host does not provide tensor creation.'); + }, + getBackendTensorStorage() { + return null; + }, +}; + +export function configureOnnxProviderHost(host: OnnxProviderHost): void { + if (fallbackEnvironment.backends.onnx) { + host.env.backends.onnx = fallbackEnvironment.backends.onnx; + } + configuredHost = host; +} + +export function getOnnxProviderHost(): OnnxProviderHost { + return configuredHost ?? fallbackHost; +} diff --git a/packages/transformers-onnx/src/index.ts b/packages/transformers-onnx/src/index.ts new file mode 100644 index 000000000..2e297cd39 --- /dev/null +++ b/packages/transformers-onnx/src/index.ts @@ -0,0 +1,4 @@ +export { OnnxInferenceProvider } from './provider.js'; +export { configureOnnxProviderHost } from './host.js'; +export { OnnxTensorOpRegistry } from './tensor-ops.js'; +export type { BackendTensorStorage, OnnxProviderHost } from './host.js'; diff --git a/packages/transformers-onnx/src/provider.ts b/packages/transformers-onnx/src/provider.ts new file mode 100644 index 000000000..3ce63a344 --- /dev/null +++ b/packages/transformers-onnx/src/provider.ts @@ -0,0 +1,449 @@ +import { + createInferenceSession, + deviceToExecutionProviders, + isONNXProxy, + isONNXTensor, + runInferenceSession, + Tensor as OrtTensor, +} from './runtime.js'; +import { getOnnxProviderHost } from './host.js'; + +const DATA_TYPES = Object.freeze({ + auto: 'auto', + fp32: 'fp32', + fp16: 'fp16', + q8: 'q8', + int8: 'int8', + uint8: 'uint8', + q4: 'q4', + bnb4: 'bnb4', + q4f16: 'q4f16', + q2: 'q2', + q2f16: 'q2f16', + q1: 'q1', + q1f16: 'q1f16', +}); + +const DEFAULT_DTYPE_SUFFIX_MAPPING: Record = Object.freeze({ + fp32: '', + fp16: '_fp16', + int8: '_int8', + uint8: '_uint8', + q8: '_quantized', + q4: '_q4', + q2: '_q2', + q1: '_q1', + q4f16: '_q4f16', + q2f16: '_q2f16', + q1f16: '_q1f16', + bnb4: '_bnb4', +}); + +const { apis, logger } = getOnnxProviderHost(); + +function selectDevice(value: any, fileName: string, { warn }: any = {}): string { + const fallback = apis.IS_NODE_ENV ? 'cpu' : 'wasm'; + if (!value) return fallback; + if (typeof value === 'string') return value; + if (Object.hasOwn(value, fileName)) return value[fileName]; + warn?.(`device not specified for "${fileName}". Using the default device (${fallback}).`); + return fallback; +} + +function selectDtype(value: any, fileName: string, device: string, { configDtype = null, warn }: any = {}): string { + let resolved = value; + let needsWarn = false; + if (value && typeof value !== 'string') { + resolved = Object.hasOwn(value, fileName) ? value[fileName] : null; + needsWarn = resolved === null; + } + if (resolved === 'auto') { + const fallback = typeof configDtype === 'string' ? configDtype : configDtype?.[fileName]; + if (fallback && fallback !== 'auto' && Object.hasOwn(DATA_TYPES, fallback)) return fallback; + } + const result = + resolved && resolved !== 'auto' && Object.hasOwn(DATA_TYPES, resolved) + ? resolved + : device === 'wasm' + ? 'q8' + : 'fp32'; + if (needsWarn) + warn?.( + `dtype not specified for "${fileName}". Using the default dtype (${result}) for this device (${device}).`, + ); + return result; +} + +let fp16Supported: boolean | undefined; +async function isWebGpuFp16Supported(): Promise { + if (fp16Supported === undefined) { + try { + const adapter = await navigator.gpu.requestAdapter(); + fp16Supported = !!adapter?.features.has('shader-f16'); + } catch { + fp16Supported = false; + } + } + return fp16Supported; +} + +/** + * ONNX Runtime adapter used for string model IDs. + */ +export class OnnxInferenceProvider { + /** + * @param {string} modelId + * @param {typeof import('../../models/modeling_utils.js').PreTrainedModel} [modelClass] + */ + static from_modelId(modelId: string): OnnxInferenceProvider { + return new OnnxInferenceProvider(modelId); + } + + static listModelArtifacts({ + sessions, + optionalConfigs, + config, + dtype: overrideDtype = null, + device: overrideDevice = null, + }: any): string[] { + const files = ['config.json']; + const customConfig = config?.['transformers.js_config'] ?? {}; + const rawDevice = overrideDevice ?? customConfig.device; + const dtype = overrideDtype ?? customConfig.dtype; + for (const [sessionName, baseName] of Object.entries(sessions) as [string, string][]) { + const device = selectDevice(rawDevice, sessionName); + const selectedDtype = selectDtype(dtype, sessionName, device); + const suffix = DEFAULT_DTYPE_SUFFIX_MAPPING[selectedDtype] ?? ''; + const fullName = `${baseName}${suffix}.onnx`; + files.push(`onnx/${fullName}`); + + const externalConfig = customConfig.use_external_data_format; + const count = + typeof externalConfig === 'object' && externalConfig !== null + ? +(externalConfig[fullName] ?? externalConfig[sessionName] ?? 0) + : +(externalConfig ?? 0); + files.push(...externalDataChunkNames(fullName, count).map((name) => `onnx/${name}`)); + } + if (optionalConfigs) files.push(...(Object.values(optionalConfigs) as string[])); + return files; + } + + static async getAvailableDtypes({ modelId, sessions, getFileMetadata, metadataOptions }: any): Promise { + const results = await Promise.all( + Object.entries(DEFAULT_DTYPE_SUFFIX_MAPPING).map(async ([dtype, suffix]) => ({ + dtype, + available: ( + await Promise.all( + (Object.values(sessions) as string[]).map( + async (baseName) => + (await getFileMetadata(modelId, `onnx/${baseName}${suffix}.onnx`, metadataOptions)) + .exists, + ), + ) + ).every(Boolean), + })), + ); + return results.filter((result) => result.available).map((result) => result.dtype); + } + + static filterModelArtifacts(files: string[], sessions: Record): string[] { + const allowedPrefixes = Object.values(sessions).map((name) => `onnx/${name}`); + return files.filter( + (file) => !file.startsWith('onnx/') || allowedPrefixes.some((prefix) => file.startsWith(prefix)), + ); + } + + readonly providerType = 'onnx'; + modelClass: any; + + constructor( + public readonly modelId: string, + modelClass: any = undefined, + ) { + this.modelClass = modelClass; + } + + /** + * Load a Transformers.js model class with this backend. + * + * @param {import('../../utils/hub.js').PretrainedModelOptions} options + */ + async load(options: any) { + const modelClass = options.modelClass ?? this.modelClass; + if (!modelClass) { + throw new Error('OnnxInferenceProvider requires a Transformers.js model class before it can load a model.'); + } + const { modelClass: _modelClass, ...loadOptions } = options; + return modelClass._from_pretrained(this.modelId, { ...loadOptions, inferenceProvider: this }); + } + + async getSession( + fileName: string, + options: any, + cache_config = false, + session_name: string | undefined = undefined, + ) { + let custom_config = options.config?.['transformers.js_config'] ?? {}; + const selectedDevice = /** @type {import('../../utils/devices.js').DeviceType} */ selectDevice( + options.device ?? custom_config.device, + fileName, + { + warn: (msg) => logger.info(msg), + }, + ); + const executionProviders = deviceToExecutionProviders(selectedDevice); + + const device_config = custom_config.device_config ?? {}; + if (Object.hasOwn(device_config, selectedDevice)) { + custom_config = { ...custom_config, ...device_config[selectedDevice] }; + } + + const selectedDtype = selectDtype(options.dtype ?? custom_config.dtype, fileName, selectedDevice, { + configDtype: custom_config.dtype, + warn: (msg) => logger.info(msg), + }); + if (!Object.hasOwn(DEFAULT_DTYPE_SUFFIX_MAPPING, selectedDtype)) { + throw new Error(`Invalid dtype: ${selectedDtype}. Should be one of: ${Object.keys(DATA_TYPES).join(', ')}`); + } + if ( + selectedDevice === 'webgpu' && + !apis.IS_NODE_ENV && + selectedDtype === DATA_TYPES.fp16 && + !(await isWebGpuFp16Supported()) + ) { + throw new Error(`The device (${selectedDevice}) does not support fp16.`); + } + + const suffix = DEFAULT_DTYPE_SUFFIX_MAPPING[selectedDtype]; + const session_options = { ...options.session_options }; + session_options.executionProviders ??= executionProviders; + + const free_dimension_overrides = custom_config.free_dimension_overrides; + if (free_dimension_overrides) { + session_options.freeDimensionOverrides ??= free_dimension_overrides; + } else if (selectedDevice.startsWith('webnn') && !session_options.freeDimensionOverrides) { + logger.warn( + `WebNN does not currently support dynamic shapes and requires 'free_dimension_overrides' to be set in config.json, preferably as a field within config["transformers.js_config"]["device_config"]["${selectedDevice}"]. ` + + `When 'free_dimension_overrides' is not set, you may experience significant performance degradation.`, + ); + } + + const bufferOrPathPromise = getCoreModelFile(this.modelId, fileName, options, suffix); + const use_external_data_format = options.use_external_data_format ?? custom_config.use_external_data_format; + const externalData = await getModelDataFiles( + this.modelId, + fileName, + suffix, + options, + use_external_data_format, + session_options, + ); + if (externalData.length > 0 && (!apis.IS_NODE_ENV || externalData.some((data) => typeof data !== 'string'))) { + session_options.externalData = externalData; + } + + if (cache_config && selectedDevice === 'webgpu') { + const names = getOnnxProviderHost().getCacheNames(options.config, { prefix: 'present', session_name }); + if (names.size > 0 && !isONNXProxy()) { + const preferredOutputLocation = {}; + for (const key of names) preferredOutputLocation[key] = 'gpu-buffer'; + session_options.preferredOutputLocation = preferredOutputLocation; + } + } + + return { + buffer_or_path: await bufferOrPathPromise, + session_options, + session_config: { dtype: selectedDtype, device: selectedDevice }, + }; + } + + async constructSessions(names: Record, options: any, cache_sessions: any = undefined) { + return Object.fromEntries( + await Promise.all( + Object.keys(names).map(async (name) => { + const sessionInfo = await this.getSession( + names[name], + options, + cache_sessions?.[name] ?? false, + name, + ); + const ortSession = await createInferenceSession( + sessionInfo.buffer_or_path, + sessionInfo.session_options, + sessionInfo.session_config, + ); + const session = { + inputNames: ortSession.inputNames, + outputNames: ortSession.outputNames, + inputMetadata: ortSession.inputMetadata, + outputMetadata: ortSession.outputMetadata, + config: (ortSession as any).config, + run: (inputs: any) => this.run(ortSession, inputs), + release: () => ortSession.release(), + }; + return [name, session]; + }), + ), + ); + } + + async run(session: any, inputs: Record) { + const checkedInputs = validateInputs(session, inputs); + try { + const ortFeed = Object.fromEntries( + Object.entries(checkedInputs).map(([key, value]) => { + const storage = getOnnxProviderHost().getBackendTensorStorage(value); + const tensor: any = + storage?.backend === 'onnx' + ? storage.handle + : new OrtTensor(value.type, value.data, value.dims); + if ( + apis.IS_NODE_ENV && + typeof Float16Array !== 'undefined' && + tensor.cpuData instanceof Float16Array + ) { + tensor.cpuData = new Uint16Array(tensor.cpuData.buffer); + } + return [key, tensor]; + }), + ); + return replaceTensors(await runInferenceSession(session, ortFeed)); + } catch (error) { + const formatted = Object.fromEntries( + Object.entries(checkedInputs).map(([key, tensor]) => { + const unpacked: any = { type: tensor.type, dims: tensor.dims, location: tensor.location }; + if (unpacked.location !== 'gpu-buffer') unpacked.data = tensor.data; + return [key, unpacked]; + }), + ); + logger.error(`An error occurred during model execution: "${error}".`); + logger.error('Inputs given to model:', formatted); + throw error; + } + } +} + +function replaceTensors(value: any): any { + for (const property in value) { + if (isONNXTensor(value[property])) { + const tensor = value[property]; + value[property] = getOnnxProviderHost().createBackendTensor({ + backend: 'onnx', + handle: tensor, + get type() { + return tensor.type; + }, + get dims() { + return tensor.dims; + }, + set dims(value) { + tensor.dims = value; + }, + get data() { + return tensor.data; + }, + get size() { + return tensor.size; + }, + get location() { + return tensor.location; + }, + dispose() { + tensor.dispose(); + }, + }); + } else if (value[property] && typeof value[property] === 'object') { + replaceTensors(value[property]); + } + } + return value; +} + +function validateInputs(session: any, inputs: Record): Record { + const checkedInputs = Object.create(null); + const missingInputs = []; + for (const inputName of session.inputNames) { + const tensor = inputs[inputName]; + if (!tensor || !Array.isArray(tensor.dims) || typeof tensor.type !== 'string') { + missingInputs.push(inputName); + continue; + } + checkedInputs[inputName] = isONNXProxy() ? tensor.clone() : tensor; + } + if (missingInputs.length > 0) { + throw new Error( + `An error occurred during model execution: "Missing the following inputs: ${missingInputs.join(', ')}.`, + ); + } + + const inputNames = Object.keys(inputs); + if (inputNames.length > session.inputNames.length) { + const ignored = inputNames.filter((inputName) => !session.inputNames.includes(inputName)); + logger.warn( + `WARNING: Too many inputs were provided (${inputNames.length} > ${session.inputNames.length}). The following inputs will be ignored: "${ignored.join(', ')}".`, + ); + } + return checkedInputs; +} + +async function getCoreModelFile(modelId: string, fileName: string, options: any, suffix: string): Promise { + const baseName = `${fileName}${suffix}.onnx`; + const subfolder = options.subfolder ?? 'onnx'; + return getOnnxProviderHost().getModelFile( + modelId, + subfolder ? `${subfolder}/${baseName}` : baseName, + true, + options, + apis.IS_NODE_ENV, + ); +} + +function externalDataChunkNames(fullName: string, count: number): string[] { + return Array.from({ length: count }, (_, index) => `${fullName}_data${index === 0 ? '' : `_${index}`}`); +} + +async function getModelDataFiles( + modelId: string, + fileName: string, + suffix: string, + options: any, + externalConfig: any, + sessionOptions: any, +): Promise { + const baseName = `${fileName}${suffix}.onnx`; + let count = 0; + if (typeof externalConfig === 'object' && externalConfig !== null) { + count = +(externalConfig[baseName] ?? externalConfig[fileName] ?? 0); + } else if (externalConfig) { + count = +externalConfig; + } + if (count > 1024) + throw new Error(`The number of external data chunks (${count}) exceeds the maximum allowed value (1024).`); + + if (count > 0) { + const subfolder = options.subfolder ?? 'onnx'; + return Promise.all( + externalDataChunkNames(baseName, count).map(async (path) => { + const data = await getOnnxProviderHost().getModelFile( + modelId, + subfolder ? `${subfolder}/${path}` : path, + true, + options, + apis.IS_NODE_ENV, + ); + return data instanceof Uint8Array ? { path, data } : path; + }), + ); + } + if (sessionOptions.externalData !== undefined) { + return Promise.all( + sessionOptions.externalData.map(async (item: any) => + typeof item.data === 'string' + ? { ...item, data: await getOnnxProviderHost().getModelFile(modelId, item.data, true, options) } + : item, + ), + ); + } + return []; +} diff --git a/packages/transformers/src/backends/onnx.js b/packages/transformers-onnx/src/runtime.ts similarity index 97% rename from packages/transformers/src/backends/onnx.js rename to packages/transformers-onnx/src/runtime.ts index 963f6fffe..df1fe5c58 100644 --- a/packages/transformers/src/backends/onnx.js +++ b/packages/transformers-onnx/src/runtime.ts @@ -16,17 +16,26 @@ * @module backends/onnx */ -import { env, apis, LogLevel } from '../env.js'; +import { getOnnxProviderHost } from './host.js'; // NOTE: Import order matters here. We need to import `onnxruntime-node` before `onnxruntime-web`. // In either case, we select the default export if it exists, otherwise we use the named export. import * as ONNX_NODE from 'onnxruntime-node'; import * as ONNX_WEB from 'onnxruntime-web/webgpu'; -import { loadWasmBinary, loadWasmFactory } from './utils/cacheWasm.js'; -import { isBlobURL, toAbsoluteURL } from '../utils/hub/utils.js'; -import { logger } from '../utils/logger.js'; +import { loadWasmBinary, loadWasmFactory } from './wasm-cache.js'; export { Tensor } from 'onnxruntime-common'; +const { env, apis, logger } = getOnnxProviderHost(); +const LogLevel = { DEBUG: 10, INFO: 20, WARNING: 30, ERROR: 40, NONE: 50 }; + +function isBlobURL(url: string): boolean { + return url.startsWith('blob:'); +} + +function toAbsoluteURL(url: string): string { + return new URL(url, globalThis.location?.href ?? 'file:///').href; +} + /** * @typedef {import('onnxruntime-common').InferenceSession.ExecutionProviderConfig} ONNXExecutionProviders */ @@ -228,7 +237,7 @@ async function ensureWasmLoaded() { wasmLoadPromise = (async () => { // At this point, we know wasmPaths is an object (not a string) because // shouldUseWasmCache checks for wasmPaths.wasm and wasmPaths.mjs - const urls = /** @type {{ wasm: string, mjs: string }} */ (ONNX_ENV.wasm.wasmPaths); + const urls = /** @type {{ wasm: string, mjs: string }} */ ONNX_ENV.wasm.wasmPaths; // Load both in parallel; the .mjs blob URL is only kept if wasmBinary succeeded. // ORT only sets locateFile when wasmBinary is provided (onnxruntime PR https://github.com/microsoft/onnxruntime/pull/27411), which diff --git a/packages/transformers-onnx/src/tensor-ops.ts b/packages/transformers-onnx/src/tensor-ops.ts new file mode 100644 index 000000000..2688796ea --- /dev/null +++ b/packages/transformers-onnx/src/tensor-ops.ts @@ -0,0 +1,218 @@ +import { createInferenceSession, runInferenceSession, isONNXProxy, Tensor as OrtTensor } from './runtime.js'; +import { getOnnxProviderHost } from './host.js'; + +/** + * Asynchronously creates a wrapper function for running an ONNX inference session. + * + * @param {number[]} session_bytes The session data in bytes. + * @param session_options The options for the ONNX session. + * @template {string | [string] | string[]} T + * @param {T} names The name(s) of the output tensor(s). + * + * @returns {Promise): Promise>} + * The wrapper function for running the ONNX inference session. + */ +const wrap = async (session_bytes: number[], session_options: any, names: string | string[]) => { + const session = await createInferenceSession(new Uint8Array(session_bytes), session_options, {}); + + return /** @type {any} */ async (inputs: Record) => { + const proxied = isONNXProxy(); + const ortFeed = Object.fromEntries( + Object.entries(inputs).map(([key, value]) => { + const input = proxied ? value.clone() : value; + const storage = getOnnxProviderHost().getBackendTensorStorage(input); + return [ + key, + storage?.backend === 'onnx' ? storage.handle : new OrtTensor(input.type, input.data, input.dims), + ]; + }), + ); + const outputs = await runInferenceSession(session, ortFeed); + if (Array.isArray(names)) { + return names.map((name) => wrapTensor(outputs[name])); + } else { + return wrapTensor(outputs[names]); + } + }; +}; + +// In-memory registry of initialized ONNX operators +export class OnnxTensorOpRegistry { + static _nearest_interpolate_4d: any; + static _bilinear_interpolate_4d: any; + static _bicubic_interpolate_4d: any; + static _matmul: any; + static _stft: any; + static _rfft: any; + static _top_k: any; + static _slice: any; + + static session_options = { + // TODO: Allow for multiple execution providers + // executionProviders: ['webgpu'], + }; + + static get nearest_interpolate_4d() { + if (!this._nearest_interpolate_4d) { + this._nearest_interpolate_4d = wrap( + [ + 8, 10, 18, 0, 58, 129, 1, 10, 41, 10, 1, 120, 10, 0, 10, 0, 10, 1, 115, 18, 1, 121, 34, 6, 82, 101, + 115, 105, 122, 101, 42, 18, 10, 4, 109, 111, 100, 101, 34, 7, 110, 101, 97, 114, 101, 115, 116, 160, + 1, 3, 18, 1, 114, 90, 31, 10, 1, 120, 18, 26, 10, 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, + 99, 10, 3, 18, 1, 104, 10, 3, 18, 1, 119, 90, 15, 10, 1, 115, 18, 10, 10, 8, 8, 7, 18, 4, 10, 2, 8, + 4, 98, 31, 10, 1, 121, 18, 26, 10, 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, 99, 10, 3, 18, + 1, 104, 10, 3, 18, 1, 119, 66, 2, 16, 21, + ], + this.session_options, + 'y', + ); + } + return this._nearest_interpolate_4d; + } + static get bilinear_interpolate_4d() { + if (!this._bilinear_interpolate_4d) { + this._bilinear_interpolate_4d = wrap( + [ + 8, 9, 18, 0, 58, 128, 1, 10, 40, 10, 1, 120, 10, 0, 10, 0, 10, 1, 115, 18, 1, 121, 34, 6, 82, 101, + 115, 105, 122, 101, 42, 17, 10, 4, 109, 111, 100, 101, 34, 6, 108, 105, 110, 101, 97, 114, 160, 1, + 3, 18, 1, 114, 90, 31, 10, 1, 120, 18, 26, 10, 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, 99, + 10, 3, 18, 1, 104, 10, 3, 18, 1, 119, 90, 15, 10, 1, 115, 18, 10, 10, 8, 8, 7, 18, 4, 10, 2, 8, 4, + 98, 31, 10, 1, 121, 18, 26, 10, 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, 99, 10, 3, 18, 1, + 104, 10, 3, 18, 1, 119, 66, 2, 16, 20, + ], + this.session_options, + 'y', + ); + } + return this._bilinear_interpolate_4d; + } + + static get bicubic_interpolate_4d() { + if (!this._bicubic_interpolate_4d) { + this._bicubic_interpolate_4d = wrap( + [ + 8, 9, 18, 0, 58, 127, 10, 39, 10, 1, 120, 10, 0, 10, 0, 10, 1, 115, 18, 1, 121, 34, 6, 82, 101, 115, + 105, 122, 101, 42, 16, 10, 4, 109, 111, 100, 101, 34, 5, 99, 117, 98, 105, 99, 160, 1, 3, 18, 1, + 114, 90, 31, 10, 1, 120, 18, 26, 10, 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, 99, 10, 3, + 18, 1, 104, 10, 3, 18, 1, 119, 90, 15, 10, 1, 115, 18, 10, 10, 8, 8, 7, 18, 4, 10, 2, 8, 4, 98, 31, + 10, 1, 121, 18, 26, 10, 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, 99, 10, 3, 18, 1, 104, 10, + 3, 18, 1, 119, 66, 2, 16, 20, + ], + this.session_options, + 'y', + ); + } + return this._bicubic_interpolate_4d; + } + + static get matmul() { + if (!this._matmul) { + this._matmul = wrap( + [ + 8, 9, 18, 0, 58, 55, 10, 17, 10, 1, 97, 10, 1, 98, 18, 1, 99, 34, 6, 77, 97, 116, 77, 117, 108, 18, + 1, 114, 90, 9, 10, 1, 97, 18, 4, 10, 2, 8, 1, 90, 9, 10, 1, 98, 18, 4, 10, 2, 8, 1, 98, 9, 10, 1, + 99, 18, 4, 10, 2, 8, 1, 66, 2, 16, 20, + ], + this.session_options, + 'c', + ); + } + return this._matmul; + } + + static get stft() { + if (!this._stft) { + this._stft = wrap( + [ + 8, 7, 18, 0, 58, 148, 1, 10, 38, 10, 1, 115, 10, 1, 106, 10, 1, 119, 10, 1, 108, 18, 1, 111, 34, 4, + 83, 84, 70, 84, 42, 15, 10, 8, 111, 110, 101, 115, 105, 100, 101, 100, 24, 1, 160, 1, 2, 18, 1, 115, + 90, 26, 10, 1, 115, 18, 21, 10, 19, 8, 1, 18, 15, 10, 3, 18, 1, 98, 10, 3, 18, 1, 115, 10, 3, 18, 1, + 99, 90, 11, 10, 1, 106, 18, 6, 10, 4, 8, 7, 18, 0, 90, 16, 10, 1, 119, 18, 11, 10, 9, 8, 1, 18, 5, + 10, 3, 18, 1, 119, 90, 11, 10, 1, 108, 18, 6, 10, 4, 8, 7, 18, 0, 98, 31, 10, 1, 111, 18, 26, 10, + 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, 102, 10, 3, 18, 1, 100, 10, 3, 18, 1, 99, 66, 2, + 16, 17, + ], + this.session_options, + 'o', + ); + } + return this._stft; + } + + static get rfft() { + if (!this._rfft) { + this._rfft = wrap( + [ + 8, 9, 18, 0, 58, 97, 10, 33, 10, 1, 120, 10, 0, 10, 1, 97, 18, 1, 121, 34, 3, 68, 70, 84, 42, 15, + 10, 8, 111, 110, 101, 115, 105, 100, 101, 100, 24, 1, 160, 1, 2, 18, 1, 100, 90, 21, 10, 1, 120, 18, + 16, 10, 14, 8, 1, 18, 10, 10, 3, 18, 1, 115, 10, 3, 18, 1, 99, 90, 11, 10, 1, 97, 18, 6, 10, 4, 8, + 7, 18, 0, 98, 21, 10, 1, 121, 18, 16, 10, 14, 8, 1, 18, 10, 10, 3, 18, 1, 115, 10, 3, 18, 1, 99, 66, + 2, 16, 20, + ], + this.session_options, + 'y', + ); + } + return this._rfft; + } + + static get top_k() { + if (!this._top_k) { + this._top_k = wrap( + [ + 8, 10, 18, 0, 58, 73, 10, 18, 10, 1, 120, 10, 1, 107, 18, 1, 118, 18, 1, 105, 34, 4, 84, 111, 112, + 75, 18, 1, 116, 90, 9, 10, 1, 120, 18, 4, 10, 2, 8, 1, 90, 15, 10, 1, 107, 18, 10, 10, 8, 8, 7, 18, + 4, 10, 2, 8, 1, 98, 9, 10, 1, 118, 18, 4, 10, 2, 8, 1, 98, 9, 10, 1, 105, 18, 4, 10, 2, 8, 7, 66, 2, + 16, 21, + ], + this.session_options, + [/* Values */ 'v', /* Indices */ 'i'], + ); + } + return this._top_k; + } + + static get slice() { + if (!this._slice) { + this._slice = wrap( + [ + 8, 7, 18, 0, 58, 96, 10, 25, 10, 1, 120, 10, 1, 115, 10, 1, 101, 10, 1, 97, 10, 1, 116, 18, 1, 121, + 34, 5, 83, 108, 105, 99, 101, 18, 1, 114, 90, 9, 10, 1, 120, 18, 4, 10, 2, 8, 1, 90, 9, 10, 1, 115, + 18, 4, 10, 2, 8, 7, 90, 9, 10, 1, 101, 18, 4, 10, 2, 8, 7, 90, 9, 10, 1, 97, 18, 4, 10, 2, 8, 7, 90, + 9, 10, 1, 116, 18, 4, 10, 2, 8, 7, 98, 9, 10, 1, 121, 18, 4, 10, 2, 8, 1, 66, 2, 16, 13, + ], + this.session_options, + 'y', + ); + } + return this._slice; + } +} + +function wrapTensor(tensor: any): any { + return getOnnxProviderHost().createBackendTensor({ + backend: 'onnx', + handle: tensor, + get type() { + return tensor.type; + }, + get dims() { + return tensor.dims; + }, + set dims(value) { + tensor.dims = value; + }, + get data() { + return tensor.data; + }, + get size() { + return tensor.size; + }, + get location() { + return tensor.location; + }, + dispose() { + tensor.dispose(); + }, + }); +} diff --git a/packages/transformers-onnx/src/testing.ts b/packages/transformers-onnx/src/testing.ts new file mode 100644 index 000000000..b8189b367 --- /dev/null +++ b/packages/transformers-onnx/src/testing.ts @@ -0,0 +1,44 @@ +import * as types from 'node:util/types'; +import { onnxruntimeBackend } from 'onnxruntime-node/dist/backend'; +import * as ONNX_COMMON from 'onnxruntime-common'; + +/** Register ONNX Runtime's Node backend in Jest's VM context. */ +export function initOnnxTestBackend(): void { + ONNX_COMMON.env.wasm.numThreads = 1; + const originalMethod = onnxruntimeBackend.init; + onnxruntimeBackend.init = function (...args: any[]) { + Array.isArray = (value: any): value is any[] => + typeof value === 'object' && + value !== null && + typeof value.length === 'number' && + value?.constructor.toString() === Array.toString(); + + const constructors = [ + 'Int8Array', + 'Int16Array', + 'Int32Array', + 'BigInt64Array', + 'Uint8Array', + 'Uint8ClampedArray', + 'Uint16Array', + 'Uint32Array', + 'BigUint64Array', + 'Float16Array', + 'Float32Array', + 'Float64Array', + ]; + for (const name of constructors) { + const constructor = (globalThis as any)[name]; + const check = (types as any)[`is${name}`]; + if (!constructor || !check) continue; + Object.defineProperty(constructor, Symbol.hasInstance, { + value: check.bind(types), + writable: true, + configurable: false, + enumerable: false, + }); + } + return originalMethod.apply(this, args as any); + }; + ONNX_COMMON.registerBackend('test', onnxruntimeBackend, Number.POSITIVE_INFINITY); +} diff --git a/packages/transformers-onnx/src/wasm-cache.ts b/packages/transformers-onnx/src/wasm-cache.ts new file mode 100644 index 000000000..5ff3185c2 --- /dev/null +++ b/packages/transformers-onnx/src/wasm-cache.ts @@ -0,0 +1,50 @@ +import { getOnnxProviderHost } from './host.js'; + +async function loadAndCacheFile(url: string): Promise { + const { env, logger, getCache } = getOnnxProviderHost(); + const fileName = url.split('/').pop(); + let cache: any; + try { + cache = await getCache?.(); + const cached = await cache?.match(url); + if (cached) return cached; + } catch (error) { + logger.warn(`Failed to load ${fileName} from cache:`, error); + } + + const response = await env.fetch(url); + if (!response.ok) throw new Error(`Failed to fetch ${fileName}: ${response.status} ${response.statusText}`); + if (cache) { + try { + await cache.put(url, response.clone()); + } catch (error) { + logger.warn(`Failed to cache ${fileName}:`, error); + } + } + return response; +} + +export async function loadWasmBinary(url: string): Promise { + const response = await loadAndCacheFile(url); + if (!response || typeof response === 'string') return null; + try { + return await response.arrayBuffer(); + } catch (error) { + getOnnxProviderHost().logger.warn('Failed to read WASM binary:', error); + return null; + } +} + +export async function loadWasmFactory(url: string): Promise { + const { apis, logger } = getOnnxProviderHost(); + if (apis.IS_SERVICE_WORKER_ENV || apis.IS_CHROME_AVAILABLE) return url; + const response = await loadAndCacheFile(url); + if (!response || typeof response === 'string') return null; + try { + const code = (await response.text()).replaceAll('globalThis.process?.versions?.node', 'false'); + return URL.createObjectURL(new Blob([code], { type: 'text/javascript' })); + } catch (error) { + logger.warn('Failed to read WASM factory:', error); + return null; + } +} diff --git a/packages/transformers-onnx/tests/provider.test.js b/packages/transformers-onnx/tests/provider.test.js new file mode 100644 index 000000000..fcceb7647 --- /dev/null +++ b/packages/transformers-onnx/tests/provider.test.js @@ -0,0 +1,11 @@ +import { OnnxInferenceProvider } from "@huggingface/transformers-onnx"; + +describe("OnnxInferenceProvider", () => { + it("creates providers from model IDs", () => { + const provider = OnnxInferenceProvider.from_modelId("onnx-community/test-model"); + + expect(provider.modelId).toBe("onnx-community/test-model"); + expect(provider.providerType).toBe("onnx"); + expect(typeof provider.constructSessions).toBe("function"); + }); +}); diff --git a/packages/transformers-onnx/tsconfig.json b/packages/transformers-onnx/tsconfig.json new file mode 100644 index 000000000..0caf6d078 --- /dev/null +++ b/packages/transformers-onnx/tsconfig.json @@ -0,0 +1,18 @@ +{ + "include": ["src/**/*"], + "compilerOptions": { + "target": "esnext", + "module": "esnext", + "moduleResolution": "bundler", + "outDir": "types", + "rootDir": "src", + "strict": false, + "skipLibCheck": true, + "declaration": true, + "declarationMap": true, + "emitDeclarationOnly": true, + "esModuleInterop": true, + "composite": true, + "types": ["node", "@webgpu/types"] + } +} diff --git a/packages/transformers/package.json b/packages/transformers/package.json index 8ea694b11..74d9c29a4 100644 --- a/packages/transformers/package.json +++ b/packages/transformers/package.json @@ -57,8 +57,7 @@ "dependencies": { "@huggingface/jinja": "^0.5.6", "@huggingface/tokenizers": "^0.1.3", - "onnxruntime-node": "1.24.3", - "onnxruntime-web": "1.26.0-dev.20260416-b7804b056c", + "@huggingface/transformers-onnx": "workspace:^", "sharp": "^0.34.5" }, "devDependencies": { diff --git a/packages/transformers/scripts/build/buildAll.mjs b/packages/transformers/scripts/build/buildAll.mjs index 8cfc89561..01b0e3c3d 100644 --- a/packages/transformers/scripts/build/buildAll.mjs +++ b/packages/transformers/scripts/build/buildAll.mjs @@ -2,7 +2,6 @@ import { build as esbuild } from "esbuild"; import path from "node:path"; import { stripNodePrefixPlugin } from "./plugins/stripNodePrefixPlugin.mjs"; import { ignoreModulesPlugin } from "./plugins/ignoreModulesPlugin.mjs"; -import { postBuildPlugin } from "./plugins/postBuildPlugin.mjs"; import { externalNodeBuiltinsPlugin } from "./plugins/externalNodeBuiltinsPlugin.mjs"; import { OUT_DIR, ROOT_DIR, getEsbuildProdConfig } from "./constants.mjs"; import { reportSize } from "../../../../scripts/reportSize.mjs"; @@ -19,7 +18,6 @@ async function buildTarget( format = "esm", // 'esm' | 'cjs' ignoreModules = [], externalModules = [], - usePostBuild = false, }, log, ) { @@ -35,9 +33,6 @@ async function buildTarget( } plugins.push(stripNodePrefixPlugin()); plugins.push(externalNodeBuiltinsPlugin()); - if (usePostBuild) { - plugins.push(postBuildPlugin(OUT_DIR, ROOT_DIR)); - } log.build(`Building ${colors.bright}${regularFile}${colors.reset}...`); await esbuild({ diff --git a/packages/transformers/scripts/build/buildAllWithWatch.mjs b/packages/transformers/scripts/build/buildAllWithWatch.mjs index 05f6aa8e8..fcb9ef4ce 100644 --- a/packages/transformers/scripts/build/buildAllWithWatch.mjs +++ b/packages/transformers/scripts/build/buildAllWithWatch.mjs @@ -2,7 +2,6 @@ import { context } from "esbuild"; import path from "node:path"; import { stripNodePrefixPlugin } from "./plugins/stripNodePrefixPlugin.mjs"; import { ignoreModulesPlugin } from "./plugins/ignoreModulesPlugin.mjs"; -import { postBuildPlugin } from "./plugins/postBuildPlugin.mjs"; import { externalNodeBuiltinsPlugin } from "./plugins/externalNodeBuiltinsPlugin.mjs"; import { rebuildPlugin } from "../../../../scripts/rebuildPlugin.mjs"; import { OUT_DIR, ROOT_DIR, getEsbuildDevConfig } from "./constants.mjs"; @@ -12,7 +11,7 @@ import { BUILD_TARGETS } from "./targets.mjs"; * Create an esbuild context for a single build target */ async function createBuildContext(targetName, targetConfig, log) { - const { name, suffix, format, ignoreModules, externalModules, usePostBuild } = targetConfig; + const { name, suffix, format, ignoreModules, externalModules } = targetConfig; const platform = format === "cjs" ? "node" : "neutral"; const outputFile = `transformers${name}${suffix}`; @@ -23,9 +22,6 @@ async function createBuildContext(targetName, targetConfig, log) { } plugins.push(stripNodePrefixPlugin()); plugins.push(externalNodeBuiltinsPlugin()); - if (usePostBuild) { - plugins.push(postBuildPlugin(OUT_DIR, ROOT_DIR)); - } plugins.push(rebuildPlugin(targetName, log)); return context({ diff --git a/packages/transformers/scripts/build/constants.mjs b/packages/transformers/scripts/build/constants.mjs index f8b2a3373..0d10113a7 100644 --- a/packages/transformers/scripts/build/constants.mjs +++ b/packages/transformers/scripts/build/constants.mjs @@ -2,16 +2,15 @@ import path from "node:path"; import { fileURLToPath } from "node:url"; export const DIST_FOLDER = "dist"; -export const NODE_IGNORE_MODULES = ["onnxruntime-web"]; +export const NODE_IGNORE_MODULES = []; export const NODE_EXTERNAL_MODULES = [ - "onnxruntime-common", - "onnxruntime-node", + "@huggingface/transformers-onnx", "sharp", // node:* modules are handled by externalNodeBuiltinsPlugin ]; -export const WEB_IGNORE_MODULES = ["onnxruntime-node", "sharp", "fs", "path", "url", "stream", "stream/promises"]; -export const WEB_EXTERNAL_MODULES = ["onnxruntime-common", "onnxruntime-web"]; +export const WEB_IGNORE_MODULES = ["sharp", "fs", "path", "url", "stream", "stream/promises"]; +export const WEB_EXTERNAL_MODULES = ["@huggingface/transformers-onnx"]; const __dirname = path.dirname(fileURLToPath(import.meta.url)); export const ROOT_DIR = path.join(__dirname, "../.."); diff --git a/packages/transformers/scripts/build/plugins/ignoreModulesPlugin.mjs b/packages/transformers/scripts/build/plugins/ignoreModulesPlugin.mjs index 4900cf39b..f8dd9f0d6 100644 --- a/packages/transformers/scripts/build/plugins/ignoreModulesPlugin.mjs +++ b/packages/transformers/scripts/build/plugins/ignoreModulesPlugin.mjs @@ -38,7 +38,6 @@ export const ignoreModulesPlugin = (modules = []) => ({ case "node:path": case "node:url": case "sharp": - case "onnxruntime-node": default: return { contents: `export default {};`, diff --git a/packages/transformers/scripts/build/plugins/postBuildPlugin.mjs b/packages/transformers/scripts/build/plugins/postBuildPlugin.mjs deleted file mode 100644 index 769b02ea7..000000000 --- a/packages/transformers/scripts/build/plugins/postBuildPlugin.mjs +++ /dev/null @@ -1,43 +0,0 @@ -import path from "node:path"; -import { copyFileSync, unlinkSync, existsSync } from "node:fs"; -import { colors, createLogger } from "../../../../../scripts/logger.mjs"; - -const log = createLogger("transformers"); - -/** - * Plugin to post-process build files. - * Equivalent to webpack's PostBuildPlugin. - */ -export const postBuildPlugin = (distDir, rootDir) => { - // it should copy the files only once. In watch mode for example it should not rerun every time - let completed = false; - - return { - name: "post-build", - setup(build) { - build.onEnd(() => { - if (completed) return; - completed = true; - - const ORT_JSEP_FILE = "ort-wasm-simd-threaded.jsep.mjs"; - const ORT_BUNDLE_FILE = "ort.bundle.min.mjs"; - - // 1. Remove unnecessary files - const file = path.join(distDir, ORT_BUNDLE_FILE); - if (existsSync(file)) unlinkSync(file); - - // 2. Copy unbundled JSEP file - const ORT_SOURCE_DIR = path.join(rootDir, "node_modules/onnxruntime-web/dist"); - const src = path.join(ORT_SOURCE_DIR, ORT_JSEP_FILE); - - if (existsSync(src)) { - const dest = path.join(distDir, ORT_JSEP_FILE); - copyFileSync(src, dest); - log.success(`${colors.gray}Copied ${ORT_JSEP_FILE}${colors.reset}`); - } else { - log.warning(`Could not find ${ORT_JSEP_FILE} in node_modules`); - } - }); - }, - }; -}; diff --git a/packages/transformers/scripts/build/targets.mjs b/packages/transformers/scripts/build/targets.mjs index fcbf4383d..0d12e7532 100644 --- a/packages/transformers/scripts/build/targets.mjs +++ b/packages/transformers/scripts/build/targets.mjs @@ -12,7 +12,6 @@ export const BUILD_TARGETS = [ suffix: ".js", format: "esm", ignoreModules: WEB_IGNORE_MODULES, - usePostBuild: true, }, }, { @@ -23,7 +22,6 @@ export const BUILD_TARGETS = [ format: "esm", ignoreModules: WEB_IGNORE_MODULES, externalModules: WEB_EXTERNAL_MODULES, - usePostBuild: false, }, }, { @@ -34,7 +32,6 @@ export const BUILD_TARGETS = [ format: "esm", ignoreModules: NODE_IGNORE_MODULES, externalModules: NODE_EXTERNAL_MODULES, - usePostBuild: false, }, }, { @@ -45,7 +42,6 @@ export const BUILD_TARGETS = [ format: "cjs", ignoreModules: NODE_IGNORE_MODULES, externalModules: NODE_EXTERNAL_MODULES, - usePostBuild: false, }, }, ]; diff --git a/packages/transformers/src/backends/artifacts.js b/packages/transformers/src/backends/artifacts.js new file mode 100644 index 000000000..0c9290e2c --- /dev/null +++ b/packages/transformers/src/backends/artifacts.js @@ -0,0 +1,26 @@ +/** + * @file Runtime-neutral random-access artifact provider contracts. + * @module backends/artifacts + */ + +/** + * @typedef {Object} ArtifactProgressEvent + * @property {string} file + * @property {number} loaded + * @property {number} [total] + */ + +/** + * @typedef {Object} RandomAccessByteSource + * @property {number} [size] Stable byte length when known. It may initially be undefined and become defined after transport metadata arrives. + * @property {(begin: number, end: number, options?: {signal?: AbortSignal}) => Promise} read Read an independent half-open byte range `[begin, end)`. The returned array is owned by the caller and remains valid after later reads and close. + * @property {() => Promise} close Idempotently reject new reads, drain reads already in progress, and release the source. + */ + +/** + * @typedef {Object} InferenceArtifactProvider + * @property {(file: string, options?: {signal?: AbortSignal}) => Promise} readJson + * @property {(file: string, options?: {signal?: AbortSignal, onProgress?: (event: ArtifactProgressEvent) => void}) => Promise} openByteSource Open a source supporting concurrent, independently positioned reads. + */ + +export {}; diff --git a/packages/transformers/src/backends/default.js b/packages/transformers/src/backends/default.js new file mode 100644 index 000000000..31f668192 --- /dev/null +++ b/packages/transformers/src/backends/default.js @@ -0,0 +1,24 @@ +import { OnnxInferenceProvider, OnnxTensorOpRegistry, configureOnnxProviderHost } from '@huggingface/transformers-onnx'; + +import { env, apis } from '../env.js'; +import { logger } from '../utils/logger.js'; +import { getModelFile } from '../utils/hub.js'; +import { getCacheNames } from '../configs.js'; +import { Tensor } from '../utils/tensor.js'; +import { TensorOpRegistry } from '../ops/registry.js'; +import { getCache } from '../utils/cache.js'; + +configureOnnxProviderHost({ + env, + apis, + logger, + getModelFile, + getCacheNames, + createBackendTensor: (storage) => Tensor.fromBackendStorage(storage), + getBackendTensorStorage: (tensor) => tensor?.getBackendStorage?.() ?? null, + getCache, +}); + +TensorOpRegistry.register(OnnxTensorOpRegistry); + +export { OnnxInferenceProvider }; diff --git a/packages/transformers/src/backends/inference.js b/packages/transformers/src/backends/inference.js new file mode 100644 index 000000000..4e6fe15a3 --- /dev/null +++ b/packages/transformers/src/backends/inference.js @@ -0,0 +1,135 @@ +/** + * @file Runtime-neutral inference backend helpers. + * + * An inference backend is a model factory with a shared pretrained model ID: + * + * ```js + * const backend = { + * modelId: 'organization/model', + * async load(options) { + * return model; + * }, + * }; + * ``` + * + * The returned model may be callable, or may expose a `forward(inputs)` method. + * Generation models expose `createAutoregressiveSession(options)`; Transformers.js installs their public `generate()`. + * + * @module backends/inference + */ + +import { installGenerationRuntime } from '../generation/runtime.js'; + +/** + * @typedef {Object} InferenceModel + * @property {(inputs: Record) => Promise>} [forward] + * @property {(options: Object) => Promise} [generate] + * @property {import('../generation/runtime.js').GenerationCapabilitiesV1} [generationCapabilities] + * @property {(options: Object) => Promise} [createAutoregressiveSession] + * @property {Object} [config] + * @property {() => Promise|unknown} dispose + */ + +/** + * @typedef {Object} InferenceBackend + * @property {string} modelId Model ID or local path used for shared config, tokenizer, and processor assets. + * @property {(options: any) => Promise} load + * @property {(names: Record, options: Object, cacheSessions?: Object) => Promise>} [constructSessions] + */ + +/** + * Returns whether a value implements the custom inference backend contract. + * Classes with static `modelId` and `load` members are supported too. + * + * @param {unknown} value + * @returns {value is InferenceBackend} + */ +export function isInferenceBackend(value) { + const backend = /** @type {any} */ (value); + return ( + (typeof value === 'object' || typeof value === 'function') && + value !== null && + typeof backend.modelId === 'string' && + typeof backend.load === 'function' + ); +} + +/** + * Resolve a string model ID from either a string or an inference backend. + * + * @param {string|InferenceBackend} model + * @returns {string} + */ +export function getModelId(model) { + if (typeof model === 'string') return model; + if (isInferenceBackend(model)) return model.modelId; + throw new TypeError('Model must be a model ID string or an inference backend with `modelId` and `load(options)`.'); +} + +/** + * Make a plain model with `forward()` callable, matching the model contract used by pipelines. + * + * @param {InferenceModel|Function} model + * @returns {InferenceModel|Function} + */ +export function normalizeInferenceModel(model) { + const implementation = /** @type {any} */ (model); + if ((typeof model !== 'object' && typeof model !== 'function') || model === null) { + throw new TypeError('Inference backend `load()` must return a model.'); + } + if (typeof implementation.dispose !== 'function') { + throw new TypeError('Inference backend models must implement `dispose()`.'); + } + if (typeof model === 'function') return model; + if ( + typeof implementation.forward !== 'function' && + typeof implementation.createAutoregressiveSession !== 'function' + ) { + throw new TypeError( + 'Inference backend models must be callable, implement `forward(inputs)`, or implement `createAutoregressiveSession(options)`.', + ); + } + + const callable = (...args) => { + if (typeof implementation.forward !== 'function') { + throw new Error('This inference model does not implement `forward(inputs)`.'); + } + return implementation.forward(...args); + }; + return new Proxy(callable, { + get(target, property, receiver) { + return property in implementation + ? Reflect.get(implementation, property, implementation) + : Reflect.get(target, property, receiver); + }, + set(_target, property, value) { + return Reflect.set(implementation, property, value, implementation); + }, + has(target, property) { + return property in implementation || property in target; + }, + }); +} + +/** + * Load and normalize a custom inference model. + * + * @param {InferenceBackend} backend + * @param {Object} options + * @returns {Promise} + */ +export async function loadInferenceModel(backend, options) { + const loadOptions = { ...options, modelId: backend.modelId }; + if (loadOptions.device === null) loadOptions.device = undefined; + if (loadOptions.dtype === null) loadOptions.dtype = undefined; + const model = /** @type {any} */ ( + installGenerationRuntime(normalizeInferenceModel(await backend.load(loadOptions))) + ); + if (model.config == null && options.config != null) { + model.config = options.config; + } + if (model.generation_config == null && options.generation_config != null) { + model.generation_config = options.generation_config; + } + return model; +} diff --git a/packages/transformers/src/backends/utils/cacheWasm.js b/packages/transformers/src/backends/utils/cacheWasm.js deleted file mode 100644 index a0c096567..000000000 --- a/packages/transformers/src/backends/utils/cacheWasm.js +++ /dev/null @@ -1,101 +0,0 @@ -import { apis, env } from '../../env.js'; -import { getCache } from '../../utils/cache.js'; -import { logger } from '../../utils/logger.js'; - -/** - * Loads and caches a file from the given URL. - * @param {string} url The URL of the file to load. - * @returns {Promise} The response object, or null if loading failed. - */ -async function loadAndCacheFile(url) { - const fileName = url.split('/').pop(); - - /** @type {import('../../utils/cache.js').CacheInterface|undefined} */ - let cache; - try { - cache = await getCache(); - - // Try to get from cache first - if (cache) { - const result = await cache.match(url); - if (result) { - return result; - } - } - } catch (error) { - logger.warn(`Failed to load ${fileName} from cache:`, error); - } - - // If not in cache, fetch it - const response = await env.fetch(url); - - if (!response.ok) { - throw new Error(`Failed to fetch ${fileName}: ${response.status} ${response.statusText}`); - } - - // Cache the response for future use - if (cache) { - try { - await cache.put(url, response.clone()); - } catch (e) { - logger.warn(`Failed to cache ${fileName}:`, e); - } - } - - return response; -} - -/** - * Loads and caches the WASM binary for ONNX Runtime. - * @param {string} wasmURL The URL of the WASM file to load. - * @returns {Promise} The WASM binary as an ArrayBuffer, or null if loading failed. - */ - -export async function loadWasmBinary(wasmURL) { - const response = await loadAndCacheFile(wasmURL); - if (!response || typeof response === 'string') return null; - - try { - return await response.arrayBuffer(); - } catch (error) { - logger.warn('Failed to read WASM binary:', error); - return null; - } -} - -/** - * Loads and caches the WASM Factory (.mjs file) for ONNX Runtime. - * Creates a blob URL from cached content (when safe) to bridge Cache API with dynamic imports used in ORT. - * @param {string} libURL The URL of the WASM Factory to load. - * @returns {Promise} The blob URL (if enabled), original URL (if disabled), or null if loading failed. - */ -export async function loadWasmFactory(libURL) { - // We can't use Blob URLs in some environments (Service Workers, Chrome extensions) due to security restrictions on dynamic import() of blob URLs. - // In such cases, just return the original URL and don't bother caching since dynamic import() won't use the Cache API anyway. - // See https://github.com/huggingface/transformers.js/issues/1532. - if (apis.IS_SERVICE_WORKER_ENV || apis.IS_CHROME_AVAILABLE) { - return libURL; - } - - // Fetch from cache or network, then create blob URL - const response = await loadAndCacheFile(libURL); - if (!response || typeof response === 'string') return null; - - try { - let code = await response.text(); - - // Handle the case where we are importing the bundled version of the library in Deno (e.g., via CDN or local file), - // where we need to patch out Node.js detection in the factory. Without this, Deno (which exposes globalThis.process.versions.node) - // would enter the Node.js branch and try to use Node.js APIs (worker_threads, fs, etc.) that aren't used in the bundled web version. - // Only needed for the asyncify (single-threaded) variant loaded via blob URL. The module-level pthread auto-start code is unreachable since asyncify never spawns workers. - // See https://github.com/huggingface/transformers.js/pull/1546/ for more information. - // - // NOTE: This does not affect default usage via Deno (i.e., imported via npm: prefix), since we'll be using onnxruntime-node (Native) instead of onnxruntime-web (WASM). - code = code.replaceAll('globalThis.process?.versions?.node', 'false'); - const blob = new Blob([code], { type: 'text/javascript' }); - return URL.createObjectURL(blob); - } catch (error) { - logger.warn('Failed to read WASM factory:', error); - return null; - } -} diff --git a/packages/transformers/src/env.js b/packages/transformers/src/env.js index bea3ef0dc..c44b8785e 100644 --- a/packages/transformers/src/env.js +++ b/packages/transformers/src/env.js @@ -205,7 +205,7 @@ export const LogLevel = Object.freeze({ * Global variable given visible to users to control execution. This provides users a simple way to configure Transformers.js. * @typedef {Object} TransformersEnvironment * @property {string} version This version of Transformers.js. - * @property {{onnx: Partial & { setLogLevel?: (logLevel: number) => void }}} backends Expose environment variables of different backends, + * @property {Record void }>} backends Expose environment variables of different inference providers, * allowing users to set these variables if they want to. * @property {number} logLevel The logging level. Use LogLevel enum values. Defaults to LogLevel.ERROR. * @property {boolean} allowRemoteModels Whether to allow loading of remote files, defaults to `true`. @@ -240,10 +240,7 @@ export const env = { /////////////////// Backends settings /////////////////// // NOTE: These will be populated later by the backends themselves. - backends: { - // onnxruntime-web/onnxruntime-node - onnx: {}, - }, + backends: {}, /////////////////// Logging settings /////////////////// get logLevel() { @@ -252,8 +249,9 @@ export const env = { set logLevel(level) { logLevel = level; - // invoke hook to set ONNX Runtime log level when Transformers.js log level changes - env.backends.onnx?.setLogLevel?.(level); + for (const backend of Object.values(env.backends)) { + backend.setLogLevel?.(level); + } }, /////////////////// Model settings /////////////////// allowRemoteModels: true, diff --git a/packages/transformers/src/generation/controller.js b/packages/transformers/src/generation/controller.js new file mode 100644 index 000000000..2b49fa438 --- /dev/null +++ b/packages/transformers/src/generation/controller.js @@ -0,0 +1,357 @@ +import { Tensor } from '../utils/tensor.js'; +import { pick } from '../utils/core.js'; +import { logger } from '../utils/logger.js'; +import { GenerationConfig } from './configuration_utils.js'; +import { + LogitsProcessorList, + ForcedBOSTokenLogitsProcessor, + ForcedEOSTokenLogitsProcessor, + SuppressTokensLogitsProcessor, + SuppressTokensAtBeginLogitsProcessor, + NoRepeatNGramLogitsProcessor, + RepetitionPenaltyLogitsProcessor, + NoBadWordsLogitsProcessor, + MinLengthLogitsProcessor, + MinNewTokensLengthLogitsProcessor, + TemperatureLogitsWarper, + ClassifierFreeGuidanceLogitsProcessor, +} from './logits_process.js'; +import { EosTokenCriteria, MaxLengthCriteria, StoppingCriteriaList } from './stopping_criteria.js'; +import { LogitsSampler } from './logits_sampler.js'; + +/** @typedef {'greedy'|'multinomial'|'top-k'|'top-p'|'beam-search'} GenerationMode */ + +/** + * Resolve the generation configuration independently of an inference runtime. + * + * @param {Object} options + * @param {Object} options.modelConfig + * @param {Object|null} [options.modelGenerationConfig] + * @param {Object|null} [options.generationConfig] + * @param {Object|null} [options.kwargs] + * @param {typeof GenerationConfig} [options.configClass] + */ +export function prepareGenerationConfig({ + modelConfig, + modelGenerationConfig = null, + generationConfig = null, + kwargs = null, + configClass = GenerationConfig, +}) { + const config = { ...modelConfig }; + for (const key of ['decoder', 'generator', 'text_config']) { + if (key in config) Object.assign(config, config[key]); + } + + const result = new configClass(config); + Object.assign(result, modelGenerationConfig ?? {}); + if (generationConfig) Object.assign(result, generationConfig); + if (kwargs) Object.assign(result, pick(kwargs, Object.getOwnPropertyNames(result))); + return result; +} + +/** + * Build the complete ordered logits processor list. + * + * @param {GenerationConfig} generationConfig + * @param {number} inputLength + * @param {import('./logits_process.js').LogitsProcessorList|import('./logits_process.js').LogitsProcessor[]|null} [userProcessors] + */ +export function createLogitsProcessorList(generationConfig, inputLength, userProcessors = null) { + const processors = new LogitsProcessorList(); + + if (generationConfig.repetition_penalty !== null && generationConfig.repetition_penalty !== 1.0) { + processors.push(new RepetitionPenaltyLogitsProcessor(generationConfig.repetition_penalty)); + } + if (generationConfig.no_repeat_ngram_size !== null && generationConfig.no_repeat_ngram_size > 0) { + processors.push(new NoRepeatNGramLogitsProcessor(generationConfig.no_repeat_ngram_size)); + } + if (generationConfig.bad_words_ids !== null) { + processors.push(new NoBadWordsLogitsProcessor(generationConfig.bad_words_ids, generationConfig.eos_token_id)); + } + if ( + generationConfig.min_length !== null && + generationConfig.eos_token_id !== null && + generationConfig.min_length > 0 + ) { + processors.push(new MinLengthLogitsProcessor(generationConfig.min_length, generationConfig.eos_token_id)); + } + if ( + generationConfig.min_new_tokens !== null && + generationConfig.eos_token_id !== null && + generationConfig.min_new_tokens > 0 + ) { + processors.push( + new MinNewTokensLengthLogitsProcessor( + inputLength, + generationConfig.min_new_tokens, + generationConfig.eos_token_id, + ), + ); + } + if (generationConfig.forced_bos_token_id !== null) { + processors.push(new ForcedBOSTokenLogitsProcessor(generationConfig.forced_bos_token_id)); + } + if (generationConfig.forced_eos_token_id !== null) { + processors.push( + new ForcedEOSTokenLogitsProcessor(generationConfig.max_length, generationConfig.forced_eos_token_id), + ); + } + if (generationConfig.suppress_tokens !== null) { + processors.push(new SuppressTokensLogitsProcessor(generationConfig.suppress_tokens)); + } + if (generationConfig.begin_suppress_tokens !== null) { + const beginIndex = + inputLength > 1 || generationConfig.forced_bos_token_id === null ? inputLength : inputLength + 1; + processors.push(new SuppressTokensAtBeginLogitsProcessor(generationConfig.begin_suppress_tokens, beginIndex)); + } + if (generationConfig.guidance_scale !== null && generationConfig.guidance_scale > 1) { + processors.push(new ClassifierFreeGuidanceLogitsProcessor(generationConfig.guidance_scale)); + } + if (generationConfig.temperature === 0 && generationConfig.do_sample) { + logger.warn( + '`do_sample` changed to false because `temperature: 0` implies greedy sampling (always selecting the most likely token), which is incompatible with `do_sample: true`.', + ); + generationConfig.do_sample = false; + } + if (generationConfig.do_sample && generationConfig.temperature !== null && generationConfig.temperature !== 1.0) { + processors.push(new TemperatureLogitsWarper(generationConfig.temperature)); + } + if (userProcessors instanceof LogitsProcessorList) { + processors.extend(userProcessors.processors); + } else if (userProcessors !== null) { + processors.extend(userProcessors); + } + return processors; +} + +/** + * Build the complete stopping criteria list. + * + * @param {GenerationConfig} generationConfig + * @param {Object} modelConfig + * @param {import('./stopping_criteria.js').StoppingCriteria|import('./stopping_criteria.js').StoppingCriteria[]|StoppingCriteriaList|null} [userCriteria] + */ +export function createStoppingCriteriaList(generationConfig, modelConfig, userCriteria = null) { + const criteria = new StoppingCriteriaList(); + if (generationConfig.max_length !== null) { + criteria.push(new MaxLengthCriteria(generationConfig.max_length, modelConfig.max_position_embeddings ?? null)); + } + if (generationConfig.eos_token_id !== null) { + criteria.push(new EosTokenCriteria(generationConfig.eos_token_id)); + } + if (userCriteria) criteria.extend(userCriteria); + return criteria; +} + +/** + * Stateful, inference-runtime-neutral generation policy. + */ +export class GenerationController { + version = 1; + + /** + * @param {Object} options + * @param {Tensor} options.inputIds + * @param {GenerationConfig} options.generationConfig + * @param {LogitsProcessorList} options.logitsProcessor + * @param {StoppingCriteriaList} options.stoppingCriteria + * @param {import('./streamers.js').BaseStreamer|null} [options.streamer] + * @param {(outputs: Object, controller: GenerationController) => void} [options.collectOutputs] + */ + constructor({ + inputIds, + generationConfig, + logitsProcessor, + stoppingCriteria, + streamer = null, + collectOutputs = null, + }) { + if (!(inputIds instanceof Tensor) || inputIds.dims.length !== 2) { + throw new TypeError('GenerationController requires a rank-2 input IDs Tensor.'); + } + + this.generationConfig = generationConfig; + this.logitsProcessor = logitsProcessor; + this.stoppingCriteria = stoppingCriteria; + this.streamer = streamer; + this.collectOutputs = collectOutputs; + this.batchSize = inputIds.dims[0]; + this.inputLength = inputIds.dims[1]; + /** @type {bigint[][]} */ + this.sequences = inputIds.tolist(); + this.scores = new Array(this.batchSize).fill(0); + this.done = new Array(this.batchSize).fill(false); + this.terminal = false; + this.finalized = false; + this.aborted = false; + + if (generationConfig.max_new_tokens !== null) { + generationConfig.max_length = this.inputLength + generationConfig.max_new_tokens; + } + this.terminal = + generationConfig.max_new_tokens === 0 || + (generationConfig.max_length !== null && this.inputLength >= generationConfig.max_length); + if (this.terminal) this.done.fill(true); + this.sampler = LogitsSampler.getSampler(generationConfig); + if (streamer) streamer.put(this.sequences.map((tokens) => [...tokens])); + } + + get allDone() { + return this.terminal; + } + + get maxSequenceLength() { + return this.generationConfig.max_length; + } + + /** + * Process CPU-visible logits and commit the selected token. + * + * @param {Tensor|{logits: Tensor, outputs?: Object}} input + */ + async step(input) { + this.#assertActive(); + const logitsInput = input instanceof Tensor ? input : input.logits; + const outputs = input instanceof Tensor ? null : (input.outputs ?? null); + if (!(logitsInput instanceof Tensor)) throw new TypeError('Generation step logits must be a Tensor.'); + if (outputs && this.collectOutputs) this.collectOutputs(outputs, this); + + let logits; + if (logitsInput.dims.length === 3) { + logits = logitsInput.slice(null, -1, null).to('float32'); + } else if (logitsInput.dims.length === 2) { + logits = logitsInput.to('float32'); + } else { + throw new Error(`Generation logits must have rank 2 or 3, received rank ${logitsInput.dims.length}.`); + } + if (logits.dims[0] !== this.batchSize) { + throw new Error(`Generation logits batch size ${logits.dims[0]} does not match ${this.batchSize}.`); + } + + const processed = this.logitsProcessor(this.sequences, logits); + const tokenIds = new Uint32Array(this.batchSize); + const tokenScores = new Float64Array(this.batchSize); + for (let batchIndex = 0; batchIndex < this.batchSize; ++batchIndex) { + const sampled = await this.sampler(processed[batchIndex]); + const [tokenId, score] = sampled[0]; + tokenIds[batchIndex] = Number(tokenId); + tokenScores[batchIndex] = score; + } + return this.commit({ tokenIds, scores: tokenScores }); + } + + /** + * Commit tokens selected by an approved runtime generation plan. + * + * @param {{tokenIds: Uint32Array, processedScores?: Float32Array, scores?: Float64Array}} decision + */ + commit(decision) { + this.#assertActive(); + if (!(decision.tokenIds instanceof Uint32Array) || decision.tokenIds.length !== this.batchSize) { + throw new TypeError(`Generation decisions must contain ${this.batchSize} uint32 token IDs.`); + } + + const generatedInputIds = []; + for (let index = 0; index < this.batchSize; ++index) { + const tokenId = BigInt(decision.tokenIds[index]); + this.sequences[index].push(tokenId); + this.scores[index] += decision.scores?.[index] ?? 0; + generatedInputIds.push([tokenId]); + } + if (this.streamer) this.streamer.put(generatedInputIds); + + this.done = this.stoppingCriteria(this.sequences, decision.processedScores); + this.terminal = this.done.every(Boolean); + const nextTokenIds = new Tensor('int64', generatedInputIds.flat(), [this.batchSize, 1]); + return { + nextTokenIds, + generatedInputIds, + done: [...this.done], + allDone: this.terminal, + }; + } + + /** + * Compile the currently safe V1 GPU plan. More operations can be added as runtimes advertise them. + * + * @param {Object} capabilities + * @returns {Object|null} + */ + compileRuntimePlan(capabilities) { + if (!capabilities?.declarativePlans?.includes('argmax')) return null; + if (!capabilities?.planModes?.includes('greedy')) return null; + if (this.generationConfig.do_sample || this.generationConfig.num_beams > 1) return null; + if (this.logitsProcessor.processors.length !== 0) return null; + return { + version: 1, + processors: [], + sampler: { op: 'argmax' }, + maxNewTokens: Math.max(0, this.maxSequenceLength - this.inputLength), + pipelineDepth: capabilities.tokenPipeline?.defaultDepth, + }; + } + + /** + * @param {Object} [extra] + */ + finalize(extra = {}) { + if (this.aborted) throw new Error('Cannot finalize an aborted generation controller.'); + if (this.finalized) throw new Error('Generation controller has already been finalized.'); + if (!this.terminal) throw new Error('Cannot finalize generation before all sequences are done.'); + this.finalized = true; + if (this.streamer) this.streamer.end(); + + // V1 generation is synchronous across rows, so all sequences have equal length. + const sequences = new Tensor('int64', this.sequences.flat(), [this.sequences.length, this.sequences[0].length]); + return this.generationConfig.return_dict_in_generate ? { sequences, ...extra } : sequences; + } + + abort(_reason = undefined) { + if (this.finalized || this.aborted) return; + this.aborted = true; + this.terminal = true; + if (this.streamer) this.streamer.end(); + } + + #assertActive() { + if (this.aborted) throw new Error('Generation controller has been aborted.'); + if (this.finalized) throw new Error('Generation controller has already been finalized.'); + if (this.terminal) throw new Error('Generation controller is already complete.'); + } +} + +/** + * Create a controller from model and user generation options. + * + * @param {Object} model + * @param {Tensor} inputIds + * @param {Object} options + * @param {(outputs: Object, controller: GenerationController) => void} [collectOutputs] + */ +export function createGenerationController(model, inputIds, options, collectOutputs = null) { + const { + generation_config = null, + logits_processor = null, + stopping_criteria = null, + streamer = null, + ...kwargs + } = options; + const generationConfig = prepareGenerationConfig({ + modelConfig: model.config ?? {}, + modelGenerationConfig: model.generation_config ?? null, + generationConfig: generation_config, + kwargs, + }); + if (generationConfig.max_new_tokens !== null) { + generationConfig.max_length = inputIds.dims.at(-1) + generationConfig.max_new_tokens; + } + return new GenerationController({ + inputIds, + generationConfig, + logitsProcessor: createLogitsProcessorList(generationConfig, inputIds.dims.at(-1), logits_processor), + stoppingCriteria: createStoppingCriteriaList(generationConfig, model.config ?? {}, stopping_criteria), + streamer, + collectOutputs, + }); +} diff --git a/packages/transformers/src/generation/runtime.js b/packages/transformers/src/generation/runtime.js new file mode 100644 index 000000000..8269d8e1d --- /dev/null +++ b/packages/transformers/src/generation/runtime.js @@ -0,0 +1,261 @@ +import { Tensor } from '../utils/tensor.js'; +import { createGenerationController } from './controller.js'; + +/** + * @typedef {Object} GenerationCapabilitiesV1 + * @property {1} sessionVersion + * @property {number} maxBatchSize + * @property {string[]} cpuModes + * @property {string[]} planModes + * @property {boolean} cpuLogits + * @property {string[]} declarativePlans + * @property {{defaultDepth: number, maxDepth: number}} tokenPipeline + * @property {boolean} customJavaScriptStoppingCriteria + * @property {false} cacheReorder + * @property {false} cacheExpand + */ + +/** + * @typedef {Object} LogitsLeaseV1 + * @property {1} version + * @property {'float32'} dtype + * @property {[number, number]} shape + * @property {() => Promise} read + * @property {(plan: Object) => Promise<{tokenIds: Uint32Array, processedScores?: Float32Array}>} [select] + * @property {() => void} release + */ + +/** + * @typedef {Object} AutoregressiveSessionV1 + * @property {1} version + * @property {number} batchSize + * @property {number} maxSequenceLength + * @property {(inputs: Object) => Promise} prefill + * @property {(inputs: Object) => Promise} decode + * @property {(inputs: Object, plan: Object) => AsyncIterable<{tokenIds: Uint32Array, processedScores?: Float32Array}>} [generateWithPlan] + * @property {() => Promise} dispose + */ + +/** + * Install the Transformers.js-owned public generation method on a custom model. + * + * @param {Object|Function} model + */ +export function installGenerationRuntime(model) { + const implementation = /** @type {any} */ (model); + if (typeof implementation.createAutoregressiveSession === 'function') { + implementation.generate = (options) => generateWithAutoregressiveSession(implementation, options); + } + return model; +} + +/** + * Run decoder-only generation through a pull-based custom runtime session. + * + * @param {Object} model + * @param {Object} options + */ +export async function generateWithAutoregressiveSession(model, options) { + const { input_ids, attention_mask = null, signal = undefined } = options; + if (!(input_ids instanceof Tensor)) { + throw new TypeError('Custom autoregressive generation requires an `input_ids` Tensor.'); + } + if (model.config?.is_encoder_decoder) { + throw new Error('Autoregressive session protocol version 1 only supports decoder-only models.'); + } + + const controller = createGenerationController(model, input_ids, options); + if (controller.allDone) return controller.finalize(); + + const capabilities = model.generationCapabilities; + try { + validateCapabilities(capabilities, controller, attention_mask); + throwIfAborted(signal); + } catch (error) { + controller.abort(error); + throw error; + } + const plan = controller.compileRuntimePlan(capabilities); + if (!plan && !capabilities.cpuLogits) { + const error = new Error( + 'This generation request requires CPU-visible logits, but the runtime does not support them.', + ); + controller.abort(error); + throw error; + } + + /** @type {AutoregressiveSessionV1|null} */ + let session = null; + /** @type {LogitsLeaseV1|null} */ + let lease = null; + try { + session = await model.createAutoregressiveSession({ + batchSize: controller.batchSize, + maxSequenceLength: controller.maxSequenceLength, + signal, + }); + validateSession(session, controller); + + const prefillInputs = { + inputIds: tensorToTokenBatch(input_ids), + attentionMask: attention_mask ? tensorToAttentionMask(attention_mask) : undefined, + signal, + }; + if (plan && typeof session.generateWithPlan === 'function') { + const decisions = session.generateWithPlan(prefillInputs, plan)[Symbol.asyncIterator](); + try { + while (true) { + const item = await decisions.next(); + if (item.done) break; + throwIfAborted(signal); + if (controller.commit(item.value).allDone) break; + } + } finally { + await decisions.return?.(); + } + if (!controller.allDone) { + throw new Error('Autoregressive runtime ended its generation plan before generation completed.'); + } + return controller.finalize(); + } + + if (!capabilities.cpuLogits) { + throw new Error( + 'This generation request requires CPU-visible logits, but the runtime does not support them.', + ); + } + + lease = await session.prefill(prefillInputs); + while (!controller.allDone) { + throwIfAborted(signal); + const currentLease = lease; + lease = null; + validateLease(currentLease, controller.batchSize); + + let values; + try { + values = await currentLease.read(); + } finally { + currentLease.release(); + } + if (!(values instanceof Float32Array)) { + throw new TypeError('Logits lease `read()` must return a Float32Array.'); + } + + const step = await controller.step(new Tensor('float32', values, currentLease.shape)); + if (step.allDone) break; + lease = await session.decode({ + tokenIds: tensorToTokenBatch(step.nextTokenIds), + signal, + }); + } + return controller.finalize(); + } catch (error) { + controller.abort(error); + throw error; + } finally { + lease?.release(); + await session?.dispose(); + } +} + +function validateCapabilities(capabilities, controller, attentionMask) { + if (!capabilities || capabilities.sessionVersion !== 1) { + throw new Error('Custom generation models must declare generation capabilities with `sessionVersion: 1`.'); + } + if (controller.batchSize > capabilities.maxBatchSize) { + throw new Error( + `Runtime supports batch size ${capabilities.maxBatchSize}, but generation received ${controller.batchSize}.`, + ); + } + if (controller.generationConfig.num_return_sequences > 1) { + throw new Error('Autoregressive session protocol version 1 does not support multiple return sequences.'); + } + if (controller.generationConfig.num_beams > 1) { + throw new Error('Autoregressive session protocol version 1 does not support beam search.'); + } + if (controller.generationConfig.guidance_scale > 1) { + throw new Error('Autoregressive session protocol version 1 does not support classifier-free guidance.'); + } + if (controller.generationConfig.output_attentions || controller.generationConfig.output_hidden_states) { + throw new Error('Autoregressive session protocol version 1 does not support attentions or hidden states.'); + } + if (attentionMask && !isAllOnes(attentionMask)) { + throw new Error('Autoregressive session protocol version 1 only supports absent or all-ones attention masks.'); + } + + const mode = controller.generationConfig.do_sample ? 'multinomial' : 'greedy'; + if (!capabilities.cpuModes?.includes(mode) && !capabilities.planModes?.includes(mode)) { + throw new Error(`Runtime does not support ${mode} generation.`); + } + const tokenPipeline = capabilities.tokenPipeline; + if ( + tokenPipeline && + (!Number.isInteger(tokenPipeline.defaultDepth) || + !Number.isInteger(tokenPipeline.maxDepth) || + tokenPipeline.defaultDepth < 1 || + tokenPipeline.defaultDepth > tokenPipeline.maxDepth) + ) { + throw new Error('Runtime declared an invalid token pipeline depth.'); + } +} + +function validateSession(session, controller) { + if (!session || session.version !== 1) throw new Error('Runtime returned an unsupported autoregressive session.'); + if (session.batchSize !== controller.batchSize) { + throw new Error(`Runtime session batch size ${session.batchSize} does not match ${controller.batchSize}.`); + } + if (session.maxSequenceLength < controller.maxSequenceLength) { + throw new Error( + `Runtime session length ${session.maxSequenceLength} is less than required length ${controller.maxSequenceLength}.`, + ); + } + if (typeof session.dispose !== 'function') throw new Error('Autoregressive sessions must implement `dispose()`.'); +} + +function validateLease(lease, batchSize) { + if (!lease || lease.version !== 1 || lease.dtype !== 'float32') { + throw new Error('Runtime returned an unsupported logits lease.'); + } + if (!Array.isArray(lease.shape) || lease.shape.length !== 2 || lease.shape[0] !== batchSize) { + throw new Error('Logits lease must have shape [batch, vocabularySize].'); + } + if (typeof lease.read !== 'function' || typeof lease.release !== 'function') { + throw new Error('Logits lease must implement `read()` and `release()`.'); + } +} + +function tensorToTokenBatch(tensor) { + if (!(tensor instanceof Tensor) || tensor.dims.length !== 2) { + throw new TypeError('Token IDs must be a rank-2 Tensor.'); + } + const data = new Uint32Array(tensor.size); + for (let index = 0; index < tensor.size; ++index) { + const value = Number(tensor.data[index]); + if (!Number.isInteger(value) || value < 0 || value > 0xffffffff) { + throw new RangeError(`Token ID at index ${index} is outside the uint32 range.`); + } + data[index] = value; + } + return { data, shape: /** @type {[number, number]} */ ([tensor.dims[0], tensor.dims[1]]) }; +} + +function tensorToAttentionMask(tensor) { + if (!(tensor instanceof Tensor) || tensor.dims.length !== 2) { + throw new TypeError('Attention mask must be a rank-2 Tensor.'); + } + return { + data: Uint8Array.from(tensor.data, Number), + shape: /** @type {[number, number]} */ ([tensor.dims[0], tensor.dims[1]]), + }; +} + +function isAllOnes(tensor) { + return Array.from(tensor.data).every((value) => Number(value) === 1); +} + +function throwIfAborted(signal) { + if (!signal?.aborted) return; + if (typeof signal.throwIfAborted === 'function') signal.throwIfAborted(); + throw signal.reason ?? new Error('Generation aborted.'); +} diff --git a/packages/transformers/src/models/auto/modeling_auto.js b/packages/transformers/src/models/auto/modeling_auto.js index 030c51acd..60e1d4298 100644 --- a/packages/transformers/src/models/auto/modeling_auto.js +++ b/packages/transformers/src/models/auto/modeling_auto.js @@ -44,6 +44,7 @@ import { CUSTOM_ARCHITECTURES, MODEL_CLASS_TYPE_MAPPING, MODEL_MAPPINGS } from ' import * as ALL_MODEL_FILES from '../models.js'; import { logger } from '../../utils/logger.js'; +import { getModelId, isInferenceBackend } from '../../backends/inference.js'; /** * Base class of all AutoModels. Contains the `from_pretrained` function @@ -85,13 +86,35 @@ class PretrainedMixin { local_files_only = false, revision = 'main', model_file_name = null, - subfolder = 'onnx', + subfolder = null, device = null, dtype = null, use_external_data_format = null, session_options = {}, + signal = undefined, + artifactProvider = undefined, } = {}, ) { + if ( + isInferenceBackend(pretrained_model_name_or_path) && + typeof pretrained_model_name_or_path.constructSessions !== 'function' + ) { + return PreTrainedModel.from_pretrained(pretrained_model_name_or_path, { + progress_callback, + config, + cache_dir, + local_files_only, + revision, + model_file_name, + subfolder, + device, + dtype, + use_external_data_format, + session_options, + signal, + artifactProvider, + }); + } const options = { progress_callback, config, @@ -104,8 +127,10 @@ class PretrainedMixin { dtype, use_external_data_format, session_options, + signal, + artifactProvider, }; - options.config = await AutoConfig.from_pretrained(pretrained_model_name_or_path, options); + options.config = await AutoConfig.from_pretrained(getModelId(pretrained_model_name_or_path), options); if (!this.MODEL_CLASS_MAPPINGS) { throw new Error('`MODEL_CLASS_MAPPINGS` not implemented for this type of `AutoClass`: ' + this.name); diff --git a/packages/transformers/src/models/modeling_utils.js b/packages/transformers/src/models/modeling_utils.js index ec9f17487..e57dacd54 100644 --- a/packages/transformers/src/models/modeling_utils.js +++ b/packages/transformers/src/models/modeling_utils.js @@ -33,7 +33,7 @@ import { } from '../generation/logits_process.js'; import { GenerationConfig } from '../generation/configuration_utils.js'; import { EosTokenCriteria, MaxLengthCriteria, StoppingCriteriaList } from '../generation/stopping_criteria.js'; -import { LogitsSampler } from '../generation/logits_sampler.js'; +import { GenerationController } from '../generation/controller.js'; import { DefaultProgressCallback, pick } from '../utils/core.js'; import { ModelOutput } from './modeling_outputs.js'; import { logger } from '../utils/logger.js'; @@ -41,6 +41,8 @@ import { DynamicCache } from '../cache_utils.js'; import { get_model_files } from '../utils/model_registry/get_model_files.js'; import { get_file_metadata } from '../utils/model_registry/get_file_metadata.js'; import { MODEL_SESSION_CONFIG, MODEL_TYPES } from './session_config.js'; +import { getModelId, isInferenceBackend, loadInferenceModel } from '../backends/inference.js'; +import { OnnxInferenceProvider } from '../backends/default.js'; /** * Converts an array or Tensor of integers to an int64 Tensor. @@ -252,16 +254,41 @@ export class PreTrainedModel extends Callable { * The model class to instantiate is selected based on the `model_type` property of the config object * (either passed as an argument or loaded from `pretrained_model_name_or_path` if possible) * - * @param {string} pretrained_model_name_or_path The name or path of the pretrained model. Can be either: + * @param {string|import('../backends/inference.js').InferenceBackend} pretrained_model_name_or_path The model backend, name, or path. A string selects the ONNX backend. It can be: * - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co. * Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced under a * user or organization name, like `dbmdz/bert-base-german-cased`. * - A path to a *directory* containing model weights, e.g., `./my_model_directory/`. * @param {import('../utils/hub.js').PretrainedModelOptions} options Additional options for loading the model. * - * @returns {Promise} A new instance of the `PreTrainedModel` class. + * @returns {Promise} A loaded model. */ - static async from_pretrained( + static async from_pretrained(pretrained_model_name_or_path, options = {}) { + if (typeof pretrained_model_name_or_path === 'string') { + return OnnxInferenceProvider.from_modelId(pretrained_model_name_or_path).load({ + ...options, + modelClass: this, + }); + } + if (typeof pretrained_model_name_or_path?.constructSessions === 'function') { + return /** @type {any} */ (pretrained_model_name_or_path.load({ ...options, modelClass: this })); + } + if (isInferenceBackend(pretrained_model_name_or_path)) { + const modelId = getModelId(pretrained_model_name_or_path); + const resolvedOptions = { ...options }; + resolvedOptions.config = + resolvedOptions.config ?? (await AutoConfig.from_pretrained(modelId, resolvedOptions)); + // Custom models are duck-typed to the same runtime contract as PreTrainedModel. + const model = /** @type {any} */ (await loadInferenceModel(pretrained_model_name_or_path, resolvedOptions)); + if (typeof model.createAutoregressiveSession === 'function' && model.generation_config == null) { + model.generation_config = await getModelJSON(modelId, 'generation_config.json', false, resolvedOptions); + } + return model; + } + throw new TypeError('Unsupported pretrained model source.'); + } + + static async _from_pretrained( pretrained_model_name_or_path, { progress_callback = null, @@ -270,11 +297,14 @@ export class PreTrainedModel extends Callable { local_files_only = false, revision = 'main', model_file_name = null, - subfolder = 'onnx', + subfolder = null, device = null, dtype = null, use_external_data_format = null, session_options = {}, + signal = undefined, + artifactProvider = undefined, + inferenceProvider = undefined, } = {}, ) { const options = { @@ -289,6 +319,9 @@ export class PreTrainedModel extends Callable { dtype, use_external_data_format, session_options, + signal, + artifactProvider, + inferenceProvider, }; const modelName = MODEL_CLASS_TO_NAME_MAPPING.get(this); @@ -347,9 +380,7 @@ export class PreTrainedModel extends Callable { } const sessions = typeConfig.sessions(config, options, textOnly); - const promises = [ - constructSessions(pretrained_model_name_or_path, sessions, options, typeConfig.cache_sessions), - ]; + const promises = [constructSessions(sessions, options, typeConfig.cache_sessions)]; if (typeConfig.optional_configs) { promises.push(get_optional_configs(pretrained_model_name_or_path, typeConfig.optional_configs, options)); } @@ -926,114 +957,54 @@ export class PreTrainedModel extends Callable { // eos_token_ids = [eos_token_ids]; // } - const numInputs = model_inputs[model_input_name].dims.at(0); - - // TODO: - // done is a list of booleans to keep track of which inputs are done - // const done = new Array(numInputs).fill(false); - // For efficiency purposes, we remove completed rows from model_inputs - // when the beam is complete, and we keep track of the row index - // const rowIndexToBatchIndex = new Map(); - - const sampler = LogitsSampler.getSampler(generation_config); - - // TODO make > numInputs - const scores = new Array(numInputs).fill(0); - /** @type {bigint[][]} */ - const all_input_ids = input_ids.tolist(); - if (streamer) { - streamer.put(all_input_ids); - } - // const all_generated_input_ids = Array.from({ length: numInputs }, () => []); - - // NOTE: For now, we don't support spawning new beams - // TODO: when we do, we simply copy past key values and accumulate into single large tensor - - //////////////////////////////////////////////////// - // Generic search which handles 4 generation modes: - // - GenerationMode.GREEDY_SEARCH - // - GenerationMode.SAMPLE - // - GenerationMode.BEAM_SEARCH - // - GenerationMode.BEAM_SAMPLE - //////////////////////////////////////////////////// - let outputs; let attentions = {}; let return_dict_items = {}; - while (true) { - // prepare model inputs - model_inputs = this.prepare_inputs_for_generation(all_input_ids, model_inputs, generation_config); - outputs = await this.forward(model_inputs); - - if (generation_config.return_dict_in_generate) { + const controller = new GenerationController({ + inputIds: input_ids, + generationConfig: generation_config, + logitsProcessor: prepared_logits_processor, + stoppingCriteria: prepared_stopping_criteria, + streamer, + collectOutputs: (stepOutputs) => { + if (!generation_config.return_dict_in_generate) return; if (generation_config.output_attentions) { - // Get attentions if they are present - const token_attentions = getAttentions(outputs); + const token_attentions = getAttentions(stepOutputs); for (const key in token_attentions) { - if (!(key in attentions)) { - attentions[key] = []; - } - attentions[key].push(token_attentions[key]); + (attentions[key] ??= []).push(token_attentions[key]); } } else if (this._return_dict_in_generate_keys) { - Object.assign(return_dict_items, pick(outputs, this._return_dict_in_generate_keys)); + Object.assign(return_dict_items, pick(stepOutputs, this._return_dict_in_generate_keys)); } - } + }, + }); - // Logits are of the form [batch_size, out_seq_length, vocab_size] - // In most cases, this will be [batch_size, 1, vocab_size] - // So, we select the last token's logits: - // (equivalent to `logits = outputs.logits[:, -1, :]`) - // The `.to('float32')` is necessary for models with float16 logits, - // and is a no-op for float32 logits. - // TODO: Support float16 sampling in the sampler directly - const logits = outputs.logits.slice(null, -1, null).to('float32'); - - const next_tokens_scores = prepared_logits_processor(all_input_ids, logits); - - /** @type {[bigint][]} */ - const generated_input_ids = []; - // const new_kv_cache = [];// NOTE: Only used for beam search when concatenating new kv - // Loop over each batch - for (let batch_idx = 0; batch_idx < next_tokens_scores.dims.at(0); ++batch_idx) { - const logs = next_tokens_scores[batch_idx]; - - const sampledTokens = await sampler(logs); - for (const [newTokenId, logProb] of sampledTokens) { - const bigint = BigInt(newTokenId); - // TODO: If branching, use previous beam as a starting point - // update generated ids, model inputs, and length for next step - scores[batch_idx] += logProb; - all_input_ids[batch_idx].push(bigint); - generated_input_ids.push([bigint]); - - // TODO: Support beam search - break; - } - } - if (streamer) { - streamer.put(generated_input_ids); - } + if (controller.allDone) return controller.finalize(); - const stop = prepared_stopping_criteria(all_input_ids); - if (stop.every((x) => x)) { - break; + let outputs; + try { + while (!controller.allDone) { + // prepare model inputs + model_inputs = this.prepare_inputs_for_generation( + controller.sequences, + model_inputs, + generation_config, + ); + outputs = await this.forward(model_inputs); + const step = await controller.step({ logits: outputs.logits, outputs }); + if (step.allDone) break; + + model_inputs = this._update_model_kwargs_for_generation({ + generated_input_ids: step.generatedInputIds, + outputs, + model_inputs, + is_encoder_decoder, + }); } - - model_inputs = this._update_model_kwargs_for_generation({ - generated_input_ids, - outputs, - model_inputs, - is_encoder_decoder, - }); - } - - if (streamer) { - streamer.end(); + } catch (error) { + controller.abort(error); + throw error; } - // TODO: ensure all_input_ids is padded correctly... - const sequences = new Tensor('int64', all_input_ids.flat(), [all_input_ids.length, all_input_ids[0].length]); - // Update past key values from the final forward pass const past_key_values = getPastKeyValues(outputs, model_inputs.past_key_values); @@ -1052,17 +1023,16 @@ export class PreTrainedModel extends Callable { } if (generation_config.return_dict_in_generate) { - return { - sequences, + return controller.finalize({ past_key_values, ...attentions, ...return_dict_items, // TODO: // scores, // logits, - }; + }); } - return sequences; + return controller.finalize(); } /** diff --git a/packages/transformers/src/models/session.js b/packages/transformers/src/models/session.js index 04f6c4c28..c32d4eb49 100644 --- a/packages/transformers/src/models/session.js +++ b/packages/transformers/src/models/session.js @@ -1,285 +1,16 @@ -import { - createInferenceSession, - deviceToExecutionProviders, - isONNXProxy, - isONNXTensor, - runInferenceSession, -} from '../backends/onnx.js'; -import { getCacheNames } from '../configs.js'; -import { DATA_TYPES, DEFAULT_DTYPE_SUFFIX_MAPPING, isWebGpuFp16Supported, selectDtype } from '../utils/dtypes.js'; -import { selectDevice } from '../utils/devices.js'; -import { apis } from '../env.js'; -import { getCoreModelFile, getModelDataFiles } from '../utils/model-loader.js'; -import { Tensor } from '../utils/tensor.js'; -import { logger } from '../utils/logger.js'; - -/** - * Constructs an InferenceSession using a model file located at the specified path. - * @param {string} pretrained_model_name_or_path The path to the directory containing the model file. - * @param {string} fileName The name of the model file. - * @param {import('../utils/hub.js').PretrainedModelOptions} options Additional options for loading the model. - * @param {boolean} [cache_config=false] Whether to compute cache shapes for GPU-pinned outputs. - * @param {string} [session_name] The name of the session (used to determine cache shapes). - * @returns {Promise<{buffer_or_path: Uint8Array|string, session_options: Object, session_config: Object}>} A Promise that resolves to the data needed to create an InferenceSession object. - * @private - */ -async function getSession( - pretrained_model_name_or_path, - fileName, - options, - cache_config = false, - session_name = undefined, -) { - let custom_config = options.config?.['transformers.js_config'] ?? {}; - - // If the device is not specified, we use the default (supported) execution providers. - const selectedDevice = /** @type {import("../utils/devices.js").DeviceType} */ ( - selectDevice(options.device ?? custom_config.device, fileName, { - warn: (msg) => logger.info(msg), - }) - ); - - const executionProviders = deviceToExecutionProviders(selectedDevice); - - // Update custom config with the selected device's config, if it exists - const device_config = custom_config.device_config ?? {}; - if (device_config.hasOwnProperty(selectedDevice)) { - custom_config = { - ...custom_config, - ...device_config[selectedDevice], - }; - } - - // If options.dtype is specified, we use it to choose the suffix for the model file. - // Otherwise, we use the default dtype for the device. - const selectedDtype = /** @type {import("../utils/dtypes.js").DataType} */ ( - selectDtype(options.dtype ?? custom_config.dtype, fileName, selectedDevice, { - configDtype: custom_config.dtype, - warn: (msg) => logger.info(msg), - }) - ); - - if (!DEFAULT_DTYPE_SUFFIX_MAPPING.hasOwnProperty(selectedDtype)) { - throw new Error(`Invalid dtype: ${selectedDtype}. Should be one of: ${Object.keys(DATA_TYPES).join(', ')}`); - } else if ( - selectedDevice === 'webgpu' && - // NOTE: Currently, we assume that the Native WebGPU EP always supports fp16. In future, we will add a check for this. - !apis.IS_NODE_ENV && - selectedDtype === DATA_TYPES.fp16 && - !(await isWebGpuFp16Supported()) - ) { - throw new Error(`The device (${selectedDevice}) does not support fp16.`); - } - - // Construct the model file suffix - const suffix = DEFAULT_DTYPE_SUFFIX_MAPPING[selectedDtype]; - - const session_options = { ...options.session_options }; - - // Overwrite `executionProviders` if not specified - session_options.executionProviders ??= executionProviders; - - // Overwrite `freeDimensionOverrides` if specified in config and not set in session options - const free_dimension_overrides = custom_config.free_dimension_overrides; - if (free_dimension_overrides) { - session_options.freeDimensionOverrides ??= free_dimension_overrides; - } else if (selectedDevice.startsWith('webnn') && !session_options.freeDimensionOverrides) { - logger.warn( - `WebNN does not currently support dynamic shapes and requires 'free_dimension_overrides' to be set in config.json, preferably as a field within config["transformers.js_config"]["device_config"]["${selectedDevice}"]. ` + - `When 'free_dimension_overrides' is not set, you may experience significant performance degradation.`, - ); - } - - const bufferOrPathPromise = getCoreModelFile(pretrained_model_name_or_path, fileName, options, suffix); - - // Handle onnx external data files - const use_external_data_format = options.use_external_data_format ?? custom_config.use_external_data_format; - const externalData = await getModelDataFiles( - pretrained_model_name_or_path, - fileName, - suffix, - options, - use_external_data_format, - session_options, - ); - - if (externalData.length > 0 && (!apis.IS_NODE_ENV || externalData.some((data) => typeof data !== 'string'))) { - session_options.externalData = externalData; - } - - if (cache_config && selectedDevice === 'webgpu') { - const names = getCacheNames(options.config, { - prefix: 'present', - session_name, - }); - if (names.size > 0 && !isONNXProxy()) { - // Only set preferredOutputLocation if names are present and we aren't proxying ONNX - /** @type {Record} */ - const preferredOutputLocation = {}; - for (const key of names) { - preferredOutputLocation[key] = 'gpu-buffer'; - } - session_options.preferredOutputLocation = preferredOutputLocation; - } - } - - const buffer_or_path = await bufferOrPathPromise; - const session_config = { - dtype: selectedDtype, - device: selectedDevice, - }; - return { buffer_or_path, session_options, session_config }; -} - -/** - * Helper function to create multiple InferenceSession objects. - * - * @param {string} pretrained_model_name_or_path The path to the directory containing the model file. - * @param {Record} names The names of the model files to load. - * @param {import('../utils/hub.js').PretrainedModelOptions} options Additional options for loading the model. - * @param {Record} [cache_sessions] A map from session name to `true`, indicating which - * sessions should have GPU-pinned KV cache outputs. - * @returns {Promise>} A Promise that resolves to a dictionary of InferenceSession objects. - * @private - */ -export async function constructSessions(pretrained_model_name_or_path, names, options, cache_sessions = undefined) { - return Object.fromEntries( - await Promise.all( - Object.keys(names).map(async (name) => { - const cache_config = cache_sessions?.[name] ?? false; - const { buffer_or_path, session_options, session_config } = await getSession( - pretrained_model_name_or_path, - names[name], - options, - cache_config, - name, - ); - const session = await createInferenceSession(buffer_or_path, session_options, session_config); - return [name, session]; - }), - ), - ); -} - /** - * Replaces ONNX Tensor objects with custom Tensor objects to support additional functions. - * @param {Object} obj The object to replace tensor objects in. - * @returns {Object} The object with tensor objects replaced by custom Tensor objects. - * @private + * Construct the normalized sessions used by built-in Transformers.js models. */ -function replaceTensors(obj) { - for (let prop in obj) { - if (isONNXTensor(obj[prop])) { - obj[prop] = new Tensor(obj[prop]); - } else if (typeof obj[prop] === 'object') { - replaceTensors(obj[prop]); - } +export async function constructSessions(names, options, cache_sessions = undefined) { + if (!options.inferenceProvider?.constructSessions) { + throw new Error('The selected inference provider does not support built-in model sessions.'); } - return obj; + return options.inferenceProvider.constructSessions(names, options, cache_sessions); } /** - * Executes an InferenceSession using the specified inputs. - * NOTE: `inputs` must contain at least the input names of the model. - * - If additional inputs are passed, they will be ignored. - * - If inputs are missing, an error will be thrown. - * - * @param {Object} session The InferenceSession object to run. - * @param {Object} inputs An object that maps input names to input tensors. - * @returns {Promise} A Promise that resolves to an object that maps output names to output tensors. - * @private + * Run a normalized built-in model session. */ export async function sessionRun(session, inputs) { - const checkedInputs = validateInputs(session, inputs); - try { - // pass the original ort tensor - const ortFeed = Object.fromEntries( - Object.entries(checkedInputs).map(([k, v]) => { - const tensor = /** @type {any} */ (v.ort_tensor); - if (apis.IS_NODE_ENV) { - // In recent versions of Node.js, which support Float16Array, we need to convert - // the Float16Array to Uint16Array for ONNX Runtime to accept it. - if (typeof Float16Array !== 'undefined' && tensor.cpuData instanceof Float16Array) { - tensor.cpuData = new Uint16Array(tensor.cpuData.buffer); // reinterpret as Uint16Array - } - } - return [k, tensor]; - }), - ); - - const output = await runInferenceSession(session, ortFeed); - return replaceTensors(output); - } catch (e) { - // Error messages can be long (nested) and uninformative. For this reason, - // we apply minor formatting to show the most important information - const formatted = Object.fromEntries( - Object.entries(checkedInputs).map(([k, tensor]) => { - // Extract these properties from the underlying ORT tensor - const unpacked = { - type: tensor.type, - dims: tensor.dims, - location: tensor.location, - }; - if (unpacked.location !== 'gpu-buffer') { - // Only return the data if it's not a GPU buffer - unpacked.data = tensor.data; - } - return [k, unpacked]; - }), - ); - - // This usually occurs when the inputs are of the wrong type. - logger.error(`An error occurred during model execution: "${e}".`); - logger.error('Inputs given to model:', formatted); - throw e; - } -} - -/** - * Validate model inputs - * @param {Object} session The InferenceSession object that will be run. - * @param {Object} inputs The inputs to check. - * @returns {Record} The checked inputs. - * @throws {Error} If any inputs are missing. - * @private - */ -function validateInputs(session, inputs) { - /** - * NOTE: Create either a shallow or deep copy based on `onnx.wasm.proxy` - * @type {Record} - */ - const checkedInputs = Object.create(null); - const missingInputs = []; - for (const inputName of session.inputNames) { - const tensor = inputs[inputName]; - // Rare case where one of the model's input names corresponds to a built-in - // object name (e.g., toString), which would cause a simple (!tensor) check to fail, - // because it's not undefined but a function. - if (!(tensor instanceof Tensor)) { - missingInputs.push(inputName); - continue; - } - // NOTE: When `env.wasm.proxy is true` the tensor is moved across the Worker - // boundary, transferring ownership to the worker and invalidating the tensor. - // So, in this case, we simply sacrifice a clone for it. - checkedInputs[inputName] = isONNXProxy() ? tensor.clone() : tensor; - } - if (missingInputs.length > 0) { - throw new Error( - `An error occurred during model execution: "Missing the following inputs: ${missingInputs.join(', ')}.`, - ); - } - - const numInputsProvided = Object.keys(inputs).length; - const numInputsNeeded = session.inputNames.length; - if (numInputsProvided > numInputsNeeded) { - // No missing inputs, but too many inputs were provided. - // Warn the user and ignore the extra inputs. - let ignored = Object.keys(inputs).filter((inputName) => !session.inputNames.includes(inputName)); - logger.warn( - `WARNING: Too many inputs were provided (${numInputsProvided} > ${numInputsNeeded}). The following inputs will be ignored: "${ignored.join(', ')}".`, - ); - } - - return checkedInputs; + return session.run(inputs); } diff --git a/packages/transformers/src/models/voxtral_realtime/modeling_voxtral_realtime.js b/packages/transformers/src/models/voxtral_realtime/modeling_voxtral_realtime.js index 75f242235..dbf0d05f6 100644 --- a/packages/transformers/src/models/voxtral_realtime/modeling_voxtral_realtime.js +++ b/packages/transformers/src/models/voxtral_realtime/modeling_voxtral_realtime.js @@ -38,7 +38,7 @@ function createEncoderState(model, input_features) { const enc_kv_cache = new DynamicCache(); const enc_names = getCacheNames(audio_config); const enc_symbols = { batch_size: 1 }; - /** @type {import('onnxruntime-common').Tensor.Type} */ + /** @type {import('../../utils/tensor.js').DataType} */ let padding_type = 'float32'; for (const meta of encoder_session.inputMetadata) { if (meta.name === 'past_padding_cache') { diff --git a/packages/transformers/src/ops/registry.js b/packages/transformers/src/ops/registry.js index 42508267d..2d06cdddd 100644 --- a/packages/transformers/src/ops/registry.js +++ b/packages/transformers/src/ops/registry.js @@ -1,176 +1,45 @@ -import { createInferenceSession, runInferenceSession, isONNXProxy } from '../backends/onnx.js'; -import { Tensor } from '../utils/tensor.js'; +let implementation = null; /** - * Asynchronously creates a wrapper function for running an ONNX inference session. - * - * @param {number[]} session_bytes The session data in bytes. - * @param {import('onnxruntime-common').InferenceSession.SessionOptions} session_options The options for the ONNX session. - * @template {string | [string] | string[]} T - * @param {T} names The name(s) of the output tensor(s). - * - * @returns {Promise): Promise>} - * The wrapper function for running the ONNX inference session. + * Runtime-neutral tensor operation registry. Inference providers install optimized implementations. */ -const wrap = async (session_bytes, session_options, names) => { - const session = await createInferenceSession(new Uint8Array(session_bytes), session_options); - - return /** @type {any} */ ( - async (/** @type {Record} */ inputs) => { - const proxied = isONNXProxy(); - const ortFeed = Object.fromEntries( - Object.entries(inputs).map(([k, v]) => [k, (proxied ? v.clone() : v).ort_tensor]), - ); - const outputs = await runInferenceSession(session, ortFeed); - if (Array.isArray(names)) { - return names.map((n) => new Tensor(outputs[n])); - } else { - return new Tensor(outputs[/** @type {string} */ (names)]); - } - } - ); -}; - -// In-memory registry of initialized ONNX operators export class TensorOpRegistry { - static session_options = { - // TODO: Allow for multiple execution providers - // executionProviders: ['webgpu'], - }; + static register(provider) { + implementation = provider; + } static get nearest_interpolate_4d() { - if (!this._nearest_interpolate_4d) { - this._nearest_interpolate_4d = wrap( - [ - 8, 10, 18, 0, 58, 129, 1, 10, 41, 10, 1, 120, 10, 0, 10, 0, 10, 1, 115, 18, 1, 121, 34, 6, 82, 101, - 115, 105, 122, 101, 42, 18, 10, 4, 109, 111, 100, 101, 34, 7, 110, 101, 97, 114, 101, 115, 116, 160, - 1, 3, 18, 1, 114, 90, 31, 10, 1, 120, 18, 26, 10, 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, - 99, 10, 3, 18, 1, 104, 10, 3, 18, 1, 119, 90, 15, 10, 1, 115, 18, 10, 10, 8, 8, 7, 18, 4, 10, 2, 8, - 4, 98, 31, 10, 1, 121, 18, 26, 10, 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, 99, 10, 3, 18, - 1, 104, 10, 3, 18, 1, 119, 66, 2, 16, 21, - ], - this.session_options, - 'y', - ); - } - return this._nearest_interpolate_4d; + return getOperation('nearest_interpolate_4d'); } static get bilinear_interpolate_4d() { - if (!this._bilinear_interpolate_4d) { - this._bilinear_interpolate_4d = wrap( - [ - 8, 9, 18, 0, 58, 128, 1, 10, 40, 10, 1, 120, 10, 0, 10, 0, 10, 1, 115, 18, 1, 121, 34, 6, 82, 101, - 115, 105, 122, 101, 42, 17, 10, 4, 109, 111, 100, 101, 34, 6, 108, 105, 110, 101, 97, 114, 160, 1, - 3, 18, 1, 114, 90, 31, 10, 1, 120, 18, 26, 10, 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, 99, - 10, 3, 18, 1, 104, 10, 3, 18, 1, 119, 90, 15, 10, 1, 115, 18, 10, 10, 8, 8, 7, 18, 4, 10, 2, 8, 4, - 98, 31, 10, 1, 121, 18, 26, 10, 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, 99, 10, 3, 18, 1, - 104, 10, 3, 18, 1, 119, 66, 2, 16, 20, - ], - this.session_options, - 'y', - ); - } - return this._bilinear_interpolate_4d; + return getOperation('bilinear_interpolate_4d'); } - static get bicubic_interpolate_4d() { - if (!this._bicubic_interpolate_4d) { - this._bicubic_interpolate_4d = wrap( - [ - 8, 9, 18, 0, 58, 127, 10, 39, 10, 1, 120, 10, 0, 10, 0, 10, 1, 115, 18, 1, 121, 34, 6, 82, 101, 115, - 105, 122, 101, 42, 16, 10, 4, 109, 111, 100, 101, 34, 5, 99, 117, 98, 105, 99, 160, 1, 3, 18, 1, - 114, 90, 31, 10, 1, 120, 18, 26, 10, 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, 99, 10, 3, - 18, 1, 104, 10, 3, 18, 1, 119, 90, 15, 10, 1, 115, 18, 10, 10, 8, 8, 7, 18, 4, 10, 2, 8, 4, 98, 31, - 10, 1, 121, 18, 26, 10, 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, 99, 10, 3, 18, 1, 104, 10, - 3, 18, 1, 119, 66, 2, 16, 20, - ], - this.session_options, - 'y', - ); - } - return this._bicubic_interpolate_4d; + return getOperation('bicubic_interpolate_4d'); } - static get matmul() { - if (!this._matmul) { - this._matmul = wrap( - [ - 8, 9, 18, 0, 58, 55, 10, 17, 10, 1, 97, 10, 1, 98, 18, 1, 99, 34, 6, 77, 97, 116, 77, 117, 108, 18, - 1, 114, 90, 9, 10, 1, 97, 18, 4, 10, 2, 8, 1, 90, 9, 10, 1, 98, 18, 4, 10, 2, 8, 1, 98, 9, 10, 1, - 99, 18, 4, 10, 2, 8, 1, 66, 2, 16, 20, - ], - this.session_options, - 'c', - ); - } - return this._matmul; + return getOperation('matmul'); } - static get stft() { - if (!this._stft) { - this._stft = wrap( - [ - 8, 7, 18, 0, 58, 148, 1, 10, 38, 10, 1, 115, 10, 1, 106, 10, 1, 119, 10, 1, 108, 18, 1, 111, 34, 4, - 83, 84, 70, 84, 42, 15, 10, 8, 111, 110, 101, 115, 105, 100, 101, 100, 24, 1, 160, 1, 2, 18, 1, 115, - 90, 26, 10, 1, 115, 18, 21, 10, 19, 8, 1, 18, 15, 10, 3, 18, 1, 98, 10, 3, 18, 1, 115, 10, 3, 18, 1, - 99, 90, 11, 10, 1, 106, 18, 6, 10, 4, 8, 7, 18, 0, 90, 16, 10, 1, 119, 18, 11, 10, 9, 8, 1, 18, 5, - 10, 3, 18, 1, 119, 90, 11, 10, 1, 108, 18, 6, 10, 4, 8, 7, 18, 0, 98, 31, 10, 1, 111, 18, 26, 10, - 24, 8, 1, 18, 20, 10, 3, 18, 1, 98, 10, 3, 18, 1, 102, 10, 3, 18, 1, 100, 10, 3, 18, 1, 99, 66, 2, - 16, 17, - ], - this.session_options, - 'o', - ); - } - return this._stft; + return getOperation('stft'); } - static get rfft() { - if (!this._rfft) { - this._rfft = wrap( - [ - 8, 9, 18, 0, 58, 97, 10, 33, 10, 1, 120, 10, 0, 10, 1, 97, 18, 1, 121, 34, 3, 68, 70, 84, 42, 15, - 10, 8, 111, 110, 101, 115, 105, 100, 101, 100, 24, 1, 160, 1, 2, 18, 1, 100, 90, 21, 10, 1, 120, 18, - 16, 10, 14, 8, 1, 18, 10, 10, 3, 18, 1, 115, 10, 3, 18, 1, 99, 90, 11, 10, 1, 97, 18, 6, 10, 4, 8, - 7, 18, 0, 98, 21, 10, 1, 121, 18, 16, 10, 14, 8, 1, 18, 10, 10, 3, 18, 1, 115, 10, 3, 18, 1, 99, 66, - 2, 16, 20, - ], - this.session_options, - 'y', - ); - } - return this._rfft; + return getOperation('rfft'); } - static get top_k() { - if (!this._top_k) { - this._top_k = wrap( - [ - 8, 10, 18, 0, 58, 73, 10, 18, 10, 1, 120, 10, 1, 107, 18, 1, 118, 18, 1, 105, 34, 4, 84, 111, 112, - 75, 18, 1, 116, 90, 9, 10, 1, 120, 18, 4, 10, 2, 8, 1, 90, 15, 10, 1, 107, 18, 10, 10, 8, 8, 7, 18, - 4, 10, 2, 8, 1, 98, 9, 10, 1, 118, 18, 4, 10, 2, 8, 1, 98, 9, 10, 1, 105, 18, 4, 10, 2, 8, 7, 66, 2, - 16, 21, - ], - this.session_options, - [/* Values */ 'v', /* Indices */ 'i'], - ); - } - return this._top_k; + return getOperation('top_k'); } - static get slice() { - if (!this._slice) { - this._slice = wrap( - [ - 8, 7, 18, 0, 58, 96, 10, 25, 10, 1, 120, 10, 1, 115, 10, 1, 101, 10, 1, 97, 10, 1, 116, 18, 1, 121, - 34, 5, 83, 108, 105, 99, 101, 18, 1, 114, 90, 9, 10, 1, 120, 18, 4, 10, 2, 8, 1, 90, 9, 10, 1, 115, - 18, 4, 10, 2, 8, 7, 90, 9, 10, 1, 101, 18, 4, 10, 2, 8, 7, 90, 9, 10, 1, 97, 18, 4, 10, 2, 8, 7, 90, - 9, 10, 1, 116, 18, 4, 10, 2, 8, 7, 98, 9, 10, 1, 121, 18, 4, 10, 2, 8, 1, 66, 2, 16, 13, - ], - this.session_options, - 'y', - ); - } - return this._slice; + return getOperation('slice'); + } +} + +async function getOperation(name) { + if (!implementation) { + await import('../backends/default.js'); + } + if (!implementation) { + throw new Error(`Tensor operation "${name}" requires an installed inference provider.`); } + return implementation[name]; } diff --git a/packages/transformers/src/pipelines.js b/packages/transformers/src/pipelines.js index 98c10c692..99bdf180e 100644 --- a/packages/transformers/src/pipelines.js +++ b/packages/transformers/src/pipelines.js @@ -51,6 +51,8 @@ import { } from './pipelines/index.js'; import { get_pipeline_files } from './utils/model_registry/get_pipeline_files.js'; import { get_file_metadata } from './utils/model_registry/get_file_metadata.js'; +import { getModelId, isInferenceBackend, loadInferenceModel } from './backends/inference.js'; +import { getModelJSON } from './utils/hub.js'; /** * @typedef {keyof typeof SUPPORTED_TASKS} TaskType @@ -89,7 +91,7 @@ import { get_file_metadata } from './utils/model_registry/get_file_metadata.js'; * - `"zero-shot-audio-classification"`: will return a `ZeroShotAudioClassificationPipeline`. * - `"zero-shot-image-classification"`: will return a `ZeroShotImageClassificationPipeline`. * - `"zero-shot-object-detection"`: will return a `ZeroShotObjectDetectionPipeline`. - * @param {string} [model=null] The name of the pre-trained model to use. If not specified, the default model for the task will be used. + * @param {string|import('./backends/inference.js').InferenceBackend} [model=null] The model ID or custom inference backend to use. If not specified, the default model for the task will be used. * @param {import('./utils/hub.js').PretrainedModelOptions} [options] Optional parameters for the pipeline. * @returns {Promise} A Pipeline object for the specified task. * @throws {Error} If an unsupported pipeline is requested. @@ -105,10 +107,12 @@ export async function pipeline( revision = 'main', device = null, dtype = null, - subfolder = 'onnx', + subfolder = null, use_external_data_format = null, model_file_name = null, session_options = {}, + signal = undefined, + artifactProvider = undefined, } = {}, ) { // Apply aliases @@ -130,17 +134,22 @@ export async function pipeline( } } + const customBackend = isInferenceBackend(model) && typeof model.constructSessions !== 'function'; + const modelId = getModelId(model); + // Determine which files the model needs - const expected_files = await get_pipeline_files(task, model, { + const expected_files = await get_pipeline_files(task, modelId, { device, dtype, + config, + include_model: !customBackend, }); /** @type {import('./utils/core.js').FilesLoadingMap} */ let files_loading = {}; if (progress_callback) { /** @type {Array<{exists: boolean, size?: number, contentType?: string, fromCache?: boolean}>} */ - const metadata = await Promise.all(expected_files.map(async (file) => get_file_metadata(model, file))); + const metadata = await Promise.all(expected_files.map(async (file) => get_file_metadata(modelId, file))); metadata.forEach((m, i) => { if (m.exists) { files_loading[expected_files[i]] = { @@ -165,6 +174,9 @@ export async function pipeline( use_external_data_format, model_file_name, session_options, + generation_config: null, + signal, + artifactProvider, }; // Determine which components to load based on the expected files @@ -174,8 +186,22 @@ export async function pipeline( // Resolve the correct model class (needs config when multiple candidates exist) const modelClasses = pipelineInfo.model; let modelPromise; - if (Array.isArray(modelClasses)) { - const resolvedConfig = config ?? (await AutoConfig.from_pretrained(model, pretrainedOptions)); + if (customBackend) { + pretrainedOptions.config = config ?? (await AutoConfig.from_pretrained(modelId, pretrainedOptions)); + if (task === 'text-generation') { + pretrainedOptions.generation_config = await getModelJSON( + modelId, + 'generation_config.json', + false, + pretrainedOptions, + ); + } + modelPromise = loadInferenceModel(/** @type {import('./backends/inference.js').InferenceBackend} */ (model), { + ...pretrainedOptions, + task, + }); + } else if (Array.isArray(modelClasses)) { + const resolvedConfig = config ?? (await AutoConfig.from_pretrained(modelId, pretrainedOptions)); const { model_type } = resolvedConfig; const matchedClass = modelClasses.find((cls) => cls.supports(model_type)); if (!matchedClass) { @@ -184,15 +210,15 @@ export async function pipeline( `None of the candidate model classes support this type.`, ); } - modelPromise = matchedClass.from_pretrained(model, { ...pretrainedOptions, config: resolvedConfig }); + modelPromise = matchedClass.from_pretrained(modelId, { ...pretrainedOptions, config: resolvedConfig }); } else { - modelPromise = modelClasses.from_pretrained(model, pretrainedOptions); + modelPromise = modelClasses.from_pretrained(modelId, pretrainedOptions); } // Load all components in parallel const [tokenizer, processor, model_loaded] = await Promise.all([ - hasTokenizer ? AutoTokenizer.from_pretrained(model, pretrainedOptions) : null, - hasProcessor ? AutoProcessor.from_pretrained(model, pretrainedOptions) : null, + hasTokenizer ? AutoTokenizer.from_pretrained(modelId, pretrainedOptions) : null, + hasProcessor ? AutoProcessor.from_pretrained(modelId, pretrainedOptions) : null, modelPromise, ]); @@ -203,7 +229,7 @@ export async function pipeline( dispatchCallback(progress_callback, { status: 'ready', task: task, - model: model, + model: modelId, }); const pipelineClass = pipelineInfo.pipeline; diff --git a/packages/transformers/src/transformers.js b/packages/transformers/src/transformers.js index ef01569ef..587d1e47a 100644 --- a/packages/transformers/src/transformers.js +++ b/packages/transformers/src/transformers.js @@ -45,6 +45,7 @@ export { PretrainedConfig, AutoConfig } from './configs.js'; export * from './generation/streamers.js'; export * from './generation/stopping_criteria.js'; export * from './generation/logits_process.js'; +export { GenerationController, createGenerationController } from './generation/controller.js'; export { load_audio, read_audio, RawAudio } from './utils/audio.js'; export { load_image, RawImage } from './utils/image.js'; @@ -58,6 +59,10 @@ export { DynamicCache } from './cache_utils.js'; // Cache and file management export { ModelRegistry } from './utils/model_registry/ModelRegistry.js'; +// Inference backends +export { getModelId, isInferenceBackend } from './backends/inference.js'; +export { OnnxInferenceProvider } from './backends/default.js'; + // Expose common types used across the library for developers to access /** * @typedef {import('./utils/hub.js').PretrainedModelOptions} PretrainedModelOptions @@ -68,4 +73,10 @@ export { ModelRegistry } from './utils/model_registry/ModelRegistry.js'; * @typedef {import('./utils/devices.js').DeviceType} DeviceType * @typedef {import('./utils/core.js').ProgressCallback} ProgressCallback * @typedef {import('./utils/core.js').ProgressInfo} ProgressInfo + * @typedef {import('./generation/runtime.js').GenerationCapabilitiesV1} GenerationCapabilitiesV1 + * @typedef {import('./generation/runtime.js').AutoregressiveSessionV1} AutoregressiveSessionV1 + * @typedef {import('./generation/runtime.js').LogitsLeaseV1} LogitsLeaseV1 + * @typedef {import('./backends/artifacts.js').InferenceArtifactProvider} InferenceArtifactProvider + * @typedef {import('./backends/artifacts.js').RandomAccessByteSource} RandomAccessByteSource + * @typedef {import('./backends/artifacts.js').ArtifactProgressEvent} ArtifactProgressEvent */ diff --git a/packages/transformers/src/utils/dtypes.js b/packages/transformers/src/utils/dtypes.js index ff026f747..e4cf2ca36 100644 --- a/packages/transformers/src/utils/dtypes.js +++ b/packages/transformers/src/utils/dtypes.js @@ -1,39 +1,5 @@ -/// - -import { apis } from '../env.js'; - -import { DEVICE_TYPES } from './devices.js'; - -// TODO: Use the adapter from `env.backends.onnx.webgpu.adapter` to check for `shader-f16` support, -// when available in https://github.com/microsoft/onnxruntime/pull/19940. -// For more information, see https://github.com/microsoft/onnxruntime/pull/19857#issuecomment-1999984753 - -/** - * Checks if WebGPU fp16 support is available in the current environment. - */ -export const isWebGpuFp16Supported = (function () { - /** @type {boolean} */ - let cachedResult; - - return async function () { - if (cachedResult === undefined) { - if (!apis.IS_WEBGPU_AVAILABLE) { - cachedResult = false; - } else { - try { - const adapter = await navigator.gpu.requestAdapter(); - cachedResult = adapter.features.has('shader-f16'); - } catch (e) { - cachedResult = false; - } - } - } - return cachedResult; - }; -})(); - export const DATA_TYPES = Object.freeze({ - auto: 'auto', // Auto-detect based on environment + auto: 'auto', fp32: 'fp32', fp16: 'fp16', q8: 'q8', @@ -41,97 +7,21 @@ export const DATA_TYPES = Object.freeze({ uint8: 'uint8', q4: 'q4', bnb4: 'bnb4', - q4f16: 'q4f16', // fp16 model with 4-bit block weight quantization + q4f16: 'q4f16', q2: 'q2', - q2f16: 'q2f16', // fp16 model with 2-bit block weight quantization + q2f16: 'q2f16', q1: 'q1', - q1f16: 'q1f16', // fp16 model with 1-bit block weight quantization + q1f16: 'q1f16', }); -/** @typedef {keyof typeof DATA_TYPES} DataType */ -export const DEFAULT_DEVICE_DTYPE = DATA_TYPES.fp32; -export const DEFAULT_DEVICE_DTYPE_MAPPING = Object.freeze({ - // NOTE: If not specified, will default to fp32 - [DEVICE_TYPES.wasm]: DATA_TYPES.q8, -}); - -/** @type {Record, string>} */ -export const DEFAULT_DTYPE_SUFFIX_MAPPING = Object.freeze({ - [DATA_TYPES.fp32]: '', - [DATA_TYPES.fp16]: '_fp16', - [DATA_TYPES.int8]: '_int8', - [DATA_TYPES.uint8]: '_uint8', - [DATA_TYPES.q8]: '_quantized', - [DATA_TYPES.q4]: '_q4', - [DATA_TYPES.q2]: '_q2', - [DATA_TYPES.q1]: '_q1', - [DATA_TYPES.q4f16]: '_q4f16', - [DATA_TYPES.q2f16]: '_q2f16', - [DATA_TYPES.q1f16]: '_q1f16', - [DATA_TYPES.bnb4]: '_bnb4', -}); - -/** - * Resolves a dtype configuration value to a concrete dtype string. - * Handles string, per-file object, and "auto" forms with device-based fallback. - * @param {DataType|Record|null|undefined} dtype The dtype config value. - * @param {string} fileName The model file name to look up if dtype is an object. - * @param {string} selectedDevice The resolved device string for fallback. - * @param {Object} [options] - * @param {DataType|Record|null} [options.configDtype=null] Config dtype used as fallback when dtype is "auto" (supports device_config overlay in session.js). - * @param {(message: string) => void} [options.warn] Optional callback invoked when dtype is a per-file object but fileName is not found. - * @returns {DataType} The resolved dtype string. - */ -export function selectDtype(dtype, fileName, selectedDevice, { configDtype = null, warn } = {}) { - /** @type {string|null|undefined} */ - let resolved; - let needsWarn = false; - if (dtype && typeof dtype !== 'string') { - if (dtype.hasOwnProperty(fileName)) { - resolved = dtype[fileName]; - } else { - resolved = null; - needsWarn = true; - } - } else { - resolved = /** @type {string|null|undefined} */ (dtype); - } - - /** @type {DataType} */ - let result; - - // Handle 'auto': try configDtype fallback - if (resolved === DATA_TYPES.auto) { - if (configDtype) { - const fallback = typeof configDtype === 'string' ? configDtype : configDtype?.[fileName]; - if (fallback && fallback !== DATA_TYPES.auto && DATA_TYPES.hasOwnProperty(fallback)) { - return /** @type {DataType} */ (fallback); - } - } - result = DEFAULT_DEVICE_DTYPE_MAPPING[selectedDevice] ?? DEFAULT_DEVICE_DTYPE; - } else if (resolved && DATA_TYPES.hasOwnProperty(resolved)) { - // Valid known dtype - result = /** @type {DataType} */ (resolved); - } else { - // Fallback to device default - result = DEFAULT_DEVICE_DTYPE_MAPPING[selectedDevice] ?? DEFAULT_DEVICE_DTYPE; - } - - if (needsWarn && warn) { - warn( - `dtype not specified for "${fileName}". Using the default dtype (${result}) for this device (${selectedDevice}).`, - ); - } - return result; -} +/** @typedef {keyof typeof DATA_TYPES} DataType */ export const DataTypeMap = Object.freeze({ float32: Float32Array, - // @ts-ignore ts(2552) Limited availability of Float16Array across browsers: - // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Float16Array + // @ts-ignore Limited availability of Float16Array across browsers. float16: typeof Float16Array !== 'undefined' ? Float16Array : Uint16Array, float64: Float64Array, - string: Array, // string[] + string: Array, int8: Int8Array, uint8: Uint8Array, int16: Int16Array, diff --git a/packages/transformers/src/utils/hub.js b/packages/transformers/src/utils/hub.js index 09f42e168..f594beedf 100755 --- a/packages/transformers/src/utils/hub.js +++ b/packages/transformers/src/utils/hub.js @@ -19,6 +19,7 @@ import { import { getCache, tryCache } from './cache.js'; import { get_file_metadata } from './model_registry/get_file_metadata.js'; import { logger } from './logger.js'; +import { getModelId } from '../backends/inference.js'; export { MAX_EXTERNAL_DATA_CHUNKS } from './hub/constants.js'; @@ -39,19 +40,21 @@ export { MAX_EXTERNAL_DATA_CHUNKS } from './hub/constants.js'; * @property {string} [cache_dir=null] Path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used. * @property {boolean} [local_files_only=false] Whether or not to only look at local files (e.g., not try downloading the model). * @property {string} [revision='main'] The specific model version to use. It can be a branch name, a tag name, or a commit id, + * @property {AbortSignal} [signal] Signal used to cancel backend-owned model loading. + * @property {import('../backends/artifacts.js').InferenceArtifactProvider} [artifactProvider] Optional random-access artifact provider supplied to custom inference backends. * since we use a git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any identifier allowed by git. * NOTE: This setting is ignored for local requests. */ /** * @typedef {Object} ModelSpecificPretrainedOptions Options for loading a pretrained model. - * @property {string} [subfolder='onnx'] In case the relevant files are located inside a subfolder of the model repo on huggingface.co, + * @property {string} [subfolder=null] In case the provider artifacts are located inside a subfolder of the model repo on huggingface.co, * you can specify the folder name here. * @property {string} [model_file_name=null] If specified, load the model with this name (excluding the dtype and .onnx suffixes). Currently only valid for encoder- or decoder-only models. * @property {import("./devices.js").DeviceType|Record} [device=null] The device to run the model on. If not specified, the device will be chosen from the environment settings. * @property {import("./dtypes.js").DataType|Record} [dtype=null] The data type to use for the model. If not specified, the data type will be chosen from the environment settings. * @property {ExternalData|Record} [use_external_data_format=false] Whether to load the model using the external data format (used for models >= 2GB in size). - * @property {import('onnxruntime-common').InferenceSession.SessionOptions} [session_options] (Optional) User-specified session options passed to the runtime. If not provided, suitable defaults will be chosen. + * @property {Object} [session_options] Compatibility options passed to providers that use session-based inference. */ /** @@ -130,6 +133,7 @@ export function getFetchHeaders(urlOrPath) { * An object containing all the paths and URLs for the resource. */ export function buildResourcePaths(path_or_repo_id, filename, options = {}, cache = null) { + path_or_repo_id = getModelId(path_or_repo_id); const revision = options.revision ?? 'main'; const requestURL = pathJoin(path_or_repo_id, filename); @@ -499,6 +503,7 @@ const INFLIGHT_LOADS = new Map(); * @returns {Promise} A Promise that resolves with the file content as a Uint8Array if `return_path` is false, or the file path as a string if `return_path` is true. */ export async function getModelFile(path_or_repo_id, filename, fatal = true, options = {}, return_path = false) { + path_or_repo_id = getModelId(path_or_repo_id); if (!env.allowLocalModels) { // User has disabled local models, so we just make sure other settings are correct. diff --git a/packages/transformers/src/utils/model-loader.js b/packages/transformers/src/utils/model-loader.js deleted file mode 100644 index 599aefe9d..000000000 --- a/packages/transformers/src/utils/model-loader.js +++ /dev/null @@ -1,111 +0,0 @@ -import { getModelFile, MAX_EXTERNAL_DATA_CHUNKS } from './hub.js'; -import { apis } from '../env.js'; - -/** - * Resolves an `use_external_data_format` config value to the number of data chunks for a given file. - * @param {import('./hub.js').ExternalData|Record|null|undefined} config The external data format configuration. - * @param {string} fullName The full ONNX file name (e.g., "model_quantized.onnx"). - * @param {string} fileName The base file name (e.g., "model"). - * @returns {number} The number of external data chunks (0 if none). - */ -export function resolveExternalDataFormat(config, fullName, fileName) { - if (!config) return 0; - if (typeof config === 'object' && config !== null) { - if (config.hasOwnProperty(fullName)) return +config[fullName]; - if (config.hasOwnProperty(fileName)) return +config[fileName]; - return 0; - } - return +config; // (false=0, true=1, number remains the same) -} - -/** - * Generates the file names for external data chunks. - * @param {string} fullName The full ONNX file name (e.g., "model_quantized.onnx"). - * @param {number} numChunks The number of external data chunks. - * @returns {string[]} Array of external data file names. - */ -export function getExternalDataChunkNames(fullName, numChunks) { - const names = []; - for (let i = 0; i < numChunks; ++i) { - names.push(`${fullName}_data${i === 0 ? '' : '_' + i}`); - } - return names; -} - -/** - * Loads the core model file. - * - * @param {string} pretrained_model_name_or_path The path to the directory containing the model file. - * @param {string} fileName The base name of the model file (without suffix or extension). - * @param {import('./hub.js').PretrainedModelOptions} options Additional options for loading the model. - * @param {string} suffix The suffix to append to the file name (e.g., '_q4', '_quantized'). - * @returns {Promise} A Promise that resolves to the model file buffer or path. - */ -export async function getCoreModelFile(pretrained_model_name_or_path, fileName, options, suffix) { - const baseName = `${fileName}${suffix}.onnx`; - const fullPath = `${options.subfolder ?? ''}/${baseName}`; - - return await getModelFile(pretrained_model_name_or_path, fullPath, true, options, apis.IS_NODE_ENV); -} - -/** - * Loads external data files for a model. - * - * @param {string} pretrained_model_name_or_path The path to the directory containing the model files. - * @param {string} fileName The base name of the model file (without suffix or extension). - * @param {string} suffix The suffix to append to the file name (e.g., '_q4'). - * @param {import('./hub.js').PretrainedModelOptions} options Additional options for loading the model. - * @param {import('./hub.js').ExternalData|Record|undefined} use_external_data_format External data format configuration. - * @param {any} [session_options] Optional session options that may contain externalData configuration. - * @returns {Promise>} A Promise that resolves to an array of external data files. - */ -export async function getModelDataFiles( - pretrained_model_name_or_path, - fileName, - suffix, - options, - use_external_data_format, - session_options = {}, -) { - const baseName = `${fileName}${suffix}.onnx`; - const return_path = apis.IS_NODE_ENV; - - /** @type {Promise[]} */ - let externalDataPromises = []; - - const num_chunks = resolveExternalDataFormat(use_external_data_format, baseName, fileName); - if (num_chunks > 0) { - if (num_chunks > MAX_EXTERNAL_DATA_CHUNKS) { - throw new Error( - `The number of external data chunks (${num_chunks}) exceeds the maximum allowed value (${MAX_EXTERNAL_DATA_CHUNKS}).`, - ); - } - const chunkNames = getExternalDataChunkNames(baseName, num_chunks); - for (const path of chunkNames) { - const fullPath = `${options.subfolder ?? ''}/${path}`; - externalDataPromises.push( - new Promise(async (resolve, reject) => { - const data = await getModelFile( - pretrained_model_name_or_path, - fullPath, - true, - options, - return_path, - ); - resolve(data instanceof Uint8Array ? { path, data } : path); - }), - ); - } - } else if (session_options.externalData !== undefined) { - externalDataPromises = session_options.externalData.map(async (ext) => { - // if the external data is a string, fetch the file and replace the string with its content - if (typeof ext.data === 'string') { - const ext_buffer = await getModelFile(pretrained_model_name_or_path, ext.data, true, options); - return { ...ext, data: ext_buffer }; - } - return ext; - }); - } - - return Promise.all(externalDataPromises); -} diff --git a/packages/transformers/src/utils/model_registry/get_available_dtypes.js b/packages/transformers/src/utils/model_registry/get_available_dtypes.js index fad61da76..869a22bbc 100644 --- a/packages/transformers/src/utils/model_registry/get_available_dtypes.js +++ b/packages/transformers/src/utils/model_registry/get_available_dtypes.js @@ -1,8 +1,8 @@ import { getSessionsConfig } from '../../models/session_config.js'; -import { DEFAULT_DTYPE_SUFFIX_MAPPING } from '../dtypes.js'; import { get_file_metadata } from './get_file_metadata.js'; import { get_config } from './get_model_files.js'; import { resolve_model_type } from './resolve_model_type.js'; +import { OnnxInferenceProvider } from '../../backends/default.js'; /** * @typedef {import('../../configs.js').PretrainedConfig} PretrainedConfig @@ -12,8 +12,6 @@ import { resolve_model_type } from './resolve_model_type.js'; * The dtypes to probe for availability (excludes 'auto' which is not a concrete dtype). * @type {string[]} */ -const CONCRETE_DTYPES = Object.keys(DEFAULT_DTYPE_SUFFIX_MAPPING); - /** * Detects which quantization levels (dtypes) are available for a model * by checking which ONNX files exist on the hub or locally. @@ -37,32 +35,13 @@ export async function get_available_dtypes( ) { config = await get_config(modelId, { config, cache_dir, local_files_only, revision }); - const subfolder = 'onnx'; - const modelType = resolve_model_type(config); - const { sessions } = getSessionsConfig(modelType, config, { model_file_name }); - - // Get all base names for model session files - const baseNames = Object.values(sessions); - - // For each dtype, check if all session files exist const metadataOptions = { revision, cache_dir, local_files_only }; - - // Probe all (dtype, baseName) combinations in parallel - const probeResults = await Promise.all( - CONCRETE_DTYPES.map(async (dtype) => { - const suffix = DEFAULT_DTYPE_SUFFIX_MAPPING[dtype] ?? ''; - const allExist = await Promise.all( - baseNames.map(async (baseName) => { - const filename = `${subfolder}/${baseName}${suffix}.onnx`; - const metadata = await get_file_metadata(modelId, filename, metadataOptions); - return metadata.exists; - }), - ); - return { dtype, available: allExist.every(Boolean) }; - }), - ); - - return probeResults.filter((r) => r.available).map((r) => r.dtype); + return OnnxInferenceProvider.getAvailableDtypes({ + modelId, + sessions, + getFileMetadata: get_file_metadata, + metadataOptions, + }); } diff --git a/packages/transformers/src/utils/model_registry/get_file_metadata.js b/packages/transformers/src/utils/model_registry/get_file_metadata.js index 2ee8a509c..2d69344f6 100644 --- a/packages/transformers/src/utils/model_registry/get_file_metadata.js +++ b/packages/transformers/src/utils/model_registry/get_file_metadata.js @@ -8,6 +8,7 @@ import { buildResourcePaths, checkCachedResource, getFetchHeaders, getFile } fro import { isValidUrl, makePretrainedOptionsKey } from '../hub/utils.js'; import { logger } from '../logger.js'; import { memoizePromise } from '../memoize_promise.js'; +import { getModelId } from '../../backends/inference.js'; /** * @typedef {import('../hub.js').PretrainedOptions} PretrainedOptions @@ -51,6 +52,7 @@ async function fetch_file_head(urlOrPath) { * @returns {Promise<{exists: boolean, size?: number, contentType?: string, fromCache?: boolean}>} A Promise that resolves to file metadata. */ export function get_file_metadata(path_or_repo_id, filename, options = {}) { + path_or_repo_id = getModelId(path_or_repo_id); const key = makePretrainedOptionsKey(path_or_repo_id, options, filename); return memoizePromise(key, () => _get_file_metadata(path_or_repo_id, filename, options)); } diff --git a/packages/transformers/src/utils/model_registry/get_files.js b/packages/transformers/src/utils/model_registry/get_files.js index eeda2891b..e13d0d291 100644 --- a/packages/transformers/src/utils/model_registry/get_files.js +++ b/packages/transformers/src/utils/model_registry/get_files.js @@ -14,6 +14,7 @@ import { get_processor_files } from './get_processor_files.js'; * @param {string|null} [options.model_file_name=null|null] Override the model file name (excluding .onnx suffix) * @param {boolean} [options.include_tokenizer=true] Whether to check for tokenizer files (set to false for vision-only models) * @param {boolean} [options.include_processor=true] Whether to check for processor files + * @param {boolean} [options.include_model=true] Whether to include built-in ONNX model files * @returns {Promise} Array of file paths that will be loaded */ export async function get_files( @@ -25,9 +26,10 @@ export async function get_files( model_file_name = null, include_tokenizer = true, include_processor = true, + include_model = true, } = {}, ) { - const files = await get_model_files(modelId, { config, dtype, device, model_file_name }); + const files = include_model ? await get_model_files(modelId, { config, dtype, device, model_file_name }) : []; if (include_tokenizer) { const tokenizerFiles = await get_tokenizer_files(modelId); diff --git a/packages/transformers/src/utils/model_registry/get_model_files.js b/packages/transformers/src/utils/model_registry/get_model_files.js index 467642421..3f267d3bc 100644 --- a/packages/transformers/src/utils/model_registry/get_model_files.js +++ b/packages/transformers/src/utils/model_registry/get_model_files.js @@ -1,11 +1,9 @@ -import { DEFAULT_DTYPE_SUFFIX_MAPPING, selectDtype } from '../dtypes.js'; -import { selectDevice } from '../devices.js'; -import { resolveExternalDataFormat, getExternalDataChunkNames } from '../model-loader.js'; import { getSessionsConfig } from '../../models/session_config.js'; import { AutoConfig } from '../../configs.js'; import { makePretrainedOptionsKey } from '../hub/utils.js'; import { memoizePromise } from '../memoize_promise.js'; import { resolve_model_type } from './resolve_model_type.js'; +import { OnnxInferenceProvider } from '../../backends/default.js'; /** * @typedef {import('../../configs.js').PretrainedConfig} PretrainedConfig @@ -63,53 +61,14 @@ export async function get_model_files( ) { config = await get_config(modelId, { config }); - const files = [ - // Add config.json (always loaded) - 'config.json', - ]; - const custom_config = config['transformers.js_config'] ?? {}; - - const use_external_data_format = custom_config.use_external_data_format; - const subfolder = 'onnx'; // Always 'onnx' as per the default in from_pretrained - - const rawDevice = overrideDevice ?? custom_config.device; - let dtype = overrideDtype ?? custom_config.dtype; - // Infer model type from config const modelType = resolve_model_type(config); - - const add_model_file = (fileName, baseName = null) => { - baseName = baseName ?? fileName; - const selectedDevice = selectDevice(rawDevice, fileName); - const selectedDtype = selectDtype(dtype, fileName, selectedDevice); - - const suffix = DEFAULT_DTYPE_SUFFIX_MAPPING[selectedDtype] ?? ''; - const fullName = `${baseName}${suffix}.onnx`; - const fullPath = subfolder ? `${subfolder}/${fullName}` : fullName; - files.push(fullPath); - - // Check for external data files - const num_chunks = resolveExternalDataFormat(use_external_data_format, fullName, fileName); - for (const dataFileName of getExternalDataChunkNames(fullName, num_chunks)) { - const dataFilePath = subfolder ? `${subfolder}/${dataFileName}` : dataFileName; - files.push(dataFilePath); - } - }; - - // Get session configuration from the shared source of truth const { sessions, optional_configs } = getSessionsConfig(modelType, config, { model_file_name }); - - // Add model files based on sessions - for (const [sessionKey, baseName] of Object.entries(sessions)) { - add_model_file(sessionKey, baseName); - } - - // Add optional config files - if (optional_configs) { - for (const configFile of Object.values(optional_configs)) { - files.push(configFile); - } - } - - return files; + return OnnxInferenceProvider.listModelArtifacts({ + sessions, + optionalConfigs: optional_configs, + config, + dtype: overrideDtype, + device: overrideDevice, + }); } diff --git a/packages/transformers/src/utils/model_registry/get_pipeline_files.js b/packages/transformers/src/utils/model_registry/get_pipeline_files.js index c6d4e3a7d..e22eddfd3 100644 --- a/packages/transformers/src/utils/model_registry/get_pipeline_files.js +++ b/packages/transformers/src/utils/model_registry/get_pipeline_files.js @@ -3,6 +3,7 @@ import { get_config } from './get_model_files.js'; import { resolve_model_type } from './resolve_model_type.js'; import { getTextOnlySessions } from '../../models/session_config.js'; import { SUPPORTED_TASKS, TASK_ALIASES } from '../../pipelines/index.js'; +import { OnnxInferenceProvider } from '../../backends/default.js'; /** * Get all files needed for a specific pipeline task. @@ -16,6 +17,7 @@ import { SUPPORTED_TASKS, TASK_ALIASES } from '../../pipelines/index.js'; * @param {import('../dtypes.js').DataType|Record} [options.dtype=null] - Override dtype * @param {import('../devices.js').DeviceType|Record} [options.device=null] - Override device * @param {string} [options.model_file_name=null] - Override the model file name (excluding .onnx suffix) + * @param {boolean} [options.include_model=true] - Whether to include built-in ONNX model files * @returns {Promise} Array of file paths that will be loaded * @throws {Error} If the task is not supported */ @@ -53,8 +55,7 @@ export async function get_pipeline_files(task, modelId, options = {}) { const textOnlySessions = getTextOnlySessions(modelType); if (textOnlySessions) { - const allowedPrefixes = Object.values(textOnlySessions).map((s) => `onnx/${s}`); - return files.filter((f) => !f.startsWith('onnx/') || allowedPrefixes.some((p) => f.startsWith(p))); + return OnnxInferenceProvider.filterModelArtifacts(files, textOnlySessions); } } diff --git a/packages/transformers/src/utils/tensor.js b/packages/transformers/src/utils/tensor.js index 85da403ef..330eaf64a 100644 --- a/packages/transformers/src/utils/tensor.js +++ b/packages/transformers/src/utils/tensor.js @@ -9,8 +9,6 @@ import { interpolate_data, max, min, permute_data, uint16_to_float32 } from './maths.js'; -import { Tensor as ONNXTensor, isONNXTensor } from '../backends/onnx.js'; - import { TensorOpRegistry } from '../ops/registry.js'; import { DataTypeMap } from './dtypes.js'; @@ -28,13 +26,10 @@ export class Tensor { * @type {number[]} */ get dims() { - // @ts-ignore - return this.ort_tensor.dims; + return this._storage.dims; } set dims(value) { - // FIXME: ONNXTensor declares dims as readonly so one needs to use the constructor() if dims change. - // @ts-ignore - this.ort_tensor.dims = value; + this._storage.dims = value; } /** @@ -42,7 +37,7 @@ export class Tensor { * @type {DataType} */ get type() { - return this.ort_tensor.type; + return this._storage.type; } /** @@ -50,7 +45,7 @@ export class Tensor { * @type {DataArray} */ get data() { - return this.ort_tensor.data; + return this._storage.data; } /** @@ -58,7 +53,7 @@ export class Tensor { * @type {number} */ get size() { - return this.ort_tensor.size; + return this._storage.size; } /** @@ -66,27 +61,36 @@ export class Tensor { * @type {string} */ get location() { - return this.ort_tensor.location; + return this._storage.location; } - ort_tensor; + _storage; /** * Create a new Tensor or copy an existing Tensor. - * @param {[DataType, DataArray, number[]]|[ONNXTensor]} args + * @param {[DataType, DataArray, number[]]} args */ constructor(...args) { - if (isONNXTensor(args[0])) { - this.ort_tensor = /** @type {ONNXTensor} */ (args[0]); - } else { - // Create new tensor - this.ort_tensor = new ONNXTensor( - /** @type {DataType} */ (args[0]), - // @ts-expect-error ts(2769) Type 'number' is not assignable to type 'bigint'. - /** @type {Exclude} */ (args[1]), - args[2], - ); + const [type, inputData, dims] = args; + let data = inputData; + if (Array.isArray(data) && type !== 'string') { + const Constructor = /** @type {any} */ (DataTypeMap[type]); + if (type === 'int64' || type === 'uint64') { + data = Constructor.from(data, (value) => BigInt(value)); + } else { + data = Constructor.from(data); + } } + this._storage = { + backend: 'cpu', + handle: null, + type, + data, + dims, + size: dims.reduce((product, dimension) => product * dimension, 1), + location: 'cpu', + dispose() {}, + }; return new Proxy(this, { get: (obj, key) => { @@ -110,8 +114,33 @@ export class Tensor { } dispose() { - this.ort_tensor.dispose(); - // this.ort_tensor = undefined; + this._storage.dispose?.(); + } + + /** + * Construct a Tensor around provider-owned storage. + * @param {Object} storage + * @returns {Tensor} + * @internal + */ + static fromBackendStorage(storage) { + const tensor = Object.create(Tensor.prototype); + tensor._storage = storage; + return new Proxy(tensor, { + get: (obj, key) => { + if (typeof key === 'string') { + const index = Number(key); + if (Number.isInteger(index)) return obj._getitem(index); + } + return obj[key]; + }, + set: (obj, key, value) => (obj[key] = value), + }); + } + + /** @internal */ + getBackendStorage() { + return this._storage; } /** diff --git a/packages/transformers/tests/generation_controller.test.js b/packages/transformers/tests/generation_controller.test.js new file mode 100644 index 000000000..adcc2a27d --- /dev/null +++ b/packages/transformers/tests/generation_controller.test.js @@ -0,0 +1,214 @@ +import { jest } from "@jest/globals"; + +import { AutoModel } from "../src/models/auto/modeling_auto.js"; +import { GenerationController, createGenerationController } from "../src/generation/controller.js"; +import { LogitsProcessor, LogitsProcessorList } from "../src/generation/logits_process.js"; +import { StoppingCriteria, StoppingCriteriaList } from "../src/generation/stopping_criteria.js"; +import { Tensor } from "../src/utils/tensor.js"; + +class ForceTokenProcessor extends LogitsProcessor { + constructor(tokenId) { + super(); + this.tokenId = tokenId; + } + + _call(_inputIds, logits) { + logits.data.fill(-Infinity); + for (let batchIndex = 0; batchIndex < logits.dims[0]; ++batchIndex) { + logits.data[batchIndex * logits.dims.at(-1) + this.tokenId] = 0; + } + return logits; + } +} + +class TokenStoppingCriteria extends StoppingCriteria { + constructor(tokenId) { + super(); + this.tokenId = BigInt(tokenId); + } + + _call(inputIds) { + return inputIds.map((tokens) => tokens.at(-1) === this.tokenId); + } +} + +function int64Tensor(values) { + return new Tensor("int64", BigInt64Array.from(values.flat().map(BigInt)), [values.length, values[0].length]); +} + +function createLease(values, release = jest.fn()) { + return { + version: 1, + dtype: "float32", + shape: [1, values.length], + read: jest.fn(async () => Float32Array.from(values)), + release, + }; +} + +describe("GenerationController", () => { + it("owns processing, sampling, stopping, streaming, and finalization", async () => { + const processors = new LogitsProcessorList(); + processors.push(new ForceTokenProcessor(2)); + const criteria = new StoppingCriteriaList(); + criteria.push(new TokenStoppingCriteria(2)); + const streamer = { put: jest.fn(), end: jest.fn() }; + const model = { + config: { eos_token_id: null }, + generation_config: null, + }; + const controller = createGenerationController(model, int64Tensor([[1]]), { max_new_tokens: 4, logits_processor: processors, stopping_criteria: criteria, streamer }); + + const step = await controller.step(new Tensor("float32", [10, 9, 0, 8], [1, 4])); + + expect(step.nextTokenIds.tolist()).toEqual([[2n]]); + expect(step.allDone).toBe(true); + expect(streamer.put.mock.calls).toEqual([[[[1n]]], [[[2n]]]]); + expect(controller.finalize().tolist()).toEqual([[1n, 2n]]); + expect(streamer.end).toHaveBeenCalledTimes(1); + }); + + it("supports zero-token generation without a model step", () => { + const streamer = { put: jest.fn(), end: jest.fn() }; + const controller = createGenerationController({ config: {}, generation_config: null }, int64Tensor([[1, 2]]), { max_new_tokens: 0, streamer }); + + expect(controller.allDone).toBe(true); + expect(controller.finalize().tolist()).toEqual([[1n, 2n]]); + expect(streamer.put).toHaveBeenCalledWith([[1n, 2n]]); + expect(streamer.end).toHaveBeenCalledTimes(1); + }); +}); + +describe("custom autoregressive sessions", () => { + it("uses leased CPU logits for arbitrary Transformers.js callbacks", async () => { + const releases = [jest.fn(), jest.fn()]; + const leases = [createLease([0, 1, 5, 2], releases[0]), createLease([0, 1, 2, 6], releases[1])]; + const session = { + version: 1, + batchSize: 1, + maxSequenceLength: 3, + prefill: jest.fn(async () => leases[0]), + decode: jest.fn(async () => leases[1]), + dispose: jest.fn(async () => {}), + }; + const backend = { + modelId: "test/controller-model", + load: jest.fn(async () => ({ + generation_config: {}, + generationCapabilities: { + sessionVersion: 1, + maxBatchSize: 1, + cpuModes: ["greedy", "multinomial"], + planModes: [], + cpuLogits: true, + declarativePlans: [], + tokenPipeline: { defaultDepth: 1, maxDepth: 1 }, + }, + createAutoregressiveSession: jest.fn(async () => session), + async forward(inputs) { + return inputs; + }, + async dispose() {}, + })), + }; + const model = await AutoModel.from_pretrained(backend, { + config: { model_type: "custom", is_encoder_decoder: false, eos_token_id: 3 }, + }); + + const output = await model.generate({ + input_ids: int64Tensor([[1]]), + attention_mask: int64Tensor([[1]]), + max_new_tokens: 2, + }); + + expect(output.tolist()).toEqual([[1n, 2n, 3n]]); + expect(session.prefill).toHaveBeenCalledWith(expect.objectContaining({ inputIds: { data: new Uint32Array([1]), shape: [1, 1] } })); + expect(session.decode).toHaveBeenCalledWith(expect.objectContaining({ tokenIds: { data: new Uint32Array([2]), shape: [1, 1] } })); + expect(releases.every((release) => release.mock.calls.length === 1)).toBe(true); + expect(session.dispose).toHaveBeenCalledTimes(1); + }); + + it("commits runtime plan decisions without reading full logits", async () => { + const iteratorClosed = jest.fn(); + const session = { + version: 1, + batchSize: 1, + maxSequenceLength: 3, + prefill: jest.fn(), + decode: jest.fn(), + async *generateWithPlan(_inputs, plan) { + try { + expect(plan.sampler).toEqual({ op: "argmax" }); + expect(plan.pipelineDepth).toBe(4); + yield { tokenIds: new Uint32Array([2]) }; + yield { tokenIds: new Uint32Array([3]) }; + yield { tokenIds: new Uint32Array([0]) }; + } finally { + iteratorClosed(); + } + }, + dispose: jest.fn(async () => {}), + }; + const backend = { + modelId: "test/fast-controller-model", + load: jest.fn(async () => ({ + generation_config: {}, + generationCapabilities: { + sessionVersion: 1, + maxBatchSize: 1, + cpuModes: ["greedy"], + planModes: ["greedy"], + cpuLogits: false, + declarativePlans: ["argmax"], + tokenPipeline: { defaultDepth: 4, maxDepth: 4 }, + }, + createAutoregressiveSession: jest.fn(async () => session), + async forward(inputs) { + return inputs; + }, + async dispose() {}, + })), + }; + const model = await AutoModel.from_pretrained(backend, { + config: { model_type: "custom", is_encoder_decoder: false, eos_token_id: 3 }, + }); + + const output = await model.generate({ input_ids: int64Tensor([[1]]), max_new_tokens: 2 }); + + expect(output.tolist()).toEqual([[1n, 2n, 3n]]); + expect(session.prefill).not.toHaveBeenCalled(); + expect(session.decode).not.toHaveBeenCalled(); + expect(iteratorClosed).toHaveBeenCalledTimes(1); + expect(session.dispose).toHaveBeenCalledTimes(1); + }); + + it("rejects unsupported batches before creating a runtime session", async () => { + const createAutoregressiveSession = jest.fn(); + const backend = { + modelId: "test/batch-controller-model", + load: jest.fn(async () => ({ + generation_config: {}, + generationCapabilities: { + sessionVersion: 1, + maxBatchSize: 1, + cpuModes: ["greedy"], + planModes: [], + cpuLogits: true, + declarativePlans: [], + tokenPipeline: { defaultDepth: 1, maxDepth: 1 }, + }, + createAutoregressiveSession, + async forward(inputs) { + return inputs; + }, + async dispose() {}, + })), + }; + const model = await AutoModel.from_pretrained(backend, { + config: { model_type: "custom", is_encoder_decoder: false }, + }); + + await expect(model.generate({ input_ids: int64Tensor([[1], [2]]), max_new_tokens: 1 })).rejects.toThrow("supports batch size 1"); + expect(createAutoregressiveSession).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/transformers/tests/inference_backends.test.js b/packages/transformers/tests/inference_backends.test.js new file mode 100644 index 000000000..ad601e02a --- /dev/null +++ b/packages/transformers/tests/inference_backends.test.js @@ -0,0 +1,126 @@ +import { jest } from "@jest/globals"; + +import { getModelId, isInferenceBackend, loadInferenceModel, normalizeInferenceModel } from "../src/backends/inference.js"; +import { OnnxInferenceProvider } from "../src/backends/default.js"; +import { AutoModel } from "../src/models/auto/modeling_auto.js"; +import { PreTrainedModel } from "../src/models/modeling_utils.js"; +import { buildResourcePaths } from "../src/utils/hub.js"; + +describe("inference backends", () => { + it("recognizes object and class backends", () => { + const objectBackend = { modelId: "test/object", load() {} }; + class ClassBackend { + static modelId = "test/class"; + static load() {} + } + + expect(isInferenceBackend(objectBackend)).toBe(true); + expect(isInferenceBackend(ClassBackend)).toBe(true); + expect(getModelId(objectBackend)).toBe("test/object"); + expect(getModelId(ClassBackend)).toBe("test/class"); + expect(getModelId("test/string")).toBe("test/string"); + }); + + it("normalizes a model with forward into a callable model", async () => { + const implementation = { + value: 42, + async forward(inputs) { + return { inputs, value: this.value }; + }, + async dispose() {}, + }; + const model = normalizeInferenceModel(implementation); + + await expect(model({ input_ids: "input" })).resolves.toEqual({ + inputs: { input_ids: "input" }, + value: 42, + }); + expect(model.value).toBe(42); + }); + + it("passes shared loading options to a custom backend", async () => { + const load = jest.fn(async () => ({ + async forward(inputs) { + return inputs; + }, + async dispose() {}, + })); + const backend = { modelId: "test/model", load }; + const config = { model_type: "custom" }; + + const model = await loadInferenceModel(backend, { config, dtype: "q4f16" }); + + expect(load).toHaveBeenCalledWith({ config, dtype: "q4f16", modelId: "test/model" }); + expect(model.config).toBe(config); + }); + + it("normalizes absent custom device and dtype options", async () => { + const backend = { + modelId: "test/model", + load: jest.fn(async () => ({ + async forward(inputs) { + return inputs; + }, + async dispose() {}, + })), + }; + + await loadInferenceModel(backend, { device: null, dtype: null }); + + expect(backend.load).toHaveBeenCalledWith(expect.objectContaining({ modelId: "test/model", device: undefined, dtype: undefined })); + }); + + it("allows custom backends through AutoModel.from_pretrained", async () => { + const loaded = { + async forward(inputs) { + return inputs; + }, + async dispose() {}, + }; + const backend = { + modelId: "test/model", + load: jest.fn(async () => loaded), + }; + const signal = new AbortController().signal; + const artifactProvider = { readJson() {}, openByteSource() {} }; + + const model = await AutoModel.from_pretrained(backend, { + config: { model_type: "custom" }, + device: "webgpu", + signal, + artifactProvider, + }); + + expect(backend.load).toHaveBeenCalledWith(expect.objectContaining({ modelId: "test/model", device: "webgpu", signal, artifactProvider })); + await expect(model({ value: 1 })).resolves.toEqual({ value: 1 }); + }); + + it("represents string IDs with the ONNX fallback backend", async () => { + const modelClass = { _from_pretrained: jest.fn(async () => "model") }; + const backend = OnnxInferenceProvider.from_modelId("test/model"); + expect(backend.modelId).toBe("test/model"); + expect(isInferenceBackend(backend)).toBe(true); + await expect(backend.load({ dtype: "q4", modelClass })).resolves.toBe("model"); + expect(modelClass._from_pretrained).toHaveBeenCalledWith("test/model", expect.objectContaining({ dtype: "q4", inferenceProvider: backend })); + }); + + it("resolves string model IDs through OnnxInferenceProvider.from_modelId", async () => { + class TestModel extends PreTrainedModel {} + TestModel._from_pretrained = jest.fn(async () => "loaded-model"); + const factory = jest.spyOn(OnnxInferenceProvider, "from_modelId"); + + await expect(TestModel.from_pretrained("test/string-model")).resolves.toBe("loaded-model"); + + expect(factory).toHaveBeenCalledWith("test/string-model"); + expect(TestModel._from_pretrained).toHaveBeenCalledWith("test/string-model", expect.objectContaining({ inferenceProvider: expect.any(OnnxInferenceProvider) })); + factory.mockRestore(); + }); + + it("uses backend model IDs for shared asset paths", () => { + const backend = { modelId: "test/model", load() {} }; + const paths = buildResourcePaths(backend, "tokenizer.json"); + + expect(paths.requestURL).toBe("test/model/tokenizer.json"); + expect(paths.remoteURL).toContain("/test/model/resolve/main/tokenizer.json"); + }); +}); diff --git a/packages/transformers/tests/init.js b/packages/transformers/tests/init.js index 9da694d72..ec30e4b87 100644 --- a/packages/transformers/tests/init.js +++ b/packages/transformers/tests/init.js @@ -1,11 +1,6 @@ // Helper functions used when initialising the testing environment. -// Import Node typing utilities -import * as types from "node:util/types"; - -// Import onnxruntime-node's default backend -import { onnxruntimeBackend } from "onnxruntime-node/dist/backend"; -import * as ONNX_COMMON from "onnxruntime-common"; +import { initOnnxTestBackend } from "@huggingface/transformers-onnx/testing"; /** * A workaround to define a new backend for onnxruntime, which @@ -13,49 +8,7 @@ import * as ONNX_COMMON from "onnxruntime-common"; * For more information, see: https://github.com/jestjs/jest/issues/11864#issuecomment-1261468011 */ export function init() { - // In rare cases (specifically when running unit tests with GitHub actions), possibly due to - // a large number of concurrent executions, onnxruntime might fallback to use the WASM backend. - // In this case, we set the number of threads to 1 to avoid errors like: - // - `TypeError: The worker script or module filename must be an absolute path or a relative path starting with './' or '../'. Received "blob:nodedata:..."` - ONNX_COMMON.env.wasm.numThreads = 1; - - let registerBackend = ONNX_COMMON.registerBackend; - - // Define the constructors to monkey-patch - const TYPED_ARRAYS_CONSTRUCTOR_NAMES = ["Int8Array", "Int16Array", "Int32Array", "BigInt64Array", "Uint8Array", "Uint8ClampedArray", "Uint16Array", "Uint32Array", "BigUint64Array", "Float16Array", "Float32Array", "Float64Array"]; - - // Keep a reference to the original initialization method - const originalMethod = onnxruntimeBackend.init; - - // Monkey-patch the initialization function - onnxruntimeBackend.init = function (...args) { - // There is probably a better way to do this - Array.isArray = (x) => typeof x === "object" && x !== null && typeof x.length === "number" && x?.constructor.toString() === Array.toString(); - - // For each typed array constructor - for (const ctorName of TYPED_ARRAYS_CONSTRUCTOR_NAMES) { - // Get the constructor from the current context - const ctor = globalThis[ctorName]; - if (ctor === undefined) continue; // If unavailable, skip the patching - - // Get the corresponding test function from the `util` module - const value = types[`is${ctorName}`].bind(types); - - // Monkey-patch the constructor so "x instanceof ctor" returns "types[`is${ctorName}`](x)" - Object.defineProperty(ctor, Symbol.hasInstance, { - value, - writable: true, // writable=true is necessary to overwrite the default implementation (and allow subsequent overwrites) - configurable: false, - enumerable: false, - }); - } - - // Call the original method - return originalMethod.apply(this, args); - }; - - // Register the backend with the highest priority, so it is used instead of the default one - registerBackend("test", onnxruntimeBackend, Number.POSITIVE_INFINITY); + initOnnxTestBackend(); } export const MAX_TOKENIZER_LOAD_TIME = 32_000; // 32 seconds diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index bea34792e..2a126f332 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -23,12 +23,9 @@ importers: '@huggingface/tokenizers': specifier: ^0.1.3 version: 0.1.3 - onnxruntime-node: - specifier: 1.24.3 - version: 1.24.3 - onnxruntime-web: - specifier: 1.26.0-dev.20260416-b7804b056c - version: 1.26.0-dev.20260416-b7804b056c + '@huggingface/transformers-onnx': + specifier: workspace:^ + version: link:../transformers-onnx sharp: specifier: ^0.34.5 version: 0.34.5 @@ -58,6 +55,34 @@ importers: specifier: 5.9.3 version: 5.9.3 + packages/transformers-onnx: + dependencies: + onnxruntime-common: + specifier: 1.24.3 + version: 1.24.3 + onnxruntime-node: + specifier: 1.24.3 + version: 1.24.3 + onnxruntime-web: + specifier: 1.26.0-dev.20260416-b7804b056c + version: 1.26.0-dev.20260416-b7804b056c + devDependencies: + '@types/node': + specifier: ^24.1.0 + version: 24.10.9 + '@webgpu/types': + specifier: ^0.1.69 + version: 0.1.69 + esbuild: + specifier: ^0.27.2 + version: 0.27.2 + jest: + specifier: ^30.2.0 + version: 30.2.0(@types/node@24.10.9) + typescript: + specifier: 5.9.3 + version: 5.9.3 + packages: '@babel/code-frame@7.28.6': diff --git a/types/webgpu-kernels.local.demo.d.ts b/types/webgpu-kernels.local.demo.d.ts new file mode 100644 index 000000000..6f6b469b6 --- /dev/null +++ b/types/webgpu-kernels.local.demo.d.ts @@ -0,0 +1,14 @@ +import type { WebGPUKernelsForwardModel, WebGPUKernelsTensorMap, WebGPUKernelsTextGenerationModel } from './webgpu-kernels.local'; +type PipelineForwardModel = { + (inputs: WebGPUKernelsTensorMap): Promise; + readonly config?: Record; + forward(inputs: WebGPUKernelsTensorMap): Promise; + dispose(): void | Promise; +}; +type PipelineTextGenerationModel = PipelineForwardModel & { + generate: WebGPUKernelsTextGenerationModel['generate']; +}; +export declare function adaptWebGPUKernelsModel(model: WebGPUKernelsTextGenerationModel): PipelineTextGenerationModel; +export declare function adaptWebGPUKernelsModel(model: WebGPUKernelsForwardModel): PipelineForwardModel; +export {}; +//# sourceMappingURL=webgpu-kernels.local.demo.d.ts.map \ No newline at end of file diff --git a/types/webgpu-kernels.local.demo.d.ts.map b/types/webgpu-kernels.local.demo.d.ts.map new file mode 100644 index 000000000..67f999ab9 --- /dev/null +++ b/types/webgpu-kernels.local.demo.d.ts.map @@ -0,0 +1 @@ +{"version":3,"file":"webgpu-kernels.local.demo.d.ts","sourceRoot":"","sources":["../webgpu-kernels.local.demo.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EACR,yBAAyB,EACzB,sBAAsB,EAEtB,gCAAgC,EACnC,MAAM,wBAAwB,CAAC;AAEhC,KAAK,oBAAoB,GAAG;IACxB,CAAC,MAAM,EAAE,sBAAsB,GAAG,OAAO,CAAC,sBAAsB,CAAC,CAAC;IAClE,QAAQ,CAAC,MAAM,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;IAC1C,OAAO,CAAC,MAAM,EAAE,sBAAsB,GAAG,OAAO,CAAC,sBAAsB,CAAC,CAAC;IACzE,OAAO,IAAI,IAAI,GAAG,OAAO,CAAC,IAAI,CAAC,CAAC;CACnC,CAAC;AAEF,KAAK,2BAA2B,GAAG,oBAAoB,GAAG;IACtD,QAAQ,EAAE,gCAAgC,CAAC,UAAU,CAAC,CAAC;CAC1D,CAAC;AAEF,wBAAgB,uBAAuB,CAAC,KAAK,EAAE,gCAAgC,GAAG,2BAA2B,CAAC;AAC9G,wBAAgB,uBAAuB,CAAC,KAAK,EAAE,yBAAyB,GAAG,oBAAoB,CAAC"} \ No newline at end of file diff --git a/webgpu-compat.local.md b/webgpu-compat.local.md new file mode 100644 index 000000000..4466d5d76 --- /dev/null +++ b/webgpu-compat.local.md @@ -0,0 +1,362 @@ +# WebGPU runtime integration contract + +## Resolved Architecture + +Transformers.js and the WebGPU runtime have agreed on: + +- pull and plan session modes; +- terminal fast-path iteration with awaited iterator `return()`; +- leased CPU logits with synchronous idempotent release; +- prefill and decode token semantics; +- separate CPU and plan generation capabilities; +- bounded token-pipeline depth capabilities; +- batch-one, all-ones-mask V1 inputs; +- uint32 runtime token batches and Transformers.js-owned int64 sequences; +- Transformers.js-owned generation policy, callbacks, stopping, streaming, and finalization; +- runtime-owned KV cache and native inference state; +- optional generation-only models without `forward()`; +- load/session cancellation and defensive model-side session cleanup; +- native QAT dtype behavior: omitted/`auto` is supported and `q4f16` throws. + +## Default Artifact Provider + +Option A is accepted for the first adapter. + +`Gemma4E2B.load()` may use the existing runtime `ResourceRootIo` when Transformers.js does not supply an `artifactProvider`. + +Initial behavior: + +- `revision`, injected fetch/auth, abort, weight progress, and runtime browser caching are mapped where supported; +- `local_files_only` is rejected before metadata or weight loading rather than ignored; +- an explicitly supplied `artifactProvider` takes precedence over runtime IO; +- unsupported cache-policy fields are rejected or documented, never silently treated as equivalent; +- exact Transformers.js cache identity, offline policy, and shared Hub caching are deferred to a separate ranged-provider project. + +This does not block a later move to a Transformers.js-provided ranged implementation. The adapter's model-loading code consumes a small internal random-access source abstraction so runtime IO and a future Transformers.js provider share the same safetensors path. + +## Provider Contract + +```ts +interface InferenceArtifactProvider { + readJson(file: string, options?: { signal?: AbortSignal }): Promise; + + openByteSource( + file: string, + options?: { + signal?: AbortSignal; + onProgress?: (event: ArtifactProgressEvent) => void; + }, + ): Promise; +} + +interface RandomAccessByteSource { + /** May become defined after transport metadata arrives. */ + readonly size?: number; + + /** Read bytes in the half-open interval [begin, end). */ + read( + begin: number, + end: number, + options?: { signal?: AbortSignal }, + ): Promise; + + /** Idempotently wait for pending reads, then release the source. */ + close(): Promise; +} +``` + +### Range Convention + +`read(begin, end)` is half-open: `[begin, end)`. + +Both values must be non-negative safe integers and `end >= begin`. An empty range returns an owned zero-length `Uint8Array`. + +This matches the runtime `ByteSource.readRange()` contract and safetensors offsets. + +### Concurrency + +Reads may execute concurrently and complete out of order. + +The provider must not rely on one global mutable seek cursor. HTTP range requests naturally satisfy this. Filesystem implementations must use positional reads. + +A provider may serialize internally when required by its backing store, but it must preserve each read's independent range and result. + +### Returned Bytes + +Every successful `read()` returns an owned `Uint8Array` whose contents remain valid and unchanged after: + +- later reads; +- out-of-order completion of other reads; +- cache writes; +- `close()`. + +The runtime may retain returned bytes until the consuming upload, decode, or transcode operation completes. Providers must not return a view into a reused scratch buffer. + +### Closing with Pending Reads + +`close()` waits for reads that were already started, then releases the source. It does not implicitly abort pending reads. + +Once closing starts: + +- new reads reject; +- existing reads settle normally or through their supplied abort signal; +- `close()` observes all settlements before resolving; +- read failures remain owned by the corresponding read promises and do not become unhandled rejections. + +A provider may internally abort transport work during process-wide teardown, but `close()` still waits for that work to settle. + +`close()` is idempotent. Concurrent calls share the same close operation. Calls after closure resolve without repeating resource release. + +### Unknown Size + +An unknown `size` is acceptable, including at open time. + +Safetensors header parsing may read an initial probe and derive tensor ranges without total file length. When size is unknown: + +- complete-file length validation is unavailable; +- progress totals may be unknown; +- size-dependent cache metadata remains unavailable until size is discovered. + +If transport metadata later reveals the size, `size` may be a getter that changes once from `undefined` to a stable non-negative safe integer. Once defined, it must not change. + +For ranges wholly inside the file, `read()` returns exactly `end - begin` bytes. When size is unknown and an initial probe extends past EOF, returning the available shorter prefix is acceptable. Arbitrary short reads for known-valid tensor ranges are errors. + +### Validation and Limits + +The runtime imposes no alignment requirement on `begin`, `end`, or range length. + +Provider validation requirements: + +- `begin` and `end` are safe integers; +- `begin >= 0`; +- `end >= begin`; +- when `size` is known, `end <= size`; +- a successful known-valid range contains exactly `end - begin` bytes. + +There is no protocol-level maximum range size. The runtime weight planner chunks and coalesces reads under its own memory and concurrency policy. A provider with implementation limits rejects oversized requests clearly instead of truncating them. + +HTTP-specific alignment, multipart behavior, and minimum request sizes do not leak into this interface. + +### Node Filesystem Sources + +The same interface supports Node filesystem-backed sources. + +Node implementations use positional reads and normally expose a known file size. They follow the same requirements: + +- half-open ranges; +- concurrent independent positional reads; +- owned returned arrays; +- per-read abort checks; +- idempotent close that drains pending reads. + +## Adapter Precedence and Errors + +The first `Gemma4E2B` adapter resolves artifacts in this order: + +1. Use `options.artifactProvider` when supplied. +2. Otherwise reject `local_files_only: true` because runtime HTTP IO cannot guarantee Transformers.js local-only semantics. +3. Otherwise resolve the fixed `Gemma4E2B.modelId` and `revision` through runtime IO. +4. Map auth/fetch, abort, progress, and supported cache options explicitly. +5. Throw on every supplied option that would otherwise imply unsupported offline or cache behavior. + +An artifact-provider failure is propagated as-is after source cleanup. The adapter does not retry through runtime IO after an explicit provider fails because that could violate local-only, auth, revision, or cache policy. + +## Generation Capabilities + +The loaded model exposes final device-dependent capabilities. The exported backend class may also expose advisory static capabilities. + +```ts +interface GenerationCapabilitiesV1 { + readonly sessionVersion: 1; + readonly maxBatchSize: 1; + + readonly cpuModes: readonly ["greedy", "multinomial"]; + readonly planModes: readonly ["greedy"]; + readonly declarativePlans: readonly ["argmax"]; + + readonly cpuLogits: true; + readonly tokenPipeline: { + readonly defaultDepth: 4; + readonly maxDepth: 4; + }; + + readonly customJavaScriptLogitsProcessors: "cpu-fallback"; + readonly customJavaScriptStoppingCriteria: true; + readonly streamers: true; + + readonly cacheReorder: false; + readonly cacheExpand: false; + readonly returnScores: "cpu-fallback"; + readonly returnLogits: "cpu-fallback"; + readonly returnAttentions: false; + readonly returnHiddenStates: false; +} +``` + +## Generation Session Contract + +```ts +interface TokenBatch { + readonly data: Uint32Array; + readonly shape: readonly [batch: 1, sequenceLength: number]; +} + +interface PrefillInputs { + readonly inputIds: TokenBatch; + readonly attentionMask?: { + readonly data: Uint8Array; + readonly shape: readonly [batch: 1, sequenceLength: number]; + }; + readonly signal?: AbortSignal; +} + +interface DecodeInputs { + readonly tokenIds: { + readonly data: Uint32Array; + readonly shape: readonly [batch: 1, one: 1]; + }; + readonly signal?: AbortSignal; +} + +interface LogitsLeaseV1 { + readonly version: 1; + readonly dtype: "float32"; + readonly shape: readonly [batch: 1, vocabularySize: number]; + + /** Exactly one call; caller owns the returned array. */ + read(): Promise; + + /** Synchronous, idempotent borrow release. */ + release(): void; +} + +interface RuntimeGenerationPlanV1 { + readonly version: 1; + readonly processors: readonly []; + readonly sampler: { readonly op: "argmax" }; + readonly maxNewTokens: number; + readonly pipelineDepth?: number; +} + +interface RuntimeTokenDecision { + /** Owned, retainable host copy. */ + readonly tokenIds: Uint32Array; +} + +interface AutoregressiveSessionV1 { + readonly version: 1; + readonly batchSize: 1; + readonly maxSequenceLength: number; + readonly consumedTokens: number; + + prefill(inputs: PrefillInputs): Promise; + decode(inputs: DecodeInputs): Promise; + + /** + * Alternative terminal execution mode. Performs prefill internally. + * Iterator completion or return makes the session terminal. + */ + generateWithPlan( + inputs: PrefillInputs, + plan: RuntimeGenerationPlanV1, + ): AsyncIterable; + + dispose(): Promise; +} +``` + +Runtime validation rejects concurrent operations, pull/plan mode mixing, decode before prefill, decode with an active lease, repeated lease reads, and operations after terminal state. + +## Fast-Path Lifecycle + +`generateWithPlan()` performs prefill internally and is an alternative to pull mode. + +Each yield represents exactly one selected token in logical generation order. Transformers.js may stop after any yielded token. Breaking iteration invokes and awaits iterator `return()`. + +After `return()` resolves: + +- no new GPU work is submitted; +- every submitted result has an attached rejection observer; +- all submitted result promises have settled; +- mapped and staging buffers are released or destroyed; +- the session is terminal; +- `session.dispose()` completes deterministically and idempotently. + +Speculative KV writes need no rollback because the terminal session is disposed rather than reused. The runtime owns decisions and staging resources that were never yielded. Yielded `tokenIds` are owned host copies and need no release. + +## Pull-Mode Lifecycle + +`prefill()` consumes the complete prompt and returns last-position logits for the first token decision. It does not sample. + +`decode()` consumes exactly the selected token from the preceding logits and returns logits for the following decision. EOS is committed and streamed by Transformers.js but is not decoded after it terminates generation. + +The runtime owns cache position, KV writes, RoPE state, causal and sliding masks, and architecture-specific positions. Its consumed-token count is authoritative for native inference state. + +Logits remain unchanged until lease release. No overwriting session operation begins while a lease is active. `read()` may be called exactly once and returns an owned row-major `Float32Array`. Transformers.js waits for `read()` to settle and then calls `release()` in `finally` before the next decode operation. + +## Cancellation and Disposal + +Session creation, prefill, and decode accept `AbortSignal`. + +Before submission, abort rejects immediately with `signal.reason`. After submission, the session becomes terminal, stops subsequent submissions, drains submitted GPU and readback work, and rejects at a cleanup-safe boundary. + +Device loss rejects active operations and makes the loaded model unusable. Ordinary abort, controller callback errors, stopping criteria, and session input errors terminate only that session. + +`session.dispose()` is safe and idempotent after partial initialization, operation failure, cancellation, and device loss. It waits for tracked work and mapped staging buffers. + +Normal disposal is session-first. Model disposal defensively marks remaining sessions terminal, awaits their disposal, and then releases model weights and adapter-owned runtime resources. + +## Gemma4E2B Adapter + +```ts +interface Gemma4E2BLoadOptions extends InferenceBackendLoadOptions { + readonly device?: "webgpu"; + readonly dtype?: "auto"; + readonly signal?: AbortSignal; + readonly artifactProvider?: InferenceArtifactProvider; + readonly progress_callback?: (event: ArtifactProgressEvent) => void; +} + +interface Gemma4E2BLoadedModel { + readonly config: PretrainedConfig; + readonly generationCapabilities: GenerationCapabilitiesV1; + + createAutoregressiveSession(options: { + readonly batchSize: 1; + readonly maxSequenceLength: number; + readonly signal?: AbortSignal; + }): Promise; + + dispose(): Promise; +} + +export class Gemma4E2B { + static readonly modelId = "google/gemma-4-E2B-it-qat-mobile-transformers"; + static readonly generationCapabilities: StaticGenerationCapabilitiesV1; + + static load(options: Gemma4E2BLoadOptions): Promise; +} +``` + +The adapter accepts `device: 'webgpu'` and omitted/`dtype: 'auto'`. Unsupported devices and dtypes, including `q4f16`, throw without fallback. + +The adapter omits demo-owned tokenization, chat templates, text decoding, history, streamers, stopping, sampling policy, and final output formatting. Transformers.js owns those concerns. + +The Transformers.js shared config and runtime config are checked for identity-critical agreement, including vocabulary size and maximum positions. + +## Implementation Status + +No runtime contract decision is outstanding. + +Transformers.js has implemented the controller, pull-session driver, CPU logits fallback, fast argmax plan commit, iterator cleanup, capability validation, option passthrough, and public provider types. + +Implementation remains in the WebGPU package for: + +- the `Gemma4E2B` adapter; +- autoregressive session lifecycle; +- CPU logits leases; +- fast argmax iteration; +- cancellation and device-loss propagation; +- model/session cleanup tracking; +- runtime IO option mapping; +- artifact-provider adaptation. diff --git a/webgpu-kernels.local.md b/webgpu-kernels.local.md new file mode 100644 index 000000000..34ed027ca --- /dev/null +++ b/webgpu-kernels.local.md @@ -0,0 +1,310 @@ +# Custom inference backends in Transformers.js + +## Generation update + +The original model-level backend decision remains active, but the generation portion of this document is superseded by the runtime-reviewed V1 protocol in `webgpu-compat.local.md`. + +Custom generation models now expose `generationCapabilities` and `createAutoregressiveSession()`. Transformers.js installs the public `generate()`, owns generation policy and finalization, and drives either leased CPU logits or an approved declarative runtime plan. Custom runtimes should not implement public generation policy themselves. + +The initial artifact-loading agreement is also finalized there: an injected random-access provider takes precedence, while the first Gemma4 adapter may otherwise use runtime IO and must reject unsupported local-only or cache semantics explicitly. + +## Decision + +Transformers.js should treat the value passed as `model` as one of two model sources: + +1. A string model ID or local path. Transformers.js calls `OnnxInferenceProvider.from_modelId(modelId)` from `@huggingface/transformers-onnx` and keeps the existing ONNX Runtime behavior. +2. An inference backend object or class. Transformers.js loads shared assets from its `modelId`, calls its `load(options)` method, and never creates an ONNX Runtime session for the model. + +The backend boundary is at the model level, not the session level. A custom runtime may have a very different execution model, tensor representation, cache layout, or generation loop, so requiring it to imitate an ORT `InferenceSession` would leak ORT assumptions into the public contract. + +The implemented public API is: + +```js +import { pipeline } from "@huggingface/transformers"; +import { Gemma4E2B } from "@huggingface/webgpu-models"; + +const pipe = await pipeline("text-generation", Gemma4E2B, { + dtype: "auto", +}); +``` + +The imported value can be an object or a class with static members. Classes are useful for packages that export one named value per model. + +## Backend contract + +```ts +interface InferenceBackend { + /** Hub model ID or local path for config/tokenizer/processor assets. */ + readonly modelId: string; + + /** Load weights, initialize the runtime, and return a model. */ + load(options: InferenceBackendLoadOptions): Promise; +} + +interface InferenceBackendLoadOptions extends PretrainedModelOptions { + /** Always supplied by Transformers.js. */ + modelId: string; + + /** Supplied when loading through pipeline(). */ + task?: string; + + /** Resolved PretrainedConfig when loading through pipeline() or AutoModel. */ + config?: PretrainedConfig; +} +``` + +Transformers.js recognizes the contract structurally: `modelId` must be a string and `load` must be a function. No inheritance, registration, global backend selection, or dependency on an internal base class is required. + +An illustrative external model definition is: + +```js +export class Gemma4E2B { + static modelId = "google/gemma-4-e-2b"; + + static async load({ dtype, device, progress_callback, config }) { + const runtime = await WebGPUGemma.load({ + modelId: this.modelId, + dtype, + device, + progress_callback, + }); + + return { + config, + forward: (inputs) => runtime.forward(inputs), + generate: (options) => runtime.generate(options), + dispose: () => runtime.dispose(), + }; + } +} +``` + +`load()` receives a copy of the options. A backend must not rely on mutating the caller's options object. + +## Model contract + +The object returned by `load()` must implement: + +```ts +interface InferenceModel { + config?: PretrainedConfig; + + forward?(inputs: Record): Promise>; + + generate?( + options: Record, + ): Promise; + + dispose(): Promise | unknown; +} +``` + +The returned value may instead be directly callable. If it is a plain object with `forward()`, Transformers.js wraps it in a callable proxy so existing pipelines can continue to invoke `model(inputs)`. Other properties and methods, including `generate`, `config`, and `dispose`, are forwarded to the original object. + +`dispose()` is required because `Pipeline.dispose()` unconditionally delegates resource cleanup to the model. A backend owns and must release its pipelines, GPU buffers, shader modules, mapped buffers, and device resources. + +If a model does not expose `config`, Transformers.js assigns the resolved shared config after `load()`. A backend may supply its own compatible config when necessary. + +## Tensor boundary + +Pipeline inputs are Transformers.js `Tensor` objects. Model outputs consumed by existing pipelines must also be Transformers.js `Tensor` objects. + +This is the remaining shared data-plane contract. `Tensor` is currently backed by an ONNX Runtime tensor internally, so a zero-copy custom WebGPU implementation is not yet possible through every generic tensor operation. The initial custom backend should therefore do one of the following: + +1. Convert input tensors to its native representation and return Transformers.js tensors at pipeline-visible boundaries. +2. Own the complete operation, especially generation, and only return the final token IDs or task output tensors. + +A later tensor refactor can replace the `ort_tensor` field with a backend-owned native handle. That change is independent of model selection and should preserve the public `Tensor` API. + +## Task-specific requirements + +The base contract is intentionally small. Each pipeline already has a task-specific model protocol. + +### Text generation + +The model must implement `generate(options)`. Transformers.js passes tokenizer outputs and user generation options in one object: + +```js +const sequences = await model.generate({ + input_ids, + attention_mask, + max_new_tokens: 256, + ...userOptions, +}); +``` + +For decoder-only generation, return an integer `Tensor` shaped `[batch * num_return_sequences, sequence_length]` containing both prompt and generated token IDs. The text-generation pipeline decodes the complete returned sequence. + +A custom runtime should normally own its generation loop. Reusing `PreTrainedModel.generate()` currently requires ORT-style session metadata, cache input/output names, and `prepare_inputs_for_generation()` behavior, which is a much larger and less stable interface. + +If requested features are supported, `generate()` must honor streamers, stopping criteria, logits processors, sampling options, return dictionaries, and timestamp output. Unsupported options should fail clearly instead of being silently ignored. + +### Feature extraction + +The model is called with tokenizer output: + +```js +const output = await model({ input_ids, attention_mask, ...inputs }); +``` + +It must return at least one of these tensor properties: + +```ts +{ + last_hidden_state?: Tensor; + logits?: Tensor; + token_embeddings?: Tensor; +} +``` + +The selected output participates in pooling, slicing, normalization, and quantization in Transformers.js. Mean pooling also uses the tokenizer's `attention_mask`. + +### Other pipelines + +Existing pipeline classes remain authoritative. Examples: + +| Pipeline family | Required model behavior | +| ----------------------------- | ------------------------------------------------------------------------------------------- | +| Classification, QA, detection | Callable model returning the output names expected by that pipeline | +| Seq2seq generation | `generate()` plus compatible `config.prefix` and `config.task_specific_params` when present | +| Image/audio generation | `generate()` accepting processor tensors under the names used by the pipeline | +| Image feature extraction | Callable model returning `pooler_output`, `last_hidden_state`, `logits`, or `image_embeds` | +| Whisper ASR timestamps | `generate()` returning `{ sequences, token_timestamps }` when timestamps are requested | + +Supporting a task means implementing that task's existing model protocol; the backend interface does not claim that every backend supports every task. + +## Asset resolution + +`modelId` separates shared pretrained assets from inference implementation. + +Transformers.js uses it for: + +- `config.json` through `AutoConfig` +- tokenizer discovery and `AutoTokenizer` +- processor discovery and `AutoProcessor` +- Hub URL construction +- local model paths +- cache keys and file metadata +- revision, cache directory, local-only, and remote/local environment policies +- `ModelRegistry` operations that ultimately resolve model files + +The Hub and metadata boundaries normalize backend values to `modelId`, so direct calls can also reuse the descriptor: + +```js +const tokenizer = await AutoTokenizer.from_pretrained(AllMiniLML6v2); +const config = await AutoConfig.from_pretrained(AllMiniLML6v2); +const model = await AutoModel.from_pretrained(AllMiniLML6v2); +``` + +During custom pipeline construction, ONNX model files are excluded from expected-file discovery. Tokenizer and processor files are still auto-detected. The backend is responsible for discovering and downloading its own weight and kernel artifacts. + +The first version deliberately requires one shared `modelId`. If weights and tokenizer live in different repositories, the external backend can use its own weight repository internally while setting `modelId` to the repository containing the Transformers-compatible config and tokenizer. Separate `tokenizerId` or `processorId` fields should only be added when a concrete use case requires them. + +## Option ownership + +Custom `load()` receives the existing pretrained/pipeline options: + +| Option | Custom backend expectation | +| -------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------- | +| `dtype` | Backend-defined weight format. Unsupported values must throw. `Gemma4E2B` initially accepts only omitted/`auto` for its native QAT checkpoint. | +| `device` | Select an available target. A backend may support a subset and should reject unsupported values. | +| `progress_callback` | Report backend-owned weight and initialization progress using existing progress event shapes. | +| `config` | Use the resolved shared config, or return a compatible replacement on the model. | +| `cache_dir` | Reuse where the backend's platform and artifact loader support it. | +| `local_files_only` | Do not perform network access when true. | +| `revision` | Resolve backend artifacts from the requested revision. | +| `subfolder` | May be reused for backend artifacts, but defaults to `onnx` for historical compatibility. A custom backend should not assume it is meaningful. | +| `model_file_name` | Optional artifact basename override; backend-defined outside ONNX. | +| `use_external_data_format` | ONNX-specific and normally ignored by custom runtimes. | +| `session_options` | ORT-specific today. It is passed through for compatibility but custom backends should not interpret arbitrary ORT settings. | +| `task` | Pipeline task hint, present only when called through `pipeline()`. | + +Global `env` policy remains available through the normal Transformers.js export. Custom backends should honor relevant fetch/cache/offline policy rather than introducing conflicting globals. + +No generic `backend_options` was added yet. Existing options cover the immediate use case, and adding an untyped escape hatch before two runtimes need the same extension would make the contract less precise. + +## Loading sequence + +For a custom pipeline: + +```text +pipeline(task, backend, options) + -> validate backend.modelId and backend.load + -> resolve modelId + -> discover shared tokenizer/processor files (not ONNX files) + -> resolve config from modelId + -> in parallel: + AutoTokenizer.from_pretrained(modelId, options) + AutoProcessor.from_pretrained(modelId, options) + backend.load({ ...options, task, modelId, config }) + -> normalize returned model to the callable model protocol + -> construct the existing task pipeline +``` + +For a string: + +```text +pipeline(task, modelId, options) + -> existing AutoModel class selection + -> PreTrainedModel.from_pretrained(modelId, options) + -> OnnxInferenceProvider.from_modelId(modelId).load({ ...options, modelClass }) + -> existing session topology + -> ONNX adapter resolves artifacts and constructs ORT sessions +``` + +## ONNX Runtime adapter + +The ONNX-specific implementation is in the TypeScript package `packages/transformers-onnx`. `packages/transformers/src/models/session.js` remains a small runtime-neutral compatibility facade so existing model implementations do not change. + +The adapter owns: + +- device-to-ORT execution-provider mapping +- dtype-to-ONNX filename suffix selection +- WebGPU fp16 capability checks +- ONNX model and external-data artifact loading +- ORT session options and free-dimension overrides +- WebGPU preferred output locations for KV caches +- ORT session construction and execution +- WASM proxy input cloning +- Transformers.js tensor to ORT tensor conversion +- ORT output wrapping as Transformers.js tensors +- ORT-specific execution diagnostics + +`packages/transformers-onnx/src/runtime.ts` owns Node/web ORT selection, WASM loading, ORT environment defaults, and serialized browser session creation/execution. Core Transformers.js does not import ONNX Runtime packages or expose raw ORT tensors and sessions. + +Built-in model forward functions still call `sessionRun()`, and built-in model construction still calls `constructSessions()`. Those compatibility functions now delegate to normalized sessions created by `OnnxInferenceProvider`, preserving existing model implementations and ONNX behavior. + +## Why not a session contract? + +Current built-in generation reads ORT session details directly: + +- `inputNames` +- `inputMetadata` +- symbolic cache shapes +- cache input/output names +- `preferredOutputLocation` +- native tensor locations + +Making these public requirements would force a fused WebGPU runtime to expose fake sessions and fake ORT cache metadata. It would also prevent a backend from implementing a faster backend-owned generation loop. A model-level boundary keeps those details private while retaining the high-level pipeline API. + +## Errors and validation + +Transformers.js rejects malformed backends early: + +- no string `modelId` +- no `load(options)` function +- `load()` returns no model +- returned model is neither callable nor has `forward()` +- returned model has no `dispose()` + +Task-specific failures, such as a text generation model without `generate()`, surface when the corresponding pipeline invokes that operation. A future task capability declaration could move those errors to pipeline construction, but it is not required for the initial interface. + +## Current limitations and follow-ups + +1. `Tensor` is still internally coupled to ORT. A backend-neutral native tensor handle is the next major architectural step for zero-copy WebGPU interoperation. +2. Backend-owned weight files are not included in `ModelRegistry.get_pipeline_files()` because Transformers.js cannot infer an external runtime's artifact graph. The backend owns its progress and cache reporting. +3. `session_options`, `subfolder`, and external-data options retain ONNX-oriented names for compatibility. Custom backends should only reuse options with meaningful semantics. +4. Generic generation remains coupled to built-in session metadata. Custom generation backends should implement `generate()`. +5. Pipeline task compatibility is duck-typed. Capability metadata can be added later if early validation becomes valuable. + +These limitations do not block the proposed `Gemma4E2B` and `AllMiniLML6v2` API. They keep the initial integration small while establishing a clean ownership boundary between Transformers.js preprocessing/postprocessing and runtime-specific inference. From 8a992e6923872399025ea278e6e0a4b2ae4be746 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Tue, 28 Jul 2026 13:45:47 +0200 Subject: [PATCH 2/5] updated gitignore --- .gitignore | 2 + webgpu-compat.local.md | 362 ---------------------------------------- webgpu-kernels.local.md | 310 ---------------------------------- 3 files changed, 2 insertions(+), 672 deletions(-) delete mode 100644 webgpu-compat.local.md delete mode 100644 webgpu-kernels.local.md diff --git a/.gitignore b/.gitignore index 21c721c4a..964c79e0d 100644 --- a/.gitignore +++ b/.gitignore @@ -14,3 +14,5 @@ packages/*/types # Do not track coverage reports packages/*/coverage + +*.local.* \ No newline at end of file diff --git a/webgpu-compat.local.md b/webgpu-compat.local.md deleted file mode 100644 index 4466d5d76..000000000 --- a/webgpu-compat.local.md +++ /dev/null @@ -1,362 +0,0 @@ -# WebGPU runtime integration contract - -## Resolved Architecture - -Transformers.js and the WebGPU runtime have agreed on: - -- pull and plan session modes; -- terminal fast-path iteration with awaited iterator `return()`; -- leased CPU logits with synchronous idempotent release; -- prefill and decode token semantics; -- separate CPU and plan generation capabilities; -- bounded token-pipeline depth capabilities; -- batch-one, all-ones-mask V1 inputs; -- uint32 runtime token batches and Transformers.js-owned int64 sequences; -- Transformers.js-owned generation policy, callbacks, stopping, streaming, and finalization; -- runtime-owned KV cache and native inference state; -- optional generation-only models without `forward()`; -- load/session cancellation and defensive model-side session cleanup; -- native QAT dtype behavior: omitted/`auto` is supported and `q4f16` throws. - -## Default Artifact Provider - -Option A is accepted for the first adapter. - -`Gemma4E2B.load()` may use the existing runtime `ResourceRootIo` when Transformers.js does not supply an `artifactProvider`. - -Initial behavior: - -- `revision`, injected fetch/auth, abort, weight progress, and runtime browser caching are mapped where supported; -- `local_files_only` is rejected before metadata or weight loading rather than ignored; -- an explicitly supplied `artifactProvider` takes precedence over runtime IO; -- unsupported cache-policy fields are rejected or documented, never silently treated as equivalent; -- exact Transformers.js cache identity, offline policy, and shared Hub caching are deferred to a separate ranged-provider project. - -This does not block a later move to a Transformers.js-provided ranged implementation. The adapter's model-loading code consumes a small internal random-access source abstraction so runtime IO and a future Transformers.js provider share the same safetensors path. - -## Provider Contract - -```ts -interface InferenceArtifactProvider { - readJson(file: string, options?: { signal?: AbortSignal }): Promise; - - openByteSource( - file: string, - options?: { - signal?: AbortSignal; - onProgress?: (event: ArtifactProgressEvent) => void; - }, - ): Promise; -} - -interface RandomAccessByteSource { - /** May become defined after transport metadata arrives. */ - readonly size?: number; - - /** Read bytes in the half-open interval [begin, end). */ - read( - begin: number, - end: number, - options?: { signal?: AbortSignal }, - ): Promise; - - /** Idempotently wait for pending reads, then release the source. */ - close(): Promise; -} -``` - -### Range Convention - -`read(begin, end)` is half-open: `[begin, end)`. - -Both values must be non-negative safe integers and `end >= begin`. An empty range returns an owned zero-length `Uint8Array`. - -This matches the runtime `ByteSource.readRange()` contract and safetensors offsets. - -### Concurrency - -Reads may execute concurrently and complete out of order. - -The provider must not rely on one global mutable seek cursor. HTTP range requests naturally satisfy this. Filesystem implementations must use positional reads. - -A provider may serialize internally when required by its backing store, but it must preserve each read's independent range and result. - -### Returned Bytes - -Every successful `read()` returns an owned `Uint8Array` whose contents remain valid and unchanged after: - -- later reads; -- out-of-order completion of other reads; -- cache writes; -- `close()`. - -The runtime may retain returned bytes until the consuming upload, decode, or transcode operation completes. Providers must not return a view into a reused scratch buffer. - -### Closing with Pending Reads - -`close()` waits for reads that were already started, then releases the source. It does not implicitly abort pending reads. - -Once closing starts: - -- new reads reject; -- existing reads settle normally or through their supplied abort signal; -- `close()` observes all settlements before resolving; -- read failures remain owned by the corresponding read promises and do not become unhandled rejections. - -A provider may internally abort transport work during process-wide teardown, but `close()` still waits for that work to settle. - -`close()` is idempotent. Concurrent calls share the same close operation. Calls after closure resolve without repeating resource release. - -### Unknown Size - -An unknown `size` is acceptable, including at open time. - -Safetensors header parsing may read an initial probe and derive tensor ranges without total file length. When size is unknown: - -- complete-file length validation is unavailable; -- progress totals may be unknown; -- size-dependent cache metadata remains unavailable until size is discovered. - -If transport metadata later reveals the size, `size` may be a getter that changes once from `undefined` to a stable non-negative safe integer. Once defined, it must not change. - -For ranges wholly inside the file, `read()` returns exactly `end - begin` bytes. When size is unknown and an initial probe extends past EOF, returning the available shorter prefix is acceptable. Arbitrary short reads for known-valid tensor ranges are errors. - -### Validation and Limits - -The runtime imposes no alignment requirement on `begin`, `end`, or range length. - -Provider validation requirements: - -- `begin` and `end` are safe integers; -- `begin >= 0`; -- `end >= begin`; -- when `size` is known, `end <= size`; -- a successful known-valid range contains exactly `end - begin` bytes. - -There is no protocol-level maximum range size. The runtime weight planner chunks and coalesces reads under its own memory and concurrency policy. A provider with implementation limits rejects oversized requests clearly instead of truncating them. - -HTTP-specific alignment, multipart behavior, and minimum request sizes do not leak into this interface. - -### Node Filesystem Sources - -The same interface supports Node filesystem-backed sources. - -Node implementations use positional reads and normally expose a known file size. They follow the same requirements: - -- half-open ranges; -- concurrent independent positional reads; -- owned returned arrays; -- per-read abort checks; -- idempotent close that drains pending reads. - -## Adapter Precedence and Errors - -The first `Gemma4E2B` adapter resolves artifacts in this order: - -1. Use `options.artifactProvider` when supplied. -2. Otherwise reject `local_files_only: true` because runtime HTTP IO cannot guarantee Transformers.js local-only semantics. -3. Otherwise resolve the fixed `Gemma4E2B.modelId` and `revision` through runtime IO. -4. Map auth/fetch, abort, progress, and supported cache options explicitly. -5. Throw on every supplied option that would otherwise imply unsupported offline or cache behavior. - -An artifact-provider failure is propagated as-is after source cleanup. The adapter does not retry through runtime IO after an explicit provider fails because that could violate local-only, auth, revision, or cache policy. - -## Generation Capabilities - -The loaded model exposes final device-dependent capabilities. The exported backend class may also expose advisory static capabilities. - -```ts -interface GenerationCapabilitiesV1 { - readonly sessionVersion: 1; - readonly maxBatchSize: 1; - - readonly cpuModes: readonly ["greedy", "multinomial"]; - readonly planModes: readonly ["greedy"]; - readonly declarativePlans: readonly ["argmax"]; - - readonly cpuLogits: true; - readonly tokenPipeline: { - readonly defaultDepth: 4; - readonly maxDepth: 4; - }; - - readonly customJavaScriptLogitsProcessors: "cpu-fallback"; - readonly customJavaScriptStoppingCriteria: true; - readonly streamers: true; - - readonly cacheReorder: false; - readonly cacheExpand: false; - readonly returnScores: "cpu-fallback"; - readonly returnLogits: "cpu-fallback"; - readonly returnAttentions: false; - readonly returnHiddenStates: false; -} -``` - -## Generation Session Contract - -```ts -interface TokenBatch { - readonly data: Uint32Array; - readonly shape: readonly [batch: 1, sequenceLength: number]; -} - -interface PrefillInputs { - readonly inputIds: TokenBatch; - readonly attentionMask?: { - readonly data: Uint8Array; - readonly shape: readonly [batch: 1, sequenceLength: number]; - }; - readonly signal?: AbortSignal; -} - -interface DecodeInputs { - readonly tokenIds: { - readonly data: Uint32Array; - readonly shape: readonly [batch: 1, one: 1]; - }; - readonly signal?: AbortSignal; -} - -interface LogitsLeaseV1 { - readonly version: 1; - readonly dtype: "float32"; - readonly shape: readonly [batch: 1, vocabularySize: number]; - - /** Exactly one call; caller owns the returned array. */ - read(): Promise; - - /** Synchronous, idempotent borrow release. */ - release(): void; -} - -interface RuntimeGenerationPlanV1 { - readonly version: 1; - readonly processors: readonly []; - readonly sampler: { readonly op: "argmax" }; - readonly maxNewTokens: number; - readonly pipelineDepth?: number; -} - -interface RuntimeTokenDecision { - /** Owned, retainable host copy. */ - readonly tokenIds: Uint32Array; -} - -interface AutoregressiveSessionV1 { - readonly version: 1; - readonly batchSize: 1; - readonly maxSequenceLength: number; - readonly consumedTokens: number; - - prefill(inputs: PrefillInputs): Promise; - decode(inputs: DecodeInputs): Promise; - - /** - * Alternative terminal execution mode. Performs prefill internally. - * Iterator completion or return makes the session terminal. - */ - generateWithPlan( - inputs: PrefillInputs, - plan: RuntimeGenerationPlanV1, - ): AsyncIterable; - - dispose(): Promise; -} -``` - -Runtime validation rejects concurrent operations, pull/plan mode mixing, decode before prefill, decode with an active lease, repeated lease reads, and operations after terminal state. - -## Fast-Path Lifecycle - -`generateWithPlan()` performs prefill internally and is an alternative to pull mode. - -Each yield represents exactly one selected token in logical generation order. Transformers.js may stop after any yielded token. Breaking iteration invokes and awaits iterator `return()`. - -After `return()` resolves: - -- no new GPU work is submitted; -- every submitted result has an attached rejection observer; -- all submitted result promises have settled; -- mapped and staging buffers are released or destroyed; -- the session is terminal; -- `session.dispose()` completes deterministically and idempotently. - -Speculative KV writes need no rollback because the terminal session is disposed rather than reused. The runtime owns decisions and staging resources that were never yielded. Yielded `tokenIds` are owned host copies and need no release. - -## Pull-Mode Lifecycle - -`prefill()` consumes the complete prompt and returns last-position logits for the first token decision. It does not sample. - -`decode()` consumes exactly the selected token from the preceding logits and returns logits for the following decision. EOS is committed and streamed by Transformers.js but is not decoded after it terminates generation. - -The runtime owns cache position, KV writes, RoPE state, causal and sliding masks, and architecture-specific positions. Its consumed-token count is authoritative for native inference state. - -Logits remain unchanged until lease release. No overwriting session operation begins while a lease is active. `read()` may be called exactly once and returns an owned row-major `Float32Array`. Transformers.js waits for `read()` to settle and then calls `release()` in `finally` before the next decode operation. - -## Cancellation and Disposal - -Session creation, prefill, and decode accept `AbortSignal`. - -Before submission, abort rejects immediately with `signal.reason`. After submission, the session becomes terminal, stops subsequent submissions, drains submitted GPU and readback work, and rejects at a cleanup-safe boundary. - -Device loss rejects active operations and makes the loaded model unusable. Ordinary abort, controller callback errors, stopping criteria, and session input errors terminate only that session. - -`session.dispose()` is safe and idempotent after partial initialization, operation failure, cancellation, and device loss. It waits for tracked work and mapped staging buffers. - -Normal disposal is session-first. Model disposal defensively marks remaining sessions terminal, awaits their disposal, and then releases model weights and adapter-owned runtime resources. - -## Gemma4E2B Adapter - -```ts -interface Gemma4E2BLoadOptions extends InferenceBackendLoadOptions { - readonly device?: "webgpu"; - readonly dtype?: "auto"; - readonly signal?: AbortSignal; - readonly artifactProvider?: InferenceArtifactProvider; - readonly progress_callback?: (event: ArtifactProgressEvent) => void; -} - -interface Gemma4E2BLoadedModel { - readonly config: PretrainedConfig; - readonly generationCapabilities: GenerationCapabilitiesV1; - - createAutoregressiveSession(options: { - readonly batchSize: 1; - readonly maxSequenceLength: number; - readonly signal?: AbortSignal; - }): Promise; - - dispose(): Promise; -} - -export class Gemma4E2B { - static readonly modelId = "google/gemma-4-E2B-it-qat-mobile-transformers"; - static readonly generationCapabilities: StaticGenerationCapabilitiesV1; - - static load(options: Gemma4E2BLoadOptions): Promise; -} -``` - -The adapter accepts `device: 'webgpu'` and omitted/`dtype: 'auto'`. Unsupported devices and dtypes, including `q4f16`, throw without fallback. - -The adapter omits demo-owned tokenization, chat templates, text decoding, history, streamers, stopping, sampling policy, and final output formatting. Transformers.js owns those concerns. - -The Transformers.js shared config and runtime config are checked for identity-critical agreement, including vocabulary size and maximum positions. - -## Implementation Status - -No runtime contract decision is outstanding. - -Transformers.js has implemented the controller, pull-session driver, CPU logits fallback, fast argmax plan commit, iterator cleanup, capability validation, option passthrough, and public provider types. - -Implementation remains in the WebGPU package for: - -- the `Gemma4E2B` adapter; -- autoregressive session lifecycle; -- CPU logits leases; -- fast argmax iteration; -- cancellation and device-loss propagation; -- model/session cleanup tracking; -- runtime IO option mapping; -- artifact-provider adaptation. diff --git a/webgpu-kernels.local.md b/webgpu-kernels.local.md deleted file mode 100644 index 34ed027ca..000000000 --- a/webgpu-kernels.local.md +++ /dev/null @@ -1,310 +0,0 @@ -# Custom inference backends in Transformers.js - -## Generation update - -The original model-level backend decision remains active, but the generation portion of this document is superseded by the runtime-reviewed V1 protocol in `webgpu-compat.local.md`. - -Custom generation models now expose `generationCapabilities` and `createAutoregressiveSession()`. Transformers.js installs the public `generate()`, owns generation policy and finalization, and drives either leased CPU logits or an approved declarative runtime plan. Custom runtimes should not implement public generation policy themselves. - -The initial artifact-loading agreement is also finalized there: an injected random-access provider takes precedence, while the first Gemma4 adapter may otherwise use runtime IO and must reject unsupported local-only or cache semantics explicitly. - -## Decision - -Transformers.js should treat the value passed as `model` as one of two model sources: - -1. A string model ID or local path. Transformers.js calls `OnnxInferenceProvider.from_modelId(modelId)` from `@huggingface/transformers-onnx` and keeps the existing ONNX Runtime behavior. -2. An inference backend object or class. Transformers.js loads shared assets from its `modelId`, calls its `load(options)` method, and never creates an ONNX Runtime session for the model. - -The backend boundary is at the model level, not the session level. A custom runtime may have a very different execution model, tensor representation, cache layout, or generation loop, so requiring it to imitate an ORT `InferenceSession` would leak ORT assumptions into the public contract. - -The implemented public API is: - -```js -import { pipeline } from "@huggingface/transformers"; -import { Gemma4E2B } from "@huggingface/webgpu-models"; - -const pipe = await pipeline("text-generation", Gemma4E2B, { - dtype: "auto", -}); -``` - -The imported value can be an object or a class with static members. Classes are useful for packages that export one named value per model. - -## Backend contract - -```ts -interface InferenceBackend { - /** Hub model ID or local path for config/tokenizer/processor assets. */ - readonly modelId: string; - - /** Load weights, initialize the runtime, and return a model. */ - load(options: InferenceBackendLoadOptions): Promise; -} - -interface InferenceBackendLoadOptions extends PretrainedModelOptions { - /** Always supplied by Transformers.js. */ - modelId: string; - - /** Supplied when loading through pipeline(). */ - task?: string; - - /** Resolved PretrainedConfig when loading through pipeline() or AutoModel. */ - config?: PretrainedConfig; -} -``` - -Transformers.js recognizes the contract structurally: `modelId` must be a string and `load` must be a function. No inheritance, registration, global backend selection, or dependency on an internal base class is required. - -An illustrative external model definition is: - -```js -export class Gemma4E2B { - static modelId = "google/gemma-4-e-2b"; - - static async load({ dtype, device, progress_callback, config }) { - const runtime = await WebGPUGemma.load({ - modelId: this.modelId, - dtype, - device, - progress_callback, - }); - - return { - config, - forward: (inputs) => runtime.forward(inputs), - generate: (options) => runtime.generate(options), - dispose: () => runtime.dispose(), - }; - } -} -``` - -`load()` receives a copy of the options. A backend must not rely on mutating the caller's options object. - -## Model contract - -The object returned by `load()` must implement: - -```ts -interface InferenceModel { - config?: PretrainedConfig; - - forward?(inputs: Record): Promise>; - - generate?( - options: Record, - ): Promise; - - dispose(): Promise | unknown; -} -``` - -The returned value may instead be directly callable. If it is a plain object with `forward()`, Transformers.js wraps it in a callable proxy so existing pipelines can continue to invoke `model(inputs)`. Other properties and methods, including `generate`, `config`, and `dispose`, are forwarded to the original object. - -`dispose()` is required because `Pipeline.dispose()` unconditionally delegates resource cleanup to the model. A backend owns and must release its pipelines, GPU buffers, shader modules, mapped buffers, and device resources. - -If a model does not expose `config`, Transformers.js assigns the resolved shared config after `load()`. A backend may supply its own compatible config when necessary. - -## Tensor boundary - -Pipeline inputs are Transformers.js `Tensor` objects. Model outputs consumed by existing pipelines must also be Transformers.js `Tensor` objects. - -This is the remaining shared data-plane contract. `Tensor` is currently backed by an ONNX Runtime tensor internally, so a zero-copy custom WebGPU implementation is not yet possible through every generic tensor operation. The initial custom backend should therefore do one of the following: - -1. Convert input tensors to its native representation and return Transformers.js tensors at pipeline-visible boundaries. -2. Own the complete operation, especially generation, and only return the final token IDs or task output tensors. - -A later tensor refactor can replace the `ort_tensor` field with a backend-owned native handle. That change is independent of model selection and should preserve the public `Tensor` API. - -## Task-specific requirements - -The base contract is intentionally small. Each pipeline already has a task-specific model protocol. - -### Text generation - -The model must implement `generate(options)`. Transformers.js passes tokenizer outputs and user generation options in one object: - -```js -const sequences = await model.generate({ - input_ids, - attention_mask, - max_new_tokens: 256, - ...userOptions, -}); -``` - -For decoder-only generation, return an integer `Tensor` shaped `[batch * num_return_sequences, sequence_length]` containing both prompt and generated token IDs. The text-generation pipeline decodes the complete returned sequence. - -A custom runtime should normally own its generation loop. Reusing `PreTrainedModel.generate()` currently requires ORT-style session metadata, cache input/output names, and `prepare_inputs_for_generation()` behavior, which is a much larger and less stable interface. - -If requested features are supported, `generate()` must honor streamers, stopping criteria, logits processors, sampling options, return dictionaries, and timestamp output. Unsupported options should fail clearly instead of being silently ignored. - -### Feature extraction - -The model is called with tokenizer output: - -```js -const output = await model({ input_ids, attention_mask, ...inputs }); -``` - -It must return at least one of these tensor properties: - -```ts -{ - last_hidden_state?: Tensor; - logits?: Tensor; - token_embeddings?: Tensor; -} -``` - -The selected output participates in pooling, slicing, normalization, and quantization in Transformers.js. Mean pooling also uses the tokenizer's `attention_mask`. - -### Other pipelines - -Existing pipeline classes remain authoritative. Examples: - -| Pipeline family | Required model behavior | -| ----------------------------- | ------------------------------------------------------------------------------------------- | -| Classification, QA, detection | Callable model returning the output names expected by that pipeline | -| Seq2seq generation | `generate()` plus compatible `config.prefix` and `config.task_specific_params` when present | -| Image/audio generation | `generate()` accepting processor tensors under the names used by the pipeline | -| Image feature extraction | Callable model returning `pooler_output`, `last_hidden_state`, `logits`, or `image_embeds` | -| Whisper ASR timestamps | `generate()` returning `{ sequences, token_timestamps }` when timestamps are requested | - -Supporting a task means implementing that task's existing model protocol; the backend interface does not claim that every backend supports every task. - -## Asset resolution - -`modelId` separates shared pretrained assets from inference implementation. - -Transformers.js uses it for: - -- `config.json` through `AutoConfig` -- tokenizer discovery and `AutoTokenizer` -- processor discovery and `AutoProcessor` -- Hub URL construction -- local model paths -- cache keys and file metadata -- revision, cache directory, local-only, and remote/local environment policies -- `ModelRegistry` operations that ultimately resolve model files - -The Hub and metadata boundaries normalize backend values to `modelId`, so direct calls can also reuse the descriptor: - -```js -const tokenizer = await AutoTokenizer.from_pretrained(AllMiniLML6v2); -const config = await AutoConfig.from_pretrained(AllMiniLML6v2); -const model = await AutoModel.from_pretrained(AllMiniLML6v2); -``` - -During custom pipeline construction, ONNX model files are excluded from expected-file discovery. Tokenizer and processor files are still auto-detected. The backend is responsible for discovering and downloading its own weight and kernel artifacts. - -The first version deliberately requires one shared `modelId`. If weights and tokenizer live in different repositories, the external backend can use its own weight repository internally while setting `modelId` to the repository containing the Transformers-compatible config and tokenizer. Separate `tokenizerId` or `processorId` fields should only be added when a concrete use case requires them. - -## Option ownership - -Custom `load()` receives the existing pretrained/pipeline options: - -| Option | Custom backend expectation | -| -------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------- | -| `dtype` | Backend-defined weight format. Unsupported values must throw. `Gemma4E2B` initially accepts only omitted/`auto` for its native QAT checkpoint. | -| `device` | Select an available target. A backend may support a subset and should reject unsupported values. | -| `progress_callback` | Report backend-owned weight and initialization progress using existing progress event shapes. | -| `config` | Use the resolved shared config, or return a compatible replacement on the model. | -| `cache_dir` | Reuse where the backend's platform and artifact loader support it. | -| `local_files_only` | Do not perform network access when true. | -| `revision` | Resolve backend artifacts from the requested revision. | -| `subfolder` | May be reused for backend artifacts, but defaults to `onnx` for historical compatibility. A custom backend should not assume it is meaningful. | -| `model_file_name` | Optional artifact basename override; backend-defined outside ONNX. | -| `use_external_data_format` | ONNX-specific and normally ignored by custom runtimes. | -| `session_options` | ORT-specific today. It is passed through for compatibility but custom backends should not interpret arbitrary ORT settings. | -| `task` | Pipeline task hint, present only when called through `pipeline()`. | - -Global `env` policy remains available through the normal Transformers.js export. Custom backends should honor relevant fetch/cache/offline policy rather than introducing conflicting globals. - -No generic `backend_options` was added yet. Existing options cover the immediate use case, and adding an untyped escape hatch before two runtimes need the same extension would make the contract less precise. - -## Loading sequence - -For a custom pipeline: - -```text -pipeline(task, backend, options) - -> validate backend.modelId and backend.load - -> resolve modelId - -> discover shared tokenizer/processor files (not ONNX files) - -> resolve config from modelId - -> in parallel: - AutoTokenizer.from_pretrained(modelId, options) - AutoProcessor.from_pretrained(modelId, options) - backend.load({ ...options, task, modelId, config }) - -> normalize returned model to the callable model protocol - -> construct the existing task pipeline -``` - -For a string: - -```text -pipeline(task, modelId, options) - -> existing AutoModel class selection - -> PreTrainedModel.from_pretrained(modelId, options) - -> OnnxInferenceProvider.from_modelId(modelId).load({ ...options, modelClass }) - -> existing session topology - -> ONNX adapter resolves artifacts and constructs ORT sessions -``` - -## ONNX Runtime adapter - -The ONNX-specific implementation is in the TypeScript package `packages/transformers-onnx`. `packages/transformers/src/models/session.js` remains a small runtime-neutral compatibility facade so existing model implementations do not change. - -The adapter owns: - -- device-to-ORT execution-provider mapping -- dtype-to-ONNX filename suffix selection -- WebGPU fp16 capability checks -- ONNX model and external-data artifact loading -- ORT session options and free-dimension overrides -- WebGPU preferred output locations for KV caches -- ORT session construction and execution -- WASM proxy input cloning -- Transformers.js tensor to ORT tensor conversion -- ORT output wrapping as Transformers.js tensors -- ORT-specific execution diagnostics - -`packages/transformers-onnx/src/runtime.ts` owns Node/web ORT selection, WASM loading, ORT environment defaults, and serialized browser session creation/execution. Core Transformers.js does not import ONNX Runtime packages or expose raw ORT tensors and sessions. - -Built-in model forward functions still call `sessionRun()`, and built-in model construction still calls `constructSessions()`. Those compatibility functions now delegate to normalized sessions created by `OnnxInferenceProvider`, preserving existing model implementations and ONNX behavior. - -## Why not a session contract? - -Current built-in generation reads ORT session details directly: - -- `inputNames` -- `inputMetadata` -- symbolic cache shapes -- cache input/output names -- `preferredOutputLocation` -- native tensor locations - -Making these public requirements would force a fused WebGPU runtime to expose fake sessions and fake ORT cache metadata. It would also prevent a backend from implementing a faster backend-owned generation loop. A model-level boundary keeps those details private while retaining the high-level pipeline API. - -## Errors and validation - -Transformers.js rejects malformed backends early: - -- no string `modelId` -- no `load(options)` function -- `load()` returns no model -- returned model is neither callable nor has `forward()` -- returned model has no `dispose()` - -Task-specific failures, such as a text generation model without `generate()`, surface when the corresponding pipeline invokes that operation. A future task capability declaration could move those errors to pipeline construction, but it is not required for the initial interface. - -## Current limitations and follow-ups - -1. `Tensor` is still internally coupled to ORT. A backend-neutral native tensor handle is the next major architectural step for zero-copy WebGPU interoperation. -2. Backend-owned weight files are not included in `ModelRegistry.get_pipeline_files()` because Transformers.js cannot infer an external runtime's artifact graph. The backend owns its progress and cache reporting. -3. `session_options`, `subfolder`, and external-data options retain ONNX-oriented names for compatibility. Custom backends should only reuse options with meaningful semantics. -4. Generic generation remains coupled to built-in session metadata. Custom generation backends should implement `generate()`. -5. Pipeline task compatibility is duck-typed. Capability metadata can be added later if early validation becomes valuable. - -These limitations do not block the proposed `Gemma4E2B` and `AllMiniLML6v2` API. They keep the initial integration small while establishing a clean ownership boundary between Transformers.js preprocessing/postprocessing and runtime-specific inference. From a453c9824f0d0ea98dbbdbe22306798ab6f3a734 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Tue, 28 Jul 2026 16:00:48 +0200 Subject: [PATCH 3/5] some fixes --- packages/transformers-onnx/src/host.ts | 81 +++++-- packages/transformers-onnx/src/provider.ts | 115 +++++---- packages/transformers-onnx/src/runtime.ts | 87 ++++--- packages/transformers-onnx/src/tensor-ops.ts | 5 +- .../transformers-onnx/tests/provider.test.js | 36 ++- packages/transformers-onnx/tsconfig.json | 2 +- .../transformers/docs/source/_toctree.yml | 2 + .../docs/source/guides/custom-backends.md | 74 ++++++ .../transformers/src/backends/artifacts.js | 24 +- packages/transformers/src/backends/default.js | 29 ++- .../transformers/src/backends/inference.js | 107 ++++++++- .../src/backends/model_registry.js | 22 ++ .../transformers/src/generation/controller.js | 10 +- .../transformers/src/generation/runtime.js | 221 +++++++++++++++--- .../transformers/src/models/modeling_utils.js | 16 +- packages/transformers/src/ops/registry.js | 3 +- packages/transformers/src/pipelines.js | 50 +++- .../transformers/src/tokenization_utils.js | 2 +- packages/transformers/src/transformers.js | 18 +- packages/transformers/src/utils/hub.js | 18 +- .../src/utils/model_registry/ModelRegistry.js | 10 +- .../model_registry/get_available_dtypes.js | 19 +- .../utils/model_registry/get_file_metadata.js | 22 +- .../src/utils/model_registry/get_files.js | 24 +- .../utils/model_registry/get_model_files.js | 22 +- .../model_registry/get_pipeline_files.js | 11 +- .../model_registry/get_processor_files.js | 5 +- .../model_registry/get_tokenizer_files.js | 5 +- packages/transformers/src/utils/tensor.js | 10 +- .../tests/generation_controller.test.js | 193 ++++++++++++++- .../tests/inference_backends.test.js | 31 ++- .../tests/onnx_provider_object.test.js | 33 +++ .../transformers/tests/onnx_wiring.test.js | 30 +++ packages/transformers/tests/types.test.js | 66 ++++++ .../tests/utils/generation.test.js | 15 ++ .../tests/utils/hub_abort.test.js | 22 ++ .../tests/utils/model_registry.test.js | 15 ++ .../transformers/tests/utils/tensor.test.js | 4 + types/webgpu-kernels.local.demo.d.ts | 14 -- types/webgpu-kernels.local.demo.d.ts.map | 1 - 40 files changed, 1225 insertions(+), 249 deletions(-) create mode 100644 packages/transformers/docs/source/guides/custom-backends.md create mode 100644 packages/transformers/src/backends/model_registry.js create mode 100644 packages/transformers/tests/onnx_provider_object.test.js create mode 100644 packages/transformers/tests/onnx_wiring.test.js create mode 100644 packages/transformers/tests/utils/hub_abort.test.js delete mode 100644 types/webgpu-kernels.local.demo.d.ts delete mode 100644 types/webgpu-kernels.local.demo.d.ts.map diff --git a/packages/transformers-onnx/src/host.ts b/packages/transformers-onnx/src/host.ts index c8edc0ce5..1cfa1ecdd 100644 --- a/packages/transformers-onnx/src/host.ts +++ b/packages/transformers-onnx/src/host.ts @@ -9,37 +9,74 @@ export interface BackendTensorStorage { dispose(): void; } +export interface OnnxProviderEnvironment { + backends: Record; + logLevel?: number; + useWasmCache?: boolean; + fetch: typeof globalThis.fetch; +} + +export interface OnnxProviderApis { + readonly IS_NODE_ENV: boolean; + readonly IS_WEB_ENV: boolean; + readonly IS_WEBGPU_AVAILABLE: boolean; + readonly IS_WEBNN_AVAILABLE: boolean; + readonly IS_DENO_WEB_RUNTIME: boolean; + readonly IS_SAFARI_BELOW_26: boolean; + readonly IS_SERVICE_WORKER_ENV: boolean; + readonly IS_CHROME_AVAILABLE: boolean; +} + +export interface OnnxProviderLogger { + info(...data: unknown[]): void; + warn(...data: unknown[]): void; + error(...data: unknown[]): void; +} + +export interface OnnxProviderCache { + match(request: string): Promise } | undefined>; + put(request: string, response: Response): Promise; +} + export interface OnnxProviderHost { - readonly env: any; - readonly apis: any; - readonly logger: any; - getModelFile(modelId: string, file: string, fatal: boolean, options: any, returnPath?: boolean): Promise; - getCacheNames(config: any, options: any): Set; - createBackendTensor(storage: BackendTensorStorage): any; - getBackendTensorStorage(tensor: any): BackendTensorStorage | null; - getCache?(): Promise; + readonly env: OnnxProviderEnvironment; + readonly apis: OnnxProviderApis; + readonly logger: OnnxProviderLogger; + getModelFile( + modelId: string, + file: string, + fatal: boolean, + options: Record, + returnPath?: boolean, + ): Promise; + getCacheNames(config: unknown, options: unknown): Set; + createBackendTensor(storage: BackendTensorStorage): unknown; + getBackendTensorStorage(tensor: unknown): BackendTensorStorage | null; + getCache?(): Promise; + readonly maxExternalDataChunks: number; } -let configuredHost: OnnxProviderHost | null = null; +const ONNX_HOST_SYMBOL = Symbol.for('transformers.js.onnxProviderHost'); +let configuredHost: OnnxProviderHost | null = + ((globalThis as any)[ONNX_HOST_SYMBOL] as OnnxProviderHost | undefined) ?? null; -const fallbackEnvironment: any = { +const fallbackEnvironment: OnnxProviderEnvironment = { backends: { onnx: {} }, logLevel: 30, useWasmCache: typeof caches !== 'undefined', - fetch: (...args: any[]) => (globalThis.fetch as any)(...args), + fetch: (...args) => globalThis.fetch(...args), }; const environment = new Proxy(fallbackEnvironment, { get(target, property) { - return (configuredHost?.env ?? target)[property]; + return Reflect.get(configuredHost?.env ?? target, property); }, set(target, property, value) { - (configuredHost?.env ?? target)[property] = value; - return true; + return Reflect.set(configuredHost?.env ?? target, property, value); }, }); -const apis = { +const fallbackApis = { IS_NODE_ENV: typeof process !== 'undefined' && process?.release?.name === 'node', IS_WEB_ENV: typeof window !== 'undefined' || typeof self !== 'undefined', IS_WEBGPU_AVAILABLE: typeof navigator !== 'undefined' && !!navigator.gpu, @@ -51,9 +88,15 @@ const apis = { IS_CHROME_AVAILABLE: 'chrome' in globalThis, }; +const apis = new Proxy(fallbackApis, { + get(target, property) { + return Reflect.get(configuredHost?.apis ?? target, property); + }, +}); + const logger = new Proxy(console, { get(target, property) { - return (configuredHost?.logger ?? target)[property]; + return Reflect.get(configuredHost?.logger ?? target, property); }, }); @@ -73,10 +116,14 @@ const fallbackHost: OnnxProviderHost = { getBackendTensorStorage() { return null; }, + maxExternalDataChunks: 100, }; export function configureOnnxProviderHost(host: OnnxProviderHost): void { - if (fallbackEnvironment.backends.onnx) { + // Direct package imports initialize ORT against the fallback environment, so migrate those + // settings when a host arrives later. A symbol-registered host was already configured before + // module evaluation and contains the authoritative ORT environment. + if (configuredHost === null && fallbackEnvironment.backends.onnx) { host.env.backends.onnx = fallbackEnvironment.backends.onnx; } configuredHost = host; diff --git a/packages/transformers-onnx/src/provider.ts b/packages/transformers-onnx/src/provider.ts index 3ce63a344..fbc2ea738 100644 --- a/packages/transformers-onnx/src/provider.ts +++ b/packages/transformers-onnx/src/provider.ts @@ -91,10 +91,6 @@ async function isWebGpuFp16Supported(): Promise { * ONNX Runtime adapter used for string model IDs. */ export class OnnxInferenceProvider { - /** - * @param {string} modelId - * @param {typeof import('../../models/modeling_utils.js').PreTrainedModel} [modelClass] - */ static from_modelId(modelId: string): OnnxInferenceProvider { return new OnnxInferenceProvider(modelId); } @@ -163,12 +159,9 @@ export class OnnxInferenceProvider { this.modelClass = modelClass; } - /** - * Load a Transformers.js model class with this backend. - * - * @param {import('../../utils/hub.js').PretrainedModelOptions} options - */ + /** Load a Transformers.js model class with this backend. */ async load(options: any) { + throwIfAborted(options.signal); const modelClass = options.modelClass ?? this.modelClass; if (!modelClass) { throw new Error('OnnxInferenceProvider requires a Transformers.js model class before it can load a model.'); @@ -183,14 +176,11 @@ export class OnnxInferenceProvider { cache_config = false, session_name: string | undefined = undefined, ) { + throwIfAborted(options.signal); let custom_config = options.config?.['transformers.js_config'] ?? {}; - const selectedDevice = /** @type {import('../../utils/devices.js').DeviceType} */ selectDevice( - options.device ?? custom_config.device, - fileName, - { - warn: (msg) => logger.info(msg), - }, - ); + const selectedDevice = selectDevice(options.device ?? custom_config.device, fileName, { + warn: (msg: string) => logger.info(msg), + }); const executionProviders = deviceToExecutionProviders(selectedDevice); const device_config = custom_config.device_config ?? {}; @@ -200,7 +190,7 @@ export class OnnxInferenceProvider { const selectedDtype = selectDtype(options.dtype ?? custom_config.dtype, fileName, selectedDevice, { configDtype: custom_config.dtype, - warn: (msg) => logger.info(msg), + warn: (msg: string) => logger.info(msg), }); if (!Object.hasOwn(DEFAULT_DTYPE_SUFFIX_MAPPING, selectedDtype)) { throw new Error(`Invalid dtype: ${selectedDtype}. Should be one of: ${Object.keys(DATA_TYPES).join(', ')}`); @@ -238,6 +228,7 @@ export class OnnxInferenceProvider { use_external_data_format, session_options, ); + throwIfAborted(options.signal); if (externalData.length > 0 && (!apis.IS_NODE_ENV || externalData.some((data) => typeof data !== 'string'))) { session_options.externalData = externalData; } @@ -245,47 +236,61 @@ export class OnnxInferenceProvider { if (cache_config && selectedDevice === 'webgpu') { const names = getOnnxProviderHost().getCacheNames(options.config, { prefix: 'present', session_name }); if (names.size > 0 && !isONNXProxy()) { - const preferredOutputLocation = {}; + const preferredOutputLocation: Record = {}; for (const key of names) preferredOutputLocation[key] = 'gpu-buffer'; session_options.preferredOutputLocation = preferredOutputLocation; } } + const buffer_or_path = await bufferOrPathPromise; + throwIfAborted(options.signal); return { - buffer_or_path: await bufferOrPathPromise, + buffer_or_path, session_options, session_config: { dtype: selectedDtype, device: selectedDevice }, }; } async constructSessions(names: Record, options: any, cache_sessions: any = undefined) { - return Object.fromEntries( - await Promise.all( - Object.keys(names).map(async (name) => { - const sessionInfo = await this.getSession( - names[name], - options, - cache_sessions?.[name] ?? false, - name, - ); - const ortSession = await createInferenceSession( - sessionInfo.buffer_or_path, - sessionInfo.session_options, - sessionInfo.session_config, - ); - const session = { - inputNames: ortSession.inputNames, - outputNames: ortSession.outputNames, - inputMetadata: ortSession.inputMetadata, - outputMetadata: ortSession.outputMetadata, - config: (ortSession as any).config, - run: (inputs: any) => this.run(ortSession, inputs), - release: () => ortSession.release(), - }; - return [name, session]; - }), - ), - ); + const tasks = Object.keys(names).map(async (name) => { + const sessionInfo = await this.getSession(names[name], options, cache_sessions?.[name] ?? false, name); + const ortSession = await createInferenceSession( + sessionInfo.buffer_or_path, + sessionInfo.session_options, + sessionInfo.session_config, + ); + try { + throwIfAborted(options.signal); + } catch (error) { + await ortSession.release(); + throw error; + } + return [ + name, + { + inputNames: ortSession.inputNames, + outputNames: ortSession.outputNames, + inputMetadata: ortSession.inputMetadata, + outputMetadata: ortSession.outputMetadata, + config: (ortSession as any).config, + run: (inputs: any) => this.run(ortSession, inputs), + release: () => ortSession.release(), + }, + ] as const; + }); + const results = await Promise.allSettled(tasks); + const sessions = results + .filter( + (result): result is PromiseFulfilledResult> => + result.status === 'fulfilled', + ) + .map((result) => result.value); + const failure = results.find((result): result is PromiseRejectedResult => result.status === 'rejected'); + if (failure) { + await Promise.allSettled(sessions.map(([, session]) => session.release())); + throw failure.reason; + } + return Object.fromEntries(sessions); } async run(session: any, inputs: Record) { @@ -335,10 +340,10 @@ function replaceTensors(value: any): any { return tensor.type; }, get dims() { - return tensor.dims; + return tensor.dims as number[]; }, set dims(value) { - tensor.dims = value; + (tensor as unknown as { dims: number[] }).dims = value; }, get data() { return tensor.data; @@ -403,6 +408,12 @@ function externalDataChunkNames(fullName: string, count: number): string[] { return Array.from({ length: count }, (_, index) => `${fullName}_data${index === 0 ? '' : `_${index}`}`); } +function throwIfAborted(signal: AbortSignal | undefined): void { + if (!signal?.aborted) return; + if (typeof signal.throwIfAborted === 'function') signal.throwIfAborted(); + throw signal.reason ?? new Error('Model loading aborted.'); +} + async function getModelDataFiles( modelId: string, fileName: string, @@ -418,8 +429,12 @@ async function getModelDataFiles( } else if (externalConfig) { count = +externalConfig; } - if (count > 1024) - throw new Error(`The number of external data chunks (${count}) exceeds the maximum allowed value (1024).`); + const maxChunks = getOnnxProviderHost().maxExternalDataChunks; + if (count > maxChunks) { + throw new Error( + `The number of external data chunks (${count}) exceeds the maximum allowed value (${maxChunks}).`, + ); + } if (count > 0) { const subfolder = options.subfolder ?? 'onnx'; diff --git a/packages/transformers-onnx/src/runtime.ts b/packages/transformers-onnx/src/runtime.ts index df1fe5c58..05c9c233e 100644 --- a/packages/transformers-onnx/src/runtime.ts +++ b/packages/transformers-onnx/src/runtime.ts @@ -22,6 +22,7 @@ import { getOnnxProviderHost } from './host.js'; // In either case, we select the default export if it exists, otherwise we use the named export. import * as ONNX_NODE from 'onnxruntime-node'; import * as ONNX_WEB from 'onnxruntime-web/webgpu'; +import type { Env, InferenceSession as OrtInferenceSession, Tensor as OrtTensor } from 'onnxruntime-common'; import { loadWasmBinary, loadWasmFactory } from './wasm-cache.js'; export { Tensor } from 'onnxruntime-common'; @@ -36,12 +37,9 @@ function toAbsoluteURL(url: string): string { return new URL(url, globalThis.location?.href ?? 'file:///').href; } -/** - * @typedef {import('onnxruntime-common').InferenceSession.ExecutionProviderConfig} ONNXExecutionProviders - */ +type ExecutionProvider = OrtInferenceSession.ExecutionProviderConfig; -/** @type {Record} */ -const DEVICE_TO_EXECUTION_PROVIDER_MAPPING = Object.freeze({ +const DEVICE_TO_EXECUTION_PROVIDER_MAPPING: Readonly> = Object.freeze({ auto: null, // Auto-detect based on device and environment gpu: null, // Auto-detect GPU cpu: 'cpu', // CPU @@ -64,7 +62,7 @@ const DEVICE_TO_EXECUTION_PROVIDER_MAPPING = Object.freeze({ * @param {number} logLevel - The LogLevel value to convert * @returns {number} ONNX Runtime severity level (0-4) */ -function getOnnxLogSeverityLevel(logLevel) { +function getOnnxLogSeverityLevel(logLevel: number): 0 | 1 | 2 | 3 | 4 { // ONNX Runtime's log severity levels are defined as follows: // (0) ORT_LOGGING_LEVEL_VERBOSE: Print all log messages. // (1) ORT_LOGGING_LEVEL_INFO: Print info and higher level log messages. @@ -88,11 +86,7 @@ function getOnnxLogSeverityLevel(logLevel) { } } -/** - * Maps ONNX Runtime numeric severity levels to string log levels. - * @type {Record<0 | 1 | 2 | 3 | 4, 'verbose' | 'info' | 'warning' | 'error' | 'fatal'>} - */ -const ONNX_LOG_LEVEL_NAMES = { +const ONNX_LOG_LEVEL_NAMES: Record<0 | 1 | 2 | 3 | 4, 'verbose' | 'info' | 'warning' | 'error' | 'fatal'> = { 0: 'verbose', 1: 'info', 2: 'warning', @@ -100,20 +94,16 @@ const ONNX_LOG_LEVEL_NAMES = { 4: 'fatal', }; -/** - * The list of supported devices, sorted by priority/performance. - * @type {import("../utils/devices.js").DeviceType[]} - */ -const supportedDevices = []; +// The list of supported devices, sorted by priority/performance. +const supportedDevices: string[] = []; -/** @type {ONNXExecutionProviders[]} */ -let defaultDevices; -let ONNX; +let defaultDevices: ExecutionProvider[]; +let ONNX: typeof ONNX_NODE; const ORT_SYMBOL = Symbol.for('onnxruntime'); if (ORT_SYMBOL in globalThis) { // If the JS runtime exposes their own ONNX runtime, use it - ONNX = globalThis[ORT_SYMBOL]; + ONNX = (globalThis as any)[ORT_SYMBOL] as typeof ONNX_NODE; } else if (apis.IS_NODE_ENV) { ONNX = ONNX_NODE; @@ -144,7 +134,7 @@ if (ORT_SYMBOL in globalThis) { supportedDevices.push('cpu'); defaultDevices = ['cpu']; } else { - ONNX = ONNX_WEB; + ONNX = ONNX_WEB as unknown as typeof ONNX_NODE; if (apis.IS_WEBNN_AVAILABLE) { // TODO: Only push supported providers (depending on available hardware) @@ -159,15 +149,10 @@ if (ORT_SYMBOL in globalThis) { defaultDevices = ['wasm']; } -// @ts-ignore const InferenceSession = ONNX.InferenceSession; -/** - * Map a device to the execution providers to use for the given device. - * @param {import("../utils/devices.js").DeviceType|"auto"|null} [device=null] (Optional) The device to run the inference on. - * @returns {ONNXExecutionProviders[]} The execution providers to use for the given device. - */ -export function deviceToExecutionProviders(device = null) { +/** Map a device to the execution providers to use for the given device. */ +export function deviceToExecutionProviders(device: string | null = null): ExecutionProvider[] { // Use the default execution providers if the user hasn't specified anything if (!device) return defaultDevices; @@ -189,16 +174,14 @@ export function deviceToExecutionProviders(device = null) { /** * Currently, Transformers.js doesn't support simultaneous loading of sessions in WASM/WebGPU. * For this reason, we need to chain the loading calls. - * @type {Promise} */ -let webInitChain = Promise.resolve(); +let webInitChain: Promise = Promise.resolve(); /** * Promise that resolves when WASM binary has been loaded (if caching is enabled). * This ensures we only attempt to load the WASM binary once. - * @type {Promise|null} */ -let wasmLoadPromise = null; +let wasmLoadPromise: Promise | null = null; /** * Ensures the WASM binary is loaded and cached before creating an inference session. @@ -237,7 +220,7 @@ async function ensureWasmLoaded() { wasmLoadPromise = (async () => { // At this point, we know wasmPaths is an object (not a string) because // shouldUseWasmCache checks for wasmPaths.wasm and wasmPaths.mjs - const urls = /** @type {{ wasm: string, mjs: string }} */ ONNX_ENV.wasm.wasmPaths; + const urls = ONNX_ENV.wasm.wasmPaths as { wasm: string; mjs: string }; // Load both in parallel; the .mjs blob URL is only kept if wasmBinary succeeded. // ORT only sets locateFile when wasmBinary is provided (onnxruntime PR https://github.com/microsoft/onnxruntime/pull/27411), which @@ -292,26 +275,34 @@ async function ensureWasmLoaded() { * @param {Object} session_config ONNX inference session configuration. * @returns {Promise} The ONNX inference session. */ -export async function createInferenceSession(buffer_or_path, session_options, session_config) { +export async function createInferenceSession( + buffer_or_path: Uint8Array | string, + session_options: OrtInferenceSession.SessionOptions, + session_config: Record, +): Promise }> { await ensureWasmLoaded(); const logSeverityLevel = getOnnxLogSeverityLevel(env.logLevel ?? LogLevel.WARNING); - const load = () => - InferenceSession.create(buffer_or_path, { + const load = () => { + const options = { // Set default log severity level, but allow overriding through session options logSeverityLevel, ...session_options, - }); + }; + return typeof buffer_or_path === 'string' + ? InferenceSession.create(buffer_or_path, options) + : InferenceSession.create(buffer_or_path, options); + }; const session = await (apis.IS_WEB_ENV ? (webInitChain = webInitChain.then(load)) : load()); - session.config = session_config; - return session; + const configuredSession = session as OrtInferenceSession & { config: Record }; + configuredSession.config = session_config; + return configuredSession; } /** * Currently, Transformers.js doesn't support simultaneous execution of sessions in WASM/WebGPU. * For this reason, we need to chain the inference calls (otherwise we get "Error: Session already started"). - * @type {Promise} */ -let webInferenceChain = Promise.resolve(); +let webInferenceChain: Promise> = Promise.resolve({}); /** * Run an inference session. @@ -319,21 +310,23 @@ let webInferenceChain = Promise.resolve(); * @param {Record} ortFeed The input tensors. * @returns {Promise>} The output tensors. */ -export async function runInferenceSession(session, ortFeed) { +export async function runInferenceSession( + session: OrtInferenceSession, + ortFeed: Record, +): Promise> { const run = () => session.run(ortFeed); return apis.IS_WEB_ENV ? (webInferenceChain = webInferenceChain.then(run)) : run(); } /** * Check if an object is an ONNX tensor. - * @param {any} x The object to check + * @param x The object to check * @returns {boolean} Whether the object is an ONNX tensor. */ -export function isONNXTensor(x) { +export function isONNXTensor(x: unknown): x is OrtTensor { return x instanceof ONNX.Tensor; } -/** @type {import('onnxruntime-common').Env} */ -const ONNX_ENV = ONNX?.env; +const ONNX_ENV: Env = ONNX.env; /** * Check if ONNX's WASM backend is being proxied. @@ -384,7 +377,7 @@ if (ONNX_ENV) { * levels, and set the log level environment variable in ONNX Runtime. * @param {number} logLevel The log level to set. */ - function setLogLevel(logLevel) { + function setLogLevel(logLevel: number) { const severityLevel = getOnnxLogSeverityLevel(logLevel); ONNX_ENV.logLevel = ONNX_LOG_LEVEL_NAMES[severityLevel]; } diff --git a/packages/transformers-onnx/src/tensor-ops.ts b/packages/transformers-onnx/src/tensor-ops.ts index 2688796ea..1f0c0ba7d 100644 --- a/packages/transformers-onnx/src/tensor-ops.ts +++ b/packages/transformers-onnx/src/tensor-ops.ts @@ -1,4 +1,5 @@ import { createInferenceSession, runInferenceSession, isONNXProxy, Tensor as OrtTensor } from './runtime.js'; +import type { Tensor as OrtTensorType } from 'onnxruntime-common'; import { getOnnxProviderHost } from './host.js'; /** @@ -15,7 +16,7 @@ import { getOnnxProviderHost } from './host.js'; const wrap = async (session_bytes: number[], session_options: any, names: string | string[]) => { const session = await createInferenceSession(new Uint8Array(session_bytes), session_options, {}); - return /** @type {any} */ async (inputs: Record) => { + return async (inputs: Record) => { const proxied = isONNXProxy(); const ortFeed = Object.fromEntries( Object.entries(inputs).map(([key, value]) => { @@ -26,7 +27,7 @@ const wrap = async (session_bytes: number[], session_options: any, names: string storage?.backend === 'onnx' ? storage.handle : new OrtTensor(input.type, input.data, input.dims), ]; }), - ); + ) as Record; const outputs = await runInferenceSession(session, ortFeed); if (Array.isArray(names)) { return names.map((name) => wrapTensor(outputs[name])); diff --git a/packages/transformers-onnx/tests/provider.test.js b/packages/transformers-onnx/tests/provider.test.js index fcceb7647..f8462a31d 100644 --- a/packages/transformers-onnx/tests/provider.test.js +++ b/packages/transformers-onnx/tests/provider.test.js @@ -1,4 +1,4 @@ -import { OnnxInferenceProvider } from "@huggingface/transformers-onnx"; +import { configureOnnxProviderHost, OnnxInferenceProvider } from "@huggingface/transformers-onnx"; describe("OnnxInferenceProvider", () => { it("creates providers from model IDs", () => { @@ -8,4 +8,38 @@ describe("OnnxInferenceProvider", () => { expect(provider.providerType).toBe("onnx"); expect(typeof provider.constructSessions).toBe("function"); }); + + it("uses the host external-data chunk limit", async () => { + configureOnnxProviderHost({ + env: { + backends: { onnx: {} }, + logLevel: 30, + useWasmCache: false, + fetch: globalThis.fetch, + }, + apis: { + IS_NODE_ENV: true, + IS_WEB_ENV: false, + IS_WEBGPU_AVAILABLE: false, + IS_WEBNN_AVAILABLE: false, + IS_DENO_WEB_RUNTIME: false, + IS_SAFARI_BELOW_26: false, + IS_SERVICE_WORKER_ENV: false, + IS_CHROME_AVAILABLE: false, + }, + logger: console, + getModelFile: async () => new Uint8Array(), + getCacheNames: () => new Set(), + createBackendTensor: () => null, + getBackendTensorStorage: () => null, + maxExternalDataChunks: 2, + }); + + const provider = OnnxInferenceProvider.from_modelId("onnx-community/test-model"); + await expect( + provider.getSession("model", { + config: { "transformers.js_config": { use_external_data_format: 3 } }, + }), + ).rejects.toThrow("exceeds the maximum allowed value (2)"); + }); }); diff --git a/packages/transformers-onnx/tsconfig.json b/packages/transformers-onnx/tsconfig.json index 0caf6d078..ede70dc64 100644 --- a/packages/transformers-onnx/tsconfig.json +++ b/packages/transformers-onnx/tsconfig.json @@ -6,7 +6,7 @@ "moduleResolution": "bundler", "outDir": "types", "rootDir": "src", - "strict": false, + "strict": true, "skipLibCheck": true, "declaration": true, "declarationMap": true, diff --git a/packages/transformers/docs/source/_toctree.yml b/packages/transformers/docs/source/_toctree.yml index 76d779249..b712b8f9c 100644 --- a/packages/transformers/docs/source/_toctree.yml +++ b/packages/transformers/docs/source/_toctree.yml @@ -27,6 +27,8 @@ - sections: - local: guides/webgpu title: Running models on WebGPU + - local: guides/custom-backends + title: Custom inference backends - local: guides/dtypes title: Using quantized models (dtypes) - local: guides/private diff --git a/packages/transformers/docs/source/guides/custom-backends.md b/packages/transformers/docs/source/guides/custom-backends.md new file mode 100644 index 000000000..e75ea7c93 --- /dev/null +++ b/packages/transformers/docs/source/guides/custom-backends.md @@ -0,0 +1,74 @@ +# Custom inference backends + +Transformers.js can load curated model backends that use a runtime other than the default ONNX provider. A pipeline-facing backend supplies a fixed model ID and a `load()` function: + +```ts +import type { + CausalGenerationCapabilitiesV1, + InferenceBackend, + InferenceModel, + PlanAutoregressiveSessionV1, +} from '@huggingface/transformers'; + +const causalGeneration = { + sessionVersion: 1, + maxBatchSize: 1, + cpuLogits: false, + cpuModes: [], + planModes: ['greedy'], + declarativePlans: ['argmax'], + tokenPipeline: { defaultDepth: 4, maxDepth: 4 }, +} as const satisfies CausalGenerationCapabilitiesV1; + +export const backend: InferenceBackend = { + modelId: 'organization/curated-model', + capabilities: { + devices: ['webgpu'], + dtypes: ['auto'], + tasks: ['text-generation'], + }, + async load(options): Promise { + // Load weights and create runtime-owned model state here. + return { + capabilities: { causalGeneration }, + async createAutoregressiveSession(sessionOptions): Promise { + throw new Error('Example only'); + }, + async dispose() {}, + }; + }, +}; +``` + +Pass the backend anywhere a model ID is accepted: + +```js +const generator = await pipeline('text-generation', backend, { + device: 'webgpu', + dtype: 'auto', + signal, + artifactProvider, +}); +``` + +## Capabilities + +Static backend capabilities are advisory and allow `pipeline()` to reject unsupported tasks early. The loaded model's `capabilities` are authoritative. Execution families are optional and task-specific: decoder-only generation uses `causalGeneration`; forward-only, encoder-decoder, diffusion, and audio runtimes should not imitate the causal protocol. + +Plan-only causal runtimes may declare `cpuLogits: false`, an empty `cpuModes`, and a supported greedy plan. Requests requiring JavaScript logits processors, sampling, or returned scores cannot use that plan and fail before a session is created. Pull sessions expose CPU logits leases for full Transformers.js policy compatibility. + +## Session lifecycle + +Transformers.js owns generation policy, stopping, callbacks, streaming, and final output construction. The runtime owns inference state and the KV cache. + +- `generateWithPlan()` yields ordered token decisions. Transformers.js awaits the iterator's `return()` when generation stops early. +- A logits lease must remain valid until its synchronous, idempotent `release()` is called. +- Session `dispose()` is always awaited after generation. +- `sessionConcurrency.maxActiveSessions` is enforced fail-fast. Callers that want queueing must serialize generation themselves. +- Abort signals are forwarded through loading and session creation. Backends must release partially created resources before propagating cancellation. + +## Artifact providers + +An `InferenceArtifactProvider` can supply JSON and random-access byte sources. Byte ranges are half-open (`[begin, end)`), independent reads may complete out of order, and returned arrays must be owned by the caller. `close()` is idempotent, rejects new reads, and waits for existing reads without implicitly aborting them. + +An explicit artifact provider takes precedence inside the custom backend. Provider failures must be propagated after cleanup rather than retried through another transport. diff --git a/packages/transformers/src/backends/artifacts.js b/packages/transformers/src/backends/artifacts.js index 0c9290e2c..0e5de5edf 100644 --- a/packages/transformers/src/backends/artifacts.js +++ b/packages/transformers/src/backends/artifacts.js @@ -13,8 +13,8 @@ /** * @typedef {Object} RandomAccessByteSource * @property {number} [size] Stable byte length when known. It may initially be undefined and become defined after transport metadata arrives. - * @property {(begin: number, end: number, options?: {signal?: AbortSignal}) => Promise} read Read an independent half-open byte range `[begin, end)`. The returned array is owned by the caller and remains valid after later reads and close. - * @property {() => Promise} close Idempotently reject new reads, drain reads already in progress, and release the source. + * @property {(begin: number, end: number, options?: {signal?: AbortSignal}) => Promise} read Read an independent half-open byte range `[begin, end)`. Bounds are non-negative safe integers with `end >= begin`. Reads may complete out of order, and every result is an owned array that remains valid after later reads and close. + * @property {() => Promise} close Idempotently reject new reads, wait for reads already in progress to settle, and release the source without implicitly aborting those reads. */ /** @@ -23,4 +23,22 @@ * @property {(file: string, options?: {signal?: AbortSignal, onProgress?: (event: ArtifactProgressEvent) => void}) => Promise} openByteSource Open a source supporting concurrent, independently positioned reads. */ -export {}; +/** + * Validate the provider shape without opening an artifact. Range and ownership conformance is + * validated by the consuming runtime, which owns the source lifecycle. + * + * @param {unknown} provider + * @returns {asserts provider is InferenceArtifactProvider|undefined} + */ +export function validateInferenceArtifactProvider(provider) { + if (provider === undefined) return; + const candidate = /** @type {any} */ (provider); + if ( + provider === null || + typeof provider !== 'object' || + typeof candidate.readJson !== 'function' || + typeof candidate.openByteSource !== 'function' + ) { + throw new TypeError('`artifactProvider` must implement `readJson()` and `openByteSource()`.'); + } +} diff --git a/packages/transformers/src/backends/default.js b/packages/transformers/src/backends/default.js index 31f668192..1cbd6851c 100644 --- a/packages/transformers/src/backends/default.js +++ b/packages/transformers/src/backends/default.js @@ -1,14 +1,13 @@ -import { OnnxInferenceProvider, OnnxTensorOpRegistry, configureOnnxProviderHost } from '@huggingface/transformers-onnx'; - import { env, apis } from '../env.js'; import { logger } from '../utils/logger.js'; -import { getModelFile } from '../utils/hub.js'; +import { getModelFile, MAX_EXTERNAL_DATA_CHUNKS } from '../utils/hub.js'; import { getCacheNames } from '../configs.js'; import { Tensor } from '../utils/tensor.js'; import { TensorOpRegistry } from '../ops/registry.js'; import { getCache } from '../utils/cache.js'; -configureOnnxProviderHost({ +const ONNX_HOST_SYMBOL = Symbol.for('transformers.js.onnxProviderHost'); +const host = { env, apis, logger, @@ -17,8 +16,24 @@ configureOnnxProviderHost({ createBackendTensor: (storage) => Tensor.fromBackendStorage(storage), getBackendTensorStorage: (tensor) => tensor?.getBackendStorage?.() ?? null, getCache, -}); + maxExternalDataChunks: MAX_EXTERNAL_DATA_CHUNKS, +}; + +let modulePromise; -TensorOpRegistry.register(OnnxTensorOpRegistry); +export function getOnnxProviderModule() { + if (!modulePromise) { + globalThis[ONNX_HOST_SYMBOL] = host; + modulePromise = import('@huggingface/transformers-onnx').then((module) => { + module.configureOnnxProviderHost(host); + TensorOpRegistry.register(module.OnnxTensorOpRegistry); + return module; + }); + } + return modulePromise; +} -export { OnnxInferenceProvider }; +export async function getDefaultInferenceProvider(modelId) { + const { OnnxInferenceProvider } = await getOnnxProviderModule(); + return OnnxInferenceProvider.from_modelId(modelId); +} diff --git a/packages/transformers/src/backends/inference.js b/packages/transformers/src/backends/inference.js index 4e6fe15a3..3a5c71fb8 100644 --- a/packages/transformers/src/backends/inference.js +++ b/packages/transformers/src/backends/inference.js @@ -18,22 +18,72 @@ * @module backends/inference */ -import { installGenerationRuntime } from '../generation/runtime.js'; +import { getCausalGenerationCapabilities, installGenerationRuntime } from '../generation/runtime.js'; +import { validateInferenceArtifactProvider } from './artifacts.js'; + +/** + * @typedef {Object} ForwardCapabilitiesV1 + * @property {1} version + */ + +/** + * Execution capabilities of a loaded inference model. Capability families without a versioned + * Transformers.js integration remain opaque until their task-specific session contract is defined. + * + * @typedef {Object} InferenceModelCapabilities + * @property {ForwardCapabilitiesV1} [forward] + * @property {import('../generation/runtime.js').CausalGenerationCapabilitiesV1} [causalGeneration] + * @property {Readonly>} [encoderDecoderGeneration] + * @property {Readonly>} [diffusion] + * @property {Readonly>} [audioGeneration] + */ + +/** + * Advisory capabilities available before a curated backend is loaded. + * `load()` remains authoritative for device and dtype validation. + * + * @typedef {Object} StaticBackendCapabilities + * @property {ReadonlyArray} devices + * @property {ReadonlyArray} dtypes + * @property {ReadonlyArray} tasks + */ + +/** + * @typedef {import('../utils/hub.js').PretrainedModelOptions & { + * modelId: string, + * task?: string, + * config?: import('../configs.js').PretrainedConfig, + * modelClass?: Function, + * generation_config?: Record, + * }} InferenceBackendLoadOptions + */ + +/** + * @typedef {import('../utils/hub.js').PretrainedModelOptions & { + * task?: string, + * config?: import('../configs.js').PretrainedConfig, + * modelClass?: Function, + * generation_config?: Record, + * }} InferenceModelLoadOptions + */ /** * @typedef {Object} InferenceModel * @property {(inputs: Record) => Promise>} [forward] - * @property {(options: Object) => Promise} [generate] - * @property {import('../generation/runtime.js').GenerationCapabilitiesV1} [generationCapabilities] - * @property {(options: Object) => Promise} [createAutoregressiveSession] - * @property {Object} [config] + * @property {(options: Record) => Promise>} [generate] + * @property {InferenceModelCapabilities} [capabilities] + * @property {import('../generation/runtime.js').GenerationCapabilitiesV1} [generationCapabilities] Deprecated flat causal-generation capabilities. + * @property {(options: import('../generation/runtime.js').AutoregressiveSessionOptionsV1) => Promise} [createAutoregressiveSession] + * @property {import('../configs.js').PretrainedConfig} [config] * @property {() => Promise|unknown} dispose */ /** * @typedef {Object} InferenceBackend * @property {string} modelId Model ID or local path used for shared config, tokenizer, and processor assets. - * @property {(options: any) => Promise} load + * @property {string} [providerType] Provider family identifier used for provider-specific host initialization. + * @property {StaticBackendCapabilities} [capabilities] + * @property {(options: InferenceBackendLoadOptions) => Promise} load * @property {(names: Record, options: Object, cacheSessions?: Object) => Promise>} [constructSessions] */ @@ -66,6 +116,40 @@ export function getModelId(model) { throw new TypeError('Model must be a model ID string or an inference backend with `modelId` and `load(options)`.'); } +/** + * Reject a task excluded by authoritative static backend metadata. + * + * @param {InferenceBackend} backend + * @param {string} task + */ +export function validateInferenceBackendTask(backend, task) { + const tasks = backend.capabilities?.tasks; + if (!tasks) return; + const canonicalTask = task.split('_', 1)[0]; + if (!tasks.includes(task) && !tasks.includes(canonicalTask)) { + throw new Error(`Inference backend "${backend.modelId}" does not support the "${task}" task.`); + } +} + +/** + * Validate the loaded execution capability required by an integrated task. + * + * @param {InferenceModel|Function} model + * @param {string} task + */ +export function validateInferenceModelTask(model, task) { + const implementation = /** @type {any} */ (model); + if (!implementation.capabilities) return; + if (task === 'text-generation') { + if ( + !getCausalGenerationCapabilities(implementation) || + typeof implementation.createAutoregressiveSession !== 'function' + ) { + throw new Error('The loaded inference model does not support causal text generation.'); + } + } +} + /** * Make a plain model with `forward()` callable, matching the model contract used by pipelines. * @@ -80,6 +164,14 @@ export function normalizeInferenceModel(model) { if (typeof implementation.dispose !== 'function') { throw new TypeError('Inference backend models must implement `dispose()`.'); } + if ( + implementation.capabilities?.causalGeneration && + typeof implementation.createAutoregressiveSession !== 'function' + ) { + throw new TypeError( + 'Models declaring `capabilities.causalGeneration` must implement `createAutoregressiveSession(options)`.', + ); + } if (typeof model === 'function') return model; if ( typeof implementation.forward !== 'function' && @@ -115,13 +207,14 @@ export function normalizeInferenceModel(model) { * Load and normalize a custom inference model. * * @param {InferenceBackend} backend - * @param {Object} options + * @param {InferenceModelLoadOptions} options * @returns {Promise} */ export async function loadInferenceModel(backend, options) { const loadOptions = { ...options, modelId: backend.modelId }; if (loadOptions.device === null) loadOptions.device = undefined; if (loadOptions.dtype === null) loadOptions.dtype = undefined; + validateInferenceArtifactProvider(loadOptions.artifactProvider); const model = /** @type {any} */ ( installGenerationRuntime(normalizeInferenceModel(await backend.load(loadOptions))) ); diff --git a/packages/transformers/src/backends/model_registry.js b/packages/transformers/src/backends/model_registry.js new file mode 100644 index 000000000..8bf10e9ad --- /dev/null +++ b/packages/transformers/src/backends/model_registry.js @@ -0,0 +1,22 @@ +import { getOnnxProviderModule } from './default.js'; + +/** + * Provider operations used by model-file discovery. + * + * @typedef {Object} ModelRegistryInferenceProvider + * @property {(options: Object) => string[]} listModelArtifacts + * @property {(options: Object) => Promise} getAvailableDtypes + * @property {(files: string[], sessions: Record) => string[]} filterModelArtifacts + */ + +/** + * Resolve an explicitly supplied registry provider or the lazy default provider. + * + * @param {ModelRegistryInferenceProvider|null} [provider] + * @returns {Promise} + */ +export async function getModelRegistryInferenceProvider(provider = null) { + if (provider) return provider; + const { OnnxInferenceProvider } = await getOnnxProviderModule(); + return OnnxInferenceProvider; +} diff --git a/packages/transformers/src/generation/controller.js b/packages/transformers/src/generation/controller.js index 2b49fa438..e959e90c6 100644 --- a/packages/transformers/src/generation/controller.js +++ b/packages/transformers/src/generation/controller.js @@ -225,11 +225,12 @@ export class GenerationController { } else { throw new Error(`Generation logits must have rank 2 or 3, received rank ${logitsInput.dims.length}.`); } - if (logits.dims[0] !== this.batchSize) { - throw new Error(`Generation logits batch size ${logits.dims[0]} does not match ${this.batchSize}.`); - } - const processed = this.logitsProcessor(this.sequences, logits); + if (processed.dims[0] !== this.batchSize) { + throw new Error( + `Processed generation logits batch size ${processed.dims[0]} does not match ${this.batchSize}.`, + ); + } const tokenIds = new Uint32Array(this.batchSize); const tokenScores = new Float64Array(this.batchSize); for (let batchIndex = 0; batchIndex < this.batchSize; ++batchIndex) { @@ -282,6 +283,7 @@ export class GenerationController { if (!capabilities?.declarativePlans?.includes('argmax')) return null; if (!capabilities?.planModes?.includes('greedy')) return null; if (this.generationConfig.do_sample || this.generationConfig.num_beams > 1) return null; + if (this.generationConfig.output_scores) return null; if (this.logitsProcessor.processors.length !== 0) return null; return { version: 1, diff --git a/packages/transformers/src/generation/runtime.js b/packages/transformers/src/generation/runtime.js index 8269d8e1d..5d835d94c 100644 --- a/packages/transformers/src/generation/runtime.js +++ b/packages/transformers/src/generation/runtime.js @@ -2,40 +2,121 @@ import { Tensor } from '../utils/tensor.js'; import { createGenerationController } from './controller.js'; /** - * @typedef {Object} GenerationCapabilitiesV1 + * @typedef {Object} SessionConcurrencyCapabilities + * @property {number} maxActiveSessions + * @property {1} concurrentOperationsPerSession + */ + +/** + * @typedef {Object} CausalGenerationCapabilitiesV1 * @property {1} sessionVersion * @property {number} maxBatchSize - * @property {string[]} cpuModes - * @property {string[]} planModes + * @property {ReadonlyArray} cpuModes + * @property {ReadonlyArray} planModes * @property {boolean} cpuLogits - * @property {string[]} declarativePlans - * @property {{defaultDepth: number, maxDepth: number}} tokenPipeline - * @property {boolean} customJavaScriptStoppingCriteria - * @property {false} cacheReorder - * @property {false} cacheExpand + * @property {ReadonlyArray<'argmax'>} declarativePlans + * @property {{readonly defaultDepth: number, readonly maxDepth: number}} tokenPipeline + * @property {boolean} [customJavaScriptStoppingCriteria] + * @property {false} [cacheReorder] + * @property {false} [cacheExpand] + * @property {SessionConcurrencyCapabilities} [sessionConcurrency] */ +/** @typedef {CausalGenerationCapabilitiesV1} GenerationCapabilitiesV1 */ + /** * @typedef {Object} LogitsLeaseV1 * @property {1} version * @property {'float32'} dtype - * @property {[number, number]} shape + * @property {readonly [number, number]} shape * @property {() => Promise} read - * @property {(plan: Object) => Promise<{tokenIds: Uint32Array, processedScores?: Float32Array}>} [select] + * @property {(plan: RuntimeGenerationPlanV1) => Promise} [select] * @property {() => void} release */ /** - * @typedef {Object} AutoregressiveSessionV1 + * @typedef {Object} RuntimeTokenBatchV1 + * @property {Uint32Array} data + * @property {readonly [number, number]} shape + */ + +/** + * @typedef {Object} RuntimeAttentionMaskV1 + * @property {Uint8Array} data + * @property {readonly [number, number]} shape + */ + +/** + * @typedef {Object} AutoregressivePrefillInputsV1 + * @property {RuntimeTokenBatchV1} inputIds + * @property {RuntimeAttentionMaskV1} [attentionMask] + * @property {AbortSignal} [signal] + */ + +/** + * @typedef {Object} AutoregressiveDecodeInputsV1 + * @property {RuntimeTokenBatchV1} tokenIds + * @property {AbortSignal} [signal] + */ + +/** + * @typedef {Object} RuntimeGenerationPlanV1 + * @property {1} version + * @property {ReadonlyArray} processors + * @property {{readonly op: 'argmax'}} sampler + * @property {number} maxNewTokens + * @property {number} [pipelineDepth] + */ + +/** + * @typedef {Object} RuntimeTokenDecisionV1 + * @property {Uint32Array} tokenIds + * @property {Float32Array} [processedScores] + * @property {Float64Array} [scores] + */ + +/** + * @typedef {Object} AutoregressiveSessionOptionsV1 + * @property {number} batchSize + * @property {number} maxSequenceLength + * @property {AbortSignal} [signal] + */ + +/** + * @typedef {Object} PlanAutoregressiveSessionV1 * @property {1} version * @property {number} batchSize * @property {number} maxSequenceLength - * @property {(inputs: Object) => Promise} prefill - * @property {(inputs: Object) => Promise} decode - * @property {(inputs: Object, plan: Object) => AsyncIterable<{tokenIds: Uint32Array, processedScores?: Float32Array}>} [generateWithPlan] + * @property {(inputs: AutoregressivePrefillInputsV1, plan: RuntimeGenerationPlanV1) => AsyncIterable} generateWithPlan * @property {() => Promise} dispose */ +/** + * @typedef {Object} PullAutoregressiveSessionV1 + * @property {1} version + * @property {number} batchSize + * @property {number} maxSequenceLength + * @property {(inputs: AutoregressivePrefillInputsV1) => Promise} prefill + * @property {(inputs: AutoregressiveDecodeInputsV1) => Promise} decode + * @property {(inputs: AutoregressivePrefillInputsV1, plan: RuntimeGenerationPlanV1) => AsyncIterable} [generateWithPlan] + * @property {() => Promise} dispose + */ + +/** @typedef {PlanAutoregressiveSessionV1|PullAutoregressiveSessionV1} AutoregressiveSessionV1 */ + +/** @type {WeakMap} */ +const activeSessionCounts = new WeakMap(); + +/** + * Resolve the causal-generation capability while accepting the original flat V1 field. + * + * @param {Object|Function} model + * @returns {CausalGenerationCapabilitiesV1|undefined} + */ +export function getCausalGenerationCapabilities(model) { + return model?.capabilities?.causalGeneration ?? model?.generationCapabilities; +} + /** * Install the Transformers.js-owned public generation method on a custom model. * @@ -67,7 +148,7 @@ export async function generateWithAutoregressiveSession(model, options) { const controller = createGenerationController(model, input_ids, options); if (controller.allDone) return controller.finalize(); - const capabilities = model.generationCapabilities; + const capabilities = getCausalGenerationCapabilities(model); try { validateCapabilities(capabilities, controller, attention_mask); throwIfAborted(signal); @@ -76,6 +157,7 @@ export async function generateWithAutoregressiveSession(model, options) { throw error; } const plan = controller.compileRuntimePlan(capabilities); + const mode = controller.generationConfig.do_sample ? 'multinomial' : 'greedy'; if (!plan && !capabilities.cpuLogits) { const error = new Error( 'This generation request requires CPU-visible logits, but the runtime does not support them.', @@ -83,6 +165,19 @@ export async function generateWithAutoregressiveSession(model, options) { controller.abort(error); throw error; } + if (!plan && !capabilities.cpuModes.includes(mode)) { + const error = new Error(`Runtime does not support ${mode} generation through its CPU logits path.`); + controller.abort(error); + throw error; + } + + let releaseSessionSlot; + try { + releaseSessionSlot = acquireSessionSlot(model, capabilities); + } catch (error) { + controller.abort(error); + throw error; + } /** @type {AutoregressiveSessionV1|null} */ let session = null; @@ -94,15 +189,16 @@ export async function generateWithAutoregressiveSession(model, options) { maxSequenceLength: controller.maxSequenceLength, signal, }); - validateSession(session, controller); + validateSession(session, controller, plan !== null); const prefillInputs = { inputIds: tensorToTokenBatch(input_ids), attentionMask: attention_mask ? tensorToAttentionMask(attention_mask) : undefined, signal, }; - if (plan && typeof session.generateWithPlan === 'function') { - const decisions = session.generateWithPlan(prefillInputs, plan)[Symbol.asyncIterator](); + if (plan) { + const planSession = /** @type {PlanAutoregressiveSessionV1} */ (session); + const decisions = planSession.generateWithPlan(prefillInputs, plan)[Symbol.asyncIterator](); try { while (true) { const item = await decisions.next(); @@ -125,26 +221,29 @@ export async function generateWithAutoregressiveSession(model, options) { ); } - lease = await session.prefill(prefillInputs); + const pullSession = /** @type {PullAutoregressiveSessionV1} */ (session); + lease = await pullSession.prefill(prefillInputs); while (!controller.allDone) { throwIfAborted(signal); const currentLease = lease; lease = null; - validateLease(currentLease, controller.batchSize); - let values; try { + validateLease(currentLease, controller.batchSize); values = await currentLease.read(); } finally { - currentLease.release(); + currentLease?.release?.(); } if (!(values instanceof Float32Array)) { throw new TypeError('Logits lease `read()` must return a Float32Array.'); } + if (values.length !== currentLease.shape[0] * currentLease.shape[1]) { + throw new Error('Logits lease data length does not match its declared shape.'); + } - const step = await controller.step(new Tensor('float32', values, currentLease.shape)); + const step = await controller.step(new Tensor('float32', values, [...currentLease.shape])); if (step.allDone) break; - lease = await session.decode({ + lease = await pullSession.decode({ tokenIds: tensorToTokenBatch(step.nextTokenIds), signal, }); @@ -154,8 +253,15 @@ export async function generateWithAutoregressiveSession(model, options) { controller.abort(error); throw error; } finally { - lease?.release(); - await session?.dispose(); + try { + lease?.release(); + } finally { + try { + await session?.dispose(); + } finally { + releaseSessionSlot(); + } + } } } @@ -163,6 +269,24 @@ function validateCapabilities(capabilities, controller, attentionMask) { if (!capabilities || capabilities.sessionVersion !== 1) { throw new Error('Custom generation models must declare generation capabilities with `sessionVersion: 1`.'); } + if (!Number.isInteger(capabilities.maxBatchSize) || capabilities.maxBatchSize < 1) { + throw new Error('Runtime must declare a positive integer `maxBatchSize`.'); + } + if (!Array.isArray(capabilities.cpuModes) || !Array.isArray(capabilities.planModes)) { + throw new Error('Runtime must declare `cpuModes` and `planModes` arrays.'); + } + if (typeof capabilities.cpuLogits !== 'boolean') { + throw new Error('Runtime must declare whether CPU-visible logits are supported.'); + } + if (!capabilities.cpuLogits && capabilities.cpuModes.length > 0) { + throw new Error('Runtime cannot declare CPU generation modes when `cpuLogits` is false.'); + } + if (!Array.isArray(capabilities.declarativePlans)) { + throw new Error('Runtime must declare a `declarativePlans` array.'); + } + if (capabilities.planModes.includes('greedy') && !capabilities.declarativePlans.includes('argmax')) { + throw new Error('Runtime greedy plan mode requires the `argmax` declarative plan.'); + } if (controller.batchSize > capabilities.maxBatchSize) { throw new Error( `Runtime supports batch size ${capabilities.maxBatchSize}, but generation received ${controller.batchSize}.`, @@ -190,17 +314,26 @@ function validateCapabilities(capabilities, controller, attentionMask) { } const tokenPipeline = capabilities.tokenPipeline; if ( - tokenPipeline && - (!Number.isInteger(tokenPipeline.defaultDepth) || - !Number.isInteger(tokenPipeline.maxDepth) || - tokenPipeline.defaultDepth < 1 || - tokenPipeline.defaultDepth > tokenPipeline.maxDepth) + !tokenPipeline || + !Number.isInteger(tokenPipeline.defaultDepth) || + !Number.isInteger(tokenPipeline.maxDepth) || + tokenPipeline.defaultDepth < 1 || + tokenPipeline.defaultDepth > tokenPipeline.maxDepth ) { throw new Error('Runtime declared an invalid token pipeline depth.'); } + const concurrency = capabilities.sessionConcurrency; + if ( + concurrency && + (!Number.isInteger(concurrency.maxActiveSessions) || + concurrency.maxActiveSessions < 1 || + concurrency.concurrentOperationsPerSession !== 1) + ) { + throw new Error('Runtime declared invalid session concurrency capabilities.'); + } } -function validateSession(session, controller) { +function validateSession(session, controller, usePlan) { if (!session || session.version !== 1) throw new Error('Runtime returned an unsupported autoregressive session.'); if (session.batchSize !== controller.batchSize) { throw new Error(`Runtime session batch size ${session.batchSize} does not match ${controller.batchSize}.`); @@ -211,6 +344,12 @@ function validateSession(session, controller) { ); } if (typeof session.dispose !== 'function') throw new Error('Autoregressive sessions must implement `dispose()`.'); + if (usePlan && typeof session.generateWithPlan !== 'function') { + throw new Error('Plan autoregressive sessions must implement `generateWithPlan()`.'); + } + if (!usePlan && (typeof session.prefill !== 'function' || typeof session.decode !== 'function')) { + throw new Error('Pull autoregressive sessions must implement `prefill()` and `decode()`.'); + } } function validateLease(lease, batchSize) { @@ -259,3 +398,21 @@ function throwIfAborted(signal) { if (typeof signal.throwIfAborted === 'function') signal.throwIfAborted(); throw signal.reason ?? new Error('Generation aborted.'); } + +function acquireSessionSlot(model, capabilities) { + const limit = capabilities.sessionConcurrency?.maxActiveSessions; + if (limit === undefined) return () => {}; + const active = activeSessionCounts.get(model) ?? 0; + if (active >= limit) { + throw new Error(`Runtime supports at most ${limit} active autoregressive session${limit === 1 ? '' : 's'}.`); + } + activeSessionCounts.set(model, active + 1); + let released = false; + return () => { + if (released) return; + released = true; + const remaining = (activeSessionCounts.get(model) ?? 1) - 1; + if (remaining > 0) activeSessionCounts.set(model, remaining); + else activeSessionCounts.delete(model); + }; +} diff --git a/packages/transformers/src/models/modeling_utils.js b/packages/transformers/src/models/modeling_utils.js index e57dacd54..80a2d4c95 100644 --- a/packages/transformers/src/models/modeling_utils.js +++ b/packages/transformers/src/models/modeling_utils.js @@ -42,7 +42,7 @@ import { get_model_files } from '../utils/model_registry/get_model_files.js'; import { get_file_metadata } from '../utils/model_registry/get_file_metadata.js'; import { MODEL_SESSION_CONFIG, MODEL_TYPES } from './session_config.js'; import { getModelId, isInferenceBackend, loadInferenceModel } from '../backends/inference.js'; -import { OnnxInferenceProvider } from '../backends/default.js'; +import { getDefaultInferenceProvider, getOnnxProviderModule } from '../backends/default.js'; /** * Converts an array or Tensor of integers to an int64 Tensor. @@ -265,13 +265,17 @@ export class PreTrainedModel extends Callable { */ static async from_pretrained(pretrained_model_name_or_path, options = {}) { if (typeof pretrained_model_name_or_path === 'string') { - return OnnxInferenceProvider.from_modelId(pretrained_model_name_or_path).load({ + const provider = await getDefaultInferenceProvider(pretrained_model_name_or_path); + return provider.load({ ...options, modelClass: this, }); } if (typeof pretrained_model_name_or_path?.constructSessions === 'function') { - return /** @type {any} */ (pretrained_model_name_or_path.load({ ...options, modelClass: this })); + if (pretrained_model_name_or_path.providerType === 'onnx') { + await getOnnxProviderModule(); + } + return /** @type {any} */ (pretrained_model_name_or_path).load({ ...options, modelClass: this }); } if (isInferenceBackend(pretrained_model_name_or_path)) { const modelId = getModelId(pretrained_model_name_or_path); @@ -978,7 +982,11 @@ export class PreTrainedModel extends Callable { }, }); - if (controller.allDone) return controller.finalize(); + if (controller.allDone) { + return generation_config.return_dict_in_generate + ? controller.finalize({ past_key_values: kwargs.past_key_values ?? new DynamicCache() }) + : controller.finalize(); + } let outputs; try { diff --git a/packages/transformers/src/ops/registry.js b/packages/transformers/src/ops/registry.js index 2d06cdddd..2c819fcbd 100644 --- a/packages/transformers/src/ops/registry.js +++ b/packages/transformers/src/ops/registry.js @@ -36,7 +36,8 @@ export class TensorOpRegistry { async function getOperation(name) { if (!implementation) { - await import('../backends/default.js'); + const { getOnnxProviderModule } = await import('../backends/default.js'); + await getOnnxProviderModule(); } if (!implementation) { throw new Error(`Tensor operation "${name}" requires an installed inference provider.`); diff --git a/packages/transformers/src/pipelines.js b/packages/transformers/src/pipelines.js index 99bdf180e..3c43b4595 100644 --- a/packages/transformers/src/pipelines.js +++ b/packages/transformers/src/pipelines.js @@ -51,7 +51,14 @@ import { } from './pipelines/index.js'; import { get_pipeline_files } from './utils/model_registry/get_pipeline_files.js'; import { get_file_metadata } from './utils/model_registry/get_file_metadata.js'; -import { getModelId, isInferenceBackend, loadInferenceModel } from './backends/inference.js'; +import { + getModelId, + isInferenceBackend, + loadInferenceModel, + validateInferenceBackendTask, + validateInferenceModelTask, +} from './backends/inference.js'; +import { validateInferenceArtifactProvider } from './backends/artifacts.js'; import { getModelJSON } from './utils/hub.js'; /** @@ -136,12 +143,20 @@ export async function pipeline( const customBackend = isInferenceBackend(model) && typeof model.constructSessions !== 'function'; const modelId = getModelId(model); + validateInferenceArtifactProvider(artifactProvider); + if (customBackend) { + validateInferenceBackendTask(/** @type {import('./backends/inference.js').InferenceBackend} */ (model), task); + } // Determine which files the model needs const expected_files = await get_pipeline_files(task, modelId, { device, dtype, config, + cache_dir, + local_files_only, + revision, + model_file_name, include_model: !customBackend, }); @@ -149,7 +164,11 @@ export async function pipeline( let files_loading = {}; if (progress_callback) { /** @type {Array<{exists: boolean, size?: number, contentType?: string, fromCache?: boolean}>} */ - const metadata = await Promise.all(expected_files.map(async (file) => get_file_metadata(modelId, file))); + const metadata = await Promise.all( + expected_files.map(async (file) => + get_file_metadata(modelId, file, { cache_dir, local_files_only, revision }), + ), + ); metadata.forEach((m, i) => { if (m.exists) { files_loading[expected_files[i]] = { @@ -215,12 +234,27 @@ export async function pipeline( modelPromise = modelClasses.from_pretrained(modelId, pretrainedOptions); } - // Load all components in parallel - const [tokenizer, processor, model_loaded] = await Promise.all([ - hasTokenizer ? AutoTokenizer.from_pretrained(modelId, pretrainedOptions) : null, - hasProcessor ? AutoProcessor.from_pretrained(modelId, pretrainedOptions) : null, - modelPromise, - ]); + let tokenizer; + let processor; + let model_loaded; + try { + // Load all components in parallel. + [tokenizer, processor, model_loaded] = await Promise.all([ + hasTokenizer ? AutoTokenizer.from_pretrained(modelId, pretrainedOptions) : null, + hasProcessor ? AutoProcessor.from_pretrained(modelId, pretrainedOptions) : null, + modelPromise, + ]); + if (customBackend) { + validateInferenceModelTask(model_loaded, task); + } + } catch (error) { + // A parallel tokenizer/processor failure may race with a successful GPU model load. + const loadedModel = model_loaded ?? (await modelPromise.catch(() => null)); + try { + await loadedModel?.dispose?.(); + } catch {} + throw error; + } const results = { task, model: model_loaded }; if (tokenizer) results.tokenizer = tokenizer; diff --git a/packages/transformers/src/tokenization_utils.js b/packages/transformers/src/tokenization_utils.js index 29de6f186..5f816078e 100644 --- a/packages/transformers/src/tokenization_utils.js +++ b/packages/transformers/src/tokenization_utils.js @@ -26,7 +26,7 @@ import { get_tokenizer_files } from './utils/model_registry/get_tokenizer_files. * @returns {Promise} A promise that resolves with information about the loaded tokenizer. */ export async function loadTokenizer(pretrained_model_name_or_path, options) { - const tokenizerFiles = await get_tokenizer_files(pretrained_model_name_or_path); + const tokenizerFiles = await get_tokenizer_files(pretrained_model_name_or_path, options); return await Promise.all( tokenizerFiles.map((file) => getModelJSON(pretrained_model_name_or_path, file, true, options)), ); diff --git a/packages/transformers/src/transformers.js b/packages/transformers/src/transformers.js index 587d1e47a..e142c7f37 100644 --- a/packages/transformers/src/transformers.js +++ b/packages/transformers/src/transformers.js @@ -61,7 +61,6 @@ export { ModelRegistry } from './utils/model_registry/ModelRegistry.js'; // Inference backends export { getModelId, isInferenceBackend } from './backends/inference.js'; -export { OnnxInferenceProvider } from './backends/default.js'; // Expose common types used across the library for developers to access /** @@ -73,8 +72,25 @@ export { OnnxInferenceProvider } from './backends/default.js'; * @typedef {import('./utils/devices.js').DeviceType} DeviceType * @typedef {import('./utils/core.js').ProgressCallback} ProgressCallback * @typedef {import('./utils/core.js').ProgressInfo} ProgressInfo + * @typedef {import('./backends/inference.js').InferenceBackend} InferenceBackend + * @typedef {import('./backends/inference.js').InferenceBackendLoadOptions} InferenceBackendLoadOptions + * @typedef {import('./backends/inference.js').InferenceModel} InferenceModel + * @typedef {import('./backends/inference.js').InferenceModelCapabilities} InferenceModelCapabilities + * @typedef {import('./backends/inference.js').StaticBackendCapabilities} StaticBackendCapabilities + * @typedef {import('./backends/inference.js').ForwardCapabilitiesV1} ForwardCapabilitiesV1 + * @typedef {import('./generation/runtime.js').CausalGenerationCapabilitiesV1} CausalGenerationCapabilitiesV1 * @typedef {import('./generation/runtime.js').GenerationCapabilitiesV1} GenerationCapabilitiesV1 * @typedef {import('./generation/runtime.js').AutoregressiveSessionV1} AutoregressiveSessionV1 + * @typedef {import('./generation/runtime.js').PlanAutoregressiveSessionV1} PlanAutoregressiveSessionV1 + * @typedef {import('./generation/runtime.js').PullAutoregressiveSessionV1} PullAutoregressiveSessionV1 + * @typedef {import('./generation/runtime.js').SessionConcurrencyCapabilities} SessionConcurrencyCapabilities + * @typedef {import('./generation/runtime.js').AutoregressiveSessionOptionsV1} AutoregressiveSessionOptionsV1 + * @typedef {import('./generation/runtime.js').AutoregressivePrefillInputsV1} AutoregressivePrefillInputsV1 + * @typedef {import('./generation/runtime.js').AutoregressiveDecodeInputsV1} AutoregressiveDecodeInputsV1 + * @typedef {import('./generation/runtime.js').RuntimeGenerationPlanV1} RuntimeGenerationPlanV1 + * @typedef {import('./generation/runtime.js').RuntimeTokenDecisionV1} RuntimeTokenDecisionV1 + * @typedef {import('./generation/runtime.js').RuntimeTokenBatchV1} RuntimeTokenBatchV1 + * @typedef {import('./generation/runtime.js').RuntimeAttentionMaskV1} RuntimeAttentionMaskV1 * @typedef {import('./generation/runtime.js').LogitsLeaseV1} LogitsLeaseV1 * @typedef {import('./backends/artifacts.js').InferenceArtifactProvider} InferenceArtifactProvider * @typedef {import('./backends/artifacts.js').RandomAccessByteSource} RandomAccessByteSource diff --git a/packages/transformers/src/utils/hub.js b/packages/transformers/src/utils/hub.js index f594beedf..e3a0071dc 100755 --- a/packages/transformers/src/utils/hub.js +++ b/packages/transformers/src/utils/hub.js @@ -65,9 +65,11 @@ export { MAX_EXTERNAL_DATA_CHUNKS } from './hub/constants.js'; * Helper function to get a file, using either the Fetch API or FileSystem API. * * @param {URL|string} urlOrPath The URL/path of the file to get. + * @param {AbortSignal} [signal] Signal used to cancel an HTTP request. * @returns {Promise} A promise that resolves to a FileResponse object (if the file is retrieved using the FileSystem API), or a Response object (if the file is retrieved using the Fetch API). */ -export async function getFile(urlOrPath) { +export async function getFile(urlOrPath, signal = undefined) { + throwIfAborted(signal); if (env.useFS && !isValidUrl(urlOrPath, ['http:', 'https:', 'blob:'])) { return new FileResponse( urlOrPath instanceof URL @@ -79,6 +81,7 @@ export async function getFile(urlOrPath) { } else { return env.fetch(urlOrPath, { headers: getFetchHeaders(urlOrPath), + signal, }); } } @@ -261,6 +264,7 @@ export async function loadResourceFile( return_path = false, cache = null, ) { + throwIfAborted(options.signal); const { requestURL, localPath, remoteURL, proposedCacheKey, validModelId } = buildResourcePaths( path_or_repo_id, filename, @@ -279,6 +283,7 @@ export async function loadResourceFile( // Check cache response = await checkCachedResource(cache, localPath, proposedCacheKey); + throwIfAborted(options.signal); const cacheHit = response !== undefined; if (cacheHit) { @@ -292,7 +297,7 @@ export async function loadResourceFile( const isURL = isValidUrl(requestURL, ['http:', 'https:']); if (!isURL) { try { - response = await getFile(localPath); + response = await getFile(localPath, options.signal); cacheKey = localPath; // Update the cache key to be the local path } catch (e) { // Something went wrong while trying to get the file locally. @@ -335,7 +340,7 @@ export async function loadResourceFile( } // File not found locally, so we try to download it from the remote server - response = await getFile(remoteURL); + response = await getFile(remoteURL, options.signal); if (response.status !== 200) { return handleError(response.status, remoteURL, fatal); @@ -425,6 +430,7 @@ export async function loadResourceFile( ); } } + throwIfAborted(options.signal); result = buffer; } @@ -484,6 +490,12 @@ export async function loadResourceFile( throw new Error('Unable to get model file path or buffer.'); } +function throwIfAborted(signal) { + if (!signal?.aborted) return; + if (typeof signal.throwIfAborted === 'function') signal.throwIfAborted(); + throw signal.reason ?? new Error('Model loading aborted.'); +} + /** @type {Map>} Pending file loads keyed by resource identity. */ const INFLIGHT_LOADS = new Map(); diff --git a/packages/transformers/src/utils/model_registry/ModelRegistry.js b/packages/transformers/src/utils/model_registry/ModelRegistry.js index 1179a92f6..04d6c6e01 100644 --- a/packages/transformers/src/utils/model_registry/ModelRegistry.js +++ b/packages/transformers/src/utils/model_registry/ModelRegistry.js @@ -181,28 +181,30 @@ export class ModelRegistry { * Get tokenizer files needed for a specific model. * * @param {string} modelId - The model id + * @param {Object} [options] - Hub metadata options * @returns {Promise} Array of tokenizer file paths * * @example * const files = await ModelRegistry.get_tokenizer_files('onnx-community/gpt2-ONNX'); * console.log(files); // ['tokenizer.json', 'tokenizer_config.json'] */ - static async get_tokenizer_files(modelId) { - return get_tokenizer_files(modelId); + static async get_tokenizer_files(modelId, options = {}) { + return get_tokenizer_files(modelId, options); } /** * Get processor files needed for a specific model. * * @param {string} modelId - The model id + * @param {Object} [options] - Hub metadata options * @returns {Promise} Array of processor file paths * * @example * const files = await ModelRegistry.get_processor_files('onnx-community/vit-base-patch16-224-ONNX'); * console.log(files); // ['preprocessor_config.json'] */ - static async get_processor_files(modelId) { - return get_processor_files(modelId); + static async get_processor_files(modelId, options = {}) { + return get_processor_files(modelId, options); } /** diff --git a/packages/transformers/src/utils/model_registry/get_available_dtypes.js b/packages/transformers/src/utils/model_registry/get_available_dtypes.js index 869a22bbc..2c534b157 100644 --- a/packages/transformers/src/utils/model_registry/get_available_dtypes.js +++ b/packages/transformers/src/utils/model_registry/get_available_dtypes.js @@ -2,16 +2,12 @@ import { getSessionsConfig } from '../../models/session_config.js'; import { get_file_metadata } from './get_file_metadata.js'; import { get_config } from './get_model_files.js'; import { resolve_model_type } from './resolve_model_type.js'; -import { OnnxInferenceProvider } from '../../backends/default.js'; +import { getModelRegistryInferenceProvider } from '../../backends/model_registry.js'; /** * @typedef {import('../../configs.js').PretrainedConfig} PretrainedConfig */ -/** - * The dtypes to probe for availability (excludes 'auto' which is not a concrete dtype). - * @type {string[]} - */ /** * Detects which quantization levels (dtypes) are available for a model * by checking which ONNX files exist on the hub or locally. @@ -27,18 +23,27 @@ import { OnnxInferenceProvider } from '../../backends/default.js'; * @param {string} [options.revision='main'] Model revision * @param {string} [options.cache_dir=null] Custom cache directory * @param {boolean} [options.local_files_only=false] Only check local files + * @param {import('../../backends/model_registry.js').ModelRegistryInferenceProvider|null} [options.inferenceProvider=null] Artifact metadata provider * @returns {Promise} Array of available dtype strings (e.g., ['fp32', 'fp16', 'q4', 'q8']) */ export async function get_available_dtypes( modelId, - { config = null, model_file_name = null, revision = 'main', cache_dir = null, local_files_only = false } = {}, + { + config = null, + model_file_name = null, + revision = 'main', + cache_dir = null, + local_files_only = false, + inferenceProvider = null, + } = {}, ) { config = await get_config(modelId, { config, cache_dir, local_files_only, revision }); const modelType = resolve_model_type(config); const { sessions } = getSessionsConfig(modelType, config, { model_file_name }); const metadataOptions = { revision, cache_dir, local_files_only }; - return OnnxInferenceProvider.getAvailableDtypes({ + const provider = await getModelRegistryInferenceProvider(inferenceProvider); + return provider.getAvailableDtypes({ modelId, sessions, getFileMetadata: get_file_metadata, diff --git a/packages/transformers/src/utils/model_registry/get_file_metadata.js b/packages/transformers/src/utils/model_registry/get_file_metadata.js index 2d69344f6..a3a827547 100644 --- a/packages/transformers/src/utils/model_registry/get_file_metadata.js +++ b/packages/transformers/src/utils/model_registry/get_file_metadata.js @@ -25,10 +25,12 @@ import { getModelId } from '../../backends/inference.js'; * 5. Range requests typically aren't compressed, and content-range header shows true uncompressed size * * @param {URL|string} urlOrPath The URL/path of the file. + * @param {AbortSignal} [signal] Signal used to cancel the request. * @returns {Promise} A promise that resolves to a Response object or null if not supported. * @private */ -async function fetch_file_head(urlOrPath) { +async function fetch_file_head(urlOrPath, signal = undefined) { + throwIfAborted(signal); // Range requests only make sense for HTTP URLs if (!isValidUrl(urlOrPath, ['http:', 'https:'])) { return null; @@ -36,7 +38,7 @@ async function fetch_file_head(urlOrPath) { const headers = getFetchHeaders(urlOrPath); headers.set('Range', 'bytes=0-0'); - return env.fetch(urlOrPath, { method: 'GET', headers, cache: 'no-store' }); + return env.fetch(urlOrPath, { method: 'GET', headers, cache: 'no-store', signal }); } /** @@ -52,12 +54,15 @@ async function fetch_file_head(urlOrPath) { * @returns {Promise<{exists: boolean, size?: number, contentType?: string, fromCache?: boolean}>} A Promise that resolves to file metadata. */ export function get_file_metadata(path_or_repo_id, filename, options = {}) { + throwIfAborted(options.signal); path_or_repo_id = getModelId(path_or_repo_id); + if (options.signal) return _get_file_metadata(path_or_repo_id, filename, options); const key = makePretrainedOptionsKey(path_or_repo_id, options, filename); return memoizePromise(key, () => _get_file_metadata(path_or_repo_id, filename, options)); } async function _get_file_metadata(path_or_repo_id, filename, options) { + throwIfAborted(options.signal); /** @type {import('../cache.js').CacheInterface | null} */ const cache = await getCache(options?.cache_dir); const { localPath, remoteURL, proposedCacheKey, validModelId } = buildResourcePaths( @@ -69,6 +74,7 @@ async function _get_file_metadata(path_or_repo_id, filename, options) { // Check cache first - if cached, we can get metadata from the cached response const cachedResponse = await checkCachedResource(cache, localPath, proposedCacheKey); + throwIfAborted(options.signal); if (cachedResponse !== undefined && typeof cachedResponse !== 'string') { const size = cachedResponse.headers.get('content-length'); const contentType = cachedResponse.headers.get('content-type'); @@ -85,7 +91,7 @@ async function _get_file_metadata(path_or_repo_id, filename, options) { const isURL = isValidUrl(localPath, ['http:', 'https:']); if (!isURL) { try { - const response = await getFile(localPath); + const response = await getFile(localPath, options.signal); if (typeof response !== 'string' && response.status !== 404) { const size = response.headers.get('content-length'); const contentType = response.headers.get('content-type'); @@ -98,6 +104,7 @@ async function _get_file_metadata(path_or_repo_id, filename, options) { }; } } catch (e) { + throwIfAborted(options.signal); // File doesn't exist locally, continue to remote check } } @@ -107,7 +114,7 @@ async function _get_file_metadata(path_or_repo_id, filename, options) { if (env.allowRemoteModels && !options.local_files_only && validModelId) { try { // Make a Range request to get metadata without downloading full content - const rangeResponse = await fetch_file_head(remoteURL); + const rangeResponse = await fetch_file_head(remoteURL, options.signal); if (rangeResponse && rangeResponse.status >= 200 && rangeResponse.status < 300) { let size; @@ -149,6 +156,7 @@ async function _get_file_metadata(path_or_repo_id, filename, options) { }; } } catch (e) { + throwIfAborted(options.signal); // Range request failed most likely because of a network error, timeout, etc. logger.warn(`Unable to fetch file metadata for "${remoteURL}": ${e}`); } @@ -156,3 +164,9 @@ async function _get_file_metadata(path_or_repo_id, filename, options) { return { exists: false, fromCache: false }; } + +function throwIfAborted(signal) { + if (!signal?.aborted) return; + if (typeof signal.throwIfAborted === 'function') signal.throwIfAborted(); + throw signal.reason ?? new Error('Metadata loading aborted.'); +} diff --git a/packages/transformers/src/utils/model_registry/get_files.js b/packages/transformers/src/utils/model_registry/get_files.js index e13d0d291..0d2c34f23 100644 --- a/packages/transformers/src/utils/model_registry/get_files.js +++ b/packages/transformers/src/utils/model_registry/get_files.js @@ -12,6 +12,10 @@ import { get_processor_files } from './get_processor_files.js'; * @param {import('../dtypes.js').DataType|Record} [options.dtype=null] Override dtype (use this if passing dtype to pipeline) * @param {import('../devices.js').DeviceType|Record} [options.device=null] Override device (use this if passing device to pipeline) * @param {string|null} [options.model_file_name=null|null] Override the model file name (excluding .onnx suffix) + * @param {string|null} [options.cache_dir=null] Custom cache directory + * @param {boolean} [options.local_files_only=false] Never hit the network if true + * @param {string} [options.revision='main'] Model revision + * @param {import('../../backends/model_registry.js').ModelRegistryInferenceProvider|null} [options.inferenceProvider=null] Artifact metadata provider * @param {boolean} [options.include_tokenizer=true] Whether to check for tokenizer files (set to false for vision-only models) * @param {boolean} [options.include_processor=true] Whether to check for processor files * @param {boolean} [options.include_model=true] Whether to include built-in ONNX model files @@ -24,19 +28,33 @@ export async function get_files( dtype = null, device = null, model_file_name = null, + cache_dir = null, + local_files_only = false, + revision = 'main', + inferenceProvider = null, include_tokenizer = true, include_processor = true, include_model = true, } = {}, ) { - const files = include_model ? await get_model_files(modelId, { config, dtype, device, model_file_name }) : []; + const metadataOptions = { cache_dir, local_files_only, revision }; + const files = include_model + ? await get_model_files(modelId, { + config, + dtype, + device, + model_file_name, + inferenceProvider, + ...metadataOptions, + }) + : []; if (include_tokenizer) { - const tokenizerFiles = await get_tokenizer_files(modelId); + const tokenizerFiles = await get_tokenizer_files(modelId, metadataOptions); files.push(...tokenizerFiles); } if (include_processor) { - const processorFiles = await get_processor_files(modelId); + const processorFiles = await get_processor_files(modelId, metadataOptions); files.push(...processorFiles); } diff --git a/packages/transformers/src/utils/model_registry/get_model_files.js b/packages/transformers/src/utils/model_registry/get_model_files.js index 3f267d3bc..a9a524dd6 100644 --- a/packages/transformers/src/utils/model_registry/get_model_files.js +++ b/packages/transformers/src/utils/model_registry/get_model_files.js @@ -3,7 +3,7 @@ import { AutoConfig } from '../../configs.js'; import { makePretrainedOptionsKey } from '../hub/utils.js'; import { memoizePromise } from '../memoize_promise.js'; import { resolve_model_type } from './resolve_model_type.js'; -import { OnnxInferenceProvider } from '../../backends/default.js'; +import { getModelRegistryInferenceProvider } from '../../backends/model_registry.js'; /** * @typedef {import('../../configs.js').PretrainedConfig} PretrainedConfig @@ -53,18 +53,32 @@ export function get_config( * @param {import('../dtypes.js').DataType|Record} [options.dtype=null] Override dtype (use this if passing dtype to pipeline) * @param {import('../devices.js').DeviceType|Record} [options.device=null] Override device (use this if passing device to pipeline) * @param {string} [options.model_file_name=null] Override the model file name (excluding .onnx suffix). + * @param {string|null} [options.cache_dir=null] Custom cache directory. + * @param {boolean} [options.local_files_only=false] Never hit the network if true. + * @param {string} [options.revision='main'] Model revision. + * @param {import('../../backends/model_registry.js').ModelRegistryInferenceProvider|null} [options.inferenceProvider=null] Artifact metadata provider. * @returns {Promise} Array of file paths that will be loaded */ export async function get_model_files( modelId, - { config = null, dtype: overrideDtype = null, device: overrideDevice = null, model_file_name = null } = {}, + { + config = null, + dtype: overrideDtype = null, + device: overrideDevice = null, + model_file_name = null, + cache_dir = null, + local_files_only = false, + revision = 'main', + inferenceProvider = null, + } = {}, ) { - config = await get_config(modelId, { config }); + config = await get_config(modelId, { config, cache_dir, local_files_only, revision }); // Infer model type from config const modelType = resolve_model_type(config); const { sessions, optional_configs } = getSessionsConfig(modelType, config, { model_file_name }); - return OnnxInferenceProvider.listModelArtifacts({ + const provider = await getModelRegistryInferenceProvider(inferenceProvider); + return provider.listModelArtifacts({ sessions, optionalConfigs: optional_configs, config, diff --git a/packages/transformers/src/utils/model_registry/get_pipeline_files.js b/packages/transformers/src/utils/model_registry/get_pipeline_files.js index e22eddfd3..0f4438ef1 100644 --- a/packages/transformers/src/utils/model_registry/get_pipeline_files.js +++ b/packages/transformers/src/utils/model_registry/get_pipeline_files.js @@ -3,7 +3,7 @@ import { get_config } from './get_model_files.js'; import { resolve_model_type } from './resolve_model_type.js'; import { getTextOnlySessions } from '../../models/session_config.js'; import { SUPPORTED_TASKS, TASK_ALIASES } from '../../pipelines/index.js'; -import { OnnxInferenceProvider } from '../../backends/default.js'; +import { getModelRegistryInferenceProvider } from '../../backends/model_registry.js'; /** * Get all files needed for a specific pipeline task. @@ -17,6 +17,10 @@ import { OnnxInferenceProvider } from '../../backends/default.js'; * @param {import('../dtypes.js').DataType|Record} [options.dtype=null] - Override dtype * @param {import('../devices.js').DeviceType|Record} [options.device=null] - Override device * @param {string} [options.model_file_name=null] - Override the model file name (excluding .onnx suffix) + * @param {string|null} [options.cache_dir=null] - Custom cache directory + * @param {boolean} [options.local_files_only=false] - Never hit the network if true + * @param {string} [options.revision='main'] - Model revision + * @param {import('../../backends/model_registry.js').ModelRegistryInferenceProvider|null} [options.inferenceProvider=null] - Artifact metadata provider * @param {boolean} [options.include_model=true] - Whether to include built-in ONNX model files * @returns {Promise} Array of file paths that will be loaded * @throws {Error} If the task is not supported @@ -49,13 +53,14 @@ export async function get_pipeline_files(task, modelId, options = {}) { // When loading multimodal models via the text-generation pipeline, // only load the sessions needed for text generation (embed_tokens, decoder_model_merged) - if (task === 'text-generation') { + if (task === 'text-generation' && options.include_model !== false) { const config = await get_config(modelId, options); const modelType = resolve_model_type(config); const textOnlySessions = getTextOnlySessions(modelType); if (textOnlySessions) { - return OnnxInferenceProvider.filterModelArtifacts(files, textOnlySessions); + const provider = await getModelRegistryInferenceProvider(options.inferenceProvider ?? null); + return provider.filterModelArtifacts(files, textOnlySessions); } } diff --git a/packages/transformers/src/utils/model_registry/get_processor_files.js b/packages/transformers/src/utils/model_registry/get_processor_files.js index 894748d84..6d294eca6 100644 --- a/packages/transformers/src/utils/model_registry/get_processor_files.js +++ b/packages/transformers/src/utils/model_registry/get_processor_files.js @@ -6,15 +6,16 @@ import { get_file_metadata } from './get_file_metadata.js'; * Auto-detects if the model has a processor by checking if preprocessor_config.json exists. * * @param {string} modelId The model id (e.g., "Xenova/detr-resnet-50") + * @param {Object} [options] Hub metadata options. * @returns {Promise} Array of processor file names (empty if no processor) */ -export async function get_processor_files(modelId) { +export async function get_processor_files(modelId, options = {}) { if (!modelId) { throw new Error('modelId is required'); } // Check if preprocessor_config.json exists - const metadata = await get_file_metadata(modelId, IMAGE_PROCESSOR_NAME, {}); + const metadata = await get_file_metadata(modelId, IMAGE_PROCESSOR_NAME, options); return metadata.exists ? [IMAGE_PROCESSOR_NAME] : []; } diff --git a/packages/transformers/src/utils/model_registry/get_tokenizer_files.js b/packages/transformers/src/utils/model_registry/get_tokenizer_files.js index 1024f9958..fc32c6de5 100644 --- a/packages/transformers/src/utils/model_registry/get_tokenizer_files.js +++ b/packages/transformers/src/utils/model_registry/get_tokenizer_files.js @@ -5,14 +5,15 @@ import { get_file_metadata } from './get_file_metadata.js'; * Automatically detects whether the model has tokenizer files. * * @param {string} modelId The model id to check for tokenizer files + * @param {Object} [options] Hub metadata options. * @returns {Promise} An array of file names that will be loaded */ -export async function get_tokenizer_files(modelId) { +export async function get_tokenizer_files(modelId, options = {}) { if (!modelId) { throw new Error('modelId is required for get_tokenizer_files'); } - const metadata = await get_file_metadata(modelId, 'tokenizer_config.json', {}); + const metadata = await get_file_metadata(modelId, 'tokenizer_config.json', options); if (metadata.exists) { return ['tokenizer.json', 'tokenizer_config.json']; } diff --git a/packages/transformers/src/utils/tensor.js b/packages/transformers/src/utils/tensor.js index 330eaf64a..585f0e5ea 100644 --- a/packages/transformers/src/utils/tensor.js +++ b/packages/transformers/src/utils/tensor.js @@ -15,6 +15,8 @@ import { DataTypeMap } from './dtypes.js'; import { random } from './random.js'; +const NOOP_DISPOSE = () => {}; + /** * @typedef {keyof typeof DataTypeMap} DataType * @typedef {import('./maths.js').AnyTypedArray | any[]} DataArray @@ -81,15 +83,19 @@ export class Tensor { data = Constructor.from(data); } } + const size = dims.reduce((product, dimension) => product * dimension, 1); + if (data.length !== size) { + throw new RangeError(`Tensor data length (${data.length}) does not match shape [${dims}] (${size}).`); + } this._storage = { backend: 'cpu', handle: null, type, data, dims, - size: dims.reduce((product, dimension) => product * dimension, 1), + size, location: 'cpu', - dispose() {}, + dispose: NOOP_DISPOSE, }; return new Proxy(this, { diff --git a/packages/transformers/tests/generation_controller.test.js b/packages/transformers/tests/generation_controller.test.js index adcc2a27d..49fc6a992 100644 --- a/packages/transformers/tests/generation_controller.test.js +++ b/packages/transformers/tests/generation_controller.test.js @@ -77,6 +77,29 @@ describe("GenerationController", () => { expect(streamer.put).toHaveBeenCalledWith([[1n, 2n]]); expect(streamer.end).toHaveBeenCalledTimes(1); }); + + it("processes classifier-free guidance before validating the output batch", async () => { + const controller = createGenerationController({ config: {}, generation_config: null }, int64Tensor([[1]]), { max_new_tokens: 1, guidance_scale: 3 }); + const logits = new Tensor( + "float32", + [ + 0, + 5, + 0, + 0, // conditional + 0, + 0, + 1, + 0, // unconditional + ], + [2, 4], + ); + + const step = await controller.step(logits); + + expect(step.nextTokenIds.tolist()).toEqual([[1n]]); + expect(step.allDone).toBe(true); + }); }); describe("custom autoregressive sessions", () => { @@ -134,8 +157,6 @@ describe("custom autoregressive sessions", () => { version: 1, batchSize: 1, maxSequenceLength: 3, - prefill: jest.fn(), - decode: jest.fn(), async *generateWithPlan(_inputs, plan) { try { expect(plan.sampler).toEqual({ op: "argmax" }); @@ -153,14 +174,16 @@ describe("custom autoregressive sessions", () => { modelId: "test/fast-controller-model", load: jest.fn(async () => ({ generation_config: {}, - generationCapabilities: { - sessionVersion: 1, - maxBatchSize: 1, - cpuModes: ["greedy"], - planModes: ["greedy"], - cpuLogits: false, - declarativePlans: ["argmax"], - tokenPipeline: { defaultDepth: 4, maxDepth: 4 }, + capabilities: { + causalGeneration: { + sessionVersion: 1, + maxBatchSize: 1, + cpuModes: [], + planModes: ["greedy"], + cpuLogits: false, + declarativePlans: ["argmax"], + tokenPipeline: { defaultDepth: 4, maxDepth: 4 }, + }, }, createAutoregressiveSession: jest.fn(async () => session), async forward(inputs) { @@ -176,12 +199,94 @@ describe("custom autoregressive sessions", () => { const output = await model.generate({ input_ids: int64Tensor([[1]]), max_new_tokens: 2 }); expect(output.tolist()).toEqual([[1n, 2n, 3n]]); - expect(session.prefill).not.toHaveBeenCalled(); - expect(session.decode).not.toHaveBeenCalled(); expect(iteratorClosed).toHaveBeenCalledTimes(1); expect(session.dispose).toHaveBeenCalledTimes(1); }); + it("rejects a plan session that omits generateWithPlan", async () => { + const session = { + version: 1, + batchSize: 1, + maxSequenceLength: 2, + dispose: jest.fn(async () => {}), + }; + const backend = { + modelId: "test/malformed-plan-model", + load: jest.fn(async () => ({ + capabilities: { + causalGeneration: { + sessionVersion: 1, + maxBatchSize: 1, + cpuModes: [], + planModes: ["greedy"], + cpuLogits: false, + declarativePlans: ["argmax"], + tokenPipeline: { defaultDepth: 1, maxDepth: 1 }, + }, + }, + createAutoregressiveSession: jest.fn(async () => session), + async dispose() {}, + })), + }; + const model = await AutoModel.from_pretrained(backend, { + config: { model_type: "custom", is_encoder_decoder: false }, + }); + + await expect(model.generate({ input_ids: int64Tensor([[1]]), max_new_tokens: 1 })).rejects.toThrow("must implement `generateWithPlan()`"); + expect(session.dispose).toHaveBeenCalledTimes(1); + }); + + it("enforces advertised active-session limits", async () => { + let releasePlan; + const planGate = new Promise((resolve) => { + releasePlan = resolve; + }); + let markStarted; + const started = new Promise((resolve) => { + markStarted = resolve; + }); + const createAutoregressiveSession = jest.fn(async () => ({ + version: 1, + batchSize: 1, + maxSequenceLength: 2, + async *generateWithPlan() { + markStarted(); + await planGate; + yield { tokenIds: new Uint32Array([2]) }; + }, + async dispose() {}, + })); + const backend = { + modelId: "test/single-session-model", + load: jest.fn(async () => ({ + capabilities: { + causalGeneration: { + sessionVersion: 1, + maxBatchSize: 1, + cpuModes: [], + planModes: ["greedy"], + cpuLogits: false, + declarativePlans: ["argmax"], + tokenPipeline: { defaultDepth: 1, maxDepth: 1 }, + sessionConcurrency: { maxActiveSessions: 1, concurrentOperationsPerSession: 1 }, + }, + }, + createAutoregressiveSession, + async dispose() {}, + })), + }; + const model = await AutoModel.from_pretrained(backend, { + config: { model_type: "custom", is_encoder_decoder: false, eos_token_id: 2 }, + }); + const first = model.generate({ input_ids: int64Tensor([[1]]), max_new_tokens: 1 }); + await started; + + await expect(model.generate({ input_ids: int64Tensor([[1]]), max_new_tokens: 1 })).rejects.toThrow("at most 1 active autoregressive session"); + releasePlan(); + await expect(first).resolves.toEqual(expect.any(Tensor)); + expect(createAutoregressiveSession).toHaveBeenCalledTimes(1); + }); + it("rejects unsupported batches before creating a runtime session", async () => { const createAutoregressiveSession = jest.fn(); const backend = { @@ -211,4 +316,68 @@ describe("custom autoregressive sessions", () => { await expect(model.generate({ input_ids: int64Tensor([[1], [2]]), max_new_tokens: 1 })).rejects.toThrow("supports batch size 1"); expect(createAutoregressiveSession).not.toHaveBeenCalled(); }); + + it("reports when a request is unsupported by the CPU mode", async () => { + const createAutoregressiveSession = jest.fn(); + const backend = { + modelId: "test/plan-only-score-model", + load: jest.fn(async () => ({ + capabilities: { + causalGeneration: { + sessionVersion: 1, + maxBatchSize: 1, + cpuModes: [], + planModes: ["greedy"], + cpuLogits: true, + declarativePlans: ["argmax"], + tokenPipeline: { defaultDepth: 1, maxDepth: 1 }, + }, + }, + createAutoregressiveSession, + async dispose() {}, + })), + }; + const model = await AutoModel.from_pretrained(backend, { + config: { model_type: "custom", is_encoder_decoder: false }, + }); + + await expect(model.generate({ input_ids: int64Tensor([[1]]), max_new_tokens: 1, output_scores: true })).rejects.toThrow("does not support greedy generation through its CPU logits path"); + expect(createAutoregressiveSession).not.toHaveBeenCalled(); + }); + + it("releases malformed logits leases", async () => { + const release = jest.fn(); + const session = { + version: 1, + batchSize: 1, + maxSequenceLength: 2, + prefill: jest.fn(async () => ({ version: 0, release })), + decode: jest.fn(), + dispose: jest.fn(async () => {}), + }; + const backend = { + modelId: "test/malformed-lease-model", + load: jest.fn(async () => ({ + generation_config: {}, + generationCapabilities: { + sessionVersion: 1, + maxBatchSize: 1, + cpuModes: ["greedy"], + planModes: [], + cpuLogits: true, + declarativePlans: [], + tokenPipeline: { defaultDepth: 1, maxDepth: 1 }, + }, + createAutoregressiveSession: jest.fn(async () => session), + async dispose() {}, + })), + }; + const model = await AutoModel.from_pretrained(backend, { + config: { model_type: "custom", is_encoder_decoder: false }, + }); + + await expect(model.generate({ input_ids: int64Tensor([[1]]), max_new_tokens: 1 })).rejects.toThrow("unsupported logits lease"); + expect(release).toHaveBeenCalledTimes(1); + expect(session.dispose).toHaveBeenCalledTimes(1); + }); }); diff --git a/packages/transformers/tests/inference_backends.test.js b/packages/transformers/tests/inference_backends.test.js index ad601e02a..341773958 100644 --- a/packages/transformers/tests/inference_backends.test.js +++ b/packages/transformers/tests/inference_backends.test.js @@ -1,7 +1,7 @@ import { jest } from "@jest/globals"; -import { getModelId, isInferenceBackend, loadInferenceModel, normalizeInferenceModel } from "../src/backends/inference.js"; -import { OnnxInferenceProvider } from "../src/backends/default.js"; +import { getModelId, isInferenceBackend, loadInferenceModel, normalizeInferenceModel, validateInferenceBackendTask, validateInferenceModelTask } from "../src/backends/inference.js"; +import { OnnxInferenceProvider } from "@huggingface/transformers-onnx"; import { AutoModel } from "../src/models/auto/modeling_auto.js"; import { PreTrainedModel } from "../src/models/modeling_utils.js"; import { buildResourcePaths } from "../src/utils/hub.js"; @@ -54,6 +54,33 @@ describe("inference backends", () => { expect(model.config).toBe(config); }); + it("rejects malformed artifact providers before backend loading", async () => { + const backend = { modelId: "test/model", load: jest.fn() }; + + await expect(loadInferenceModel(backend, { artifactProvider: { readJson() {} } })).rejects.toThrow("must implement `readJson()` and `openByteSource()`"); + expect(backend.load).not.toHaveBeenCalled(); + }); + + it("uses declared task and loaded execution capabilities for setup validation", async () => { + const backend = { + modelId: "test/model", + capabilities: { devices: ["webgpu"], dtypes: ["auto"], tasks: ["text-generation"] }, + load() {}, + }; + expect(() => validateInferenceBackendTask(backend, "feature-extraction")).toThrow('does not support the "feature-extraction" task'); + expect(() => validateInferenceBackendTask(backend, "text-generation")).not.toThrow(); + expect(() => + validateInferenceModelTask( + { + capabilities: { forward: { version: 1 } }, + async forward() {}, + async dispose() {}, + }, + "text-generation", + ), + ).toThrow("does not support causal text generation"); + }); + it("normalizes absent custom device and dtype options", async () => { const backend = { modelId: "test/model", diff --git a/packages/transformers/tests/onnx_provider_object.test.js b/packages/transformers/tests/onnx_provider_object.test.js new file mode 100644 index 000000000..6e1b4292c --- /dev/null +++ b/packages/transformers/tests/onnx_provider_object.test.js @@ -0,0 +1,33 @@ +import { jest } from "@jest/globals"; + +import { env } from "../src/env.js"; +import { AutoModelForCausalLM } from "../src/models/auto/modeling_auto.js"; +import { PreTrainedModel } from "../src/models/modeling_utils.js"; +import { DEFAULT_MODEL_OPTIONS, MAX_MODEL_LOAD_TIME } from "./init.js"; + +describe("ONNX provider objects", () => { + it("configures the Transformers.js host before provider loading", async () => { + const { OnnxInferenceProvider } = await import("@huggingface/transformers-onnx"); + const provider = OnnxInferenceProvider.from_modelId("test/provider-object"); + class TestModel extends PreTrainedModel {} + TestModel._from_pretrained = jest.fn(async () => "loaded"); + + await expect(TestModel.from_pretrained(provider, { config: { model_type: "custom" } })).resolves.toBe("loaded"); + + expect(TestModel._from_pretrained).toHaveBeenCalledWith("test/provider-object", expect.objectContaining({ inferenceProvider: provider })); + expect(typeof env.backends.onnx.setLogLevel).toBe("function"); + }); + + it( + "loads a real model through a provider object", + async () => { + const { OnnxInferenceProvider } = await import("@huggingface/transformers-onnx"); + const provider = OnnxInferenceProvider.from_modelId("hf-internal-testing/tiny-random-LlamaForCausalLM"); + const model = await AutoModelForCausalLM.from_pretrained(provider, DEFAULT_MODEL_OPTIONS); + + expect(model).toBeInstanceOf(PreTrainedModel); + await model.dispose(); + }, + MAX_MODEL_LOAD_TIME, + ); +}); diff --git a/packages/transformers/tests/onnx_wiring.test.js b/packages/transformers/tests/onnx_wiring.test.js new file mode 100644 index 000000000..5992911ef --- /dev/null +++ b/packages/transformers/tests/onnx_wiring.test.js @@ -0,0 +1,30 @@ +import { jest } from "@jest/globals"; + +import { env } from "../src/env.js"; +import { getOnnxProviderModule } from "../src/backends/default.js"; +import { PreTrainedModel } from "../src/models/modeling_utils.js"; + +describe("ONNX module wiring", () => { + it("preserves the ORT environment installed through a pre-registered host", async () => { + await getOnnxProviderModule(); + + expect(typeof env.backends.onnx.setLogLevel).toBe("function"); + expect(env.backends.onnx.wasm).toBeDefined(); + }); + + it("honors an aborted signal before built-in model loading", async () => { + class TestModel extends PreTrainedModel {} + TestModel._from_pretrained = jest.fn(); + const controller = new AbortController(); + const reason = new Error("cancel loading"); + controller.abort(reason); + + await expect( + TestModel.from_pretrained("test/aborted-model", { + config: { model_type: "custom" }, + signal: controller.signal, + }), + ).rejects.toBe(reason); + expect(TestModel._from_pretrained).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/transformers/tests/types.test.js b/packages/transformers/tests/types.test.js index d2c058c85..77509297e 100644 --- a/packages/transformers/tests/types.test.js +++ b/packages/transformers/tests/types.test.js @@ -69,6 +69,72 @@ describe("TypeScript compilation succeeds", () => { } }); } + + it("compiles a readonly plan-only inference backend", () => { + const diagnostics = getDiagnosticsFromSource(` + import type { + CausalGenerationCapabilitiesV1, + InferenceBackend, + InferenceModel, + LogitsLeaseV1, + PlanAutoregressiveSessionV1, + StaticBackendCapabilities, + } from "../../types/transformers.js"; + + const staticCapabilities = { + devices: ["webgpu"], + dtypes: ["auto"], + tasks: ["text-generation"], + } as const satisfies StaticBackendCapabilities; + + const causalGeneration = { + sessionVersion: 1, + maxBatchSize: 1, + cpuModes: [], + planModes: ["greedy"], + cpuLogits: false, + declarativePlans: ["argmax"], + tokenPipeline: { defaultDepth: 4, maxDepth: 4 }, + } as const satisfies CausalGenerationCapabilitiesV1; + + const session: PlanAutoregressiveSessionV1 = { + version: 1, + batchSize: 1, + maxSequenceLength: 128, + async *generateWithPlan(_inputs, _plan) { + yield { tokenIds: new Uint32Array([1]) }; + }, + async dispose() {}, + }; + + const lease: LogitsLeaseV1 = { + version: 1, + dtype: "float32", + shape: [1, 4] as const, + async read() { return new Float32Array(4); }, + release() {}, + }; + + const backend: InferenceBackend = { + modelId: "test/model", + capabilities: staticCapabilities, + async load(_options) { + const model: InferenceModel = { + capabilities: { causalGeneration }, + async createAutoregressiveSession() { return session; }, + async dispose() {}, + }; + return model; + }, + }; + + void backend; + void lease; + `); + if (diagnostics.length > 0) { + throw new Error(formatDiagnostics(diagnostics)); + } + }); }); describe("TypeScript expected errors", () => { diff --git a/packages/transformers/tests/utils/generation.test.js b/packages/transformers/tests/utils/generation.test.js index 231985571..aef8e167e 100644 --- a/packages/transformers/tests/utils/generation.test.js +++ b/packages/transformers/tests/utils/generation.test.js @@ -140,6 +140,21 @@ describe("Generation parameters", () => { MAX_TEST_EXECUTION_TIME, ); + it( + "max_new_tokens=0 with return_dict_in_generate", + async () => { + const inputs = tokenizer(DUMMY_TEXT); + const outputs = await model.generate({ + ...inputs, + max_new_tokens: 0, + return_dict_in_generate: true, + }); + expect(outputs.sequences.tolist()).toEqual(inputs.input_ids.tolist()); + expect(outputs.past_key_values).toBeInstanceOf(DynamicCache); + }, + MAX_TEST_EXECUTION_TIME, + ); + it( "min_length", async () => { diff --git a/packages/transformers/tests/utils/hub_abort.test.js b/packages/transformers/tests/utils/hub_abort.test.js new file mode 100644 index 000000000..0584b7c00 --- /dev/null +++ b/packages/transformers/tests/utils/hub_abort.test.js @@ -0,0 +1,22 @@ +import { jest } from "@jest/globals"; + +import { env } from "../../src/env.js"; +import { getFile } from "../../src/utils/hub.js"; + +describe("Hub cancellation", () => { + it("passes the loading signal to remote fetches", async () => { + const originalFetch = env.fetch; + const controller = new AbortController(); + env.fetch = jest.fn(async (_url, options) => { + expect(options.signal).toBe(controller.signal); + return new Response(new Uint8Array()); + }); + + try { + await getFile("https://huggingface.co/test/model", controller.signal); + expect(env.fetch).toHaveBeenCalledTimes(1); + } finally { + env.fetch = originalFetch; + } + }); +}); diff --git a/packages/transformers/tests/utils/model_registry.test.js b/packages/transformers/tests/utils/model_registry.test.js index cb9e6408a..ab9713c18 100644 --- a/packages/transformers/tests/utils/model_registry.test.js +++ b/packages/transformers/tests/utils/model_registry.test.js @@ -55,6 +55,21 @@ describe("get_available_dtypes", () => { mockGetFileMetadata.mockReset(); }); + it("uses an explicitly supplied inference provider", async () => { + const inferenceProvider = { + getAvailableDtypes: jest.fn(async () => ["native"]), + listModelArtifacts() { + return []; + }, + filterModelArtifacts(files) { + return files; + }, + }; + + await expect(get_available_dtypes("test/model", { config: ENCODER_ONLY_CONFIG, inferenceProvider })).resolves.toEqual(["native"]); + expect(inferenceProvider.getAvailableDtypes).toHaveBeenCalledWith(expect.objectContaining({ modelId: "test/model" })); + }); + it("should detect fp32 and q4 for an encoder-only model", async () => { setupExistingFiles( "onnx/model.onnx", // fp32 diff --git a/packages/transformers/tests/utils/tensor.test.js b/packages/transformers/tests/utils/tensor.test.js index 278aa8771..583b0a277 100644 --- a/packages/transformers/tests/utils/tensor.test.js +++ b/packages/transformers/tests/utils/tensor.test.js @@ -4,6 +4,10 @@ import { init } from "../init.js"; init(); describe("Tensor operations", () => { + it("rejects data that does not match the declared shape", () => { + expect(() => new Tensor("float32", [1, 2, 3], [2, 2])).toThrow("does not match shape [2,2] (4)"); + }); + describe("cat", () => { it("should concatenate on dim=0", () => { const t1 = new Tensor("float32", [1, 2, 3], [1, 3]); diff --git a/types/webgpu-kernels.local.demo.d.ts b/types/webgpu-kernels.local.demo.d.ts deleted file mode 100644 index 6f6b469b6..000000000 --- a/types/webgpu-kernels.local.demo.d.ts +++ /dev/null @@ -1,14 +0,0 @@ -import type { WebGPUKernelsForwardModel, WebGPUKernelsTensorMap, WebGPUKernelsTextGenerationModel } from './webgpu-kernels.local'; -type PipelineForwardModel = { - (inputs: WebGPUKernelsTensorMap): Promise; - readonly config?: Record; - forward(inputs: WebGPUKernelsTensorMap): Promise; - dispose(): void | Promise; -}; -type PipelineTextGenerationModel = PipelineForwardModel & { - generate: WebGPUKernelsTextGenerationModel['generate']; -}; -export declare function adaptWebGPUKernelsModel(model: WebGPUKernelsTextGenerationModel): PipelineTextGenerationModel; -export declare function adaptWebGPUKernelsModel(model: WebGPUKernelsForwardModel): PipelineForwardModel; -export {}; -//# sourceMappingURL=webgpu-kernels.local.demo.d.ts.map \ No newline at end of file diff --git a/types/webgpu-kernels.local.demo.d.ts.map b/types/webgpu-kernels.local.demo.d.ts.map deleted file mode 100644 index 67f999ab9..000000000 --- a/types/webgpu-kernels.local.demo.d.ts.map +++ /dev/null @@ -1 +0,0 @@ -{"version":3,"file":"webgpu-kernels.local.demo.d.ts","sourceRoot":"","sources":["../webgpu-kernels.local.demo.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EACR,yBAAyB,EACzB,sBAAsB,EAEtB,gCAAgC,EACnC,MAAM,wBAAwB,CAAC;AAEhC,KAAK,oBAAoB,GAAG;IACxB,CAAC,MAAM,EAAE,sBAAsB,GAAG,OAAO,CAAC,sBAAsB,CAAC,CAAC;IAClE,QAAQ,CAAC,MAAM,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;IAC1C,OAAO,CAAC,MAAM,EAAE,sBAAsB,GAAG,OAAO,CAAC,sBAAsB,CAAC,CAAC;IACzE,OAAO,IAAI,IAAI,GAAG,OAAO,CAAC,IAAI,CAAC,CAAC;CACnC,CAAC;AAEF,KAAK,2BAA2B,GAAG,oBAAoB,GAAG;IACtD,QAAQ,EAAE,gCAAgC,CAAC,UAAU,CAAC,CAAC;CAC1D,CAAC;AAEF,wBAAgB,uBAAuB,CAAC,KAAK,EAAE,gCAAgC,GAAG,2BAA2B,CAAC;AAC9G,wBAAgB,uBAAuB,CAAC,KAAK,EAAE,yBAAyB,GAAG,oBAAoB,CAAC"} \ No newline at end of file From 7432d00ec87a3f837b49780c9697435a84bb9be2 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Wed, 29 Jul 2026 15:24:25 +0200 Subject: [PATCH 4/5] updates to align with ModelRegistry caching and other improvements --- packages/transformers-onnx/README.md | 12 + packages/transformers-onnx/package.json | 10 +- packages/transformers-onnx/scripts/build.mjs | 35 ++- packages/transformers-onnx/src/host.ts | 9 +- packages/transformers-onnx/src/runtime.ts | 35 ++- .../docs/source/guides/custom-backends.md | 2 + .../transformers/scripts/build/buildAll.mjs | 10 + packages/transformers/src/backends/default.js | 16 +- .../transformers/src/backends/inference.js | 16 ++ .../src/backends/model_registry.js | 10 +- packages/transformers/src/env.js | 10 +- .../transformers/src/generation/controller.js | 50 ++-- .../transformers/src/generation/runtime.js | 45 ++- .../transformers/src/generation/streamers.js | 17 ++ .../transformers/src/models/modeling_utils.js | 267 ++++-------------- packages/transformers/src/pipelines.js | 63 ++++- packages/transformers/src/transformers.js | 1 + packages/transformers/src/utils/core.js | 51 +++- packages/transformers/src/utils/hub.js | 12 +- .../src/utils/model_registry/ModelRegistry.js | 99 ++++--- .../src/utils/model_registry/clear_cache.js | 20 +- .../model_registry/get_available_dtypes.js | 3 + .../utils/model_registry/get_file_metadata.js | 35 ++- .../src/utils/model_registry/get_files.js | 3 + .../utils/model_registry/get_model_files.js | 20 +- .../model_registry/get_pipeline_files.js | 5 +- .../src/utils/model_registry/is_cached.js | 11 +- .../tests/generation_controller.test.js | 47 ++- .../tests/inference_backends.test.js | 24 +- .../tests/progress_callbacks.test.js | 33 +++ packages/transformers/tests/types.test.js | 4 + .../tests/utils/get_file_metadata.test.js | 54 ++++ .../tests/utils/model_registry.test.js | 62 ++++ pnpm-lock.yaml | 6 +- 34 files changed, 691 insertions(+), 406 deletions(-) create mode 100644 packages/transformers/tests/utils/get_file_metadata.test.js diff --git a/packages/transformers-onnx/README.md b/packages/transformers-onnx/README.md index f0705bb07..a72266aaa 100644 --- a/packages/transformers-onnx/README.md +++ b/packages/transformers-onnx/README.md @@ -7,3 +7,15 @@ import { OnnxInferenceProvider } from '@huggingface/transformers-onnx'; const provider = OnnxInferenceProvider.from_modelId('onnx-community/model-ONNX'); ``` + +## Runtime dependencies + +Node applications using the ONNX provider must install `onnxruntime-node` alongside Transformers.js: + +```sh +npm install @huggingface/transformers onnxruntime-node +``` + +The Node runtime is an optional peer so browser-only and custom-backend-only installations do not download native ONNX binaries. + +Browser ESM builds load this package lazily through the bare `@huggingface/transformers-onnx` specifier. Direct CDN usage therefore requires an import map for this package and its `onnxruntime-web` dependencies. Keep the provider's copied `.mjs` and `.wasm` files beside `transformers-onnx.web.js`; default relative `wasmPaths` are resolved from that module URL. Alternatively, set `env.backends.onnx.wasm.wasmPaths` before loading a model. diff --git a/packages/transformers-onnx/package.json b/packages/transformers-onnx/package.json index 67c94f9f0..3e9dbab83 100644 --- a/packages/transformers-onnx/package.json +++ b/packages/transformers-onnx/package.json @@ -37,14 +37,22 @@ }, "dependencies": { "onnxruntime-common": "1.24.3", - "onnxruntime-node": "1.24.3", "onnxruntime-web": "1.26.0-dev.20260416-b7804b056c" }, + "peerDependencies": { + "onnxruntime-node": "1.24.3" + }, + "peerDependenciesMeta": { + "onnxruntime-node": { + "optional": true + } + }, "devDependencies": { "@types/node": "^24.1.0", "@webgpu/types": "^0.1.69", "esbuild": "^0.27.2", "jest": "^30.2.0", + "onnxruntime-node": "1.24.3", "typescript": "5.9.3" }, "files": [ diff --git a/packages/transformers-onnx/scripts/build.mjs b/packages/transformers-onnx/scripts/build.mjs index b5fe605cb..bdac31c1a 100644 --- a/packages/transformers-onnx/scripts/build.mjs +++ b/packages/transformers-onnx/scripts/build.mjs @@ -1,10 +1,17 @@ import { build } from "esbuild"; -import { mkdir } from "node:fs/promises"; +import { copyFile, mkdir } from "node:fs/promises"; +import { fileURLToPath } from "node:url"; await mkdir(new URL("../dist/", import.meta.url), { recursive: true }); +const filePath = (path) => fileURLToPath(new URL(path, import.meta.url)); +const entry = filePath("../src/index.ts"); +const testingEntry = filePath("../src/testing.ts"); +const empty = filePath("../src/empty.ts"); +const wasmAssets = ["ort-wasm-simd-threaded.asyncify.mjs", "ort-wasm-simd-threaded.asyncify.wasm", "ort-wasm-simd-threaded.mjs", "ort-wasm-simd-threaded.wasm"]; + const shared = { - entryPoints: [new URL("../src/index.ts", import.meta.url).pathname], + entryPoints: [entry], bundle: true, sourcemap: false, logLevel: "warning", @@ -13,46 +20,50 @@ const shared = { await Promise.all([ build({ ...shared, - outfile: new URL("../dist/transformers-onnx.node.mjs", import.meta.url).pathname, + outfile: filePath("../dist/transformers-onnx.node.mjs"), platform: "node", format: "esm", external: ["onnxruntime-common", "onnxruntime-node"], - alias: { "onnxruntime-web/webgpu": "./src/empty.ts" }, + alias: { "onnxruntime-web/webgpu": empty }, + banner: { js: "const __ONNX_MODULE_URL__ = import.meta.url;" }, }), build({ ...shared, - outfile: new URL("../dist/transformers-onnx.node.cjs", import.meta.url).pathname, + outfile: filePath("../dist/transformers-onnx.node.cjs"), platform: "node", format: "cjs", external: ["onnxruntime-common", "onnxruntime-node"], - alias: { "onnxruntime-web/webgpu": "./src/empty.ts" }, + alias: { "onnxruntime-web/webgpu": empty }, + banner: { js: 'const __ONNX_MODULE_URL__ = require("node:url").pathToFileURL(__filename).href;' }, }), build({ ...shared, - outfile: new URL("../dist/transformers-onnx.web.js", import.meta.url).pathname, + outfile: filePath("../dist/transformers-onnx.web.js"), platform: "browser", format: "esm", external: ["onnxruntime-common", "onnxruntime-web"], - alias: { "onnxruntime-node": "./src/empty.ts" }, + alias: { "onnxruntime-node": empty }, + banner: { js: "const __ONNX_MODULE_URL__ = import.meta.url;" }, }), build({ - entryPoints: [new URL("../src/testing.ts", import.meta.url).pathname], + entryPoints: [testingEntry], bundle: true, sourcemap: false, logLevel: "warning", - outfile: new URL("../dist/testing.mjs", import.meta.url).pathname, + outfile: filePath("../dist/testing.mjs"), platform: "node", format: "esm", external: ["onnxruntime-common", "onnxruntime-node"], }), build({ - entryPoints: [new URL("../src/testing.ts", import.meta.url).pathname], + entryPoints: [testingEntry], bundle: true, sourcemap: false, logLevel: "warning", - outfile: new URL("../dist/testing.cjs", import.meta.url).pathname, + outfile: filePath("../dist/testing.cjs"), platform: "node", format: "cjs", external: ["onnxruntime-common", "onnxruntime-node"], }), + ...wasmAssets.map((file) => copyFile(filePath(`../node_modules/onnxruntime-web/dist/${file}`), filePath(`../dist/${file}`))), ]); diff --git a/packages/transformers-onnx/src/host.ts b/packages/transformers-onnx/src/host.ts index 1cfa1ecdd..c39ab1f6a 100644 --- a/packages/transformers-onnx/src/host.ts +++ b/packages/transformers-onnx/src/host.ts @@ -124,7 +124,14 @@ export function configureOnnxProviderHost(host: OnnxProviderHost): void { // settings when a host arrives later. A symbol-registered host was already configured before // module evaluation and contains the authoritative ORT environment. if (configuredHost === null && fallbackEnvironment.backends.onnx) { - host.env.backends.onnx = fallbackEnvironment.backends.onnx; + const target = (host.env.backends.onnx ?? {}) as Record; + const source = fallbackEnvironment.backends.onnx as Record; + const targetWasm = target.wasm; + const targetWebgpu = target.webgpu; + Object.assign(target, source); + target.wasm = Object.assign(source.wasm ?? {}, targetWasm ?? {}); + target.webgpu = Object.assign(source.webgpu ?? {}, targetWebgpu ?? {}); + host.env.backends.onnx = target; } configuredHost = host; } diff --git a/packages/transformers-onnx/src/runtime.ts b/packages/transformers-onnx/src/runtime.ts index 05c9c233e..bb6ca077e 100644 --- a/packages/transformers-onnx/src/runtime.ts +++ b/packages/transformers-onnx/src/runtime.ts @@ -18,6 +18,8 @@ import { getOnnxProviderHost } from './host.js'; +declare const __ONNX_MODULE_URL__: string; + // NOTE: Import order matters here. We need to import `onnxruntime-node` before `onnxruntime-web`. // In either case, we select the default export if it exists, otherwise we use the named export. import * as ONNX_NODE from 'onnxruntime-node'; @@ -34,7 +36,7 @@ function isBlobURL(url: string): boolean { } function toAbsoluteURL(url: string): string { - return new URL(url, globalThis.location?.href ?? 'file:///').href; + return new URL(url, globalThis.location?.href ?? __ONNX_MODULE_URL__).href; } type ExecutionProvider = OrtInferenceSession.ExecutionProviderConfig; @@ -288,9 +290,11 @@ export async function createInferenceSession( logSeverityLevel, ...session_options, }; - return typeof buffer_or_path === 'string' - ? InferenceSession.create(buffer_or_path, options) - : InferenceSession.create(buffer_or_path, options); + const create = InferenceSession.create as ( + model: Uint8Array | string, + options: OrtInferenceSession.SessionOptions, + ) => Promise; + return create(buffer_or_path, options); }; const session = await (apis.IS_WEB_ENV ? (webInitChain = webInitChain.then(load)) : load()); const configuredSession = session as OrtInferenceSession & { config: Record }; @@ -338,7 +342,17 @@ export function isONNXProxy() { } if (ONNX_ENV) { + const configured = (env.backends.onnx ??= {}) as Record; + const configuredWasm = configured.wasm; + const configuredWebgpu = configured.webgpu; + for (const [key, value] of Object.entries(configured)) { + if (key !== 'wasm' && key !== 'webgpu' && key !== 'setLogLevel') { + (ONNX_ENV as any)[key] = value; + } + } + if (ONNX_ENV.wasm) { + Object.assign(ONNX_ENV.wasm, configuredWasm ?? {}); // Initialize wasm backend with suitable default settings. // (Optional) Set path to wasm files. This will override the default path search behavior of onnxruntime-web. @@ -365,11 +379,12 @@ if (ONNX_ENV) { // Users may wish to proxy the WASM backend to prevent the UI from freezing, // However, this is not necessary when using WebGPU, so we default to false. - ONNX_ENV.wasm.proxy = false; + ONNX_ENV.wasm.proxy ??= false; } if (ONNX_ENV.webgpu) { - ONNX_ENV.webgpu.powerPreference = 'high-performance'; + Object.assign(ONNX_ENV.webgpu, configuredWebgpu ?? {}); + ONNX_ENV.webgpu.powerPreference ??= 'high-performance'; } /** @@ -386,8 +401,8 @@ if (ONNX_ENV) { setLogLevel(env.logLevel ?? LogLevel.WARNING); // Expose ONNX environment variables to `env.backends.onnx` - env.backends.onnx = { - ...ONNX_ENV, - setLogLevel, - }; + Object.assign(configured, ONNX_ENV, { setLogLevel }); + configured.wasm = ONNX_ENV.wasm; + configured.webgpu = ONNX_ENV.webgpu; + env.backends.onnx = configured; } diff --git a/packages/transformers/docs/source/guides/custom-backends.md b/packages/transformers/docs/source/guides/custom-backends.md index e75ea7c93..918a1730a 100644 --- a/packages/transformers/docs/source/guides/custom-backends.md +++ b/packages/transformers/docs/source/guides/custom-backends.md @@ -67,6 +67,8 @@ Transformers.js owns generation policy, stopping, callbacks, streaming, and fina - `sessionConcurrency.maxActiveSessions` is enforced fail-fast. Callers that want queueing must serialize generation themselves. - Abort signals are forwarded through loading and session creation. Backends must release partially created resources before propagating cancellation. +Protocol V1 does not export runtime KV state. With `return_dict_in_generate: true`, custom sessions return an empty `DynamicCache` for result-shape compatibility and dispose their runtime-owned cache with the session. + ## Artifact providers An `InferenceArtifactProvider` can supply JSON and random-access byte sources. Byte ranges are half-open (`[begin, end)`), independent reads may complete out of order, and returned arrays must be owned by the caller. `close()` is idempotent, rejects new reads, and waits for existing reads without implicitly aborting them. diff --git a/packages/transformers/scripts/build/buildAll.mjs b/packages/transformers/scripts/build/buildAll.mjs index 01b0e3c3d..59278bbfb 100644 --- a/packages/transformers/scripts/build/buildAll.mjs +++ b/packages/transformers/scripts/build/buildAll.mjs @@ -1,4 +1,5 @@ import { build as esbuild } from "esbuild"; +import { copyFile } from "node:fs/promises"; import path from "node:path"; import { stripNodePrefixPlugin } from "./plugins/stripNodePrefixPlugin.mjs"; import { ignoreModulesPlugin } from "./plugins/ignoreModulesPlugin.mjs"; @@ -67,4 +68,13 @@ export async function buildAll(log) { log.section(target.name); await buildTarget(target.config, log); } + + const ortDist = path.resolve(ROOT_DIR, "../transformers-onnx/node_modules/onnxruntime-web/dist"); + const wasmAssets = [ + "ort-wasm-simd-threaded.asyncify.mjs", + "ort-wasm-simd-threaded.asyncify.wasm", + "ort-wasm-simd-threaded.mjs", + "ort-wasm-simd-threaded.wasm", + ]; + await Promise.all(wasmAssets.map((file) => copyFile(path.join(ortDist, file), path.join(OUT_DIR, file)))); } diff --git a/packages/transformers/src/backends/default.js b/packages/transformers/src/backends/default.js index 1cbd6851c..1248cecef 100644 --- a/packages/transformers/src/backends/default.js +++ b/packages/transformers/src/backends/default.js @@ -24,11 +24,17 @@ let modulePromise; export function getOnnxProviderModule() { if (!modulePromise) { globalThis[ONNX_HOST_SYMBOL] = host; - modulePromise = import('@huggingface/transformers-onnx').then((module) => { - module.configureOnnxProviderHost(host); - TensorOpRegistry.register(module.OnnxTensorOpRegistry); - return module; - }); + const pending = import('@huggingface/transformers-onnx') + .then((module) => { + module.configureOnnxProviderHost(host); + TensorOpRegistry.register(module.OnnxTensorOpRegistry); + return module; + }) + .catch((error) => { + if (modulePromise === pending) modulePromise = undefined; + throw error; + }); + modulePromise = pending; } return modulePromise; } diff --git a/packages/transformers/src/backends/inference.js b/packages/transformers/src/backends/inference.js index 3a5c71fb8..3b21686dc 100644 --- a/packages/transformers/src/backends/inference.js +++ b/packages/transformers/src/backends/inference.js @@ -55,6 +55,7 @@ import { validateInferenceArtifactProvider } from './artifacts.js'; * config?: import('../configs.js').PretrainedConfig, * modelClass?: Function, * generation_config?: Record, + * artifactMetadata?: Record, * }} InferenceBackendLoadOptions */ @@ -64,6 +65,7 @@ import { validateInferenceArtifactProvider } from './artifacts.js'; * config?: import('../configs.js').PretrainedConfig, * modelClass?: Function, * generation_config?: Record, + * artifactMetadata?: Record, * }} InferenceModelLoadOptions */ @@ -78,11 +80,25 @@ import { validateInferenceArtifactProvider } from './artifacts.js'; * @property {() => Promise|unknown} dispose */ +/** + * A backend-provided default chat template. File sources default to the backend model ID and + * `chat_template.jinja`; inline content is never fetched or cached. + * + * @typedef { + * | {content: string, modelId?: never, file?: never} + * | {content?: never, modelId?: string, file?: string} + * } InferenceBackendChatTemplate + */ + /** * @typedef {Object} InferenceBackend * @property {string} modelId Model ID or local path used for shared config, tokenizer, and processor assets. + * @property {InferenceBackendChatTemplate} [chatTemplate] Default chat template installed on a pipeline tokenizer. * @property {string} [providerType] Provider family identifier used for provider-specific host initialization. * @property {StaticBackendCapabilities} [capabilities] + * @property {(options: Object) => ReadonlyArray|Promise>} [listModelArtifacts] Lists backend-owned files required for the selected load options. + * @property {(file: string, options: Object) => Promise<{size?: number, fromCache?: boolean}|null>} [getModelArtifactMetadata] Returns backend-owned cache metadata for an artifact. + * @property {(file: string, options: Object) => Promise} [deleteModelArtifact] Deletes an artifact from backend-owned cache storage. * @property {(options: InferenceBackendLoadOptions) => Promise} load * @property {(names: Record, options: Object, cacheSessions?: Object) => Promise>} [constructSessions] */ diff --git a/packages/transformers/src/backends/model_registry.js b/packages/transformers/src/backends/model_registry.js index 8bf10e9ad..deeaf352f 100644 --- a/packages/transformers/src/backends/model_registry.js +++ b/packages/transformers/src/backends/model_registry.js @@ -4,9 +4,9 @@ import { getOnnxProviderModule } from './default.js'; * Provider operations used by model-file discovery. * * @typedef {Object} ModelRegistryInferenceProvider - * @property {(options: Object) => string[]} listModelArtifacts - * @property {(options: Object) => Promise} getAvailableDtypes - * @property {(files: string[], sessions: Record) => string[]} filterModelArtifacts + * @property {(options: Object) => ReadonlyArray|Promise>} listModelArtifacts + * @property {(options: Object) => Promise} [getAvailableDtypes] + * @property {(files: string[], sessions: Record) => string[]} [filterModelArtifacts] */ /** @@ -16,7 +16,9 @@ import { getOnnxProviderModule } from './default.js'; * @returns {Promise} */ export async function getModelRegistryInferenceProvider(provider = null) { - if (provider) return provider; + if (typeof provider?.listModelArtifacts === 'function') return provider; + const providerClass = /** @type {any} */ (provider?.constructor); + if (typeof providerClass?.listModelArtifacts === 'function') return providerClass; const { OnnxInferenceProvider } = await getOnnxProviderModule(); return OnnxInferenceProvider; } diff --git a/packages/transformers/src/env.js b/packages/transformers/src/env.js index c44b8785e..0c5af6f42 100644 --- a/packages/transformers/src/env.js +++ b/packages/transformers/src/env.js @@ -239,8 +239,14 @@ export const env = { version: VERSION, /////////////////// Backends settings /////////////////// - // NOTE: These will be populated later by the backends themselves. - backends: {}, + // ONNX settings are available before the lazy provider loads so existing configuration code + // can set wasmPaths, numThreads, proxy, and WebGPU options immediately after importing. + backends: { + onnx: { + wasm: {}, + webgpu: {}, + }, + }, /////////////////// Logging settings /////////////////// get logLevel() { diff --git a/packages/transformers/src/generation/controller.js b/packages/transformers/src/generation/controller.js index e959e90c6..8f251b959 100644 --- a/packages/transformers/src/generation/controller.js +++ b/packages/transformers/src/generation/controller.js @@ -144,12 +144,23 @@ export function createStoppingCriteriaList(generationConfig, modelConfig, userCr return criteria; } +/** + * Resolve a `max_new_tokens` limit before constructing processors and stopping criteria. + * + * @param {GenerationConfig} generationConfig + * @param {number} inputLength + */ +export function prepareGenerationLength(generationConfig, inputLength) { + if (generationConfig.max_new_tokens !== null) { + generationConfig.max_length = inputLength + generationConfig.max_new_tokens; + } + return generationConfig; +} + /** * Stateful, inference-runtime-neutral generation policy. */ export class GenerationController { - version = 1; - /** * @param {Object} options * @param {Tensor} options.inputIds @@ -180,18 +191,13 @@ export class GenerationController { this.inputLength = inputIds.dims[1]; /** @type {bigint[][]} */ this.sequences = inputIds.tolist(); - this.scores = new Array(this.batchSize).fill(0); this.done = new Array(this.batchSize).fill(false); this.terminal = false; this.finalized = false; this.aborted = false; + this.abortReason = undefined; - if (generationConfig.max_new_tokens !== null) { - generationConfig.max_length = this.inputLength + generationConfig.max_new_tokens; - } - this.terminal = - generationConfig.max_new_tokens === 0 || - (generationConfig.max_length !== null && this.inputLength >= generationConfig.max_length); + this.terminal = generationConfig.max_new_tokens === 0; if (this.terminal) this.done.fill(true); this.sampler = LogitsSampler.getSampler(generationConfig); if (streamer) streamer.put(this.sequences.map((tokens) => [...tokens])); @@ -202,7 +208,8 @@ export class GenerationController { } get maxSequenceLength() { - return this.generationConfig.max_length; + if (this.generationConfig.max_new_tokens === 0) return this.inputLength; + return Math.max(this.generationConfig.max_length ?? 0, this.inputLength + 1); } /** @@ -232,20 +239,18 @@ export class GenerationController { ); } const tokenIds = new Uint32Array(this.batchSize); - const tokenScores = new Float64Array(this.batchSize); for (let batchIndex = 0; batchIndex < this.batchSize; ++batchIndex) { const sampled = await this.sampler(processed[batchIndex]); - const [tokenId, score] = sampled[0]; + const [tokenId] = sampled[0]; tokenIds[batchIndex] = Number(tokenId); - tokenScores[batchIndex] = score; } - return this.commit({ tokenIds, scores: tokenScores }); + return this.commit({ tokenIds }); } /** * Commit tokens selected by an approved runtime generation plan. * - * @param {{tokenIds: Uint32Array, processedScores?: Float32Array, scores?: Float64Array}} decision + * @param {{tokenIds: Uint32Array, processedScores?: Float32Array}} decision */ commit(decision) { this.#assertActive(); @@ -257,7 +262,6 @@ export class GenerationController { for (let index = 0; index < this.batchSize; ++index) { const tokenId = BigInt(decision.tokenIds[index]); this.sequences[index].push(tokenId); - this.scores[index] += decision.scores?.[index] ?? 0; generatedInputIds.push([tokenId]); } if (this.streamer) this.streamer.put(generatedInputIds); @@ -268,7 +272,6 @@ export class GenerationController { return { nextTokenIds, generatedInputIds, - done: [...this.done], allDone: this.terminal, }; } @@ -298,7 +301,7 @@ export class GenerationController { * @param {Object} [extra] */ finalize(extra = {}) { - if (this.aborted) throw new Error('Cannot finalize an aborted generation controller.'); + if (this.aborted) throw this.abortReason; if (this.finalized) throw new Error('Generation controller has already been finalized.'); if (!this.terminal) throw new Error('Cannot finalize generation before all sequences are done.'); this.finalized = true; @@ -309,15 +312,16 @@ export class GenerationController { return this.generationConfig.return_dict_in_generate ? { sequences, ...extra } : sequences; } - abort(_reason = undefined) { + abort(reason = undefined) { if (this.finalized || this.aborted) return; this.aborted = true; this.terminal = true; - if (this.streamer) this.streamer.end(); + this.abortReason = reason ?? new Error('Generation controller has been aborted.'); + this.streamer?.abort?.(this.abortReason); } #assertActive() { - if (this.aborted) throw new Error('Generation controller has been aborted.'); + if (this.aborted) throw this.abortReason; if (this.finalized) throw new Error('Generation controller has already been finalized.'); if (this.terminal) throw new Error('Generation controller is already complete.'); } @@ -345,9 +349,7 @@ export function createGenerationController(model, inputIds, options, collectOutp generationConfig: generation_config, kwargs, }); - if (generationConfig.max_new_tokens !== null) { - generationConfig.max_length = inputIds.dims.at(-1) + generationConfig.max_new_tokens; - } + prepareGenerationLength(generationConfig, inputIds.dims.at(-1)); return new GenerationController({ inputIds, generationConfig, diff --git a/packages/transformers/src/generation/runtime.js b/packages/transformers/src/generation/runtime.js index 5d835d94c..352498de8 100644 --- a/packages/transformers/src/generation/runtime.js +++ b/packages/transformers/src/generation/runtime.js @@ -1,4 +1,6 @@ import { Tensor } from '../utils/tensor.js'; +import { DynamicCache } from '../cache_utils.js'; +import { throwIfAborted } from '../utils/core.js'; import { createGenerationController } from './controller.js'; /** @@ -16,9 +18,6 @@ import { createGenerationController } from './controller.js'; * @property {boolean} cpuLogits * @property {ReadonlyArray<'argmax'>} declarativePlans * @property {{readonly defaultDepth: number, readonly maxDepth: number}} tokenPipeline - * @property {boolean} [customJavaScriptStoppingCriteria] - * @property {false} [cacheReorder] - * @property {false} [cacheExpand] * @property {SessionConcurrencyCapabilities} [sessionConcurrency] */ @@ -30,7 +29,6 @@ import { createGenerationController } from './controller.js'; * @property {'float32'} dtype * @property {readonly [number, number]} shape * @property {() => Promise} read - * @property {(plan: RuntimeGenerationPlanV1) => Promise} [select] * @property {() => void} release */ @@ -72,7 +70,6 @@ import { createGenerationController } from './controller.js'; * @typedef {Object} RuntimeTokenDecisionV1 * @property {Uint32Array} tokenIds * @property {Float32Array} [processedScores] - * @property {Float64Array} [scores] */ /** @@ -137,16 +134,23 @@ export function installGenerationRuntime(model) { * @param {Object} options */ export async function generateWithAutoregressiveSession(model, options) { - const { input_ids, attention_mask = null, signal = undefined } = options; - if (!(input_ids instanceof Tensor)) { - throw new TypeError('Custom autoregressive generation requires an `input_ids` Tensor.'); + const { inputs = null, input_ids = null, attention_mask = null, signal = undefined } = options; + if (inputs !== null && input_ids !== null) { + throw new TypeError('Custom autoregressive generation accepts either `inputs` or `input_ids`, but not both.'); + } + const resolvedInputIds = input_ids ?? inputs; + if (!(resolvedInputIds instanceof Tensor)) { + throw new TypeError('Custom autoregressive generation requires an `inputs` or `input_ids` Tensor.'); } if (model.config?.is_encoder_decoder) { throw new Error('Autoregressive session protocol version 1 only supports decoder-only models.'); } - const controller = createGenerationController(model, input_ids, options); - if (controller.allDone) return controller.finalize(); + const controller = createGenerationController(model, resolvedInputIds, options); + const finalize = () => + controller.finalize( + controller.generationConfig.return_dict_in_generate ? { past_key_values: new DynamicCache() } : {}, + ); const capabilities = getCausalGenerationCapabilities(model); try { @@ -156,6 +160,7 @@ export async function generateWithAutoregressiveSession(model, options) { controller.abort(error); throw error; } + if (controller.allDone) return finalize(); const plan = controller.compileRuntimePlan(capabilities); const mode = controller.generationConfig.do_sample ? 'multinomial' : 'greedy'; if (!plan && !capabilities.cpuLogits) { @@ -192,7 +197,7 @@ export async function generateWithAutoregressiveSession(model, options) { validateSession(session, controller, plan !== null); const prefillInputs = { - inputIds: tensorToTokenBatch(input_ids), + inputIds: tensorToTokenBatch(resolvedInputIds), attentionMask: attention_mask ? tensorToAttentionMask(attention_mask) : undefined, signal, }; @@ -212,13 +217,7 @@ export async function generateWithAutoregressiveSession(model, options) { if (!controller.allDone) { throw new Error('Autoregressive runtime ended its generation plan before generation completed.'); } - return controller.finalize(); - } - - if (!capabilities.cpuLogits) { - throw new Error( - 'This generation request requires CPU-visible logits, but the runtime does not support them.', - ); + return finalize(); } const pullSession = /** @type {PullAutoregressiveSessionV1} */ (session); @@ -248,13 +247,13 @@ export async function generateWithAutoregressiveSession(model, options) { signal, }); } - return controller.finalize(); + return finalize(); } catch (error) { controller.abort(error); throw error; } finally { try { - lease?.release(); + lease?.release?.(); } finally { try { await session?.dispose(); @@ -393,12 +392,6 @@ function isAllOnes(tensor) { return Array.from(tensor.data).every((value) => Number(value) === 1); } -function throwIfAborted(signal) { - if (!signal?.aborted) return; - if (typeof signal.throwIfAborted === 'function') signal.throwIfAborted(); - throw signal.reason ?? new Error('Generation aborted.'); -} - function acquireSessionSlot(model, capabilities) { const limit = capabilities.sessionConcurrency?.maxActiveSessions; if (limit === undefined) return () => {}; diff --git a/packages/transformers/src/generation/streamers.js b/packages/transformers/src/generation/streamers.js index 813dd959b..0e0a9a982 100644 --- a/packages/transformers/src/generation/streamers.js +++ b/packages/transformers/src/generation/streamers.js @@ -30,6 +30,12 @@ export class BaseStreamer { end() { throw Error('Not implemented'); } + + /** + * Function that is called when generation fails. Unlike `end()`, this must not signal successful completion. + * @param {unknown} _reason The generation failure. + */ + abort(_reason) {} } const stdout_write = apis.IS_PROCESS_AVAILABLE ? (x) => process.stdout.write(x) : (x) => console.log(x); @@ -152,6 +158,12 @@ export class TextStreamer extends BaseStreamer { this.on_finalized_text(printable_text, true); } + abort(_reason) { + this.token_cache = []; + this.print_len = 0; + this.next_tokens_are_prompt = true; + } + /** * Prints the new text to stdout. If the stream is ending, also prints a newline. * @param {string} text @@ -255,4 +267,9 @@ export class WhisperTextStreamer extends TextStreamer { super.end(); this.on_finalize?.(); } + + abort(reason) { + super.abort(reason); + this.waiting_for_timestamp = false; + } } diff --git a/packages/transformers/src/models/modeling_utils.js b/packages/transformers/src/models/modeling_utils.js index 80a2d4c95..66263cb3e 100644 --- a/packages/transformers/src/models/modeling_utils.js +++ b/packages/transformers/src/models/modeling_utils.js @@ -17,23 +17,14 @@ export function registerTaskMappings(mappings) { import { GITHUB_ISSUE_URL } from '../utils/constants.js'; import { getModelJSON } from '../utils/hub.js'; import { Seq2SeqLMOutput } from './modeling_outputs.js'; -import { - LogitsProcessorList, - ForcedBOSTokenLogitsProcessor, - ForcedEOSTokenLogitsProcessor, - SuppressTokensLogitsProcessor, - SuppressTokensAtBeginLogitsProcessor, - NoRepeatNGramLogitsProcessor, - RepetitionPenaltyLogitsProcessor, - NoBadWordsLogitsProcessor, - MinLengthLogitsProcessor, - MinNewTokensLengthLogitsProcessor, - TemperatureLogitsWarper, - ClassifierFreeGuidanceLogitsProcessor, -} from '../generation/logits_process.js'; import { GenerationConfig } from '../generation/configuration_utils.js'; -import { EosTokenCriteria, MaxLengthCriteria, StoppingCriteriaList } from '../generation/stopping_criteria.js'; -import { GenerationController } from '../generation/controller.js'; +import { + GenerationController, + createLogitsProcessorList, + createStoppingCriteriaList, + prepareGenerationConfig, + prepareGenerationLength, +} from '../generation/controller.js'; import { DefaultProgressCallback, pick } from '../utils/core.js'; import { ModelOutput } from './modeling_outputs.js'; import { logger } from '../utils/logger.js'; @@ -279,9 +270,41 @@ export class PreTrainedModel extends Callable { } if (isInferenceBackend(pretrained_model_name_or_path)) { const modelId = getModelId(pretrained_model_name_or_path); + /** @type {import('../backends/inference.js').InferenceModelLoadOptions} */ const resolvedOptions = { ...options }; resolvedOptions.config = resolvedOptions.config ?? (await AutoConfig.from_pretrained(modelId, resolvedOptions)); + if ( + resolvedOptions.progress_callback && + !(resolvedOptions.progress_callback instanceof DefaultProgressCallback) && + typeof pretrained_model_name_or_path.listModelArtifacts === 'function' + ) { + const expectedFiles = await pretrained_model_name_or_path.listModelArtifacts({ + ...resolvedOptions, + modelId, + }); + const metadata = await Promise.all( + expectedFiles.map((file) => + get_file_metadata(pretrained_model_name_or_path, file, resolvedOptions), + ), + ); + /** @type {import('../utils/core.js').FilesLoadingMap} */ + const filesLoading = {}; + resolvedOptions.artifactMetadata = {}; + metadata.forEach((entry, index) => { + if (!entry.exists) return; + const file = expectedFiles[index]; + resolvedOptions.artifactMetadata[file] = { size: entry.size, fromCache: entry.fromCache }; + filesLoading[file] = { + loaded: entry.fromCache ? (entry.size ?? 0) : 0, + total: entry.size ?? 0, + }; + }); + resolvedOptions.progress_callback = new DefaultProgressCallback( + resolvedOptions.progress_callback, + filesLoading, + ); + } // Custom models are duck-typed to the same runtime contract as PreTrainedModel. const model = /** @type {any} */ (await loadInferenceModel(pretrained_model_name_or_path, resolvedOptions)); if (typeof model.createAutoregressiveSession === 'function' && model.generation_config == null) { @@ -357,6 +380,7 @@ export class PreTrainedModel extends Callable { dtype, device, model_file_name, + inferenceProvider, }); const metadata = await Promise.all( @@ -425,7 +449,7 @@ export class PreTrainedModel extends Callable { /** * @param {GenerationConfig} generation_config * @param {number} input_ids_seq_length The starting sequence length for the input ids. - * @returns {LogitsProcessorList} + * @returns {import('../generation/logits_process.js').LogitsProcessorList} * @private */ _get_logits_processor( @@ -435,154 +459,7 @@ export class PreTrainedModel extends Callable { // prefix_allowed_tokens_fn, TODO logits_processor = null, ) { - const processors = new LogitsProcessorList(); - - // if (generation_config.diversity_penalty !== null && generation_config.diversity_penalty > 0.0) { - // processors.push(new HammingDiversityLogitsProcessor( - // generation_config.diversity_penalty, - // generation_config.num_beams, - // generation_config.num_beam_groups - // )); - // } - - // if (generation_config.encoder_repetition_penalty !== null && generation_config.encoder_repetition_penalty !== 1.0) { - // processors.push(new EncoderRepetitionPenaltyLogitsProcessor( - // generation_config.encoder_repetition_penalty, - // encoder_input_ids - // )); - // } - - if (generation_config.repetition_penalty !== null && generation_config.repetition_penalty !== 1.0) { - processors.push(new RepetitionPenaltyLogitsProcessor(generation_config.repetition_penalty)); - } - - if (generation_config.no_repeat_ngram_size !== null && generation_config.no_repeat_ngram_size > 0) { - processors.push(new NoRepeatNGramLogitsProcessor(generation_config.no_repeat_ngram_size)); - } - - // if (generation_config.encoder_no_repeat_ngram_size !== null && generation_config.encoder_no_repeat_ngram_size > 0) { - // if (this.config.is_encoder_decoder) { - // processors.push(new EncoderNoRepeatNGramLogitsProcessor( - // generation_config.encoder_no_repeat_ngram_size, - // encoder_input_ids - // )); - // } else { - // throw new Error("It's impossible to use `encoder_no_repeat_ngram_size` with decoder-only architecture"); - // } - // } - - if (generation_config.bad_words_ids !== null) { - processors.push( - new NoBadWordsLogitsProcessor(generation_config.bad_words_ids, generation_config.eos_token_id), - ); - } - - if ( - generation_config.min_length !== null && - generation_config.eos_token_id !== null && - generation_config.min_length > 0 - ) { - processors.push(new MinLengthLogitsProcessor(generation_config.min_length, generation_config.eos_token_id)); - } - - if ( - generation_config.min_new_tokens !== null && - generation_config.eos_token_id !== null && - generation_config.min_new_tokens > 0 - ) { - processors.push( - new MinNewTokensLengthLogitsProcessor( - input_ids_seq_length, - generation_config.min_new_tokens, - generation_config.eos_token_id, - ), - ); - } - - // if (prefix_allowed_tokens_fn !== null) { - // processors.push(new PrefixConstrainedLogitsProcessor( - // prefix_allowed_tokens_fn, - // generation_config.num_beams / generation_config.num_beam_groups - // )); - // } - - if (generation_config.forced_bos_token_id !== null) { - processors.push(new ForcedBOSTokenLogitsProcessor(generation_config.forced_bos_token_id)); - } - - if (generation_config.forced_eos_token_id !== null) { - processors.push( - new ForcedEOSTokenLogitsProcessor(generation_config.max_length, generation_config.forced_eos_token_id), - ); - } - - // if (generation_config.remove_invalid_values === true) { - // processors.push(new InfNanRemoveLogitsProcessor()); - // } - - // if (generation_config.exponential_decay_length_penalty !== null) { - // processors.push(new ExponentialDecayLengthPenalty( - // generation_config.exponential_decay_length_penalty, - // generation_config.eos_token_id, - // input_ids_seq_length - // )); - // } - - if (generation_config.suppress_tokens !== null) { - processors.push(new SuppressTokensLogitsProcessor(generation_config.suppress_tokens)); - } - - if (generation_config.begin_suppress_tokens !== null) { - const begin_index = - input_ids_seq_length > 1 || generation_config.forced_bos_token_id === null - ? input_ids_seq_length - : input_ids_seq_length + 1; - - processors.push( - new SuppressTokensAtBeginLogitsProcessor(generation_config.begin_suppress_tokens, begin_index), - ); - } - - // DEPRECATED: https://github.com/huggingface/transformers/pull/29485 - // if (generation_config.forced_decoder_ids !== null) { - // processors.push(new ForceTokensLogitsProcessor(generation_config.forced_decoder_ids)); - // } - - // 8. prepare batched CFG externally - if (generation_config.guidance_scale !== null && generation_config.guidance_scale > 1) { - processors.push(new ClassifierFreeGuidanceLogitsProcessor(generation_config.guidance_scale)); - } - - if (generation_config.temperature === 0 && generation_config.do_sample) { - logger.warn( - '`do_sample` changed to false because `temperature: 0` implies greedy sampling (always selecting the most likely token), which is incompatible with `do_sample: true`.', - ); - generation_config.do_sample = false; - } - - if (generation_config.do_sample) { - if (generation_config.temperature !== null && generation_config.temperature !== 1.0) { - processors.push(new TemperatureLogitsWarper(generation_config.temperature)); - } - // TODO: Add TopPLogitsWarper and TopKLogitsWarper - // if (generation_config.top_k !== null && generation_config.top_k !== 0) { - // processors.push(new TopKLogitsWarper(generation_config.top_k)); - // } - // if (generation_config.top_p !== null && generation_config.top_p < 1.0) { - // processors.push(new TopPLogitsWarper(generation_config.top_p)); - // } - } - - if (logits_processor !== null) { - processors.extend(logits_processor); - } - - // `LogitNormalization` should always be the last logit processor, when present - // if (generation_config.renormalize_logits === true) { - // processors.push(new LogitNormalization()); - // } - - return processors; + return createLogitsProcessorList(generation_config, input_ids_seq_length, logits_processor); } /** @@ -593,60 +470,22 @@ export class PreTrainedModel extends Callable { * @returns {GenerationConfig} The final generation config object to be used by the model for text generation. */ _prepare_generation_config(generation_config, kwargs, cls = GenerationConfig) { - // Create empty generation config (contains defaults) - // We pass `this.config` so that if `eos_token_id` or `bos_token_id` exist in the model's config, we will use them - const config = { ...this.config }; - for (const key of ['decoder', 'generator', 'text_config']) { - // Special case: some models have generation attributes set in the decoder. - // Use them if still unset in the generation config. - if (key in config) { - Object.assign(config, config[key]); - } - } - - const gen_config = new cls(config); - - // Apply model's generation config, if it exists - Object.assign(gen_config, this.generation_config ?? {}); - - // Next, use any generation config specified by the user - // when calling `generate` - if (generation_config) { - Object.assign(gen_config, generation_config); - } - - // Finally, if any kwargs were passed, use them to overwrite - if (kwargs) { - Object.assign(gen_config, pick(kwargs, Object.getOwnPropertyNames(gen_config))); - } - - return gen_config; + return prepareGenerationConfig({ + modelConfig: this.config, + modelGenerationConfig: this.generation_config, + generationConfig: generation_config, + kwargs, + configClass: cls, + }); } /** * * @param {GenerationConfig} generation_config - * @param {import('../generation/stopping_criteria.js').StoppingCriteria|import('../generation/stopping_criteria.js').StoppingCriteria[]|StoppingCriteriaList} [stopping_criteria=null] + * @param {import('../generation/stopping_criteria.js').StoppingCriteria|import('../generation/stopping_criteria.js').StoppingCriteria[]|import('../generation/stopping_criteria.js').StoppingCriteriaList} [stopping_criteria=null] */ _get_stopping_criteria(generation_config, stopping_criteria = null) { - const criteria = new StoppingCriteriaList(); - - if (generation_config.max_length !== null) { - criteria.push( - new MaxLengthCriteria(generation_config.max_length, this.config.max_position_embeddings ?? null), - ); - } - // if (generation_config.max_time !== null) { - // criteria.push(new MaxTimeCriteria(generation_config.max_time)); - // } - if (generation_config.eos_token_id !== null) { - criteria.push(new EosTokenCriteria(generation_config.eos_token_id)); - } - - if (stopping_criteria) { - criteria.extend(stopping_criteria); - } - return criteria; + return createStoppingCriteriaList(generation_config, this.config, stopping_criteria); } /** @@ -926,9 +765,7 @@ export class PreTrainedModel extends Callable { // 6. Prepare `max_length` depending on other stopping criteria. let input_ids_length = input_ids.dims.at(-1); - if (generation_config.max_new_tokens !== null) { - generation_config.max_length = input_ids_length + generation_config.max_new_tokens; - } + prepareGenerationLength(generation_config, input_ids_length); // input_ids_length = model_inputs[model_input_name].dims.at(1); // // inputs instanceof Tensor ? : inputs.length; diff --git a/packages/transformers/src/pipelines.js b/packages/transformers/src/pipelines.js index 3c43b4595..58226a6fb 100644 --- a/packages/transformers/src/pipelines.js +++ b/packages/transformers/src/pipelines.js @@ -59,7 +59,8 @@ import { validateInferenceModelTask, } from './backends/inference.js'; import { validateInferenceArtifactProvider } from './backends/artifacts.js'; -import { getModelJSON } from './utils/hub.js'; +import { getModelJSON, getModelText } from './utils/hub.js'; +import { CHAT_TEMPLATE_NAME } from './utils/constants.js'; /** * @typedef {keyof typeof SUPPORTED_TASKS} TaskType @@ -70,6 +71,36 @@ import { getModelJSON } from './utils/hub.js'; * @typedef {SupportedTasks & AliasTasks} AllTasks A mapping from all pipeline names and aliases to their corresponding pipeline classes. */ +/** + * Resolve a custom backend's default chat template without moving rendering policy into the backend. + * + * @param {import('./backends/inference.js').InferenceBackend} backend + * @param {import('./utils/hub.js').PretrainedModelOptions} options + * @returns {Promise|null} + */ +export function loadInferenceBackendChatTemplate(backend, options) { + const source = backend.chatTemplate; + if (source == null) return null; + if (typeof source !== 'object') { + throw new TypeError('Inference backend `chatTemplate` must be an object.'); + } + if (Object.hasOwn(source, 'content')) { + if (typeof source.content !== 'string' || Object.hasOwn(source, 'modelId') || Object.hasOwn(source, 'file')) { + throw new TypeError( + 'Inference backend `chatTemplate.content` must be a string and cannot be combined with `modelId` or `file`.', + ); + } + return Promise.resolve(source.content); + } + if (source.modelId !== undefined && typeof source.modelId !== 'string') { + throw new TypeError('Inference backend `chatTemplate.modelId` must be a string.'); + } + if (source.file !== undefined && typeof source.file !== 'string') { + throw new TypeError('Inference backend `chatTemplate.file` must be a string.'); + } + return getModelText(source.modelId ?? backend.modelId, source.file ?? CHAT_TEMPLATE_NAME, true, options); +} + /** * Utility factory method to build a `Pipeline` object. * @@ -142,6 +173,10 @@ export async function pipeline( } const customBackend = isInferenceBackend(model) && typeof model.constructSessions !== 'function'; + const customRegistryProvider = + customBackend && typeof (/** @type {any} */ (model).listModelArtifacts) === 'function' + ? /** @type {import('./backends/model_registry.js').ModelRegistryInferenceProvider} */ (model) + : null; const modelId = getModelId(model); validateInferenceArtifactProvider(artifactProvider); if (customBackend) { @@ -157,22 +192,31 @@ export async function pipeline( local_files_only, revision, model_file_name, - include_model: !customBackend, + inferenceProvider: customRegistryProvider, + include_model: !customBackend || customRegistryProvider !== null, }); /** @type {import('./utils/core.js').FilesLoadingMap} */ let files_loading = {}; + /** @type {Record} */ + const artifactMetadata = {}; if (progress_callback) { /** @type {Array<{exists: boolean, size?: number, contentType?: string, fromCache?: boolean}>} */ const metadata = await Promise.all( expected_files.map(async (file) => - get_file_metadata(modelId, file, { cache_dir, local_files_only, revision }), + get_file_metadata(customBackend ? model : modelId, file, { + cache_dir, + local_files_only, + revision, + signal, + }), ), ); metadata.forEach((m, i) => { if (m.exists) { + artifactMetadata[expected_files[i]] = { size: m.size, fromCache: m.fromCache }; files_loading[expected_files[i]] = { - loaded: 0, + loaded: m.fromCache ? (m.size ?? 0) : 0, total: m.size ?? 0, }; } @@ -196,6 +240,7 @@ export async function pipeline( generation_config: null, signal, artifactProvider, + artifactMetadata, }; // Determine which components to load based on the expected files @@ -237,13 +282,21 @@ export async function pipeline( let tokenizer; let processor; let model_loaded; + let chat_template; try { // Load all components in parallel. - [tokenizer, processor, model_loaded] = await Promise.all([ + [tokenizer, processor, model_loaded, chat_template] = await Promise.all([ hasTokenizer ? AutoTokenizer.from_pretrained(modelId, pretrainedOptions) : null, hasProcessor ? AutoProcessor.from_pretrained(modelId, pretrainedOptions) : null, modelPromise, + customBackend && hasTokenizer + ? loadInferenceBackendChatTemplate( + /** @type {import('./backends/inference.js').InferenceBackend} */ (model), + pretrainedOptions, + ) + : null, ]); + if (tokenizer && chat_template != null) tokenizer.chat_template = chat_template; if (customBackend) { validateInferenceModelTask(model_loaded, task); } diff --git a/packages/transformers/src/transformers.js b/packages/transformers/src/transformers.js index e142c7f37..cb1523603 100644 --- a/packages/transformers/src/transformers.js +++ b/packages/transformers/src/transformers.js @@ -73,6 +73,7 @@ export { getModelId, isInferenceBackend } from './backends/inference.js'; * @typedef {import('./utils/core.js').ProgressCallback} ProgressCallback * @typedef {import('./utils/core.js').ProgressInfo} ProgressInfo * @typedef {import('./backends/inference.js').InferenceBackend} InferenceBackend + * @typedef {import('./backends/inference.js').InferenceBackendChatTemplate} InferenceBackendChatTemplate * @typedef {import('./backends/inference.js').InferenceBackendLoadOptions} InferenceBackendLoadOptions * @typedef {import('./backends/inference.js').InferenceModel} InferenceModel * @typedef {import('./backends/inference.js').InferenceModelCapabilities} InferenceModelCapabilities diff --git a/packages/transformers/src/utils/core.js b/packages/transformers/src/utils/core.js index 3304d5f90..f64e460aa 100644 --- a/packages/transformers/src/utils/core.js +++ b/packages/transformers/src/utils/core.js @@ -92,6 +92,18 @@ export function dispatchCallback(progress_callback, data) { if (progress_callback) progress_callback(data); } +/** + * Throw an abort signal's reason when cancellation has been requested. + * + * @param {AbortSignal|null|undefined} signal + * @param {string} [fallbackMessage] + */ +export function throwIfAborted(signal, fallbackMessage = 'Operation aborted.') { + if (!signal?.aborted) return; + if (typeof signal.throwIfAborted === 'function') signal.throwIfAborted(); + throw signal.reason ?? new Error(fallbackMessage); +} + /** * A callable progress callback that wraps an original callback and emits * aggregate `progress_total` events. Because it extends `Callable`, instances @@ -110,6 +122,8 @@ export class DefaultProgressCallback extends Callable { super(); this.callback = callback; this.files_loading = files_loading; + this.lastAggregateLoaded = -1; + this.lastAggregateTotal = -1; /** @type {Map>} Pending and completed file loads, used to deduplicate work within a single pipeline() call. */ this.loads = new Map(); } @@ -119,23 +133,34 @@ export class DefaultProgressCallback extends Callable { */ _call(info) { if (info.status === 'progress') { + const previous = this.files_loading[info.file]; + const previousLoaded = Number.isFinite(previous?.loaded) ? previous.loaded : 0; + const previousTotal = Number.isFinite(previous?.total) ? previous.total : 0; + const reportedLoaded = Number.isFinite(info.loaded) ? Math.max(0, info.loaded) : previousLoaded; + const reportedTotal = Number.isFinite(info.total) ? Math.max(0, info.total) : 0; + const total = Math.max(previousTotal, reportedTotal, reportedLoaded); + const loaded = Math.min(total, Math.max(previousLoaded, reportedLoaded)); this.files_loading[info.file] = { - loaded: info.loaded, - total: info.total, + loaded, + total, }; - const loaded = Object.values(this.files_loading).reduce((acc, curr) => acc + curr.loaded, 0); - const total = Object.values(this.files_loading).reduce((acc, curr) => acc + curr.total, 0); - const progress = total > 0 ? (loaded / total) * 100 : 0; + const aggregateLoaded = Object.values(this.files_loading).reduce((acc, curr) => acc + curr.loaded, 0); + const aggregateTotal = Object.values(this.files_loading).reduce((acc, curr) => acc + curr.total, 0); + const progress = aggregateTotal > 0 ? (aggregateLoaded / aggregateTotal) * 100 : 0; - this.callback({ - status: 'progress_total', - name: info.name, - progress, - loaded, - total, - files: structuredClone(this.files_loading), - }); + if (aggregateLoaded !== this.lastAggregateLoaded || aggregateTotal !== this.lastAggregateTotal) { + this.lastAggregateLoaded = aggregateLoaded; + this.lastAggregateTotal = aggregateTotal; + this.callback({ + status: 'progress_total', + name: info.name, + progress, + loaded: aggregateLoaded, + total: aggregateTotal, + files: structuredClone(this.files_loading), + }); + } } this.callback(info); } diff --git a/packages/transformers/src/utils/hub.js b/packages/transformers/src/utils/hub.js index e3a0071dc..cf587c23a 100755 --- a/packages/transformers/src/utils/hub.js +++ b/packages/transformers/src/utils/hub.js @@ -5,7 +5,7 @@ */ import { apis, env } from '../env.js'; -import { DefaultProgressCallback, dispatchCallback } from './core.js'; +import { DefaultProgressCallback, dispatchCallback, throwIfAborted } from './core.js'; import { FileResponse } from './hub/FileResponse.js'; import { FileCache } from './cache/FileCache.js'; import { @@ -40,10 +40,10 @@ export { MAX_EXTERNAL_DATA_CHUNKS } from './hub/constants.js'; * @property {string} [cache_dir=null] Path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used. * @property {boolean} [local_files_only=false] Whether or not to only look at local files (e.g., not try downloading the model). * @property {string} [revision='main'] The specific model version to use. It can be a branch name, a tag name, or a commit id, - * @property {AbortSignal} [signal] Signal used to cancel backend-owned model loading. - * @property {import('../backends/artifacts.js').InferenceArtifactProvider} [artifactProvider] Optional random-access artifact provider supplied to custom inference backends. * since we use a git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any identifier allowed by git. * NOTE: This setting is ignored for local requests. + * @property {AbortSignal} [signal] Signal used to cancel backend-owned model loading. + * @property {import('../backends/artifacts.js').InferenceArtifactProvider} [artifactProvider] Optional random-access artifact provider supplied to custom inference backends. */ /** @@ -490,12 +490,6 @@ export async function loadResourceFile( throw new Error('Unable to get model file path or buffer.'); } -function throwIfAborted(signal) { - if (!signal?.aborted) return; - if (typeof signal.throwIfAborted === 'function') signal.throwIfAborted(); - throw signal.reason ?? new Error('Model loading aborted.'); -} - /** @type {Map>} Pending file loads keyed by resource identity. */ const INFLIGHT_LOADS = new Map(); diff --git a/packages/transformers/src/utils/model_registry/ModelRegistry.js b/packages/transformers/src/utils/model_registry/ModelRegistry.js index 04d6c6e01..b38731ec8 100644 --- a/packages/transformers/src/utils/model_registry/ModelRegistry.js +++ b/packages/transformers/src/utils/model_registry/ModelRegistry.js @@ -110,6 +110,19 @@ import { is_cached, is_cached_files, is_pipeline_cached, is_pipeline_cached_file import { get_file_metadata } from './get_file_metadata.js'; import { clear_cache, clear_pipeline_cache } from './clear_cache.js'; import { get_available_dtypes } from './get_available_dtypes.js'; +import { getModelId, isInferenceBackend } from '../../backends/inference.js'; + +function resolveModelRegistryRequest(model, options) { + const inferenceBackend = isInferenceBackend(model) ? model : null; + const inferenceProvider = + inferenceBackend && typeof inferenceBackend.listModelArtifacts === 'function' + ? inferenceBackend + : options.inferenceProvider; + return { + modelId: getModelId(model), + options: inferenceProvider || inferenceBackend ? { ...options, inferenceProvider, inferenceBackend } : options, + }; +} /** * Static class for cache and file management operations. @@ -119,7 +132,7 @@ export class ModelRegistry { /** * Get all files (model, tokenizer, processor) needed for a model. * - * @param {string} modelId - The model id (e.g., "onnx-community/bert-base-uncased-ONNX") + * @param {string|import('../../backends/inference.js').InferenceBackend} model - The model id or custom inference backend. * @param {Object} [options] - Optional parameters * @param {import('../../configs.js').PretrainedConfig} [options.config=null] - Pre-loaded config * @param {import('../dtypes.js').DataType|Record} [options.dtype=null] - Override dtype @@ -133,8 +146,9 @@ export class ModelRegistry { * const files = await ModelRegistry.get_files('onnx-community/gpt2-ONNX'); * console.log(files); // ['config.json', 'tokenizer.json', 'onnx/model_q4.onnx', ...] */ - static async get_files(modelId, options = {}) { - return get_files(modelId, options); + static async get_files(model, options = {}) { + const request = resolveModelRegistryRequest(model, options); + return get_files(request.modelId, request.options); } /** @@ -142,7 +156,7 @@ export class ModelRegistry { * Automatically determines which components are needed based on the task. * * @param {string} task - The pipeline task (e.g., "text-generation", "background-removal") - * @param {string} modelId - The model id (e.g., "onnx-community/bert-base-uncased-ONNX") + * @param {string|import('../../backends/inference.js').InferenceBackend} model - The model id or custom inference backend. * @param {Object} [options] - Optional parameters * @param {import('../../configs.js').PretrainedConfig} [options.config=null] - Pre-loaded config * @param {import('../dtypes.js').DataType|Record} [options.dtype=null] - Override dtype @@ -154,14 +168,15 @@ export class ModelRegistry { * const files = await ModelRegistry.get_pipeline_files('text-generation', 'onnx-community/gpt2-ONNX'); * console.log(files); // ['config.json', 'tokenizer.json', 'onnx/model_q4.onnx', ...] */ - static async get_pipeline_files(task, modelId, options = {}) { - return get_pipeline_files(task, modelId, options); + static async get_pipeline_files(task, model, options = {}) { + const request = resolveModelRegistryRequest(model, options); + return get_pipeline_files(task, request.modelId, request.options); } /** * Get model files needed for a specific model. * - * @param {string} modelId - The model id + * @param {string|import('../../backends/inference.js').InferenceBackend} model - The model id or custom inference backend. * @param {Object} [options] - Optional parameters * @param {import('../../configs.js').PretrainedConfig} [options.config=null] - Pre-loaded config * @param {import('../dtypes.js').DataType|Record} [options.dtype=null] - Override dtype @@ -173,14 +188,15 @@ export class ModelRegistry { * const files = await ModelRegistry.get_model_files('onnx-community/bert-base-uncased-ONNX'); * console.log(files); // ['config.json', 'onnx/model_q4.onnx', 'generation_config.json'] */ - static async get_model_files(modelId, options = {}) { - return get_model_files(modelId, options); + static async get_model_files(model, options = {}) { + const request = resolveModelRegistryRequest(model, options); + return get_model_files(request.modelId, request.options); } /** * Get tokenizer files needed for a specific model. * - * @param {string} modelId - The model id + * @param {string|import('../../backends/inference.js').InferenceBackend} model - The model id or custom inference backend * @param {Object} [options] - Hub metadata options * @returns {Promise} Array of tokenizer file paths * @@ -188,14 +204,15 @@ export class ModelRegistry { * const files = await ModelRegistry.get_tokenizer_files('onnx-community/gpt2-ONNX'); * console.log(files); // ['tokenizer.json', 'tokenizer_config.json'] */ - static async get_tokenizer_files(modelId, options = {}) { - return get_tokenizer_files(modelId, options); + static async get_tokenizer_files(model, options = {}) { + const request = resolveModelRegistryRequest(model, options); + return get_tokenizer_files(request.modelId, request.options); } /** * Get processor files needed for a specific model. * - * @param {string} modelId - The model id + * @param {string|import('../../backends/inference.js').InferenceBackend} model - The model id or custom inference backend * @param {Object} [options] - Hub metadata options * @returns {Promise} Array of processor file paths * @@ -203,8 +220,9 @@ export class ModelRegistry { * const files = await ModelRegistry.get_processor_files('onnx-community/vit-base-patch16-224-ONNX'); * console.log(files); // ['preprocessor_config.json'] */ - static async get_processor_files(modelId, options = {}) { - return get_processor_files(modelId, options); + static async get_processor_files(model, options = {}) { + const request = resolveModelRegistryRequest(model, options); + return get_processor_files(request.modelId, request.options); } /** @@ -214,7 +232,7 @@ export class ModelRegistry { * A dtype is considered available if all required model session files * exist for that dtype. * - * @param {string} modelId - The model id (e.g., "onnx-community/all-MiniLM-L6-v2-ONNX") + * @param {string|import('../../backends/inference.js').InferenceBackend} model - The model id or custom inference backend * @param {Object} [options] - Optional parameters * @param {import('../../configs.js').PretrainedConfig} [options.config=null] - Pre-loaded config * @param {string} [options.model_file_name=null] - Override the model file name (excluding .onnx suffix) @@ -227,8 +245,9 @@ export class ModelRegistry { * const dtypes = await ModelRegistry.get_available_dtypes('onnx-community/all-MiniLM-L6-v2-ONNX'); * console.log(dtypes); // ['fp32', 'fp16', 'int8', 'uint8', 'q8', 'q4'] */ - static async get_available_dtypes(modelId, options = {}) { - return get_available_dtypes(modelId, options); + static async get_available_dtypes(model, options = {}) { + const request = resolveModelRegistryRequest(model, options); + return get_available_dtypes(request.modelId, request.options); } /** @@ -236,7 +255,7 @@ export class ModelRegistry { * then confirming all required files are cached. * Returns a plain boolean — use `is_cached_files` if you need per-file detail. * - * @param {string} modelId - The model id + * @param {string|import('../../backends/inference.js').InferenceBackend} model - The model id or custom inference backend * @param {Object} [options] - Optional parameters * @param {string} [options.cache_dir] - Custom cache directory * @param {string} [options.revision] - Model revision (default: 'main') @@ -249,15 +268,16 @@ export class ModelRegistry { * const cached = await ModelRegistry.is_cached('onnx-community/bert-base-uncased-ONNX'); * console.log(cached); // true or false */ - static async is_cached(modelId, options = {}) { - return is_cached(modelId, options); + static async is_cached(model, options = {}) { + const request = resolveModelRegistryRequest(model, options); + return is_cached(request.modelId, request.options); } /** * Checks if all files for a given model are already cached, with per-file detail. * Automatically determines which files are needed using get_files(). * - * @param {string} modelId - The model id + * @param {string|import('../../backends/inference.js').InferenceBackend} model - The model id or custom inference backend * @param {Object} [options] - Optional parameters * @param {string} [options.cache_dir] - Custom cache directory * @param {string} [options.revision] - Model revision (default: 'main') @@ -271,8 +291,9 @@ export class ModelRegistry { * console.log(status.allCached); // true or false * console.log(status.files); // [{ file: 'config.json', cached: true }, ...] */ - static async is_cached_files(modelId, options = {}) { - return is_cached_files(modelId, options); + static async is_cached_files(model, options = {}) { + const request = resolveModelRegistryRequest(model, options); + return is_cached_files(request.modelId, request.options); } /** @@ -281,7 +302,7 @@ export class ModelRegistry { * Returns a plain boolean — use `is_pipeline_cached_files` if you need per-file detail. * * @param {string} task - The pipeline task (e.g., "text-generation", "background-removal") - * @param {string} modelId - The model id + * @param {string|import('../../backends/inference.js').InferenceBackend} model - The model id or custom inference backend * @param {Object} [options] - Optional parameters * @param {string} [options.cache_dir] - Custom cache directory * @param {string} [options.revision] - Model revision (default: 'main') @@ -294,8 +315,9 @@ export class ModelRegistry { * const cached = await ModelRegistry.is_pipeline_cached('text-generation', 'onnx-community/gpt2-ONNX'); * console.log(cached); // true or false */ - static async is_pipeline_cached(task, modelId, options = {}) { - return is_pipeline_cached(task, modelId, options); + static async is_pipeline_cached(task, model, options = {}) { + const request = resolveModelRegistryRequest(model, options); + return is_pipeline_cached(task, request.modelId, request.options); } /** @@ -303,7 +325,7 @@ export class ModelRegistry { * Automatically determines which components are needed based on the task. * * @param {string} task - The pipeline task (e.g., "text-generation", "background-removal") - * @param {string} modelId - The model id + * @param {string|import('../../backends/inference.js').InferenceBackend} model - The model id or custom inference backend * @param {Object} [options] - Optional parameters * @param {string} [options.cache_dir] - Custom cache directory * @param {string} [options.revision] - Model revision (default: 'main') @@ -317,14 +339,15 @@ export class ModelRegistry { * console.log(status.allCached); // true or false * console.log(status.files); // [{ file: 'config.json', cached: true }, ...] */ - static async is_pipeline_cached_files(task, modelId, options = {}) { - return is_pipeline_cached_files(task, modelId, options); + static async is_pipeline_cached_files(task, model, options = {}) { + const request = resolveModelRegistryRequest(model, options); + return is_pipeline_cached_files(task, request.modelId, request.options); } /** * Get metadata for a specific file without downloading it. * - * @param {string} path_or_repo_id - Model id or path + * @param {string|import('../../backends/inference.js').InferenceBackend} path_or_repo_id - Model id, path, or custom backend * @param {string} filename - The file name * @param {import('../hub.js').PretrainedOptions} [options] - Optional parameters * @returns {Promise<{exists: boolean, size?: number, contentType?: string, fromCache?: boolean}>} File metadata @@ -341,7 +364,7 @@ export class ModelRegistry { * Clears all cached files for a given model. * Automatically determines which files are needed and removes them from the cache. * - * @param {string} modelId - The model id (e.g., "onnx-community/gpt2-ONNX") + * @param {string|import('../../backends/inference.js').InferenceBackend} model - The model id or custom inference backend * @param {Object} [options] - Optional parameters * @param {string} [options.cache_dir] - Custom cache directory * @param {string} [options.revision] - Model revision (default: 'main') @@ -356,8 +379,9 @@ export class ModelRegistry { * const result = await ModelRegistry.clear_cache('onnx-community/bert-base-uncased-ONNX'); * console.log(`Deleted ${result.filesDeleted} of ${result.filesCached} cached files`); */ - static async clear_cache(modelId, options = {}) { - return clear_cache(modelId, options); + static async clear_cache(model, options = {}) { + const request = resolveModelRegistryRequest(model, options); + return clear_cache(request.modelId, request.options); } /** @@ -365,7 +389,7 @@ export class ModelRegistry { * Automatically determines which components are needed based on the task. * * @param {string} task - The pipeline task (e.g., "text-generation", "image-classification") - * @param {string} modelId - The model id (e.g., "onnx-community/gpt2-ONNX") + * @param {string|import('../../backends/inference.js').InferenceBackend} model - The model id or custom inference backend * @param {Object} [options] - Optional parameters * @param {string} [options.cache_dir] - Custom cache directory * @param {string} [options.revision] - Model revision (default: 'main') @@ -378,7 +402,8 @@ export class ModelRegistry { * const result = await ModelRegistry.clear_pipeline_cache('text-generation', 'onnx-community/gpt2-ONNX'); * console.log(`Deleted ${result.filesDeleted} of ${result.filesCached} cached files`); */ - static async clear_pipeline_cache(task, modelId, options = {}) { - return clear_pipeline_cache(task, modelId, options); + static async clear_pipeline_cache(task, model, options = {}) { + const request = resolveModelRegistryRequest(model, options); + return clear_pipeline_cache(task, request.modelId, request.options); } } diff --git a/packages/transformers/src/utils/model_registry/clear_cache.js b/packages/transformers/src/utils/model_registry/clear_cache.js index fe4850e8c..f667f2359 100644 --- a/packages/transformers/src/utils/model_registry/clear_cache.js +++ b/packages/transformers/src/utils/model_registry/clear_cache.js @@ -35,21 +35,21 @@ import { get_pipeline_files } from './get_pipeline_files.js'; */ async function clear_files_from_cache(modelId, files, options = {}) { const cache = await getCache(options?.cache_dir); - - if (!cache) { - return { - filesDeleted: 0, - filesCached: 0, - files: files.map((filename) => ({ file: filename, deleted: false, wasCached: false })), - }; - } - - if (!cache.delete) { + if (cache && !cache.delete) { throw new Error('Cache does not support delete operation'); } const results = await Promise.all( files.map(async (filename) => { + const backendMetadata = await options.inferenceBackend?.getModelArtifactMetadata?.(filename, options); + if (backendMetadata) { + const wasCached = backendMetadata.fromCache === true; + const deleted = wasCached + ? (await options.inferenceBackend?.deleteModelArtifact?.(filename, options)) === true + : false; + return { file: filename, deleted, wasCached }; + } + if (!cache) return { file: filename, deleted: false, wasCached: false }; const { localPath, proposedCacheKey } = buildResourcePaths(modelId, filename, options, cache); const cached = await checkCachedResource(cache, localPath, proposedCacheKey); diff --git a/packages/transformers/src/utils/model_registry/get_available_dtypes.js b/packages/transformers/src/utils/model_registry/get_available_dtypes.js index 2c534b157..7b16e1cd0 100644 --- a/packages/transformers/src/utils/model_registry/get_available_dtypes.js +++ b/packages/transformers/src/utils/model_registry/get_available_dtypes.js @@ -43,6 +43,9 @@ export async function get_available_dtypes( const { sessions } = getSessionsConfig(modelType, config, { model_file_name }); const metadataOptions = { revision, cache_dir, local_files_only }; const provider = await getModelRegistryInferenceProvider(inferenceProvider); + if (!provider.getAvailableDtypes) { + throw new Error('The inference backend does not support dtype discovery.'); + } return provider.getAvailableDtypes({ modelId, sessions, diff --git a/packages/transformers/src/utils/model_registry/get_file_metadata.js b/packages/transformers/src/utils/model_registry/get_file_metadata.js index a3a827547..f1c2cb6dd 100644 --- a/packages/transformers/src/utils/model_registry/get_file_metadata.js +++ b/packages/transformers/src/utils/model_registry/get_file_metadata.js @@ -7,8 +7,9 @@ import { getCache } from '../cache.js'; import { buildResourcePaths, checkCachedResource, getFetchHeaders, getFile } from '../hub.js'; import { isValidUrl, makePretrainedOptionsKey } from '../hub/utils.js'; import { logger } from '../logger.js'; -import { memoizePromise } from '../memoize_promise.js'; +import { throwIfAborted } from '../core.js'; import { getModelId } from '../../backends/inference.js'; +import { isInferenceBackend } from '../../backends/inference.js'; /** * @typedef {import('../hub.js').PretrainedOptions} PretrainedOptions @@ -46,19 +47,37 @@ async function fetch_file_head(urlOrPath, signal = undefined) { * Uses Range requests for remote files to be efficient. * Can also be used as a lightweight file existence check by checking the `.exists` property. * - * @param {string} path_or_repo_id This can be either: + * @param {string|import('../../backends/inference.js').InferenceBackend} path_or_repo_id This can be either: * - a string, the *model id* of a model repo on huggingface.co. * - a path to a *directory* potentially containing the file. * @param {string} filename The name of the file to check. * @param {PretrainedOptions} [options] An object containing optional parameters. * @returns {Promise<{exists: boolean, size?: number, contentType?: string, fromCache?: boolean}>} A Promise that resolves to file metadata. */ -export function get_file_metadata(path_or_repo_id, filename, options = {}) { +const INFLIGHT_METADATA = new Map(); +const BACKEND_INFLIGHT_METADATA = new WeakMap(); + +export async function get_file_metadata(path_or_repo_id, filename, options = {}) { throwIfAborted(options.signal); + const backend = isInferenceBackend(path_or_repo_id) ? path_or_repo_id : null; path_or_repo_id = getModelId(path_or_repo_id); - if (options.signal) return _get_file_metadata(path_or_repo_id, filename, options); + const load = async () => { + const backendMetadata = await backend?.getModelArtifactMetadata?.(filename, options); + throwIfAborted(options.signal); + if (backendMetadata) return { exists: true, ...backendMetadata }; + return _get_file_metadata(path_or_repo_id, filename, options); + }; + + if (options.signal) return load(); const key = makePretrainedOptionsKey(path_or_repo_id, options, filename); - return memoizePromise(key, () => _get_file_metadata(path_or_repo_id, filename, options)); + const inflight = backend ? (BACKEND_INFLIGHT_METADATA.get(backend) ?? new Map()) : INFLIGHT_METADATA; + if (backend && !BACKEND_INFLIGHT_METADATA.has(backend)) BACKEND_INFLIGHT_METADATA.set(backend, inflight); + let pending = inflight.get(key); + if (!pending) { + pending = load().finally(() => inflight.delete(key)); + inflight.set(key, pending); + } + return pending; } async function _get_file_metadata(path_or_repo_id, filename, options) { @@ -164,9 +183,3 @@ async function _get_file_metadata(path_or_repo_id, filename, options) { return { exists: false, fromCache: false }; } - -function throwIfAborted(signal) { - if (!signal?.aborted) return; - if (typeof signal.throwIfAborted === 'function') signal.throwIfAborted(); - throw signal.reason ?? new Error('Metadata loading aborted.'); -} diff --git a/packages/transformers/src/utils/model_registry/get_files.js b/packages/transformers/src/utils/model_registry/get_files.js index 0d2c34f23..c27e4de4b 100644 --- a/packages/transformers/src/utils/model_registry/get_files.js +++ b/packages/transformers/src/utils/model_registry/get_files.js @@ -15,6 +15,7 @@ import { get_processor_files } from './get_processor_files.js'; * @param {string|null} [options.cache_dir=null] Custom cache directory * @param {boolean} [options.local_files_only=false] Never hit the network if true * @param {string} [options.revision='main'] Model revision + * @param {string|null} [options.task=null] Pipeline task requesting the artifacts * @param {import('../../backends/model_registry.js').ModelRegistryInferenceProvider|null} [options.inferenceProvider=null] Artifact metadata provider * @param {boolean} [options.include_tokenizer=true] Whether to check for tokenizer files (set to false for vision-only models) * @param {boolean} [options.include_processor=true] Whether to check for processor files @@ -31,6 +32,7 @@ export async function get_files( cache_dir = null, local_files_only = false, revision = 'main', + task = null, inferenceProvider = null, include_tokenizer = true, include_processor = true, @@ -44,6 +46,7 @@ export async function get_files( dtype, device, model_file_name, + task, inferenceProvider, ...metadataOptions, }) diff --git a/packages/transformers/src/utils/model_registry/get_model_files.js b/packages/transformers/src/utils/model_registry/get_model_files.js index a9a524dd6..9c74f65c5 100644 --- a/packages/transformers/src/utils/model_registry/get_model_files.js +++ b/packages/transformers/src/utils/model_registry/get_model_files.js @@ -56,6 +56,7 @@ export function get_config( * @param {string|null} [options.cache_dir=null] Custom cache directory. * @param {boolean} [options.local_files_only=false] Never hit the network if true. * @param {string} [options.revision='main'] Model revision. + * @param {string|null} [options.task=null] Pipeline task requesting the artifacts. * @param {import('../../backends/model_registry.js').ModelRegistryInferenceProvider|null} [options.inferenceProvider=null] Artifact metadata provider. * @returns {Promise} Array of file paths that will be loaded */ @@ -69,6 +70,7 @@ export async function get_model_files( cache_dir = null, local_files_only = false, revision = 'main', + task = null, inferenceProvider = null, } = {}, ) { @@ -78,11 +80,15 @@ export async function get_model_files( const modelType = resolve_model_type(config); const { sessions, optional_configs } = getSessionsConfig(modelType, config, { model_file_name }); const provider = await getModelRegistryInferenceProvider(inferenceProvider); - return provider.listModelArtifacts({ - sessions, - optionalConfigs: optional_configs, - config, - dtype: overrideDtype, - device: overrideDevice, - }); + return [ + ...(await provider.listModelArtifacts({ + modelId, + task, + sessions, + optionalConfigs: optional_configs, + config, + dtype: overrideDtype, + device: overrideDevice, + })), + ]; } diff --git a/packages/transformers/src/utils/model_registry/get_pipeline_files.js b/packages/transformers/src/utils/model_registry/get_pipeline_files.js index 0f4438ef1..e5b3d5ff7 100644 --- a/packages/transformers/src/utils/model_registry/get_pipeline_files.js +++ b/packages/transformers/src/utils/model_registry/get_pipeline_files.js @@ -47,6 +47,7 @@ export async function get_pipeline_files(task, modelId, options = {}) { const files = await get_files(modelId, { ...options, + task, include_tokenizer, include_processor, }); @@ -60,7 +61,9 @@ export async function get_pipeline_files(task, modelId, options = {}) { if (textOnlySessions) { const provider = await getModelRegistryInferenceProvider(options.inferenceProvider ?? null); - return provider.filterModelArtifacts(files, textOnlySessions); + if (provider.filterModelArtifacts) { + return provider.filterModelArtifacts(files, textOnlySessions); + } } } diff --git a/packages/transformers/src/utils/model_registry/is_cached.js b/packages/transformers/src/utils/model_registry/is_cached.js index 9c6530149..e2e381f07 100644 --- a/packages/transformers/src/utils/model_registry/is_cached.js +++ b/packages/transformers/src/utils/model_registry/is_cached.js @@ -26,14 +26,11 @@ import { get_pipeline_files } from './get_pipeline_files.js'; async function check_files_cache(modelId, files, options = {}) { const cache = await getCache(options?.cache_dir); - if (!cache) { - const fileStatuses = files.map((filename) => ({ file: filename, cached: false })); - // No cache available, all files considered not cached - return { allCached: false, files: fileStatuses }; - } - const fileStatuses = await Promise.all( files.map(async (filename) => { + const backendMetadata = await options.inferenceBackend?.getModelArtifactMetadata?.(filename, options); + if (backendMetadata) return { file: filename, cached: backendMetadata.fromCache === true }; + if (!cache) return { file: filename, cached: false }; const { localPath, proposedCacheKey } = buildResourcePaths(modelId, filename, options, cache); const cached = await checkCachedResource(cache, localPath, proposedCacheKey); return { file: filename, cached: !!cached }; @@ -52,6 +49,8 @@ async function check_files_cache(modelId, files, options = {}) { * @returns {Promise} */ async function is_file_cached(modelId, filename, options = {}) { + const backendMetadata = await options.inferenceBackend?.getModelArtifactMetadata?.(filename, options); + if (backendMetadata) return backendMetadata.fromCache === true; const cache = await getCache(options?.cache_dir); if (!cache) return false; const { localPath, proposedCacheKey } = buildResourcePaths(modelId, filename, options, cache); diff --git a/packages/transformers/tests/generation_controller.test.js b/packages/transformers/tests/generation_controller.test.js index 49fc6a992..8be80789d 100644 --- a/packages/transformers/tests/generation_controller.test.js +++ b/packages/transformers/tests/generation_controller.test.js @@ -78,6 +78,32 @@ describe("GenerationController", () => { expect(streamer.end).toHaveBeenCalledTimes(1); }); + it.each([2, 1])("generates one legacy token when max_length is at or below the prompt (%i)", async (max_length) => { + const controller = createGenerationController({ config: {}, generation_config: null }, int64Tensor([[1, 2]]), { max_length }); + + expect(controller.allDone).toBe(false); + expect(controller.maxSequenceLength).toBe(3); + expect(controller.compileRuntimePlan({ declarativePlans: ["argmax"], planModes: ["greedy"], tokenPipeline: { defaultDepth: 1 } })).toMatchObject({ maxNewTokens: 1 }); + + const step = await controller.step(new Tensor("float32", [0, 0, 4], [1, 3])); + expect(step.allDone).toBe(true); + expect(controller.finalize().tolist()).toEqual([[1n, 2n, 2n]]); + }); + + it("preserves abort reasons without signaling successful stream completion", async () => { + const reason = new Error("generation failed"); + const streamer = { put: jest.fn(), end: jest.fn(), abort: jest.fn() }; + const controller = createGenerationController({ config: {}, generation_config: null }, int64Tensor([[1]]), { max_new_tokens: 1, streamer }); + + controller.abort(reason); + + expect(controller.abortReason).toBe(reason); + expect(streamer.abort).toHaveBeenCalledWith(reason); + expect(streamer.end).not.toHaveBeenCalled(); + await expect(controller.step(new Tensor("float32", [0, 1], [1, 2]))).rejects.toBe(reason); + expect(() => controller.finalize()).toThrow(reason); + }); + it("processes classifier-free guidance before validating the output batch", async () => { const controller = createGenerationController({ config: {}, generation_config: null }, int64Tensor([[1]]), { max_new_tokens: 1, guidance_scale: 3 }); const logits = new Tensor( @@ -139,7 +165,7 @@ describe("custom autoregressive sessions", () => { }); const output = await model.generate({ - input_ids: int64Tensor([[1]]), + inputs: int64Tensor([[1]]), attention_mask: int64Tensor([[1]]), max_new_tokens: 2, }); @@ -196,9 +222,10 @@ describe("custom autoregressive sessions", () => { config: { model_type: "custom", is_encoder_decoder: false, eos_token_id: 3 }, }); - const output = await model.generate({ input_ids: int64Tensor([[1]]), max_new_tokens: 2 }); + const output = await model.generate({ input_ids: int64Tensor([[1]]), max_new_tokens: 2, return_dict_in_generate: true }); - expect(output.tolist()).toEqual([[1n, 2n, 3n]]); + expect(output.sequences.tolist()).toEqual([[1n, 2n, 3n]]); + expect(output.past_key_values).toBeDefined(); expect(iteratorClosed).toHaveBeenCalledTimes(1); expect(session.dispose).toHaveBeenCalledTimes(1); }); @@ -317,6 +344,20 @@ describe("custom autoregressive sessions", () => { expect(createAutoregressiveSession).not.toHaveBeenCalled(); }); + it("validates zero-token requests without creating a runtime session", async () => { + const createAutoregressiveSession = jest.fn(); + const backend = { + modelId: "test/zero-token-model", + load: jest.fn(async () => ({ createAutoregressiveSession, async dispose() {} })), + }; + const model = await AutoModel.from_pretrained(backend, { + config: { model_type: "custom", is_encoder_decoder: false }, + }); + + await expect(model.generate({ input_ids: int64Tensor([[1]]), max_new_tokens: 0 })).rejects.toThrow("must declare generation capabilities"); + expect(createAutoregressiveSession).not.toHaveBeenCalled(); + }); + it("reports when a request is unsupported by the CPU mode", async () => { const createAutoregressiveSession = jest.fn(); const backend = { diff --git a/packages/transformers/tests/inference_backends.test.js b/packages/transformers/tests/inference_backends.test.js index 341773958..e5fcf842a 100644 --- a/packages/transformers/tests/inference_backends.test.js +++ b/packages/transformers/tests/inference_backends.test.js @@ -5,6 +5,8 @@ import { OnnxInferenceProvider } from "@huggingface/transformers-onnx"; import { AutoModel } from "../src/models/auto/modeling_auto.js"; import { PreTrainedModel } from "../src/models/modeling_utils.js"; import { buildResourcePaths } from "../src/utils/hub.js"; +import { loadInferenceBackendChatTemplate } from "../src/pipelines.js"; +import { DefaultProgressCallback } from "../src/utils/core.js"; describe("inference backends", () => { it("recognizes object and class backends", () => { @@ -107,18 +109,32 @@ describe("inference backends", () => { const backend = { modelId: "test/model", load: jest.fn(async () => loaded), + listModelArtifacts: jest.fn(() => ["model.bin"]), + getModelArtifactMetadata: jest.fn(async () => ({ size: 10, fromCache: true })), }; const signal = new AbortController().signal; const artifactProvider = { readJson() {}, openByteSource() {} }; + const progress_callback = jest.fn(); const model = await AutoModel.from_pretrained(backend, { config: { model_type: "custom" }, device: "webgpu", signal, artifactProvider, + progress_callback, }); - expect(backend.load).toHaveBeenCalledWith(expect.objectContaining({ modelId: "test/model", device: "webgpu", signal, artifactProvider })); + expect(backend.getModelArtifactMetadata).toHaveBeenCalledWith("model.bin", expect.objectContaining({ signal })); + expect(backend.load).toHaveBeenCalledWith( + expect.objectContaining({ + modelId: "test/model", + device: "webgpu", + signal, + artifactProvider, + artifactMetadata: { "model.bin": { size: 10, fromCache: true } }, + progress_callback: expect.any(DefaultProgressCallback), + }), + ); await expect(model({ value: 1 })).resolves.toEqual({ value: 1 }); }); @@ -150,4 +166,10 @@ describe("inference backends", () => { expect(paths.requestURL).toBe("test/model/tokenizer.json"); expect(paths.remoteURL).toContain("/test/model/resolve/main/tokenizer.json"); }); + + it("resolves and validates inline backend chat templates", async () => { + await expect(loadInferenceBackendChatTemplate({ modelId: "test/model", load() {}, chatTemplate: { content: "{{ messages }}" } }, {})).resolves.toBe("{{ messages }}"); + expect(loadInferenceBackendChatTemplate({ modelId: "test/model", load() {} }, {})).toBeNull(); + expect(() => loadInferenceBackendChatTemplate({ modelId: "test/model", load() {}, chatTemplate: { content: "inline", file: "chat.jinja" } }, {})).toThrow("cannot be combined"); + }); }); diff --git a/packages/transformers/tests/progress_callbacks.test.js b/packages/transformers/tests/progress_callbacks.test.js index d9a2b8427..b69af41c7 100644 --- a/packages/transformers/tests/progress_callbacks.test.js +++ b/packages/transformers/tests/progress_callbacks.test.js @@ -1,6 +1,7 @@ import { pipeline, LlamaForCausalLM, AutoModelForCausalLM, WhisperForConditionalGeneration, Gemma3ForConditionalGeneration, Gemma3nForConditionalGeneration, VoxtralRealtimeForConditionalGeneration } from "../src/transformers.js"; import { init, MAX_MODEL_LOAD_TIME, MAX_MODEL_DISPOSE_TIME, DEFAULT_MODEL_OPTIONS } from "./init.js"; +import { DefaultProgressCallback } from "../src/utils/core.js"; // Initialise the testing environment init(); @@ -243,6 +244,38 @@ describe("Progress Callbacks", () => { describe("Edge cases", () => { const model_id = "hf-internal-testing/tiny-random-LlamaForCausalLM"; + it("keeps aggregate totals stable and suppresses duplicate or regressing updates", () => { + const events = []; + const callback = new DefaultProgressCallback((event) => events.push(event), { + "model.safetensors": { loaded: 0, total: 1000 }, + }); + callback({ status: "progress", name: "test/model", file: "model.safetensors", loaded: 500, total: 500 }); + callback({ status: "progress", name: "test/model", file: "model.safetensors", loaded: 1, total: 4 }); + callback({ status: "progress", name: "test/model", file: "model.safetensors", loaded: 1000, total: 1000 }); + callback({ status: "progress", name: "test/model", file: "model.safetensors", loaded: 1000, total: 1000 }); + + const totals = events.filter((event) => event.status === "progress_total"); + expect(totals.map(({ loaded }) => loaded)).toEqual([500, 1000]); + expect(totals.every(({ total }) => total === 1000)).toBe(true); + }); + + it("grows unknown totals without emitting NaN", () => { + const events = []; + const callback = new DefaultProgressCallback((event) => events.push(event), {}); + callback({ status: "progress", name: "test/model", file: "model.bin", loaded: 0, total: undefined }); + callback({ status: "progress", name: "test/model", file: "model.bin", loaded: 100, total: 100 }); + callback({ status: "progress", name: "test/model", file: "model.bin", loaded: 200, total: 200 }); + callback({ status: "progress", name: "test/model", file: "model.bin", loaded: 200, total: undefined }); + + const totals = events.filter((event) => event.status === "progress_total"); + expect(totals.map(({ loaded, total }) => [loaded, total])).toEqual([ + [0, 0], + [100, 100], + [200, 200], + ]); + expect(totals.every(({ loaded, total, progress }) => [loaded, total, progress].every(Number.isFinite))).toBe(true); + }); + it( "no progress_total without progress_callback", async () => { diff --git a/packages/transformers/tests/types.test.js b/packages/transformers/tests/types.test.js index 77509297e..f22c94bfd 100644 --- a/packages/transformers/tests/types.test.js +++ b/packages/transformers/tests/types.test.js @@ -75,6 +75,7 @@ describe("TypeScript compilation succeeds", () => { import type { CausalGenerationCapabilitiesV1, InferenceBackend, + InferenceBackendChatTemplate, InferenceModel, LogitsLeaseV1, PlanAutoregressiveSessionV1, @@ -117,6 +118,7 @@ describe("TypeScript compilation succeeds", () => { const backend: InferenceBackend = { modelId: "test/model", + chatTemplate: { modelId: "test/templates", file: "chat.jinja" }, capabilities: staticCapabilities, async load(_options) { const model: InferenceModel = { @@ -130,6 +132,8 @@ describe("TypeScript compilation succeeds", () => { void backend; void lease; + const inlineTemplate = { content: "{{ messages }}" } as const satisfies InferenceBackendChatTemplate; + void inlineTemplate; `); if (diagnostics.length > 0) { throw new Error(formatDiagnostics(diagnostics)); diff --git a/packages/transformers/tests/utils/get_file_metadata.test.js b/packages/transformers/tests/utils/get_file_metadata.test.js new file mode 100644 index 000000000..28dea3ace --- /dev/null +++ b/packages/transformers/tests/utils/get_file_metadata.test.js @@ -0,0 +1,54 @@ +import { jest } from "@jest/globals"; + +import { get_file_metadata } from "../../src/utils/model_registry/get_file_metadata.js"; + +function createBackend(getModelArtifactMetadata) { + return { + modelId: "test/metadata-backend", + load() {}, + getModelArtifactMetadata, + }; +} + +describe("get_file_metadata", () => { + it("deduplicates only in-flight backend metadata requests", async () => { + let resolveMetadata; + const pendingMetadata = new Promise((resolve) => { + resolveMetadata = resolve; + }); + const getModelArtifactMetadata = jest.fn(() => pendingMetadata); + const backend = createBackend(getModelArtifactMetadata); + + const first = get_file_metadata(backend, "model.bin"); + const second = get_file_metadata(backend, "model.bin"); + expect(getModelArtifactMetadata).toHaveBeenCalledTimes(1); + resolveMetadata({ size: 12, fromCache: true }); + + await expect(Promise.all([first, second])).resolves.toEqual([ + { exists: true, size: 12, fromCache: true }, + { exists: true, size: 12, fromCache: true }, + ]); + await get_file_metadata(backend, "model.bin"); + expect(getModelArtifactMetadata).toHaveBeenCalledTimes(2); + }); + + it("returns an asynchronously rejected promise for an already-aborted signal", async () => { + const reason = new Error("cancelled"); + const controller = new AbortController(); + controller.abort(reason); + let result; + + expect(() => { + result = get_file_metadata(createBackend(jest.fn()), "model.bin", { signal: controller.signal }); + }).not.toThrow(); + await expect(result).rejects.toBe(reason); + }); + + it("does not share signaled requests", async () => { + const getModelArtifactMetadata = jest.fn(async () => ({ size: 1 })); + const backend = createBackend(getModelArtifactMetadata); + + await Promise.all([get_file_metadata(backend, "model.bin", { signal: new AbortController().signal }), get_file_metadata(backend, "model.bin", { signal: new AbortController().signal })]); + expect(getModelArtifactMetadata).toHaveBeenCalledTimes(2); + }); +}); diff --git a/packages/transformers/tests/utils/model_registry.test.js b/packages/transformers/tests/utils/model_registry.test.js index ab9713c18..921981e9e 100644 --- a/packages/transformers/tests/utils/model_registry.test.js +++ b/packages/transformers/tests/utils/model_registry.test.js @@ -11,6 +11,7 @@ await import("../../src/models/registry.js"); // Dynamic import after mock setup (required for ESM) const { get_available_dtypes } = await import("../../src/utils/model_registry/get_available_dtypes.js"); +const { ModelRegistry } = await import("../../src/utils/model_registry/ModelRegistry.js"); // A minimal config that mimics a BERT-like encoder-only model const ENCODER_ONLY_CONFIG = { @@ -196,3 +197,64 @@ describe("get_available_dtypes", () => { } }); }); + +describe("custom inference backend artifacts", () => { + beforeEach(() => { + mockGetFileMetadata.mockReset(); + }); + + it("composes backend-owned and tokenizer files for a pipeline", async () => { + setupExistingFiles("tokenizer_config.json"); + const listModelArtifacts = jest.fn(() => ["config.json", "model.safetensors"]); + const backend = { + modelId: "test/custom-model", + load() {}, + listModelArtifacts, + }; + + await expect( + ModelRegistry.get_pipeline_files("text-generation", backend, { + config: DECODER_ONLY_CONFIG, + device: "webgpu", + dtype: "auto", + }), + ).resolves.toEqual(["config.json", "model.safetensors", "tokenizer.json", "tokenizer_config.json"]); + expect(listModelArtifacts).toHaveBeenCalledWith( + expect.objectContaining({ + modelId: "test/custom-model", + task: "text-generation", + device: "webgpu", + dtype: "auto", + }), + ); + }); + + it("uses backend-owned cache metadata and deletion hooks", async () => { + const deleteModelArtifact = jest.fn(async () => true); + const backend = { + modelId: "test/cached-custom-model", + load() {}, + listModelArtifacts: jest.fn(() => ["config.json", "model.bin"]), + getModelArtifactMetadata: jest.fn(async (file) => ({ size: 10, fromCache: file === "model.bin" })), + deleteModelArtifact, + }; + const options = { + config: ENCODER_ONLY_CONFIG, + include_tokenizer: false, + include_processor: false, + }; + + await expect(ModelRegistry.is_cached_files(backend, options)).resolves.toEqual({ + allCached: false, + files: [ + { file: "config.json", cached: false }, + { file: "model.bin", cached: true }, + ], + }); + await expect(ModelRegistry.clear_cache(backend, options)).resolves.toMatchObject({ + filesDeleted: 1, + filesCached: 1, + }); + expect(deleteModelArtifact).toHaveBeenCalledWith("model.bin", expect.objectContaining({ inferenceBackend: backend })); + }); +}); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 2a126f332..75d83120e 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -60,9 +60,6 @@ importers: onnxruntime-common: specifier: 1.24.3 version: 1.24.3 - onnxruntime-node: - specifier: 1.24.3 - version: 1.24.3 onnxruntime-web: specifier: 1.26.0-dev.20260416-b7804b056c version: 1.26.0-dev.20260416-b7804b056c @@ -79,6 +76,9 @@ importers: jest: specifier: ^30.2.0 version: 30.2.0(@types/node@24.10.9) + onnxruntime-node: + specifier: 1.24.3 + version: 1.24.3 typescript: specifier: 5.9.3 version: 5.9.3 From c40d28afae3c254396fc39182318d1d6acd7dba7 Mon Sep 17 00:00:00 2001 From: Nico Martin Date: Thu, 13 Aug 2026 11:32:38 +0200 Subject: [PATCH 5/5] added sampling --- .../transformers/src/generation/controller.js | 21 +++++++- .../src/generation/logits_process.js | 28 +++++++++++ .../transformers/src/generation/runtime.js | 7 ++- .../tests/generation_controller.test.js | 50 +++++++++++++++++++ 4 files changed, 102 insertions(+), 4 deletions(-) diff --git a/packages/transformers/src/generation/controller.js b/packages/transformers/src/generation/controller.js index 8f251b959..c5d4cd71d 100644 --- a/packages/transformers/src/generation/controller.js +++ b/packages/transformers/src/generation/controller.js @@ -264,6 +264,7 @@ export class GenerationController { this.sequences[index].push(tokenId); generatedInputIds.push([tokenId]); } + this.logitsProcessor.onTokensSampled(Array.from(decision.tokenIds), this.sequences); if (this.streamer) this.streamer.put(generatedInputIds); this.done = this.stoppingCriteria(this.sequences, decision.processedScores); @@ -283,10 +284,26 @@ export class GenerationController { * @returns {Object|null} */ compileRuntimePlan(capabilities) { + if (this.generationConfig.num_beams > 1) return null; + if (this.generationConfig.output_scores) return null; + if (this.generationConfig.do_sample) { + if (!capabilities?.declarativePlans?.includes('multinomial')) return null; + if (!capabilities?.planModes?.includes('multinomial')) return null; + if (!this.logitsProcessor.processors.every((processor) => processor instanceof TemperatureLogitsWarper)) return null; + const temperature = this.generationConfig.temperature ?? 1.0; + const topK = this.generationConfig.top_k; + if (!Number.isFinite(temperature) || temperature <= 0) return null; + if (!Number.isInteger(topK) || topK < 1 || topK > 128) return null; + return { + version: 1, + processors: [], + sampler: { op: 'multinomial', temperature, topK }, + maxNewTokens: Math.max(0, this.maxSequenceLength - this.inputLength), + pipelineDepth: capabilities.tokenPipeline?.defaultDepth, + }; + } if (!capabilities?.declarativePlans?.includes('argmax')) return null; if (!capabilities?.planModes?.includes('greedy')) return null; - if (this.generationConfig.do_sample || this.generationConfig.num_beams > 1) return null; - if (this.generationConfig.output_scores) return null; if (this.logitsProcessor.processors.length !== 0) return null; return { version: 1, diff --git a/packages/transformers/src/generation/logits_process.js b/packages/transformers/src/generation/logits_process.js index 647a30806..e2258721f 100644 --- a/packages/transformers/src/generation/logits_process.js +++ b/packages/transformers/src/generation/logits_process.js @@ -22,6 +22,14 @@ export class LogitsProcessor extends Callable { _call(input_ids, logits) { throw Error('`_call` should be implemented in a subclass'); } + + /** + * Notify the processor after tokens have been committed. + * + * @param {number[]} token_ids The sampled token ID for each batch item. + * @param {bigint[][]} input_ids The updated input IDs. + */ + onTokensSampled(token_ids, input_ids) {} } /** @@ -39,6 +47,14 @@ export class LogitsWarper extends Callable { _call(input_ids, logits) { throw Error('`_call` should be implemented in a subclass'); } + + /** + * Notify the warper after tokens have been committed. + * + * @param {number[]} token_ids The sampled token ID for each batch item. + * @param {bigint[][]} input_ids The updated input IDs. + */ + onTokensSampled(token_ids, input_ids) {} } /** @@ -88,6 +104,18 @@ export class LogitsProcessorList extends Callable { return toReturn; } + /** + * Notify all processors after tokens have been committed. + * + * @param {number[]} token_ids The sampled token ID for each batch item. + * @param {bigint[][]} input_ids The updated input IDs. + */ + onTokensSampled(token_ids, input_ids) { + for (const processor of this.processors) { + processor.onTokensSampled(token_ids, input_ids); + } + } + [Symbol.iterator]() { return this.processors.values(); } diff --git a/packages/transformers/src/generation/runtime.js b/packages/transformers/src/generation/runtime.js index 352498de8..548a5cc74 100644 --- a/packages/transformers/src/generation/runtime.js +++ b/packages/transformers/src/generation/runtime.js @@ -16,7 +16,7 @@ import { createGenerationController } from './controller.js'; * @property {ReadonlyArray} cpuModes * @property {ReadonlyArray} planModes * @property {boolean} cpuLogits - * @property {ReadonlyArray<'argmax'>} declarativePlans + * @property {ReadonlyArray<'argmax'|'multinomial'>} declarativePlans * @property {{readonly defaultDepth: number, readonly maxDepth: number}} tokenPipeline * @property {SessionConcurrencyCapabilities} [sessionConcurrency] */ @@ -61,7 +61,7 @@ import { createGenerationController } from './controller.js'; * @typedef {Object} RuntimeGenerationPlanV1 * @property {1} version * @property {ReadonlyArray} processors - * @property {{readonly op: 'argmax'}} sampler + * @property {{readonly op: 'argmax'}|{readonly op: 'multinomial', readonly temperature: number, readonly topK: number}} sampler * @property {number} maxNewTokens * @property {number} [pipelineDepth] */ @@ -286,6 +286,9 @@ function validateCapabilities(capabilities, controller, attentionMask) { if (capabilities.planModes.includes('greedy') && !capabilities.declarativePlans.includes('argmax')) { throw new Error('Runtime greedy plan mode requires the `argmax` declarative plan.'); } + if (capabilities.planModes.includes('multinomial') && !capabilities.declarativePlans.includes('multinomial')) { + throw new Error('Runtime multinomial plan mode requires the `multinomial` declarative plan.'); + } if (controller.batchSize > capabilities.maxBatchSize) { throw new Error( `Runtime supports batch size ${capabilities.maxBatchSize}, but generation received ${controller.batchSize}.`, diff --git a/packages/transformers/tests/generation_controller.test.js b/packages/transformers/tests/generation_controller.test.js index 8be80789d..9acd03f4c 100644 --- a/packages/transformers/tests/generation_controller.test.js +++ b/packages/transformers/tests/generation_controller.test.js @@ -21,6 +21,17 @@ class ForceTokenProcessor extends LogitsProcessor { } } +class TrackingTokenProcessor extends LogitsProcessor { + constructor() { + super(); + this.onTokensSampled = jest.fn(); + } + + _call(_inputIds, logits) { + return logits; + } +} + class TokenStoppingCriteria extends StoppingCriteria { constructor(tokenId) { super(); @@ -47,6 +58,45 @@ function createLease(values, release = jest.fn()) { } describe("GenerationController", () => { + it("compiles native multinomial plans with temperature", () => { + const controller = createGenerationController( + { config: {}, generation_config: null }, + int64Tensor([[1]]), + { do_sample: true, temperature: 0.7, top_k: 50, max_new_tokens: 1 }, + ); + expect(controller.compileRuntimePlan({ + declarativePlans: ["argmax", "multinomial"], + planModes: ["greedy", "multinomial"], + tokenPipeline: { defaultDepth: 1 }, + })).toMatchObject({ + sampler: { op: "multinomial", temperature: 0.7, topK: 50 }, + }); + }); + + it("compiles native multinomial plans with model sampling defaults", () => { + const controller = createGenerationController( + { config: {}, generation_config: null }, + int64Tensor([[1]]), + { do_sample: true, temperature: 1.0, top_k: 64, top_p: 0.95, max_new_tokens: 1 }, + ); + expect(controller.compileRuntimePlan({ + declarativePlans: ["argmax", "multinomial"], + planModes: ["greedy", "multinomial"], + tokenPipeline: { defaultDepth: 1 }, + })).toMatchObject({ sampler: { op: "multinomial", temperature: 1.0, topK: 64 } }); + }); + + it("notifies logits processors after committing sampled tokens", () => { + const processor = new TrackingTokenProcessor(); + const processors = new LogitsProcessorList(); + processors.push(processor); + const controller = createGenerationController({ config: {}, generation_config: null }, int64Tensor([[1]]), { max_new_tokens: 1, logits_processor: processors }); + + controller.commit({ tokenIds: Uint32Array.of(2) }); + + expect(processor.onTokensSampled).toHaveBeenCalledWith([2], [[1n, 2n]]); + }); + it("owns processing, sampling, stopping, streaming, and finalization", async () => { const processors = new LogitsProcessorList(); processors.push(new ForceTokenProcessor(2));