diff --git a/.gitignore b/.gitignore index 21c721c4a..10344ec19 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ __pycache__ node_modules deno.lock package-lock.json +*.local.* # Do not track build artifacts/generated files packages/*/dist diff --git a/packages/transformers/src/configs.js b/packages/transformers/src/configs.js index 68a29d850..b93432340 100644 --- a/packages/transformers/src/configs.js +++ b/packages/transformers/src/configs.js @@ -27,7 +27,7 @@ */ import { pick } from './utils/core.js'; -import { getModelJSON } from './utils/hub.js'; +import { getModelJSON, maybeAddDeprecatedEnvWarning } from './utils/hub.js'; /** * @typedef {import('./utils/hub.js').PretrainedOptions} PretrainedOptions @@ -531,8 +531,17 @@ export class PretrainedConfig { */ static async from_pretrained( pretrained_model_name_or_path, - { progress_callback = null, config = null, cache_dir = null, local_files_only = false, revision = 'main' } = {}, + { + progress_callback = null, + config = null, + cache_dir = null, + local_files_only = false, + revision = 'main', + env: sessionEnv = {}, + } = {}, ) { + maybeAddDeprecatedEnvWarning(cache_dir, local_files_only); + if (config && !(config instanceof PretrainedConfig)) { config = new PretrainedConfig(config); } @@ -545,6 +554,7 @@ export class PretrainedConfig { cache_dir, local_files_only, revision, + env: sessionEnv, })); return new this(data); } diff --git a/packages/transformers/src/env.js b/packages/transformers/src/env.js index bea3ef0dc..f73e13b08 100644 --- a/packages/transformers/src/env.js +++ b/packages/transformers/src/env.js @@ -161,6 +161,7 @@ if (RUNNING_LOCALLY) { // Only used for environments with access to file system const DEFAULT_CACHE_DIR = RUNNING_LOCALLY ? path.join(dirname__, '/.cache/') : null; +const DEFAULT_REMOTE_HOST = globalThis.process?.env?.HF_ENDPOINT ?? 'https://huggingface.co/'; // Set local model path, based on available APIs const DEFAULT_LOCAL_MODEL_PATH = '/models/'; @@ -168,6 +169,7 @@ const localModelPath = RUNNING_LOCALLY ? path.join(dirname__, DEFAULT_LOCAL_MODE // Ensure default fetch is called with the correct receiver in browser environments. const DEFAULT_FETCH = typeof globalThis.fetch === 'function' ? globalThis.fetch.bind(globalThis) : undefined; +const DEFAULT_HF_TOKEN = globalThis.process?.env?.HF_TOKEN ?? globalThis.process?.env?.HF_ACCESS_TOKEN; /** * Log levels for controlling output verbosity. @@ -202,12 +204,9 @@ 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, - * allowing users to set these variables if they want to. - * @property {number} logLevel The logging level. Use LogLevel enum values. Defaults to LogLevel.ERROR. + * Session-scopable variables that describe the model, task, or resource-loading contract that one pipeline or library must control without affecting others + * + * @typedef {Object} TransformersEnvironmentSession * @property {boolean} allowRemoteModels Whether to allow loading of remote files, defaults to `true`. * If set to `false`, it will have the same effect as setting `local_files_only=true` when loading pipelines, models, tokenizers, processors, etc. * @property {string} remoteHost Host URL to load models from. Defaults to the Hugging Face Hub. @@ -215,6 +214,18 @@ export const LogLevel = Object.freeze({ * @property {boolean} allowLocalModels Whether to allow loading of local files, defaults to `false` if running in-browser, and `true` otherwise. * If set to `false`, it will skip the local file check and try to load the model from the remote host. * @property {string} localModelPath Path to load local models from. Defaults to `/models/`. + * @property {(input: string | URL, init?: any) => Promise} fetch The fetch function to use. Defaults to `fetch`. + * @property {string|undefined} hfToken Hugging Face access token to use for requests to the Hugging Face Hub. + */ + +/** + * Global variables that describe runtime facts, backend singleton state, or application-owned infrastructure shared across models + * + * @typedef {Object} TransformersEnvironmentGlobal + * @property {string} version This version of Transformers.js. + * @property {{onnx: Partial & { setLogLevel?: (logLevel: number) => void }}} backends Expose environment variables of different backends, + * 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} useFS Whether to use the file system to load files. By default, it is `true` if available. * @property {boolean} useBrowserCache Whether to use Cache API to cache models. By default, it is `true` if available. * @property {boolean} useFSCache Whether to use the file system to cache files. By default, it is `true` if available. @@ -230,7 +241,12 @@ export const LogLevel = Object.freeze({ * Requires the Cross-Origin Storage Chrome extension: {@link https://chromewebstore.google.com/detail/cross-origin-storage/denpnpcgjgikjpoglpjefakmdcbmlgih}. * The `experimental_` prefix indicates that the underlying browser API is not yet standardised and may change or be * removed without a major version bump. For more information, see {@link https://github.com/WICG/cross-origin-storage}. - * @property {(input: string | URL, init?: any) => Promise} fetch The fetch function to use. Defaults to `fetch`. + */ + +/** + * Global variable given visible to users to control execution. This provides users a simple way to configure Transformers.js. + * + * @typedef {TransformersEnvironmentGlobal & TransformersEnvironmentSession} TransformersEnvironment */ let logLevel = LogLevel.WARNING; // Default log level @@ -257,7 +273,7 @@ export const env = { }, /////////////////// Model settings /////////////////// allowRemoteModels: true, - remoteHost: 'https://huggingface.co/', + remoteHost: DEFAULT_REMOTE_HOST, remotePathTemplate: '{model}/resolve/{revision}/', allowLocalModels: !(IS_BROWSER_ENV || IS_WEBWORKER_ENV || IS_DENO_WEB_RUNTIME), // Default to true for non-web environments, false for web environments @@ -280,10 +296,24 @@ export const env = { /////////////////// Custom fetch ///////////////////// fetch: DEFAULT_FETCH, + hfToken: DEFAULT_HF_TOKEN, ////////////////////////////////////////////////////// }; +/** + * Create a session environment by applying session-scopable overrides to the global environment defaults. + * + * @param {Partial} sessionEnv Session-scopable environment overrides. + * @returns {TransformersEnvironment} + */ +export function resolveEnv(sessionEnv = {}) { + return { + ...env, + ...sessionEnv, + }; +} + /** * @param {Object} obj * @private diff --git a/packages/transformers/src/models/auto/modeling_auto.js b/packages/transformers/src/models/auto/modeling_auto.js index 030c51acd..f179451e1 100644 --- a/packages/transformers/src/models/auto/modeling_auto.js +++ b/packages/transformers/src/models/auto/modeling_auto.js @@ -90,6 +90,7 @@ class PretrainedMixin { dtype = null, use_external_data_format = null, session_options = {}, + env: sessionEnv = {}, } = {}, ) { const options = { @@ -104,6 +105,7 @@ class PretrainedMixin { dtype, use_external_data_format, session_options, + env: sessionEnv, }; options.config = await AutoConfig.from_pretrained(pretrained_model_name_or_path, options); diff --git a/packages/transformers/src/models/auto/tokenization_auto.js b/packages/transformers/src/models/auto/tokenization_auto.js index 62f3c613e..4b504d041 100644 --- a/packages/transformers/src/models/auto/tokenization_auto.js +++ b/packages/transformers/src/models/auto/tokenization_auto.js @@ -40,7 +40,14 @@ export class AutoTokenizer { */ static async from_pretrained( pretrained_model_name_or_path, - { progress_callback = null, config = null, cache_dir = null, local_files_only = false, revision = 'main' } = {}, + { + progress_callback = null, + config = null, + cache_dir = null, + local_files_only = false, + revision = 'main', + env: sessionEnv = {}, + } = {}, ) { const [tokenizerJSON, tokenizerConfig] = await loadTokenizer(pretrained_model_name_or_path, { progress_callback, @@ -48,6 +55,7 @@ export class AutoTokenizer { cache_dir, local_files_only, revision, + env: sessionEnv, }); // Some tokenizers are saved with the "Fast" suffix, so we remove that if present. diff --git a/packages/transformers/src/models/modeling_utils.js b/packages/transformers/src/models/modeling_utils.js index ec9f17487..0d403361b 100644 --- a/packages/transformers/src/models/modeling_utils.js +++ b/packages/transformers/src/models/modeling_utils.js @@ -15,7 +15,7 @@ export function registerTaskMappings(mappings) { MODEL_MAPPING_NAMES = mappings; } import { GITHUB_ISSUE_URL } from '../utils/constants.js'; -import { getModelJSON } from '../utils/hub.js'; +import { getModelJSON, maybeAddDeprecatedEnvWarning } from '../utils/hub.js'; import { Seq2SeqLMOutput } from './modeling_outputs.js'; import { LogitsProcessorList, @@ -275,8 +275,11 @@ export class PreTrainedModel extends Callable { dtype = null, use_external_data_format = null, session_options = {}, + env: sessionEnv = {}, } = {}, ) { + maybeAddDeprecatedEnvWarning(cache_dir, local_files_only); + const options = { progress_callback, config, @@ -289,11 +292,19 @@ export class PreTrainedModel extends Callable { dtype, use_external_data_format, session_options, + env: sessionEnv, }; const modelName = MODEL_CLASS_TO_NAME_MAPPING.get(this); - config = options.config = await AutoConfig.from_pretrained(pretrained_model_name_or_path, options); + config = options.config = await AutoConfig.from_pretrained(pretrained_model_name_or_path, { + progress_callback, + config, + cache_dir, + local_files_only, + revision, + env: sessionEnv, + }); const { typeConfig, textOnly, modelType } = resolveTypeConfig(modelName, config); @@ -320,6 +331,10 @@ export class PreTrainedModel extends Callable { dtype, device, model_file_name, + cache_dir, + local_files_only, + revision, + env: sessionEnv, }); const metadata = await Promise.all( diff --git a/packages/transformers/src/pipelines.js b/packages/transformers/src/pipelines.js index 98c10c692..76884bcbe 100644 --- a/packages/transformers/src/pipelines.js +++ b/packages/transformers/src/pipelines.js @@ -51,6 +51,7 @@ 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 { maybeAddDeprecatedEnvWarning } from './utils/hub.js'; /** * @typedef {keyof typeof SUPPORTED_TASKS} TaskType @@ -109,8 +110,11 @@ export async function pipeline( use_external_data_format = null, model_file_name = null, session_options = {}, + env: sessionEnv = {}, } = {}, ) { + maybeAddDeprecatedEnvWarning(cache_dir, local_files_only); + // Apply aliases // @ts-ignore task = TASK_ALIASES[task] ?? task; @@ -132,15 +136,25 @@ export async function pipeline( // Determine which files the model needs const expected_files = await get_pipeline_files(task, model, { + config, + cache_dir, + local_files_only, + revision, device, dtype, + model_file_name, + env: sessionEnv, }); /** @type {import('./utils/core.js').FilesLoadingMap} */ let files_loading = {}; if (progress_callback) { /** @type {Array<{exists: boolean, size?: number, contentType?: string, fromCache?: boolean}>} */ - const metadata = await Promise.all(expected_files.map(async (file) => get_file_metadata(model, file))); + const metadata = await Promise.all( + expected_files.map(async (file) => + get_file_metadata(model, file, { cache_dir, local_files_only, revision, env: sessionEnv }), + ), + ); metadata.forEach((m, i) => { if (m.exists) { files_loading[expected_files[i]] = { @@ -165,6 +179,7 @@ export async function pipeline( use_external_data_format, model_file_name, session_options, + env: sessionEnv, }; // Determine which components to load based on the expected files @@ -196,7 +211,7 @@ export async function pipeline( modelPromise, ]); - const results = { task, model: model_loaded }; + const results = { task, model: model_loaded, sessionEnv }; if (tokenizer) results.tokenizer = tokenizer; if (processor) results.processor = processor; diff --git a/packages/transformers/src/pipelines/_base.js b/packages/transformers/src/pipelines/_base.js index 495c10c6a..744cf60b4 100644 --- a/packages/transformers/src/pipelines/_base.js +++ b/packages/transformers/src/pipelines/_base.js @@ -1,6 +1,7 @@ import { PreTrainedTokenizer } from '../tokenization_utils.js'; import { PreTrainedModel } from '../models/modeling_utils.js'; import { Processor } from '../processing_utils.js'; +import { resolveEnv } from '../env.js'; import { Callable } from '../utils/generic.js'; @@ -15,15 +16,16 @@ import { RawImage } from '../utils/image.js'; /** * Prepare images for further tasks. * @param {ImagePipelineInputs} images images to prepare. + * @param {Partial} [sessionEnv={}] Session-scopable environment overrides. * @returns {Promise} returns processed images. */ -export async function prepareImages(images) { +export async function prepareImages(images, sessionEnv = {}) { if (!Array.isArray(images)) { images = [images]; } // Possibly convert any non-images to images - return await Promise.all(images.map((x) => RawImage.read(x))); + return await Promise.all(images.map((x) => RawImage.read(x, sessionEnv))); } /** @@ -35,9 +37,10 @@ export async function prepareImages(images) { * Prepare audios for further tasks. * @param {AudioPipelineInputs} audios audios to prepare. * @param {number} sampling_rate sampling rate of the audios. + * @param {Partial} [sessionEnv={}] Session-scopable environment overrides. * @returns {Promise} The preprocessed audio data. */ -export async function prepareAudios(audios, sampling_rate) { +export async function prepareAudios(audios, sampling_rate, sessionEnv = {}) { if (!Array.isArray(audios)) { audios = [audios]; } @@ -45,7 +48,7 @@ export async function prepareAudios(audios, sampling_rate) { return await Promise.all( audios.map((x) => { if (typeof x === 'string' || x instanceof URL) { - return read_audio(x, sampling_rate); + return read_audio(x, sampling_rate, sessionEnv); } else if (x instanceof Float64Array) { return new Float32Array(x); } @@ -91,6 +94,12 @@ export function get_bounding_box(box, asInteger) { * Refer to this class for methods shared across different pipelines. */ export class Pipeline extends Callable { + /** @type {Partial} */ + sessionEnv; + + /** @type {import('../env.js').TransformersEnvironment} */ + env; + /** * Create a new Pipeline. * @param {Object} options An object containing the following properties: @@ -98,13 +107,16 @@ export class Pipeline extends Callable { * @param {PreTrainedModel} [options.model] The model used by the pipeline. * @param {PreTrainedTokenizer} [options.tokenizer=null] The tokenizer used by the pipeline (if any). * @param {Processor} [options.processor=null] The processor used by the pipeline (if any). + * @param {Partial} [options.sessionEnv={}] Session-scopable environment overrides. */ - constructor({ task, model, tokenizer = null, processor = null }) { + constructor({ task, model, tokenizer = null, processor = null, sessionEnv = {} }) { super(); this.task = task; this.model = model; this.tokenizer = tokenizer; this.processor = processor; + this.sessionEnv = sessionEnv; + this.env = resolveEnv(sessionEnv); } /** @type {DisposeType} */ @@ -118,6 +130,8 @@ export class Pipeline extends Callable { * @property {string} task The task of the pipeline. Useful for specifying subtasks. * @property {PreTrainedModel} model The model used by the pipeline. * @property {PreTrainedTokenizer} tokenizer The tokenizer used by the pipeline. + * @property {Partial} sessionEnv The session-scopable environment overrides used by the pipeline. + * @property {import('../env.js').TransformersEnvironment} env The effective environment used by the pipeline. * * @typedef {ModelTokenizerConstructorArgs} TextPipelineConstructorArgs An object used to instantiate a text-based pipeline. */ @@ -127,6 +141,8 @@ export class Pipeline extends Callable { * @property {string} task The task of the pipeline. Useful for specifying subtasks. * @property {PreTrainedModel} model The model used by the pipeline. * @property {Processor} processor The processor used by the pipeline. + * @property {Partial} sessionEnv The session-scopable environment overrides used by the pipeline. + * @property {import('../env.js').TransformersEnvironment} env The effective environment used by the pipeline. * * @typedef {ModelProcessorConstructorArgs} AudioPipelineConstructorArgs An object used to instantiate an audio-based pipeline. * @typedef {ModelProcessorConstructorArgs} ImagePipelineConstructorArgs An object used to instantiate an image-based pipeline. @@ -138,6 +154,8 @@ export class Pipeline extends Callable { * @property {PreTrainedModel} model The model used by the pipeline. * @property {PreTrainedTokenizer} tokenizer The tokenizer used by the pipeline. * @property {Processor} processor The processor used by the pipeline. + * @property {Partial} sessionEnv The session-scopable environment overrides used by the pipeline. + * @property {import('../env.js').TransformersEnvironment} env The effective environment used by the pipeline. * * @typedef {ModelTokenizerProcessorConstructorArgs} TextAudioPipelineConstructorArgs An object used to instantiate a text- and audio-based pipeline. * @typedef {ModelTokenizerProcessorConstructorArgs} TextImagePipelineConstructorArgs An object used to instantiate a text- and image-based pipeline. diff --git a/packages/transformers/src/pipelines/audio-classification.js b/packages/transformers/src/pipelines/audio-classification.js index 60cd9ff10..dc0b806f2 100644 --- a/packages/transformers/src/pipelines/audio-classification.js +++ b/packages/transformers/src/pipelines/audio-classification.js @@ -69,7 +69,7 @@ export class AudioClassificationPipeline { async _call(audio, { top_k = 5 } = {}) { const sampling_rate = this.processor.feature_extractor.config.sampling_rate; - const preparedAudios = await prepareAudios(audio, sampling_rate); + const preparedAudios = await prepareAudios(audio, sampling_rate, this.sessionEnv); // @ts-expect-error TS2339 const id2label = this.model.config.id2label; diff --git a/packages/transformers/src/pipelines/automatic-speech-recognition.js b/packages/transformers/src/pipelines/automatic-speech-recognition.js index 238bb45a4..f3bbb6962 100644 --- a/packages/transformers/src/pipelines/automatic-speech-recognition.js +++ b/packages/transformers/src/pipelines/automatic-speech-recognition.js @@ -175,7 +175,7 @@ export class AutomaticSpeechRecognitionPipeline const batchedAudio = single ? [audio] : audio; const sampling_rate = this.processor.feature_extractor.config.sampling_rate; - const preparedAudios = await prepareAudios(batchedAudio, sampling_rate); + const preparedAudios = await prepareAudios(batchedAudio, sampling_rate, this.sessionEnv); const toReturn = []; for (const aud of preparedAudios) { @@ -215,7 +215,7 @@ export class AutomaticSpeechRecognitionPipeline const hop_length = feature_extractor_config.hop_length; const sampling_rate = feature_extractor_config.sampling_rate; - const preparedAudios = await prepareAudios(batchedAudio, sampling_rate); + const preparedAudios = await prepareAudios(batchedAudio, sampling_rate, this.sessionEnv); const toReturn = []; for (const aud of preparedAudios) { @@ -314,7 +314,7 @@ export class AutomaticSpeechRecognitionPipeline const single = !Array.isArray(audio); const batchedAudio = single ? [audio] : audio; const sampling_rate = this.processor.feature_extractor.config.sampling_rate; - const preparedAudios = await prepareAudios(batchedAudio, sampling_rate); + const preparedAudios = await prepareAudios(batchedAudio, sampling_rate, this.sessionEnv); const toReturn = []; for (const aud of preparedAudios) { const inputs = await this.processor(aud); @@ -337,7 +337,7 @@ export class AutomaticSpeechRecognitionPipeline const feature_extractor = this.processor.feature_extractor; const sampling_rate = feature_extractor.config.sampling_rate; - const preparedAudios = await prepareAudios(batchedAudio, sampling_rate); + const preparedAudios = await prepareAudios(batchedAudio, sampling_rate, this.sessionEnv); const language = kwargs.language ?? 'en'; // @ts-expect-error TS2339 diff --git a/packages/transformers/src/pipelines/background-removal.js b/packages/transformers/src/pipelines/background-removal.js index 37b84706c..518a80c09 100644 --- a/packages/transformers/src/pipelines/background-removal.js +++ b/packages/transformers/src/pipelines/background-removal.js @@ -43,7 +43,7 @@ export class BackgroundRemovalPipeline ) { async _call(images, options = {}) { - const preparedImages = await prepareImages(images); + const preparedImages = await prepareImages(images, this.sessionEnv); // @ts-expect-error TS2339 const masks = await super._call(images, options); diff --git a/packages/transformers/src/pipelines/depth-estimation.js b/packages/transformers/src/pipelines/depth-estimation.js index fd60aac1a..14fad8a74 100644 --- a/packages/transformers/src/pipelines/depth-estimation.js +++ b/packages/transformers/src/pipelines/depth-estimation.js @@ -56,7 +56,7 @@ export class DepthEstimationPipeline extends /** @type {new (options: ImagePipelineConstructorArgs) => DepthEstimationPipelineType} */ (Pipeline) { async _call(images) { - const preparedImages = await prepareImages(images); + const preparedImages = await prepareImages(images, this.sessionEnv); const inputs = await this.processor(preparedImages); const { predicted_depth } = await this.model(inputs); diff --git a/packages/transformers/src/pipelines/document-question-answering.js b/packages/transformers/src/pipelines/document-question-answering.js index 091481174..6eea489b5 100644 --- a/packages/transformers/src/pipelines/document-question-answering.js +++ b/packages/transformers/src/pipelines/document-question-answering.js @@ -55,7 +55,7 @@ export class DocumentQuestionAnsweringPipeline } // Preprocess image - const preparedImage = (await prepareImages(image))[0]; + const preparedImage = (await prepareImages(image, this.sessionEnv))[0]; const { pixel_values } = await this.processor(preparedImage); // Run tokenization diff --git a/packages/transformers/src/pipelines/image-classification.js b/packages/transformers/src/pipelines/image-classification.js index f9374e896..6af0d3555 100644 --- a/packages/transformers/src/pipelines/image-classification.js +++ b/packages/transformers/src/pipelines/image-classification.js @@ -80,7 +80,7 @@ export class ImageClassificationPipeline extends /** @type {new (options: ImagePipelineConstructorArgs) => ImageClassificationPipelineType} */ (Pipeline) { async _call(images, { top_k = 5 } = {}) { - const preparedImages = await prepareImages(images); + const preparedImages = await prepareImages(images, this.sessionEnv); const { pixel_values } = await this.processor(preparedImages); const output = await this.model({ pixel_values }); diff --git a/packages/transformers/src/pipelines/image-feature-extraction.js b/packages/transformers/src/pipelines/image-feature-extraction.js index 25687f3a6..2e8d12e88 100644 --- a/packages/transformers/src/pipelines/image-feature-extraction.js +++ b/packages/transformers/src/pipelines/image-feature-extraction.js @@ -59,7 +59,7 @@ export class ImageFeatureExtractionPipeline { /** @type {ImageFeatureExtractionPipelineCallback} */ async _call(images, { pool = null } = {}) { - const preparedImages = await prepareImages(images); + const preparedImages = await prepareImages(images, this.sessionEnv); const { pixel_values } = await this.processor(preparedImages); const outputs = await this.model({ pixel_values }); diff --git a/packages/transformers/src/pipelines/image-segmentation.js b/packages/transformers/src/pipelines/image-segmentation.js index 566ced417..8bed911cb 100644 --- a/packages/transformers/src/pipelines/image-segmentation.js +++ b/packages/transformers/src/pipelines/image-segmentation.js @@ -78,7 +78,7 @@ export class ImageSegmentationPipeline throw Error('Image segmentation pipeline currently only supports a batch size of 1.'); } - const preparedImages = await prepareImages(images); + const preparedImages = await prepareImages(images, this.sessionEnv); const imageSizes = preparedImages.map((x) => [x.height, x.width]); const inputs = await this.processor(preparedImages); diff --git a/packages/transformers/src/pipelines/image-to-image.js b/packages/transformers/src/pipelines/image-to-image.js index 66324e4b6..f135516be 100644 --- a/packages/transformers/src/pipelines/image-to-image.js +++ b/packages/transformers/src/pipelines/image-to-image.js @@ -42,7 +42,7 @@ export class ImageToImagePipeline extends /** @type {new (options: ImagePipelineConstructorArgs) => ImageToImagePipelineType} */ (Pipeline) { async _call(images) { - const preparedImages = await prepareImages(images); + const preparedImages = await prepareImages(images, this.sessionEnv); const inputs = await this.processor(preparedImages); const outputs = await this.model(inputs); diff --git a/packages/transformers/src/pipelines/image-to-text.js b/packages/transformers/src/pipelines/image-to-text.js index ddd0e3d07..cd292d623 100644 --- a/packages/transformers/src/pipelines/image-to-text.js +++ b/packages/transformers/src/pipelines/image-to-text.js @@ -53,7 +53,7 @@ export class ImageToTextPipeline { async _call(images, generate_kwargs = {}) { const isBatched = Array.isArray(images); - const preparedImages = await prepareImages(images); + const preparedImages = await prepareImages(images, this.sessionEnv); const { pixel_values } = await this.processor(preparedImages); diff --git a/packages/transformers/src/pipelines/object-detection.js b/packages/transformers/src/pipelines/object-detection.js index a1fed1789..a685cd866 100644 --- a/packages/transformers/src/pipelines/object-detection.js +++ b/packages/transformers/src/pipelines/object-detection.js @@ -63,7 +63,7 @@ export class ObjectDetectionPipeline if (isBatched && images.length !== 1) { throw Error('Object detection pipeline currently only supports a batch size of 1.'); } - const preparedImages = await prepareImages(images); + const preparedImages = await prepareImages(images, this.sessionEnv); const imageSizes = percentage ? null : preparedImages.map((x) => [x.height, x.width]); diff --git a/packages/transformers/src/pipelines/text-to-audio.js b/packages/transformers/src/pipelines/text-to-audio.js index aab19b2cb..78ff440ab 100644 --- a/packages/transformers/src/pipelines/text-to-audio.js +++ b/packages/transformers/src/pipelines/text-to-audio.js @@ -5,7 +5,6 @@ import { RawAudio } from '../utils/audio.js'; import { logger } from '../utils/logger.js'; import { AutoModel } from '../models/auto/modeling_auto.js'; -import { env } from '../env.js'; /** * @typedef {import('./_base.js').TextAudioPipelineConstructorArgs} TextAudioPipelineConstructorArgs @@ -91,7 +90,7 @@ export class TextToAudioPipeline // Load speaker embeddings as Float32Array from path/URL if (typeof speaker_embeddings === 'string' || speaker_embeddings instanceof URL) { // Load from URL with fetch - speaker_embeddings = new Float32Array(await (await env.fetch(speaker_embeddings)).arrayBuffer()); + speaker_embeddings = new Float32Array(await (await this.env.fetch(speaker_embeddings)).arrayBuffer()); } if (speaker_embeddings instanceof Float32Array) { @@ -194,7 +193,10 @@ export class TextToAudioPipeline // Load vocoder, if not provided if (!this.vocoder) { logger.info('No vocoder specified, using default HifiGan vocoder.'); - this.vocoder = await AutoModel.from_pretrained(this.DEFAULT_VOCODER_ID, { dtype: 'fp32' }); + this.vocoder = await AutoModel.from_pretrained(this.DEFAULT_VOCODER_ID, { + dtype: 'fp32', + env: this.sessionEnv, + }); } // Run tokenization diff --git a/packages/transformers/src/pipelines/zero-shot-audio-classification.js b/packages/transformers/src/pipelines/zero-shot-audio-classification.js index 20a9184af..8c5ea2251 100644 --- a/packages/transformers/src/pipelines/zero-shot-audio-classification.js +++ b/packages/transformers/src/pipelines/zero-shot-audio-classification.js @@ -71,7 +71,7 @@ export class ZeroShotAudioClassificationPipeline }); const sampling_rate = this.processor.feature_extractor.config.sampling_rate; - const preparedAudios = await prepareAudios(audio, sampling_rate); + const preparedAudios = await prepareAudios(audio, sampling_rate, this.sessionEnv); const toReturn = []; for (const aud of preparedAudios) { diff --git a/packages/transformers/src/pipelines/zero-shot-image-classification.js b/packages/transformers/src/pipelines/zero-shot-image-classification.js index 80a62f197..460b1f814 100644 --- a/packages/transformers/src/pipelines/zero-shot-image-classification.js +++ b/packages/transformers/src/pipelines/zero-shot-image-classification.js @@ -58,7 +58,7 @@ export class ZeroShotImageClassificationPipeline { async _call(images, candidate_labels, { hypothesis_template = 'This is a photo of {}' } = {}) { const isBatched = Array.isArray(images); - const preparedImages = await prepareImages(images); + const preparedImages = await prepareImages(images, this.sessionEnv); // Insert label into hypothesis template const texts = candidate_labels.map((x) => hypothesis_template.replace('{}', x)); diff --git a/packages/transformers/src/pipelines/zero-shot-object-detection.js b/packages/transformers/src/pipelines/zero-shot-object-detection.js index b96c6a1cd..9d87ade56 100644 --- a/packages/transformers/src/pipelines/zero-shot-object-detection.js +++ b/packages/transformers/src/pipelines/zero-shot-object-detection.js @@ -109,7 +109,7 @@ export class ZeroShotObjectDetectionPipeline { async _call(images, candidate_labels, { threshold = 0.1, top_k = null, percentage = false } = {}) { const isBatched = Array.isArray(images); - const preparedImages = await prepareImages(images); + const preparedImages = await prepareImages(images, this.sessionEnv); // Run tokenization const text_inputs = this.tokenizer(candidate_labels, { diff --git a/packages/transformers/src/tokenization_utils.js b/packages/transformers/src/tokenization_utils.js index 29de6f186..123c0bc71 100644 --- a/packages/transformers/src/tokenization_utils.js +++ b/packages/transformers/src/tokenization_utils.js @@ -9,7 +9,7 @@ import { Template } from '@huggingface/jinja'; import { Callable } from './utils/generic.js'; import { isIntegralNumber, mergeArrays } from './utils/core.js'; -import { getModelJSON } from './utils/hub.js'; +import { getModelJSON, maybeAddDeprecatedEnvWarning } from './utils/hub.js'; import { max } from './utils/maths.js'; import { Tensor } from './utils/tensor.js'; import { logger } from './utils/logger.js'; @@ -26,7 +26,9 @@ 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); + maybeAddDeprecatedEnvWarning(options.cache_dir, options.local_files_only); + + 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)), ); @@ -302,7 +304,14 @@ export class PreTrainedTokenizer */ static async from_pretrained( pretrained_model_name_or_path, - { progress_callback = null, config = null, cache_dir = null, local_files_only = false, revision = 'main' } = {}, + { + progress_callback = null, + config = null, + cache_dir = null, + local_files_only = false, + revision = 'main', + env: sessionEnv = {}, + } = {}, ) { const info = await loadTokenizer(pretrained_model_name_or_path, { progress_callback, @@ -310,6 +319,7 @@ export class PreTrainedTokenizer cache_dir, local_files_only, revision, + env: sessionEnv, }); // @ts-ignore diff --git a/packages/transformers/src/utils/audio.js b/packages/transformers/src/utils/audio.js index c2af80753..4f9b3e1b9 100644 --- a/packages/transformers/src/utils/audio.js +++ b/packages/transformers/src/utils/audio.js @@ -8,6 +8,7 @@ */ import { getFile } from './hub.js'; +import { resolveEnv } from '../env.js'; import { FFT, max } from './maths.js'; import { calculateReflectOffset } from './core.js'; import { saveBlob } from './io.js'; @@ -18,9 +19,10 @@ import { logger } from './logger.js'; * Helper function to load audio from a path/URL. * @param {string|URL} url The path/URL to load the audio from. * @param {number} sampling_rate The sampling rate to use when decoding the audio. + * @param {Partial} [sessionEnv={}] Session-scopable environment overrides. * @returns {Promise} The decoded audio as a `Float32Array`. */ -export async function load_audio(url, sampling_rate) { +export async function load_audio(url, sampling_rate, sessionEnv = {}) { if (typeof AudioContext === 'undefined') { // Running in node or an environment without AudioContext throw Error( @@ -30,7 +32,15 @@ export async function load_audio(url, sampling_rate) { ); } - const response = await (await getFile(url)).arrayBuffer(); + const env = resolveEnv(sessionEnv); + const response = await ( + await getFile(url, { + useFS: env.useFS, + fetch: env.fetch, + version: env.version, + hfToken: env.hfToken, + }) + ).arrayBuffer(); const audioCTX = new AudioContext({ sampleRate: sampling_rate }); if (typeof sampling_rate === 'undefined') { logger.warn(`No sampling rate provided, using default of ${audioCTX.sampleRate}Hz.`); diff --git a/packages/transformers/src/utils/hub.js b/packages/transformers/src/utils/hub.js index 09f42e168..8880b1bbd 100755 --- a/packages/transformers/src/utils/hub.js +++ b/packages/transformers/src/utils/hub.js @@ -4,7 +4,7 @@ * @module utils/hub */ -import { apis, env } from '../env.js'; +import { apis, resolveEnv } from '../env.js'; import { DefaultProgressCallback, dispatchCallback } from './core.js'; import { FileResponse } from './hub/FileResponse.js'; import { FileCache } from './cache/FileCache.js'; @@ -30,14 +30,50 @@ export { MAX_EXTERNAL_DATA_CHUNKS } from './hub/constants.js'; * - `number`: Use external data format with the specified number of chunks */ +/** + * @typedef {Object} ModelLoadingOptions Options for fetching model files. + * @property {string|null} [cache_dir=null] Path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used. Deprecated: use `env.cacheDir` for the default cache directory and `options.env` for session-scopable resource loading settings. + * @property {string} [revision='main'] The specific model version to use. Ignored for local requests. + * @property {string} localModelPath Path to load local models from. + * @property {string} remoteHost Host URL to load models from. + * @property {string} remotePathTemplate Path template to fill in and append to `remoteHost` when loading models. + */ + +/** + * @typedef {Object} FetchHeadersOptions Options for generating fetch headers. + * @property {string} version This version of Transformers.js. + * @property {string|undefined} hfToken Hugging Face access token to use for requests to the Hugging Face Hub. + */ + +/** + * @typedef {Object} FetchOptions Options for fetching files. + * @property {boolean} useFS Whether to use the file system to load files. + * @property {(input: string | URL, init?: any) => Promise} fetch The fetch function to use. + * @property {string} version This version of Transformers.js. + * @property {string|undefined} hfToken Hugging Face access token to use for requests to the Hugging Face Hub. + */ + +/** + * @typedef {Object} ResourceLoadingOptions Options that control where resources may be loaded from. + * @property {boolean} allowLocalModels Whether to allow loading of local files. + * @property {boolean} allowRemoteModels Whether to allow loading of remote files. + * @property {boolean} localFilesOnly Whether to only look at local files. + */ + +/** + * @typedef {Object} ResourceProgressOptions Options for resource loading progress callbacks. + * @property {import('./core.js').ProgressCallback} [progressCallback] + */ + /** * @typedef {Object} PretrainedOptions Options for loading a pretrained model. * @property {import('./core.js').ProgressCallback} [progress_callback=null] If specified, this function will be called during model construction, to provide the user with progress updates. * @property {import('../configs.js').PretrainedConfig} [config=null] Configuration for the model to use instead of an automatically loaded configuration. Configuration can be automatically loaded when: * - The model is a model provided by the library (loaded with the *model id* string of a pretrained model). * - The model is loaded by supplying a local directory as `pretrained_model_name_or_path` and a configuration JSON file named *config.json* is found in the directory. - * @property {string} [cache_dir=null] Path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used. - * @property {boolean} [local_files_only=false] Whether or not to only look at local files (e.g., not try downloading the model). + * @property {Partial} [env={}] Session-scopable environment overrides. + * @property {string} [cache_dir=null] Path to a directory in which a downloaded pretrained model configuration should be cached if the standard cache should not be used. Deprecated: use `env.cacheDir` for the default cache directory and `options.env` for session-scopable resource loading settings. + * @property {boolean} [local_files_only=false] Whether or not to only look at local files (e.g., not try downloading the model). Deprecated: use `options.env.allowRemoteModels=false` for session-scoped remote loading control. * @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. @@ -48,8 +84,8 @@ export { MAX_EXTERNAL_DATA_CHUNKS } from './hub/constants.js'; * @property {string} [subfolder='onnx'] In case the relevant files 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 {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. */ @@ -62,10 +98,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 {FetchOptions} options Options for fetching files. * @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) { - if (env.useFS && !isValidUrl(urlOrPath, ['http:', 'https:', 'blob:'])) { +export async function getFile(urlOrPath, options) { + if (options.useFS && !isValidUrl(urlOrPath, ['http:', 'https:', 'blob:'])) { return new FileResponse( urlOrPath instanceof URL ? urlOrPath.protocol === 'file:' @@ -74,8 +111,8 @@ export async function getFile(urlOrPath) { : urlOrPath, ); } else { - return env.fetch(urlOrPath, { - headers: getFetchHeaders(urlOrPath), + return options.fetch(urlOrPath, { + headers: getFetchHeaders(urlOrPath, options), }); } } @@ -86,25 +123,21 @@ export async function getFile(urlOrPath) { * In browser environments, returns minimal headers for security. * * @param {URL|string} urlOrPath The URL or path being fetched. + * @param {FetchHeadersOptions} options Options for generating fetch headers. * @returns {Headers} A Headers object with appropriate headers for the request. */ -export function getFetchHeaders(urlOrPath) { +export function getFetchHeaders(urlOrPath, options) { const isNode = typeof process !== 'undefined' && process?.release?.name === 'node'; const headers = new Headers(); if (isNode) { const IS_CI = !!process.env?.TESTING_REMOTELY; - const version = env.version; - headers.set('User-Agent', `transformers.js/${version}; is_ci/${IS_CI};`); + headers.set('User-Agent', `transformers.js/${options.version}; is_ci/${IS_CI};`); const isHFURL = isValidUrl(urlOrPath, ['http:', 'https:'], ['huggingface.co', 'hf.co']); if (isHFURL) { - // If an access token is present in the environment variables, - // we add it to the request headers. - // NOTE: We keep `HF_ACCESS_TOKEN` for backwards compatibility (as a fallback). - const token = process.env?.HF_TOKEN ?? process.env?.HF_ACCESS_TOKEN; - if (token) { - headers.set('Authorization', `Bearer ${token}`); + if (options.hfToken) { + headers.set('Authorization', `Bearer ${options.hfToken}`); } } } else { @@ -124,20 +157,20 @@ export function getFetchHeaders(urlOrPath) { * - 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 locate. - * @param {PretrainedOptions} [options] An object containing optional parameters. + * @param {ModelLoadingOptions} [options] * @param {import('./cache.js').CacheInterface | null} [cache] The cache instance to use for determining cache keys. * @returns {{ requestURL: string, localPath: string, remoteURL: string, proposedCacheKey: string, validModelId: boolean }} * An object containing all the paths and URLs for the resource. */ -export function buildResourcePaths(path_or_repo_id, filename, options = {}, cache = null) { +export function buildResourcePaths(path_or_repo_id, filename, options, cache = null) { const revision = options.revision ?? 'main'; const requestURL = pathJoin(path_or_repo_id, filename); const validModelId = isValidHfModelId(path_or_repo_id); - const localPath = validModelId ? pathJoin(env.localModelPath, requestURL) : requestURL; + const localPath = validModelId ? pathJoin(options.localModelPath, requestURL) : requestURL; const remoteURL = pathJoin( - env.remoteHost, - env.remotePathTemplate + options.remoteHost, + options.remotePathTemplate .replaceAll('{model}', path_or_repo_id) .replaceAll('{revision}', encodeURIComponent(revision)), filename, @@ -192,7 +225,7 @@ export async function checkCachedResource(cache, localPath, proposedCacheKey) { * @param {string} cacheKey The cache key to use. * @param {Response|import('./hub/FileResponse.js').FileResponse} response The response to cache. * @param {Uint8Array} [result] The result buffer if already read. - * @param {PretrainedOptions} [options] Options containing progress callback and context for progress updates. + * @param {ResourceProgressOptions} [options] Options containing progress callback and context for progress updates. * @returns {Promise} */ export async function storeCachedResource(path_or_repo_id, filename, cache, cacheKey, response, result, options = {}) { @@ -204,9 +237,9 @@ export async function storeCachedResource(path_or_repo_id, filename, cache, cach if (!result) { // We haven't yet read the response body, so we need to do so now. // Ensure progress updates include consistent metadata. - const wrapped_progress = options.progress_callback + const wrapped_progress = options.progressCallback ? (data) => - dispatchCallback(options.progress_callback, { + dispatchCallback(options.progressCallback, { status: 'progress', name: path_or_repo_id, file: filename, @@ -241,28 +274,31 @@ export async function storeCachedResource(path_or_repo_id, filename, cache, cach * - 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 locate. - * @param {boolean} [fatal=true] Whether to throw an error if the file is not found. - * @param {PretrainedOptions} [options] An object containing optional parameters. - * @param {boolean} [return_path=false] Whether to return the path of the file instead of the file content. - * @param {import('./cache.js').CacheInterface | null} [cache] The cache instance to use. + * @param {Object} options Options for loading the resource. + * @param {boolean} [options.fatal=true] Whether to throw an error if the file is not found. + * @param {boolean} [options.returnPath=false] Whether to return the path of the file instead of the file content. + * @param {import('./cache.js').CacheInterface | null} [options.cache] The cache instance to use. + * @param {ReturnType} options.paths Resource paths and cache keys. + * @param {ResourceLoadingOptions} options.loading Options that control where resources may be loaded from. + * @param {FetchOptions} options.fetch Options for fetching files. + * @param {import('./core.js').ProgressCallback} [options.progressCallback] Progress callback. + * @param {PretrainedOptions} [options.metadataOptions] Options to use when fetching response metadata for progress totals. * * @throws Will throw an error if the file is not found and `fatal` is true. * @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 loadResourceFile( - path_or_repo_id, - filename, - fatal = true, - options = {}, - return_path = false, - cache = null, -) { - const { requestURL, localPath, remoteURL, proposedCacheKey, validModelId } = buildResourcePaths( - path_or_repo_id, - filename, - options, - cache, - ); +export async function loadResourceFile(path_or_repo_id, filename, options) { + const { + fatal = true, + returnPath = false, + cache = null, + paths, + loading, + fetch, + progressCallback, + metadataOptions = {}, + } = options; + const { requestURL, localPath, remoteURL, proposedCacheKey, validModelId } = paths; /** @type {string} */ let cacheKey; @@ -282,22 +318,22 @@ export async function loadResourceFile( } else { // Caching not available, or file is not cached, so we perform the request - if (env.allowLocalModels) { + if (loading.allowLocalModels) { // Accessing local models is enabled, so we try to get the file locally. // If request is a valid HTTP URL, we skip the local file check. Otherwise, we try to get the file locally. const isURL = isValidUrl(requestURL, ['http:', 'https:']); if (!isURL) { try { - response = await getFile(localPath); + response = await getFile(localPath, fetch); cacheKey = localPath; // Update the cache key to be the local path } catch (e) { // Something went wrong while trying to get the file locally. // NOTE: error handling is done in the next step (since `response` will be undefined) logger.warn(`Unable to load from local path "${localPath}": "${e}"`); } - } else if (options.local_files_only) { + } else if (loading.localFilesOnly) { throw new Error(`\`local_files_only=true\`, but attempted to load a remote file from: ${requestURL}.`); - } else if (!env.allowRemoteModels) { + } else if (!loading.allowRemoteModels) { throw new Error( `\`env.allowRemoteModels=false\`, but attempted to load a remote file from: ${requestURL}.`, ); @@ -310,7 +346,7 @@ export async function loadResourceFile( // - the path is a valid HTTP url (`response === undefined`) // - the path is not a valid HTTP url and the file is not present on the file system or local server (`response.status === 404`) - if (options.local_files_only || !env.allowRemoteModels) { + if (loading.localFilesOnly || !loading.allowRemoteModels) { // User requested local files only, but the file is not found locally. if (fatal) { throw Error( @@ -331,7 +367,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, fetch); if (response.status !== 200) { return handleError(response.status, remoteURL, fatal); @@ -350,14 +386,14 @@ export async function loadResourceFile( } // Start downloading - dispatchCallback(options.progress_callback, { + dispatchCallback(progressCallback, { status: 'download', name: path_or_repo_id, file: filename, }); let result; - if (apis.IS_NODE_ENV && return_path) { + if (apis.IS_NODE_ENV && returnPath) { // In Node.js with return_path, we skip the buffer read (ONNX runtime // loads from disk directly). A completion progress event is emitted // after the caching block below to ensure progress_total reaches 100%. @@ -366,7 +402,7 @@ export async function loadResourceFile( let buffer; if (typeof response !== 'string') { - if (!options.progress_callback) { + if (!progressCallback) { // If no progress callback is specified, we can use the `.arrayBuffer()` // method to read the response. buffer = new Uint8Array(await response.arrayBuffer()); @@ -380,7 +416,7 @@ export async function loadResourceFile( buffer = new Uint8Array(await response.arrayBuffer()); // For completeness, we still fire the final progress callback - dispatchCallback(options.progress_callback, { + dispatchCallback(progressCallback, { status: 'progress', name: path_or_repo_id, file: filename, @@ -398,7 +434,7 @@ export async function loadResourceFile( } else { // Try to get size from metadata (useful when content-length is missing during download) try { - const metadata = await get_file_metadata(path_or_repo_id, filename, options); + const metadata = await get_file_metadata(path_or_repo_id, filename, metadataOptions); if (metadata.size) { expectedSize = metadata.size; } @@ -410,7 +446,7 @@ export async function loadResourceFile( buffer = await readResponse( response, (data) => { - dispatchCallback(options.progress_callback, { + dispatchCallback(progressCallback, { status: 'progress', name: path_or_repo_id, file: filename, @@ -431,16 +467,16 @@ export async function loadResourceFile( cacheKey && typeof response !== 'string' ) { - await storeCachedResource(path_or_repo_id, filename, cache, cacheKey, response, result, options); + await storeCachedResource(path_or_repo_id, filename, cache, cacheKey, response, result, { progressCallback }); } // In Node.js with return_path, the buffer read is skipped so no progress // events are emitted during loading. Emit a final completion event so // that aggregate progress_total tracking reaches 100%. This is placed // after storeCachedResource so it doesn't conflict with caching progress. - if (apis.IS_NODE_ENV && return_path && options.progress_callback && typeof response !== 'string') { + if (apis.IS_NODE_ENV && returnPath && progressCallback && typeof response !== 'string') { const size = parseInt(response.headers.get('content-length'), 10) || 0; - dispatchCallback(options.progress_callback, { + dispatchCallback(progressCallback, { status: 'progress', name: path_or_repo_id, file: filename, @@ -450,14 +486,14 @@ export async function loadResourceFile( }); } - dispatchCallback(options.progress_callback, { + dispatchCallback(progressCallback, { status: 'done', name: path_or_repo_id, file: filename, }); if (result) { - if (!apis.IS_NODE_ENV && return_path) { + if (!apis.IS_NODE_ENV && returnPath) { throw new Error('Cannot return path in a browser environment.'); } return result; @@ -499,10 +535,39 @@ 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) { + const env = resolveEnv(options.env); + const revision = options.revision ?? 'main'; + const cacheDir = options.cache_dir ?? null; + const localFilesOnly = options.local_files_only ?? false; + const pathOptions = { + cache_dir: cacheDir, + revision, + localModelPath: env.localModelPath, + remoteHost: env.remoteHost, + remotePathTemplate: env.remotePathTemplate, + }; + const loadingOptions = { + allowLocalModels: env.allowLocalModels, + allowRemoteModels: env.allowRemoteModels, + localFilesOnly, + }; + const fetchOptions = { + useFS: env.useFS, + fetch: env.fetch, + version: env.version, + hfToken: env.hfToken, + }; + const metadataOptions = { + cache_dir: cacheDir, + local_files_only: localFilesOnly, + revision, + env: options.env, + }; + if (!env.allowLocalModels) { // User has disabled local models, so we just make sure other settings are correct. - if (options.local_files_only) { + if (localFilesOnly) { throw Error( 'Invalid configuration detected: local models are disabled (`env.allowLocalModels=false`) but you have requested to only use local models (`local_files_only=true`).', ); @@ -528,9 +593,19 @@ export async function getModelFile(path_or_repo_id, filename, fatal = true, opti name: path_or_repo_id, file: filename, }); - pending = getCache(options.cache_dir).then((cache) => - loadResourceFile(path_or_repo_id, filename, fatal, options, return_path, cache), - ); + pending = getCache(cacheDir).then((cache) => { + const paths = buildResourcePaths(path_or_repo_id, filename, pathOptions, cache); + return loadResourceFile(path_or_repo_id, filename, { + fatal, + returnPath: return_path, + cache, + paths, + loading: loadingOptions, + fetch: fetchOptions, + progressCallback: progress_callback, + metadataOptions, + }); + }); if (loads === INFLIGHT_LOADS) { pending = pending.finally(() => INFLIGHT_LOADS.delete(key)); } @@ -578,3 +653,23 @@ export async function getModelJSON(modelPath, fileName, fatal = true, options = return JSON.parse(text); } + +/** + * Emits deprecation warnings for legacy resource-loading options that should be + * represented by environment configuration instead. + * + * @param {string|null|undefined} cache_dir Custom cache directory option. + * @param {boolean} local_files_only Whether loading is restricted to local files. + */ +export function maybeAddDeprecatedEnvWarning(cache_dir, local_files_only) { + if (cache_dir !== null && cache_dir !== undefined) { + logger.warn( + '`cache_dir` is deprecated for environment-style configuration. Use `env.cacheDir` for the default cache directory and `options.env` for session-scopable resource loading settings.', + ); + } + if (local_files_only) { + logger.warn( + '`local_files_only` is deprecated. Use `options.env.allowRemoteModels=false` for session-scoped remote loading control.', + ); + } +} diff --git a/packages/transformers/src/utils/hub/utils.js b/packages/transformers/src/utils/hub/utils.js index 3cb1431a3..1151cc72d 100644 --- a/packages/transformers/src/utils/hub/utils.js +++ b/packages/transformers/src/utils/hub/utils.js @@ -1,6 +1,36 @@ import { ERROR_MAPPING, REPO_ID_REGEX } from './constants.js'; +import { resolveEnv } from '../../env.js'; import { logger } from '../logger.js'; +const FETCH_IDS = new WeakMap(); +const HF_TOKEN_IDS = new Map(); +let NEXT_FETCH_ID = 0; +let NEXT_HF_TOKEN_ID = 0; + +function getFetchId(fetch) { + if (typeof fetch !== 'function') { + return null; + } + let id = FETCH_IDS.get(fetch); + if (id === undefined) { + id = ++NEXT_FETCH_ID; + FETCH_IDS.set(fetch, id); + } + return id; +} + +function getHfTokenId(hfToken) { + if (!hfToken) { + return null; + } + let id = HF_TOKEN_IDS.get(hfToken); + if (id === undefined) { + id = ++NEXT_HF_TOKEN_ID; + HF_TOKEN_IDS.set(hfToken, id); + } + return id; +} + /** * Joins multiple parts of a path into a single path, while handling leading and trailing slashes. * @@ -66,15 +96,24 @@ export function isValidHfModelId(string) { * @param {string} [options.revision='main'] Model revision. * @param {string|null} [options.cache_dir=null] Custom cache directory. * @param {boolean} [options.local_files_only=false] Whether to avoid remote lookups. + * @param {Partial} [options.env={}] Session-scopable environment overrides. * @param {...unknown} parts Additional key parts for the specific operation. * @returns {string} */ export function makePretrainedOptionsKey(model_id, options = {}, ...parts) { + const env = resolveEnv(options.env); return JSON.stringify([ model_id, options.revision ?? 'main', options.cache_dir ?? null, options.local_files_only ?? false, + env.allowRemoteModels, + env.remoteHost, + env.remotePathTemplate, + env.allowLocalModels, + env.localModelPath, + getFetchId(env.fetch), + getHfTokenId(env.hfToken), ...parts, ]); } diff --git a/packages/transformers/src/utils/image.js b/packages/transformers/src/utils/image.js index d7d8c7fcd..d48550c6c 100644 --- a/packages/transformers/src/utils/image.js +++ b/packages/transformers/src/utils/image.js @@ -9,7 +9,7 @@ import { isNullishDimension } from './core.js'; import { getFile } from './hub.js'; -import { apis } from '../env.js'; +import { apis, resolveEnv } from '../env.js'; import { Tensor } from './tensor.js'; import { saveBlob } from './io.js'; @@ -97,6 +97,7 @@ export class RawImage { /** * Helper method for reading an image from a variety of input types. * @param {RawImage|string|URL|Blob|HTMLCanvasElement|OffscreenCanvas} input + * @param {Partial} [sessionEnv={}] Session-scopable environment overrides. * @returns {Promise} The image object. * * **Example:** Read image from a URL. @@ -110,11 +111,11 @@ export class RawImage { * // } * ``` */ - static async read(input) { + static async read(input, sessionEnv = {}) { if (input instanceof RawImage) { return input; } else if (typeof input === 'string' || input instanceof URL) { - return await this.fromURL(input); + return await this.fromURL(input, sessionEnv); } else if (input instanceof Blob) { return await this.fromBlob(input); } else if ( @@ -147,10 +148,17 @@ export class RawImage { /** * Read an image from a URL or file path. * @param {string|URL} url The URL or file path to read the image from. + * @param {Partial} [sessionEnv={}] Session-scopable environment overrides. * @returns {Promise} The image object. */ - static async fromURL(url) { - const response = await getFile(url); + static async fromURL(url, sessionEnv = {}) { + const env = resolveEnv(sessionEnv); + const response = await getFile(url, { + useFS: env.useFS, + fetch: env.fetch, + version: env.version, + hfToken: env.hfToken, + }); if (response.status !== 200) { throw new Error(`Unable to read image from "${url}" (${response.status} ${response.statusText})`); } diff --git a/packages/transformers/src/utils/model_registry/ModelRegistry.js b/packages/transformers/src/utils/model_registry/ModelRegistry.js index 1179a92f6..6ab88826e 100644 --- a/packages/transformers/src/utils/model_registry/ModelRegistry.js +++ b/packages/transformers/src/utils/model_registry/ModelRegistry.js @@ -181,28 +181,30 @@ export class ModelRegistry { * Get tokenizer files needed for a specific model. * * @param {string} modelId - The model id + * @param {import('../hub.js').PretrainedOptions} [options] An object containing optional parameters. * @returns {Promise} Array of tokenizer file paths * * @example * const files = await ModelRegistry.get_tokenizer_files('onnx-community/gpt2-ONNX'); * console.log(files); // ['tokenizer.json', 'tokenizer_config.json'] */ - static async get_tokenizer_files(modelId) { - return get_tokenizer_files(modelId); + static async get_tokenizer_files(modelId, options = {}) { + return get_tokenizer_files(modelId, options); } /** * Get processor files needed for a specific model. * * @param {string} modelId - The model id + * @param {import('../hub.js').PretrainedOptions} [options] An object containing optional parameters. * @returns {Promise} Array of processor file paths * * @example * const files = await ModelRegistry.get_processor_files('onnx-community/vit-base-patch16-224-ONNX'); * console.log(files); // ['preprocessor_config.json'] */ - static async get_processor_files(modelId) { - return get_processor_files(modelId); + static async get_processor_files(modelId, options = {}) { + return get_processor_files(modelId, options); } /** diff --git a/packages/transformers/src/utils/model_registry/clear_cache.js b/packages/transformers/src/utils/model_registry/clear_cache.js index fe4850e8c..b72630547 100644 --- a/packages/transformers/src/utils/model_registry/clear_cache.js +++ b/packages/transformers/src/utils/model_registry/clear_cache.js @@ -5,6 +5,7 @@ */ import { getCache } from '../cache.js'; +import { resolveEnv } from '../../env.js'; import { buildResourcePaths, checkCachedResource } from '../hub.js'; import { get_files } from './get_files.js'; import { get_pipeline_files } from './get_pipeline_files.js'; @@ -35,6 +36,14 @@ 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); + const env = resolveEnv(options.env); + const pathOptions = { + cache_dir: options.cache_dir ?? null, + revision: options.revision ?? 'main', + localModelPath: env.localModelPath, + remoteHost: env.remoteHost, + remotePathTemplate: env.remotePathTemplate, + }; if (!cache) { return { @@ -50,7 +59,7 @@ async function clear_files_from_cache(modelId, files, options = {}) { const results = await Promise.all( files.map(async (filename) => { - const { localPath, proposedCacheKey } = buildResourcePaths(modelId, filename, options, cache); + const { localPath, proposedCacheKey } = buildResourcePaths(modelId, filename, pathOptions, cache); const cached = await checkCachedResource(cache, localPath, proposedCacheKey); const wasCached = !!cached; 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..7201560d5 100644 --- a/packages/transformers/src/utils/model_registry/get_available_dtypes.js +++ b/packages/transformers/src/utils/model_registry/get_available_dtypes.js @@ -29,13 +29,21 @@ 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 {Partial} [options.env={}] Session-scopable environment overrides. * @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 = await get_config(modelId, { config, cache_dir, local_files_only, revision }); +export async function get_available_dtypes(modelId, options = {}) { + const { + config: providedConfig = null, + model_file_name = null, + revision = 'main', + cache_dir = null, + local_files_only = false, + } = options; + const sessionEnv = options.env ?? {}; + + let config = providedConfig; + config = await get_config(modelId, { config, cache_dir, local_files_only, revision, env: sessionEnv }); const subfolder = 'onnx'; @@ -47,7 +55,7 @@ export async function get_available_dtypes( const baseNames = Object.values(sessions); // For each dtype, check if all session files exist - const metadataOptions = { revision, cache_dir, local_files_only }; + const metadataOptions = { revision, cache_dir, local_files_only, env: sessionEnv }; // Probe all (dtype, baseName) combinations in parallel const probeResults = await Promise.all( 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..27b52cf80 100644 --- a/packages/transformers/src/utils/model_registry/get_file_metadata.js +++ b/packages/transformers/src/utils/model_registry/get_file_metadata.js @@ -2,7 +2,7 @@ * @file File metadata utilities for cache-aware operations */ -import { env } from '../../env.js'; +import { resolveEnv } from '../../env.js'; import { getCache } from '../cache.js'; import { buildResourcePaths, checkCachedResource, getFetchHeaders, getFile } from '../hub.js'; import { isValidUrl, makePretrainedOptionsKey } from '../hub/utils.js'; @@ -27,15 +27,15 @@ import { memoizePromise } from '../memoize_promise.js'; * @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, fetchOptions) { // Range requests only make sense for HTTP URLs if (!isValidUrl(urlOrPath, ['http:', 'https:'])) { return null; } - const headers = getFetchHeaders(urlOrPath); + const headers = getFetchHeaders(urlOrPath, fetchOptions); headers.set('Range', 'bytes=0-0'); - return env.fetch(urlOrPath, { method: 'GET', headers, cache: 'no-store' }); + return fetchOptions.fetch(urlOrPath, { method: 'GET', headers, cache: 'no-store' }); } /** @@ -56,12 +56,26 @@ export function get_file_metadata(path_or_repo_id, filename, options = {}) { } async function _get_file_metadata(path_or_repo_id, filename, options) { + const env = resolveEnv(options.env); + const fetchOptions = { + useFS: env.useFS, + fetch: env.fetch, + version: env.version, + hfToken: env.hfToken, + }; + const pathOptions = { + cache_dir: options.cache_dir ?? null, + revision: options.revision ?? 'main', + localModelPath: env.localModelPath, + remoteHost: env.remoteHost, + remotePathTemplate: env.remotePathTemplate, + }; /** @type {import('../cache.js').CacheInterface | null} */ const cache = await getCache(options?.cache_dir); const { localPath, remoteURL, proposedCacheKey, validModelId } = buildResourcePaths( path_or_repo_id, filename, - options, + pathOptions, cache, ); @@ -83,7 +97,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, fetchOptions); if (typeof response !== 'string' && response.status !== 404) { const size = response.headers.get('content-length'); const contentType = response.headers.get('content-type'); @@ -105,7 +119,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, fetchOptions); if (rangeResponse && rangeResponse.status >= 200 && rangeResponse.status < 300) { let size; diff --git a/packages/transformers/src/utils/model_registry/get_files.js b/packages/transformers/src/utils/model_registry/get_files.js index eeda2891b..658ae1da5 100644 --- a/packages/transformers/src/utils/model_registry/get_files.js +++ b/packages/transformers/src/utils/model_registry/get_files.js @@ -9,32 +9,47 @@ import { get_processor_files } from './get_processor_files.js'; * @param {string} modelId The model id (e.g., "Xenova/llama-2-7b") * @param {Object} [options] Optional parameters * @param {import('../../configs.js').PretrainedConfig} [options.config=null] Pre-loaded model config (optional, will be fetched if not provided) + * @param {string|null} [options.cache_dir=null] Custom local cache directory. + * @param {boolean} [options.local_files_only=false] Never hit the network if true. + * @param {string} [options.revision='main'] Git branch, tag, or commit SHA. * @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 {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 {Partial} [options.env={}] Session-scopable environment overrides. * @returns {Promise} Array of file paths that will be loaded */ -export async function get_files( - modelId, - { +export async function get_files(modelId, options = {}) { + const { config = null, + cache_dir = null, + local_files_only = false, + revision = 'main', dtype = null, device = null, model_file_name = null, include_tokenizer = true, include_processor = true, - } = {}, -) { - const files = await get_model_files(modelId, { config, dtype, device, model_file_name }); + } = options; + const sessionEnv = options.env ?? {}; + const files = await get_model_files(modelId, { + config, + cache_dir, + local_files_only, + revision, + dtype, + device, + model_file_name, + env: sessionEnv, + }); if (include_tokenizer) { - const tokenizerFiles = await get_tokenizer_files(modelId); + const tokenizerFiles = await get_tokenizer_files(modelId, options); files.push(...tokenizerFiles); } if (include_processor) { - const processorFiles = await get_processor_files(modelId); + const processorFiles = await get_processor_files(modelId, options); 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..ee8ff708f 100644 --- a/packages/transformers/src/utils/model_registry/get_model_files.js +++ b/packages/transformers/src/utils/model_registry/get_model_files.js @@ -25,20 +25,20 @@ import { resolve_model_type } from './resolve_model_type.js'; * @param {string|null} [options.cache_dir=null] Custom local cache directory. * @param {boolean} [options.local_files_only=false] Never hit the network if true. * @param {string} [options.revision='main'] Git branch, tag, or commit SHA. + * @param {Partial} [options.env={}] Session-scopable environment overrides. * @returns {Promise} */ -export function get_config( - modelId, - { config = null, cache_dir = null, local_files_only = false, revision = 'main' } = {}, -) { +export function get_config(modelId, options = {}) { + const { config = null, cache_dir = null, local_files_only = false, revision = 'main' } = options; + const sessionEnv = options.env ?? {}; // When a pre-loaded config is provided, skip memoization — no fetch occurs // and there is no meaningful key to deduplicate on. if (config !== null) { - return AutoConfig.from_pretrained(modelId, { config, cache_dir, local_files_only, revision }); + return AutoConfig.from_pretrained(modelId, { config, cache_dir, local_files_only, revision, env: sessionEnv }); } - const key = makePretrainedOptionsKey(modelId, { cache_dir, local_files_only, revision }); + const key = makePretrainedOptionsKey(modelId, { cache_dir, local_files_only, revision, env: sessionEnv }); return memoizePromise(key, () => - AutoConfig.from_pretrained(modelId, { config, cache_dir, local_files_only, revision }), + AutoConfig.from_pretrained(modelId, { config, cache_dir, local_files_only, revision, env: sessionEnv }), ); } @@ -52,16 +52,29 @@ export function get_config( * @param {string} modelId The model id (e.g., "onnx-community/granite-4.0-350m-ONNX-web") * @param {Object} [options] Optional parameters * @param {import('../../configs.js').PretrainedConfig} [options.config=null] Pre-loaded model config (optional, will be fetched if not provided) + * @param {string|null} [options.cache_dir=null] Custom local cache directory. + * @param {boolean} [options.local_files_only=false] Never hit the network if true. + * @param {string} [options.revision='main'] Git branch, tag, or commit SHA. * @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 {Partial} [options.env={}] Session-scopable environment overrides. * @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 = await get_config(modelId, { config }); +export async function get_model_files(modelId, options = {}) { + const { + config: providedConfig = null, + cache_dir = null, + local_files_only = false, + revision = 'main', + dtype: overrideDtype = null, + device: overrideDevice = null, + model_file_name = null, + } = options; + const sessionEnv = options.env ?? {}; + + let config = providedConfig; + config = await get_config(modelId, { config, cache_dir, local_files_only, revision, env: sessionEnv }); const files = [ // Add config.json (always loaded) 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..a1d82b008 100644 --- a/packages/transformers/src/utils/model_registry/get_pipeline_files.js +++ b/packages/transformers/src/utils/model_registry/get_pipeline_files.js @@ -13,9 +13,13 @@ import { SUPPORTED_TASKS, TASK_ALIASES } from '../../pipelines/index.js'; * @param {string} modelId - The model id (e.g., "Xenova/bert-base-uncased") * @param {Object} [options] - Optional parameters * @param {import('../../configs.js').PretrainedConfig} [options.config=null] - Pre-loaded config + * @param {string|null} [options.cache_dir=null] Custom local cache directory. + * @param {boolean} [options.local_files_only=false] Never hit the network if true. + * @param {string} [options.revision='main'] Git branch, tag, or commit SHA. * @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 {Partial} [options.env={}] Session-scopable environment overrides. * @returns {Promise} Array of file paths that will be loaded * @throws {Error} If the task is not supported */ 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..de8d26e68 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 {import('../hub.js').PretrainedOptions} [options] An object containing optional parameters. * @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..18acf15f3 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 {import('../hub.js').PretrainedOptions} [options] An object containing optional parameters. * @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..03c042e66 100644 --- a/packages/transformers/src/utils/model_registry/is_cached.js +++ b/packages/transformers/src/utils/model_registry/is_cached.js @@ -1,4 +1,5 @@ import { getCache } from '../cache.js'; +import { resolveEnv } from '../../env.js'; import { buildResourcePaths, checkCachedResource } from '../hub.js'; import { get_files } from './get_files.js'; import { get_pipeline_files } from './get_pipeline_files.js'; @@ -25,6 +26,14 @@ import { get_pipeline_files } from './get_pipeline_files.js'; */ async function check_files_cache(modelId, files, options = {}) { const cache = await getCache(options?.cache_dir); + const env = resolveEnv(options.env); + const pathOptions = { + cache_dir: options.cache_dir ?? null, + revision: options.revision ?? 'main', + localModelPath: env.localModelPath, + remoteHost: env.remoteHost, + remotePathTemplate: env.remotePathTemplate, + }; if (!cache) { const fileStatuses = files.map((filename) => ({ file: filename, cached: false })); @@ -34,7 +43,7 @@ async function check_files_cache(modelId, files, options = {}) { const fileStatuses = await Promise.all( files.map(async (filename) => { - const { localPath, proposedCacheKey } = buildResourcePaths(modelId, filename, options, cache); + const { localPath, proposedCacheKey } = buildResourcePaths(modelId, filename, pathOptions, cache); const cached = await checkCachedResource(cache, localPath, proposedCacheKey); return { file: filename, cached: !!cached }; }), @@ -54,7 +63,15 @@ async function check_files_cache(modelId, files, options = {}) { async function is_file_cached(modelId, filename, options = {}) { const cache = await getCache(options?.cache_dir); if (!cache) return false; - const { localPath, proposedCacheKey } = buildResourcePaths(modelId, filename, options, cache); + const env = resolveEnv(options.env); + const pathOptions = { + cache_dir: options.cache_dir ?? null, + revision: options.revision ?? 'main', + localModelPath: env.localModelPath, + remoteHost: env.remoteHost, + remotePathTemplate: env.remotePathTemplate, + }; + const { localPath, proposedCacheKey } = buildResourcePaths(modelId, filename, pathOptions, cache); return !!(await checkCachedResource(cache, localPath, proposedCacheKey)); } diff --git a/packages/transformers/tests/utils/audio.test.js b/packages/transformers/tests/utils/audio.test.js index 1c6c0042d..5910bad5e 100644 --- a/packages/transformers/tests/utils/audio.test.js +++ b/packages/transformers/tests/utils/audio.test.js @@ -1,4 +1,5 @@ -import { RawAudio, spectrogram, window_function, mel_filter_bank } from "../../src/utils/audio.js"; +import { jest } from "@jest/globals"; +import { load_audio, RawAudio, spectrogram, window_function, mel_filter_bank } from "../../src/utils/audio.js"; import { init } from "../init.js"; init(); @@ -67,6 +68,44 @@ function identityMelFilters(numBins) { } describe("Audio utilities", () => { + describe("Session env", () => { + let originalAudioContext; + + beforeEach(() => { + originalAudioContext = globalThis.AudioContext; + globalThis.AudioContext = class { + sampleRate; + + constructor({ sampleRate }) { + this.sampleRate = sampleRate; + } + + async decodeAudioData() { + return { + numberOfChannels: 1, + length: 2, + getChannelData: () => new Float32Array([0.25, 0.5]), + }; + } + }; + }); + + afterEach(() => { + globalThis.AudioContext = originalAudioContext; + }); + + it("should use scoped fetch when loading an audio URL", async () => { + const fetch = jest.fn(async () => ({ + arrayBuffer: async () => new ArrayBuffer(8), + })); + + const audio = await load_audio("https://example.com/audio.wav", 16000, { fetch }); + + expect(fetch).toHaveBeenCalledWith("https://example.com/audio.wav", expect.any(Object)); + expect(audio).toEqual(new Float32Array([0.25, 0.5])); + }); + }); + describe("RawAudio", () => { it("should create RawAudio from a single Float32Array", () => { const sampling_rate = 16000; diff --git a/packages/transformers/tests/utils/hub.test.js b/packages/transformers/tests/utils/hub.test.js index 96598092d..d89af42dd 100644 --- a/packages/transformers/tests/utils/hub.test.js +++ b/packages/transformers/tests/utils/hub.test.js @@ -1,4 +1,5 @@ import { AutoModel, PreTrainedModel } from "../../src/transformers.js"; +import { buildResourcePaths, getFetchHeaders } from "../../src/utils/hub.js"; import { MAX_TEST_EXECUTION_TIME, DEFAULT_MODEL_OPTIONS } from "../init.js"; import fs from "node:fs"; @@ -6,6 +7,29 @@ import fs from "node:fs"; // TODO: Set cache folder to a temp directory describe("Hub", () => { + describe("Session env", () => { + it("should use scoped resource path options", () => { + const paths = buildResourcePaths("org/model", "config.json", { + revision: "refs/pr/1", + localModelPath: "/scoped-models/", + remoteHost: "https://models.example.com/", + remotePathTemplate: "{model}/at/{revision}/", + }); + + expect(paths.localPath).toBe("/scoped-models/org/model/config.json"); + expect(paths.remoteURL).toBe("https://models.example.com/org/model/at/refs%2Fpr%2F1/config.json"); + }); + + it("should use scoped Hugging Face token for request headers", () => { + const headers = getFetchHeaders("https://huggingface.co/org/model/resolve/main/config.json", { + version: "test-version", + hfToken: "scoped-token", + }); + + expect(headers.get("Authorization")).toBe("Bearer scoped-token"); + }); + }); + describe("Loading models", () => { it( "should load a model from the local cache", diff --git a/packages/transformers/tests/utils/image.test.js b/packages/transformers/tests/utils/image.test.js index 7fc5d5b4a..b1164b410 100644 --- a/packages/transformers/tests/utils/image.test.js +++ b/packages/transformers/tests/utils/image.test.js @@ -1,3 +1,4 @@ +import { jest } from "@jest/globals"; import { RawImage, rand } from "../../src/transformers.js"; import { load_cached_image } from "../asset_cache.js"; @@ -36,6 +37,20 @@ describe("Image utilities", () => { }); }); + describe("Session env", () => { + it("should use scoped fetch when reading an image URL", async () => { + const fetch = jest.fn(async () => ({ + status: 404, + statusText: "Not Found", + })); + + await expect(RawImage.fromURL("https://example.com/image.png", { fetch })).rejects.toThrow( + 'Unable to read image from "https://example.com/image.png"', + ); + expect(fetch).toHaveBeenCalledWith("https://example.com/image.png", expect.any(Object)); + }); + }); + describe("Channel conversions", () => { it("should convert RGBA to L (grayscale)", async () => { const grayscale = TEST_IMAGES.rgba.clone().grayscale(); diff --git a/packages/transformers/tests/utils/model_registry.test.js b/packages/transformers/tests/utils/model_registry.test.js index cb9e6408a..4996062d3 100644 --- a/packages/transformers/tests/utils/model_registry.test.js +++ b/packages/transformers/tests/utils/model_registry.test.js @@ -154,19 +154,26 @@ describe("get_available_dtypes", () => { expect(dtypes).not.toContain("fp16"); }); - it("should pass revision and cache_dir to get_file_metadata", async () => { + it("should pass session env to get_file_metadata", async () => { setupExistingFiles("onnx/model.onnx"); + const fetch = jest.fn(); + const env = { + remoteHost: "https://models.example.com/", + allowRemoteModels: false, + fetch, + hfToken: "test-token", + }; + await get_available_dtypes("test/model", { config: ENCODER_ONLY_CONFIG, revision: "v2", - cache_dir: "/tmp/cache", + env, }); - // Verify that metadata calls received the correct options for (const call of mockGetFileMetadata.mock.calls) { expect(call[0]).toBe("test/model"); - expect(call[2]).toMatchObject({ revision: "v2", cache_dir: "/tmp/cache" }); + expect(call[2]).toMatchObject({ revision: "v2", env }); } }); diff --git a/packages/transformers/tests/utils/utils.test.js b/packages/transformers/tests/utils/utils.test.js index f3113c044..7f8cfc65c 100644 --- a/packages/transformers/tests/utils/utils.test.js +++ b/packages/transformers/tests/utils/utils.test.js @@ -1,4 +1,5 @@ import { AutoProcessor } from "../../src/transformers.js"; +import { resolveEnv } from "../../src/env.js"; import { hamming, hanning, mel_filter_bank } from "../../src/utils/audio.js"; import { getFile } from "../../src/utils/hub.js"; import { RawImage } from "../../src/utils/image.js"; @@ -56,7 +57,13 @@ describe("Utilities", () => { it("Read data from blob", async () => { const blob = new Blob(["Hello, world!"], { type: "text/plain" }); const blobUrl = URL.createObjectURL(blob); - const data = await getFile(blobUrl); + const env = resolveEnv(); + const data = await getFile(blobUrl, { + useFS: env.useFS, + fetch: env.fetch, + version: env.version, + hfToken: env.hfToken, + }); expect(await data.text()).toBe("Hello, world!"); }); });