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/.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/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..a72266aaa --- /dev/null +++ b/packages/transformers-onnx/README.md @@ -0,0 +1,21 @@ +# @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'); +``` + +## 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/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..3e9dbab83 --- /dev/null +++ b/packages/transformers-onnx/package.json @@ -0,0 +1,69 @@ +{ + "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-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": [ + "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..bdac31c1a --- /dev/null +++ b/packages/transformers-onnx/scripts/build.mjs @@ -0,0 +1,69 @@ +import { build } from "esbuild"; +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: [entry], + bundle: true, + sourcemap: false, + logLevel: "warning", +}; + +await Promise.all([ + build({ + ...shared, + outfile: filePath("../dist/transformers-onnx.node.mjs"), + platform: "node", + format: "esm", + external: ["onnxruntime-common", "onnxruntime-node"], + alias: { "onnxruntime-web/webgpu": empty }, + banner: { js: "const __ONNX_MODULE_URL__ = import.meta.url;" }, + }), + build({ + ...shared, + outfile: filePath("../dist/transformers-onnx.node.cjs"), + platform: "node", + format: "cjs", + external: ["onnxruntime-common", "onnxruntime-node"], + alias: { "onnxruntime-web/webgpu": empty }, + banner: { js: 'const __ONNX_MODULE_URL__ = require("node:url").pathToFileURL(__filename).href;' }, + }), + build({ + ...shared, + outfile: filePath("../dist/transformers-onnx.web.js"), + platform: "browser", + format: "esm", + external: ["onnxruntime-common", "onnxruntime-web"], + alias: { "onnxruntime-node": empty }, + banner: { js: "const __ONNX_MODULE_URL__ = import.meta.url;" }, + }), + build({ + entryPoints: [testingEntry], + bundle: true, + sourcemap: false, + logLevel: "warning", + outfile: filePath("../dist/testing.mjs"), + platform: "node", + format: "esm", + external: ["onnxruntime-common", "onnxruntime-node"], + }), + build({ + entryPoints: [testingEntry], + bundle: true, + sourcemap: false, + logLevel: "warning", + 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/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..c39ab1f6a --- /dev/null +++ b/packages/transformers-onnx/src/host.ts @@ -0,0 +1,141 @@ +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 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: 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; +} + +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: OnnxProviderEnvironment = { + backends: { onnx: {} }, + logLevel: 30, + useWasmCache: typeof caches !== 'undefined', + fetch: (...args) => globalThis.fetch(...args), +}; + +const environment = new Proxy(fallbackEnvironment, { + get(target, property) { + return Reflect.get(configuredHost?.env ?? target, property); + }, + set(target, property, value) { + return Reflect.set(configuredHost?.env ?? target, property, value); + }, +}); + +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, + 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 apis = new Proxy(fallbackApis, { + get(target, property) { + return Reflect.get(configuredHost?.apis ?? target, property); + }, +}); + +const logger = new Proxy(console, { + get(target, property) { + return Reflect.get(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; + }, + maxExternalDataChunks: 100, +}; + +export function configureOnnxProviderHost(host: OnnxProviderHost): void { + // 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) { + 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; +} + +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..fbc2ea738 --- /dev/null +++ b/packages/transformers-onnx/src/provider.ts @@ -0,0 +1,464 @@ +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 { + 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. */ + 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.'); + } + 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, + ) { + throwIfAborted(options.signal); + let custom_config = options.config?.['transformers.js_config'] ?? {}; + 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 ?? {}; + 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: 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(', ')}`); + } + 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, + ); + throwIfAborted(options.signal); + 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: 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, + session_options, + session_config: { dtype: selectedDtype, device: selectedDevice }, + }; + } + + async constructSessions(names: Record, options: any, cache_sessions: any = undefined) { + 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) { + 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 as number[]; + }, + set dims(value) { + (tensor as unknown as { dims: number[] }).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}`}`); +} + +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, + 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; + } + 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'; + 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 80% rename from packages/transformers/src/backends/onnx.js rename to packages/transformers-onnx/src/runtime.ts index 963f6fffe..bb6ca077e 100644 --- a/packages/transformers/src/backends/onnx.js +++ b/packages/transformers-onnx/src/runtime.ts @@ -16,23 +16,32 @@ * @module backends/onnx */ -import { env, apis, LogLevel } from '../env.js'; +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'; 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 type { Env, InferenceSession as OrtInferenceSession, Tensor as OrtTensor } from 'onnxruntime-common'; +import { loadWasmBinary, loadWasmFactory } from './wasm-cache.js'; export { Tensor } from 'onnxruntime-common'; -/** - * @typedef {import('onnxruntime-common').InferenceSession.ExecutionProviderConfig} ONNXExecutionProviders - */ +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:'); +} -/** @type {Record} */ -const DEVICE_TO_EXECUTION_PROVIDER_MAPPING = Object.freeze({ +function toAbsoluteURL(url: string): string { + return new URL(url, globalThis.location?.href ?? __ONNX_MODULE_URL__).href; +} + +type ExecutionProvider = OrtInferenceSession.ExecutionProviderConfig; + +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 @@ -55,7 +64,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. @@ -79,11 +88,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', @@ -91,20 +96,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; @@ -135,7 +136,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) @@ -150,15 +151,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; @@ -180,16 +176,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. @@ -228,7 +222,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 @@ -283,26 +277,36 @@ 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, - }); + }; + 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()); - 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. @@ -310,21 +314,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. @@ -336,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. @@ -363,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'; } /** @@ -375,7 +392,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]; } @@ -384,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-onnx/src/tensor-ops.ts b/packages/transformers-onnx/src/tensor-ops.ts new file mode 100644 index 000000000..1f0c0ba7d --- /dev/null +++ b/packages/transformers-onnx/src/tensor-ops.ts @@ -0,0 +1,219 @@ +import { createInferenceSession, runInferenceSession, isONNXProxy, Tensor as OrtTensor } from './runtime.js'; +import type { Tensor as OrtTensorType } from 'onnxruntime-common'; +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 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), + ]; + }), + ) as Record; + 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..f8462a31d --- /dev/null +++ b/packages/transformers-onnx/tests/provider.test.js @@ -0,0 +1,45 @@ +import { configureOnnxProviderHost, 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"); + }); + + 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 new file mode 100644 index 000000000..ede70dc64 --- /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": true, + "skipLibCheck": true, + "declaration": true, + "declarationMap": true, + "emitDeclarationOnly": true, + "esModuleInterop": true, + "composite": true, + "types": ["node", "@webgpu/types"] + } +} 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..918a1730a --- /dev/null +++ b/packages/transformers/docs/source/guides/custom-backends.md @@ -0,0 +1,76 @@ +# 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. + +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. + +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/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..59278bbfb 100644 --- a/packages/transformers/scripts/build/buildAll.mjs +++ b/packages/transformers/scripts/build/buildAll.mjs @@ -1,8 +1,8 @@ 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"; -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 +19,6 @@ async function buildTarget( format = "esm", // 'esm' | 'cjs' ignoreModules = [], externalModules = [], - usePostBuild = false, }, log, ) { @@ -35,9 +34,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({ @@ -72,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/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..0e5de5edf --- /dev/null +++ b/packages/transformers/src/backends/artifacts.js @@ -0,0 +1,44 @@ +/** + * @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)`. 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. + */ + +/** + * @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. + */ + +/** + * 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 new file mode 100644 index 000000000..1248cecef --- /dev/null +++ b/packages/transformers/src/backends/default.js @@ -0,0 +1,45 @@ +import { env, apis } from '../env.js'; +import { logger } from '../utils/logger.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'; + +const ONNX_HOST_SYMBOL = Symbol.for('transformers.js.onnxProviderHost'); +const host = { + env, + apis, + logger, + getModelFile, + getCacheNames, + createBackendTensor: (storage) => Tensor.fromBackendStorage(storage), + getBackendTensorStorage: (tensor) => tensor?.getBackendStorage?.() ?? null, + getCache, + maxExternalDataChunks: MAX_EXTERNAL_DATA_CHUNKS, +}; + +let modulePromise; + +export function getOnnxProviderModule() { + if (!modulePromise) { + globalThis[ONNX_HOST_SYMBOL] = host; + 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; +} + +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 new file mode 100644 index 000000000..3b21686dc --- /dev/null +++ b/packages/transformers/src/backends/inference.js @@ -0,0 +1,244 @@ +/** + * @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 { 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, + * artifactMetadata?: Record, + * }} InferenceBackendLoadOptions + */ + +/** + * @typedef {import('../utils/hub.js').PretrainedModelOptions & { + * task?: string, + * config?: import('../configs.js').PretrainedConfig, + * modelClass?: Function, + * generation_config?: Record, + * artifactMetadata?: Record, + * }} InferenceModelLoadOptions + */ + +/** + * @typedef {Object} InferenceModel + * @property {(inputs: Record) => Promise>} [forward] + * @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 + */ + +/** + * 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] + */ + +/** + * 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)`.'); +} + +/** + * 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. + * + * @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 ( + 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' && + 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 {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))) + ); + 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/model_registry.js b/packages/transformers/src/backends/model_registry.js new file mode 100644 index 000000000..deeaf352f --- /dev/null +++ b/packages/transformers/src/backends/model_registry.js @@ -0,0 +1,24 @@ +import { getOnnxProviderModule } from './default.js'; + +/** + * Provider operations used by model-file discovery. + * + * @typedef {Object} ModelRegistryInferenceProvider + * @property {(options: Object) => ReadonlyArray|Promise>} 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 (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/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..0c5af6f42 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`. @@ -239,10 +239,13 @@ export const env = { version: VERSION, /////////////////// Backends settings /////////////////// - // NOTE: These will be populated later by the backends themselves. + // 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: { - // onnxruntime-web/onnxruntime-node - onnx: {}, + onnx: { + wasm: {}, + webgpu: {}, + }, }, /////////////////// Logging settings /////////////////// @@ -252,8 +255,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..c5d4cd71d --- /dev/null +++ b/packages/transformers/src/generation/controller.js @@ -0,0 +1,378 @@ +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; +} + +/** + * 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 { + /** + * @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.done = new Array(this.batchSize).fill(false); + this.terminal = false; + this.finalized = false; + this.aborted = false; + this.abortReason = undefined; + + 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])); + } + + get allDone() { + return this.terminal; + } + + get maxSequenceLength() { + if (this.generationConfig.max_new_tokens === 0) return this.inputLength; + return Math.max(this.generationConfig.max_length ?? 0, this.inputLength + 1); + } + + /** + * 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}.`); + } + 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); + for (let batchIndex = 0; batchIndex < this.batchSize; ++batchIndex) { + const sampled = await this.sampler(processed[batchIndex]); + const [tokenId] = sampled[0]; + tokenIds[batchIndex] = Number(tokenId); + } + return this.commit({ tokenIds }); + } + + /** + * Commit tokens selected by an approved runtime generation plan. + * + * @param {{tokenIds: Uint32Array, processedScores?: Float32Array}} 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); + 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); + this.terminal = this.done.every(Boolean); + const nextTokenIds = new Tensor('int64', generatedInputIds.flat(), [this.batchSize, 1]); + return { + nextTokenIds, + generatedInputIds, + 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 (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.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 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; + 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; + this.abortReason = reason ?? new Error('Generation controller has been aborted.'); + this.streamer?.abort?.(this.abortReason); + } + + #assertActive() { + 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.'); + } +} + +/** + * 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, + }); + prepareGenerationLength(generationConfig, inputIds.dims.at(-1)); + 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/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 new file mode 100644 index 000000000..548a5cc74 --- /dev/null +++ b/packages/transformers/src/generation/runtime.js @@ -0,0 +1,414 @@ +import { Tensor } from '../utils/tensor.js'; +import { DynamicCache } from '../cache_utils.js'; +import { throwIfAborted } from '../utils/core.js'; +import { createGenerationController } from './controller.js'; + +/** + * @typedef {Object} SessionConcurrencyCapabilities + * @property {number} maxActiveSessions + * @property {1} concurrentOperationsPerSession + */ + +/** + * @typedef {Object} CausalGenerationCapabilitiesV1 + * @property {1} sessionVersion + * @property {number} maxBatchSize + * @property {ReadonlyArray} cpuModes + * @property {ReadonlyArray} planModes + * @property {boolean} cpuLogits + * @property {ReadonlyArray<'argmax'|'multinomial'>} declarativePlans + * @property {{readonly defaultDepth: number, readonly maxDepth: number}} tokenPipeline + * @property {SessionConcurrencyCapabilities} [sessionConcurrency] + */ + +/** @typedef {CausalGenerationCapabilitiesV1} GenerationCapabilitiesV1 */ + +/** + * @typedef {Object} LogitsLeaseV1 + * @property {1} version + * @property {'float32'} dtype + * @property {readonly [number, number]} shape + * @property {() => Promise} read + * @property {() => void} release + */ + +/** + * @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'}|{readonly op: 'multinomial', readonly temperature: number, readonly topK: number}} sampler + * @property {number} maxNewTokens + * @property {number} [pipelineDepth] + */ + +/** + * @typedef {Object} RuntimeTokenDecisionV1 + * @property {Uint32Array} tokenIds + * @property {Float32Array} [processedScores] + */ + +/** + * @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: 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. + * + * @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 { 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, resolvedInputIds, options); + const finalize = () => + controller.finalize( + controller.generationConfig.return_dict_in_generate ? { past_key_values: new DynamicCache() } : {}, + ); + + const capabilities = getCausalGenerationCapabilities(model); + try { + validateCapabilities(capabilities, controller, attention_mask); + throwIfAborted(signal); + } catch (error) { + 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) { + const error = new Error( + 'This generation request requires CPU-visible logits, but the runtime does not support them.', + ); + 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; + /** @type {LogitsLeaseV1|null} */ + let lease = null; + try { + session = await model.createAutoregressiveSession({ + batchSize: controller.batchSize, + maxSequenceLength: controller.maxSequenceLength, + signal, + }); + validateSession(session, controller, plan !== null); + + const prefillInputs = { + inputIds: tensorToTokenBatch(resolvedInputIds), + attentionMask: attention_mask ? tensorToAttentionMask(attention_mask) : undefined, + signal, + }; + if (plan) { + const planSession = /** @type {PlanAutoregressiveSessionV1} */ (session); + const decisions = planSession.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 finalize(); + } + + const pullSession = /** @type {PullAutoregressiveSessionV1} */ (session); + lease = await pullSession.prefill(prefillInputs); + while (!controller.allDone) { + throwIfAborted(signal); + const currentLease = lease; + lease = null; + let values; + try { + validateLease(currentLease, controller.batchSize); + values = await currentLease.read(); + } finally { + 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])); + if (step.allDone) break; + lease = await pullSession.decode({ + tokenIds: tensorToTokenBatch(step.nextTokenIds), + signal, + }); + } + return finalize(); + } catch (error) { + controller.abort(error); + throw error; + } finally { + try { + lease?.release?.(); + } finally { + try { + await session?.dispose(); + } finally { + releaseSessionSlot(); + } + } + } +} + +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 (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}.`, + ); + } + 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.'); + } + 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, 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}.`); + } + 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()`.'); + 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) { + 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 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/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/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..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 { LogitsSampler } from '../generation/logits_sampler.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'; @@ -41,6 +32,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 { getDefaultInferenceProvider, getOnnxProviderModule } from '../backends/default.js'; /** * Converts an array or Tensor of integers to an int64 Tensor. @@ -252,16 +245,77 @@ 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') { + const provider = await getDefaultInferenceProvider(pretrained_model_name_or_path); + return provider.load({ + ...options, + modelClass: this, + }); + } + if (typeof pretrained_model_name_or_path?.constructSessions === 'function') { + 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); + /** @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) { + 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 +324,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 +346,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); @@ -320,6 +380,7 @@ export class PreTrainedModel extends Callable { dtype, device, model_file_name, + inferenceProvider, }); const metadata = await Promise.all( @@ -347,9 +408,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)); } @@ -390,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( @@ -400,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); } /** @@ -558,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); } /** @@ -891,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; @@ -926,114 +798,58 @@ 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); - } - - const stop = prepared_stopping_criteria(all_input_ids); - if (stop.every((x) => x)) { - break; - } + }, + }); - model_inputs = this._update_model_kwargs_for_generation({ - generated_input_ids, - outputs, - model_inputs, - is_encoder_decoder, - }); + if (controller.allDone) { + return generation_config.return_dict_in_generate + ? controller.finalize({ past_key_values: kwargs.past_key_values ?? new DynamicCache() }) + : controller.finalize(); } - if (streamer) { - streamer.end(); + 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, + }); + } + } 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 +868,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..2c819fcbd 100644 --- a/packages/transformers/src/ops/registry.js +++ b/packages/transformers/src/ops/registry.js @@ -1,176 +1,46 @@ -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) { + const { getOnnxProviderModule } = await import('../backends/default.js'); + await getOnnxProviderModule(); + } + 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..58226a6fb 100644 --- a/packages/transformers/src/pipelines.js +++ b/packages/transformers/src/pipelines.js @@ -51,6 +51,16 @@ 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, + validateInferenceBackendTask, + validateInferenceModelTask, +} from './backends/inference.js'; +import { validateInferenceArtifactProvider } from './backends/artifacts.js'; +import { getModelJSON, getModelText } from './utils/hub.js'; +import { CHAT_TEMPLATE_NAME } from './utils/constants.js'; /** * @typedef {keyof typeof SUPPORTED_TASKS} TaskType @@ -61,6 +71,36 @@ import { get_file_metadata } from './utils/model_registry/get_file_metadata.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. * @@ -89,7 +129,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 +145,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,21 +172,51 @@ 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) { + validateInferenceBackendTask(/** @type {import('./backends/inference.js').InferenceBackend} */ (model), task); + } + // 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, + cache_dir, + local_files_only, + revision, + model_file_name, + 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(model, file))); + const metadata = await Promise.all( + expected_files.map(async (file) => + 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, }; } @@ -165,6 +237,10 @@ export async function pipeline( use_external_data_format, model_file_name, session_options, + generation_config: null, + signal, + artifactProvider, + artifactMetadata, }; // Determine which components to load based on the expected files @@ -174,8 +250,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,17 +274,40 @@ 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, - modelPromise, - ]); + let tokenizer; + let processor; + let model_loaded; + let chat_template; + try { + // Load all components in parallel. + [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); + } + } 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; @@ -203,7 +316,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/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 ef01569ef..cb1523603 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,9 @@ 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'; + // Expose common types used across the library for developers to access /** * @typedef {import('./utils/hub.js').PretrainedModelOptions} PretrainedModelOptions @@ -68,4 +72,28 @@ 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('./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 + * @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 + * @typedef {import('./backends/artifacts.js').ArtifactProgressEvent} ArtifactProgressEvent */ 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/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..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 { @@ -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'; @@ -41,17 +42,19 @@ export { MAX_EXTERNAL_DATA_CHUNKS } from './hub/constants.js'; * @property {string} [revision='main'] The specific model version to use. It can be a branch name, a tag name, or a commit id, * 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. */ /** * @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. */ /** @@ -62,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 @@ -76,6 +81,7 @@ export async function getFile(urlOrPath) { } else { return env.fetch(urlOrPath, { headers: getFetchHeaders(urlOrPath), + signal, }); } } @@ -130,6 +136,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); @@ -257,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, @@ -275,6 +283,7 @@ export async function loadResourceFile( // Check cache response = await checkCachedResource(cache, localPath, proposedCacheKey); + throwIfAborted(options.signal); const cacheHit = response !== undefined; if (cacheHit) { @@ -288,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. @@ -331,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); @@ -421,6 +430,7 @@ export async function loadResourceFile( ); } } + throwIfAborted(options.signal); result = buffer; } @@ -499,6 +509,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/ModelRegistry.js b/packages/transformers/src/utils/model_registry/ModelRegistry.js index 1179a92f6..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,36 +188,41 @@ 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 * * @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(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 * * @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(model, options = {}) { + const request = resolveModelRegistryRequest(model, options); + return get_processor_files(request.modelId, request.options); } /** @@ -212,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) @@ -225,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); } /** @@ -234,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') @@ -247,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') @@ -269,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); } /** @@ -279,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') @@ -292,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); } /** @@ -301,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') @@ -315,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 @@ -339,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') @@ -354,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); } /** @@ -363,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') @@ -376,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 fad61da76..7b16e1cd0 100644 --- a/packages/transformers/src/utils/model_registry/get_available_dtypes.js +++ b/packages/transformers/src/utils/model_registry/get_available_dtypes.js @@ -1,19 +1,13 @@ 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 { 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[]} - */ -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. @@ -29,40 +23,33 @@ const CONCRETE_DTYPES = Object.keys(DEFAULT_DTYPE_SUFFIX_MAPPING); * @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 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); + const provider = await getModelRegistryInferenceProvider(inferenceProvider); + if (!provider.getAvailableDtypes) { + throw new Error('The inference backend does not support dtype discovery.'); + } + return provider.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..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,7 +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 @@ -24,10 +26,12 @@ import { memoizePromise } from '../memoize_promise.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; @@ -35,7 +39,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 }); } /** @@ -43,19 +47,41 @@ async function fetch_file_head(urlOrPath) { * 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); + 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) { + throwIfAborted(options.signal); /** @type {import('../cache.js').CacheInterface | null} */ const cache = await getCache(options?.cache_dir); const { localPath, remoteURL, proposedCacheKey, validModelId } = buildResourcePaths( @@ -67,6 +93,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'); @@ -83,7 +110,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'); @@ -96,6 +123,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 } } @@ -105,7 +133,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; @@ -147,6 +175,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}`); } diff --git a/packages/transformers/src/utils/model_registry/get_files.js b/packages/transformers/src/utils/model_registry/get_files.js index eeda2891b..c27e4de4b 100644 --- a/packages/transformers/src/utils/model_registry/get_files.js +++ b/packages/transformers/src/utils/model_registry/get_files.js @@ -12,8 +12,14 @@ 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 {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 + * @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( @@ -23,18 +29,35 @@ export async function get_files( dtype = null, device = null, model_file_name = null, + cache_dir = null, + local_files_only = false, + revision = 'main', + task = null, + inferenceProvider = null, include_tokenizer = true, include_processor = true, + include_model = true, } = {}, ) { - const files = 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, + task, + 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 467642421..9c74f65c5 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 { getModelRegistryInferenceProvider } from '../../backends/model_registry.js'; /** * @typedef {import('../../configs.js').PretrainedConfig} PretrainedConfig @@ -55,61 +53,42 @@ 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 {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 */ 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', + task = null, + inferenceProvider = null, + } = {}, ) { - 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; + config = await get_config(modelId, { config, cache_dir, local_files_only, revision }); // 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; + const provider = await getModelRegistryInferenceProvider(inferenceProvider); + 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 c6d4e3a7d..e5b3d5ff7 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 { getModelRegistryInferenceProvider } from '../../backends/model_registry.js'; /** * Get all files needed for a specific pipeline task. @@ -16,6 +17,11 @@ 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 {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 */ @@ -41,20 +47,23 @@ export async function get_pipeline_files(task, modelId, options = {}) { const files = await get_files(modelId, { ...options, + task, include_tokenizer, include_processor, }); // 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) { - const allowedPrefixes = Object.values(textOnlySessions).map((s) => `onnx/${s}`); - return files.filter((f) => !f.startsWith('onnx/') || allowedPrefixes.some((p) => f.startsWith(p))); + const provider = await getModelRegistryInferenceProvider(options.inferenceProvider ?? null); + if (provider.filterModelArtifacts) { + 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/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/src/utils/tensor.js b/packages/transformers/src/utils/tensor.js index 85da403ef..585f0e5ea 100644 --- a/packages/transformers/src/utils/tensor.js +++ b/packages/transformers/src/utils/tensor.js @@ -9,14 +9,14 @@ 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'; import { random } from './random.js'; +const NOOP_DISPOSE = () => {}; + /** * @typedef {keyof typeof DataTypeMap} DataType * @typedef {import('./maths.js').AnyTypedArray | any[]} DataArray @@ -28,13 +28,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 +39,7 @@ export class Tensor { * @type {DataType} */ get type() { - return this.ort_tensor.type; + return this._storage.type; } /** @@ -50,7 +47,7 @@ export class Tensor { * @type {DataArray} */ get data() { - return this.ort_tensor.data; + return this._storage.data; } /** @@ -58,7 +55,7 @@ export class Tensor { * @type {number} */ get size() { - return this.ort_tensor.size; + return this._storage.size; } /** @@ -66,27 +63,40 @@ 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); + } + } + 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, + location: 'cpu', + dispose: NOOP_DISPOSE, + }; return new Proxy(this, { get: (obj, key) => { @@ -110,8 +120,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..9acd03f4c --- /dev/null +++ b/packages/transformers/tests/generation_controller.test.js @@ -0,0 +1,474 @@ +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 TrackingTokenProcessor extends LogitsProcessor { + constructor() { + super(); + this.onTokensSampled = jest.fn(); + } + + _call(_inputIds, logits) { + 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("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)); + 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); + }); + + 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( + "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", () => { + 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({ + inputs: 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, + 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: {}, + 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) { + 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, return_dict_in_generate: true }); + + expect(output.sequences.tolist()).toEqual([[1n, 2n, 3n]]); + expect(output.past_key_values).toBeDefined(); + 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 = { + 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(); + }); + + 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 = { + 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 new file mode 100644 index 000000000..e5fcf842a --- /dev/null +++ b/packages/transformers/tests/inference_backends.test.js @@ -0,0 +1,175 @@ +import { jest } from "@jest/globals"; + +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"; +import { loadInferenceBackendChatTemplate } from "../src/pipelines.js"; +import { DefaultProgressCallback } from "../src/utils/core.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("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", + 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), + 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.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 }); + }); + + 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"); + }); + + 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/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/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/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 d2c058c85..f22c94bfd 100644 --- a/packages/transformers/tests/types.test.js +++ b/packages/transformers/tests/types.test.js @@ -69,6 +69,76 @@ describe("TypeScript compilation succeeds", () => { } }); } + + it("compiles a readonly plan-only inference backend", () => { + const diagnostics = getDiagnosticsFromSource(` + import type { + CausalGenerationCapabilitiesV1, + InferenceBackend, + InferenceBackendChatTemplate, + 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", + chatTemplate: { modelId: "test/templates", file: "chat.jinja" }, + capabilities: staticCapabilities, + async load(_options) { + const model: InferenceModel = { + capabilities: { causalGeneration }, + async createAutoregressiveSession() { return session; }, + async dispose() {}, + }; + return model; + }, + }; + + void backend; + void lease; + const inlineTemplate = { content: "{{ messages }}" } as const satisfies InferenceBackendChatTemplate; + void inlineTemplate; + `); + 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/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/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..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 = { @@ -55,6 +56,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 @@ -181,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/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/pnpm-lock.yaml b/pnpm-lock.yaml index bea34792e..75d83120e 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-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) + onnxruntime-node: + specifier: 1.24.3 + version: 1.24.3 + typescript: + specifier: 5.9.3 + version: 5.9.3 + packages: '@babel/code-frame@7.28.6':