Add type declaration for progress_callback (#1034)

* Add type declaration for `progress_callback` * export ProgressCallback and ProgressInfo * Update JSDoc --------- Co-authored-by: Joshua Lochner <[email protected]>
huggingface · Nov 25, 2024 · d38134d · d38134d
1 parent 82bae5e
commit d38134d
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 3 deletions.
diff --git a/src/configs.js b/src/configs.js
@@ -36,6 +36,13 @@ import {
  * @typedef {import('./utils/hub.js').PretrainedOptions} PretrainedOptions
  */
 
+/**
+ * @typedef {import('./utils/core.js').ProgressCallback} ProgressCallback
+ */
+
+/**
+ * @typedef {import('./utils/core.js').ProgressInfo} ProgressInfo
+ */
 
 /**
  * Loads a config from the specified path.

diff --git a/src/utils/core.js b/src/utils/core.js
@@ -8,11 +8,30 @@
  * @module utils/core
  */
 
+/**
+ * @typedef {Object} ProgressInfo
+ * @property {'initiate' | 'download' | 'progress' | 'done'} status The status of the progress item.
+ * @property {string} name 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.
+ * @property {string} file The name of the file
+ * @property {number} [progress] A number between 0 and 100. Only available for the 'progress' status.
+ * @property {number} [loaded] The number of bytes loaded. Only available for the 'progress' status.
+ * @property {number} [total] The total number of bytes to be loaded. Only available for the 'progress' status.
+ */
+
+/**
+ * A callback function that is called with progress information.
+ * @callback ProgressCallback
+ * @param {ProgressInfo} progressInfo
+ * @returns {void}
+ */
+
 /**
  * Helper function to dispatch progress callbacks.
  *
- * @param {Function} progress_callback The progress callback function to dispatch.
- * @param {any} data The data to pass to the progress callback function.
+ * @param {ProgressCallback | null | undefined} progress_callback The progress callback function to dispatch.
+ * @param {ProgressInfo} data The data to pass to the progress callback function.
  * @returns {void}
  * @private
  */

diff --git a/src/utils/hub.js b/src/utils/hub.js
@@ -13,7 +13,7 @@ import { dispatchCallback } from './core.js';
 
 /**
  * @typedef {Object} PretrainedOptions Options for loading a pretrained model.     
- * @property {function} [progress_callback=null] If specified, this function will be called during model construction, to provide the user with progress updates.
+ * @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.
@@ -504,6 +504,7 @@ export async function getModelFile(path_or_repo_id, filename, fatal = true, opti
         file: filename
     })
 
+    /** @type {import('./core.js').ProgressInfo} */
     const progressInfo = {
         status: 'progress',
         name: path_or_repo_id,