@huggingface/transformers 3.0.0-alpha.0

package/src/generation/logits_sampler.js

@@ -0,0 +1,204 @@
+
+/**
+ * @module generation/logits_sampler
+ */
+
+import { Callable } from "../utils/generic.js";
+import { Tensor, topk } from "../utils/tensor.js";
+
+import {
+    max,
+    softmax,
+} from '../utils/maths.js';
+import { GenerationConfig } from '../generation/configuration_utils.js';
+
+/**
+ * Sampler is a base class for all sampling methods used for text generation.
+ */
+export class LogitsSampler extends Callable {
+    /**
+     * Creates a new Sampler object with the specified generation config.
+     * @param {GenerationConfig} generation_config The generation config.
+     */
+    constructor(generation_config) {
+        super();
+        this.generation_config = generation_config;
+    }
+
+    /**
+     * Executes the sampler, using the specified logits.
+     * @param {Tensor} logits
+     * @returns {Promise<[bigint, number][]>}
+     */
+    async _call(logits) {
+        // Sample from logits, of dims [batch, sequence_length, vocab_size].
+        // If index is specified, sample from [batch, index, vocab_size].
+        return this.sample(logits);
+    }
+
+    /**
+     * Abstract method for sampling the logits.
+     * @param {Tensor} logits
+     * @throws {Error} If not implemented in subclass.
+     * @returns {Promise<[bigint, number][]>}
+     */
+    async sample(logits) {
+        throw Error("sample should be implemented in subclasses.")
+    }
+
+    /**
+     * Returns the specified logits as an array, with temperature applied.
+     * @param {Tensor} logits
+     * @param {number} index
+     * @returns {Float32Array}
+     */
+    getLogits(logits, index) {
+        let vocabSize = logits.dims.at(-1);
+
+        let logs = /** @type {Float32Array} */(logits.data);
+
+        if (index === -1) {
+            logs = logs.slice(-vocabSize);
+        } else {
+            let startIndex = index * vocabSize;
+            logs = logs.slice(startIndex, startIndex + vocabSize);
+        }
+        return logs;
+    }
+
+    /**
+     * Selects an item randomly based on the specified probabilities.
+     * @param {import("../transformers.js").DataArray} probabilities An array of probabilities to use for selection.
+     * @returns {number} The index of the selected item.
+     */
+    randomSelect(probabilities) {
+        // Return index of chosen item
+        let sumProbabilities = 0;
+        for (let i = 0; i < probabilities.length; ++i) {
+            sumProbabilities += probabilities[i];
+        }
+
+        let r = Math.random() * sumProbabilities;
+        for (let i = 0; i < probabilities.length; ++i) {
+            r -= probabilities[i];
+            if (r <= 0) {
+                return i;
+            }
+        }
+        return 0; // return first (most probable) as a fallback
+    }
+
+    /**
+     * Returns a Sampler object based on the specified options.
+     * @param {GenerationConfig} generation_config An object containing options for the sampler.
+     * @returns {LogitsSampler} A Sampler object.
+     */
+    static getSampler(generation_config) {
+        // - *greedy decoding*: `num_beams=1` and `do_sample=False`
+        // - *contrastive search*: `penalty_alpha>0` and `top_k>1`
+        // - *multinomial sampling*: `num_beams=1` and `do_sample=True`
+        // - *beam-search decoding*: `num_beams>1` and `do_sample=False`
+        // - *beam-search multinomial sampling*: `num_beams>1` and `do_sample=True`
+        // - *diverse beam-search decoding*: `num_beams>1` and `num_beam_groups>1`
+        // - *constrained beam-search decoding*: `constraints!=None` or `force_words_ids!=None`
+
+        // NOTE: beam search is implemented directly into the generation function
+        if (generation_config.do_sample) {
+            return new MultinomialSampler(generation_config);
+
+        } else if (generation_config.num_beams > 1) {
+            return new BeamSearchSampler(generation_config);
+
+        } else {
+            if (generation_config.num_return_sequences > 1) {
+                throw Error(`num_return_sequences has to be 1 when doing greedy search, but is ${generation_config.num_return_sequences}.`)
+            }
+            return new GreedySampler(generation_config);
+        }
+    }
+}
+
+/**
+ * Class representing a Greedy Sampler.
+ */
+class GreedySampler extends LogitsSampler {
+    /**
+     * Sample the maximum probability of a given logits tensor.
+     * @param {Tensor} logits
+     * @returns {Promise<[bigint, number][]>} An array with a single tuple, containing the index of the maximum value and a meaningless score (since this is a greedy search).
+     */
+    async sample(logits) {
+        // NOTE: no need to do log_softmax here since we only take the maximum
+        const argmax = max(logits.data)[1];
+
+        // Note: score is meaningless in this context, since we are performing
+        // greedy search (p = 1 => log(p) = 0)
+        return [
+            [BigInt(argmax), 0]
+        ];
+    }
+}
+
+/**
+ * Class representing a MultinomialSampler.
+ */
+class MultinomialSampler extends LogitsSampler {
+
+    /**
+     * Sample from the logits.
+     * @param {Tensor} logits
+     * @returns {Promise<[bigint, number][]>}
+     */
+    async sample(logits) {
+        let k = logits.dims.at(-1); // defaults to vocab size
+        if (this.generation_config.top_k > 0) {
+            k = Math.min(this.generation_config.top_k, k);
+        }
+
+        // Get top k tokens
+        const [v, i] = await topk(logits, k);
+
+        // Compute softmax over logits
+        const probabilities = softmax(/** @type {Float32Array} */(v.data));
+
+        return Array.from({ length: this.generation_config.num_beams }, () => {
+            const sampledIndex = this.randomSelect(probabilities);
+            return [
+                i.data[sampledIndex], // token id
+                Math.log(probabilities[sampledIndex]), // score
+            ];
+        });
+    }
+}
+
+
+/**
+ * Class representing a BeamSearchSampler.
+ */
+class BeamSearchSampler extends LogitsSampler {
+
+    /**
+     * Sample from the logits.
+     * @param {Tensor} logits
+     * @returns {Promise<[bigint, number][]>}
+     */
+    async sample(logits) {
+        let k = logits.dims.at(-1); // defaults to vocab size
+        if (this.generation_config.top_k > 0) {
+            k = Math.min(this.generation_config.top_k, k);
+        }
+
+        // Get top k tokens
+        const [v, i] = await topk(logits, k);
+
+        // Compute softmax over logits
+        const probabilities = softmax(/** @type {Float32Array} */(v.data));
+
+        return Array.from({ length: this.generation_config.num_beams }, (_, x) => {
+            return [
+                i.data[x], // token id
+                Math.log(probabilities[x]), // score
+            ];
+        });
+    }
+}

package/src/generation/parameters.js

@@ -0,0 +1,35 @@
+
+/**
+ * @module generation/parameters
+ */
+
+/**
+ * @typedef {Object} GenerationFunctionParameters
+ * @property {import('../utils/tensor.js').Tensor} [inputs=null] (`Tensor` of varying shape depending on the modality, *optional*):
+ * The sequence used as a prompt for the generation or as model inputs to the encoder. If `null` the
+ * method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs`
+ * should be in the format of `input_ids`. For encoder-decoder models *inputs* can represent any of
+ * `input_ids`, `input_values`, `input_features`, or `pixel_values`.
+ * @property {import('./configuration_utils.js').GenerationConfig} [generation_config=null] (`GenerationConfig`, *optional*):
+ * The generation configuration to be used as base parametrization for the generation call.
+ * `**kwargs` passed to generate matching the attributes of `generation_config` will override them.
+ * If `generation_config` is not provided, the default will be used, which has the following loading
+ * priority:
+ * - (1) from the `generation_config.json` model file, if it exists;
+ * - (2) from the model configuration. Please note that unspecified parameters will inherit [`GenerationConfig`]'s
+ * default values, whose documentation should be checked to parameterize generation.
+ * @property {import('./logits_process.js').LogitsProcessorList} [logits_processor=null] (`LogitsProcessorList`, *optional*):
+ * Custom logits processors that complement the default logits processors built from arguments and
+ * generation config. If a logit processor is passed that is already created with the arguments or a
+ * generation config an error is thrown. This feature is intended for advanced users.
+ * @property {import('./stopping_criteria.js').StoppingCriteriaList} [stopping_criteria=null] (`StoppingCriteriaList`, *optional*):
+ * Custom stopping criteria that complements the default stopping criteria built from arguments and a
+ * generation config. If a stopping criteria is passed that is already created with the arguments or a
+ * generation config an error is thrown. This feature is intended for advanced users.
+ * @property {import('./streamers.js').BaseStreamer} [streamer=null] (`BaseStreamer`, *optional*):
+ * Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+ * through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+ * @property {number[]} [decoder_input_ids=null] (`number[]`, *optional*):
+ * If the model is an encoder-decoder model, this argument is used to pass the `decoder_input_ids`.
+ * @param {any} [kwargs] (`Dict[str, any]`, *optional*):
+ */

package/src/generation/stopping_criteria.js

@@ -0,0 +1,156 @@
+
+/**
+ * @module generation/stopping_criteria
+ */
+
+import { Callable } from "../utils/generic.js";
+
+// NOTE:
+// Stopping Criteria returns a list of `batch_size` booleans, indicating whether each sequence in the batch should be stopped.
+
+/**
+ * Abstract base class for all stopping criteria that can be applied during generation.
+ */
+export class StoppingCriteria extends Callable {
+    /**
+     * 
+     * @param {number[][]} input_ids (`number[][]` of shape `(batch_size, sequence_length)`):
+     * Indices of input sequence tokens in the vocabulary.
+     * @param {number[][]} scores scores (`number[][]` of shape `(batch_size, config.vocab_size)`):
+     * Prediction scores of a language modeling head. These can be scores for each vocabulary token before SoftMax
+     * or scores for each vocabulary token after SoftMax.
+     * @returns {boolean[]} A list of booleans indicating whether each sequence should be stopped.
+     */
+    _call(input_ids, scores) {
+        throw Error("StoppingCriteria needs to be subclassed");
+    }
+}
+/**
+ */
+export class StoppingCriteriaList extends Callable {
+    /**
+     * Constructs a new instance of `StoppingCriteriaList`.
+     */
+    constructor() {
+        super();
+        this.criteria = [];
+    }
+
+    /**
+     * Adds a new stopping criterion to the list.
+     *
+     * @param {StoppingCriteria} item The stopping criterion to add.
+     */
+    push(item) {
+        this.criteria.push(item);
+    }
+
+    /**
+     * Adds multiple stopping criteria to the list.
+     *
+     * @param {StoppingCriteria|StoppingCriteriaList|StoppingCriteria[]} items The stopping criteria to add.
+     */
+    extend(items) {
+        if (items instanceof StoppingCriteriaList) {
+            items = items.criteria;
+        } else if (items instanceof StoppingCriteria) {
+            items = [items];
+        }
+        this.criteria.push(...items);
+    }
+
+    _call(input_ids, scores) {
+        const is_done = new Array(input_ids.length).fill(false);
+        for (const criterion of this.criteria) {
+            const criterion_done = criterion(input_ids, scores);
+            for (let i = 0; i < is_done.length; ++i) {
+                is_done[i] ||= criterion_done[i];
+            }
+        }
+        return is_done;
+    }
+
+    [Symbol.iterator]() {
+        return this.criteria.values();
+    }
+}
+
+/**
+ * This class can be used to stop generation whenever the full generated number of tokens exceeds `max_length`.
+ * Keep in mind for decoder-only type of transformers, this will include the initial prompted tokens.
+ */
+export class MaxLengthCriteria extends StoppingCriteria {
+
+    /**
+     * 
+     * @param {number} max_length The maximum length that the output sequence can have in number of tokens.
+     * @param {number} [max_position_embeddings=null] The maximum model length, as defined by the model's `config.max_position_embeddings` attribute.
+     */
+    constructor(max_length, max_position_embeddings = null) {
+        super();
+        this.max_length = max_length;
+        this.max_position_embeddings = max_position_embeddings;
+    }
+
+    _call(input_ids) {
+        return input_ids.map(ids => ids.length >= this.max_length);
+    }
+}
+
+// TODO: add MaxTimeCriteria
+
+/**
+ * This class can be used to stop generation whenever the "end-of-sequence" token is generated.
+ * By default, it uses the `model.generation_config.eos_token_id`.
+ */
+export class EosTokenCriteria extends StoppingCriteria {
+
+    /**
+     * 
+     * @param {number|number[]} eos_token_id The id of the *end-of-sequence* token.
+     * Optionally, use a list to set multiple *end-of-sequence* tokens.
+     */
+    constructor(eos_token_id) {
+        super();
+        if (!Array.isArray(eos_token_id)) {
+            eos_token_id = [eos_token_id];
+        }
+        this.eos_token_id = eos_token_id;
+    }
+
+    /**
+     * 
+     * @param {number[][]} input_ids 
+     * @param {number[][]} scores 
+     * @returns {boolean[]}
+     */
+    _call(input_ids, scores) {
+        return input_ids.map(ids => {
+            const last = ids.at(-1);
+            // NOTE: We use == instead of === to allow for number/bigint comparison
+            return this.eos_token_id.some(eos_id => last == eos_id);
+        });
+    }
+}
+
+/**
+ * This class can be used to stop generation whenever the user interrupts the process.
+ */
+export class InterruptableStoppingCriteria extends StoppingCriteria {
+    constructor() {
+        super();
+        this.interrupted = false;
+    }
+
+    interrupt() {
+        this.interrupted = true;
+    }
+
+    reset() {
+        this.interrupted = false;
+    }
+
+    _call(input_ids, scores) {
+        return new Array(input_ids.length).fill(this.interrupted);
+    }
+}

package/src/generation/streamers.js

@@ -0,0 +1,212 @@
+
+/**
+ * @module generation/streamers
+ */
+
+import { mergeArrays } from '../utils/core.js';
+import { is_chinese_char } from '../tokenizers.js';
+import { apis } from '../env.js';
+
+export class BaseStreamer {
+    /**
+     * Function that is called by `.generate()` to push new tokens
+     * @param {bigint[][]} value 
+     */
+    put(value) {
+        throw Error('Not implemented');
+    }
+
+    /**
+     * Function that is called by `.generate()` to signal the end of generation
+     */
+    end() {
+        throw Error('Not implemented');
+    }
+}
+
+const stdout_write = apis.IS_PROCESS_AVAILABLE
+    ? x => process.stdout.write(x)
+    : x => console.log(x);
+
+/**
+ * Simple text streamer that prints the token(s) to stdout as soon as entire words are formed.
+ */
+export class TextStreamer extends BaseStreamer {
+    /**
+     * 
+     * @param {import('../tokenizers.js').PreTrainedTokenizer} tokenizer 
+     */
+    constructor(tokenizer, {
+        skip_prompt = false,
+        callback_function = null,
+        token_callback_function = null,
+        decode_kwargs = {},
+        ...kwargs
+    } = {}) {
+        super();
+        this.tokenizer = tokenizer;
+        this.skip_prompt = skip_prompt;
+        this.callback_function = callback_function ?? stdout_write;
+        this.token_callback_function = token_callback_function;
+        this.decode_kwargs = { ...decode_kwargs, ...kwargs };
+
+        // variables used in the streaming process
+        this.token_cache = [];
+        this.print_len = 0;
+        this.next_tokens_are_prompt = true;
+    }
+
+    /**
+     * Receives tokens, decodes them, and prints them to stdout as soon as they form entire words.
+     * @param {bigint[][]} value 
+     */
+    put(value) {
+        if (value.length > 1) {
+            throw Error('TextStreamer only supports batch size of 1');
+        }
+
+        const tokens = value[0];
+        this.token_callback_function?.(tokens)
+
+        if (this.skip_prompt && this.next_tokens_are_prompt) {
+            this.next_tokens_are_prompt = false;
+            return;
+        }
+
+        // Add the new token to the cache and decodes the entire thing.
+        this.token_cache = mergeArrays(this.token_cache, tokens);
+        const text = this.tokenizer.decode(this.token_cache, this.decode_kwargs);
+
+        let printable_text;
+        if (text.endsWith('\n')) {
+            // After the symbol for a new line, we flush the cache.
+            printable_text = text.slice(this.print_len);
+            this.token_cache = [];
+            this.print_len = 0;
+        } else if (text.length > 0 && is_chinese_char(text.charCodeAt(text.length - 1))) {
+            // If the last token is a CJK character, we print the characters.
+            printable_text = text.slice(this.print_len);
+            this.print_len += printable_text.length;
+        } else {
+            // Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words,
+            // which may change with the subsequent token -- there are probably smarter ways to do this!)
+            printable_text = text.slice(this.print_len, text.lastIndexOf(' ') + 1);
+            this.print_len += printable_text.length;
+        }
+
+        this.on_finalized_text(printable_text, false);
+    }
+
+    /**
+     * Flushes any remaining cache and prints a newline to stdout.
+     */
+    end() {
+        let printable_text;
+        if (this.token_cache.length > 0) {
+            const text = this.tokenizer.decode(this.token_cache, this.decode_kwargs);
+            printable_text = text.slice(this.print_len);
+            this.token_cache = [];
+            this.print_len = 0;
+        } else {
+            printable_text = '';
+        }
+        this.next_tokens_are_prompt = true;
+        this.on_finalized_text(printable_text, true);
+    }
+
+    /**
+     * Prints the new text to stdout. If the stream is ending, also prints a newline.
+     * @param {string} text 
+     * @param {boolean} stream_end 
+     */
+    on_finalized_text(text, stream_end) {
+        if (text.length > 0) {
+            this.callback_function?.(text);
+        }
+        if (stream_end && this.callback_function === stdout_write && apis.IS_PROCESS_AVAILABLE) {
+            this.callback_function?.('\n');
+        }
+    }
+}
+
+/**
+ * Utility class to handle streaming of tokens generated by whisper speech-to-text models.
+ * Callback functions are invoked when each of the following events occur:
+ *  - A new chunk starts (on_chunk_start)
+ *  - A new token is generated (callback_function)
+ *  - A chunk ends (on_chunk_end)
+ *  - The stream is finalized (on_finalize)
+ */
+export class WhisperTextStreamer extends TextStreamer {
+    /**
+     * @param {import('../tokenizers.js').WhisperTokenizer} tokenizer
+     * @param {Object} options
+     * @param {boolean} [options.skip_prompt=false] Whether to skip the prompt tokens
+     * @param {function(string): void} [options.callback_function=null] Function to call when a piece of text is ready to display
+     * @param {function(string): void} [options.token_callback_function=null] Function to call when a new token is generated
+     * @param {function(number): void} [options.on_chunk_start=null] Function to call when a new chunk starts
+     * @param {function(number): void} [options.on_chunk_end=null] Function to call when a chunk ends
+     * @param {function(): void} [options.on_finalize=null] Function to call when the stream is finalized
+     * @param {number} [options.time_precision=0.02] Precision of the timestamps
+     * @param {boolean} [options.skip_special_tokens=true] Whether to skip special tokens when decoding
+     * @param {Object} [options.decode_kwargs={}] Additional keyword arguments to pass to the tokenizer's decode method
+     */
+    constructor(tokenizer, {
+        skip_prompt = false,
+        callback_function = null,
+        token_callback_function = null,
+        on_chunk_start = null,
+        on_chunk_end = null,
+        on_finalize = null,
+        time_precision = 0.02,
+        skip_special_tokens = true,
+        decode_kwargs = {},
+    } = {}) {
+        super(tokenizer, {
+            skip_prompt,
+            callback_function,
+            token_callback_function,
+            decode_kwargs: { skip_special_tokens, ...decode_kwargs },
+        });
+        this.timestamp_begin = tokenizer.timestamp_begin;
+
+        this.on_chunk_start = on_chunk_start;
+        this.on_chunk_end = on_chunk_end;
+        this.on_finalize = on_finalize;
+
+        this.time_precision = time_precision;
+
+        this.waiting_for_timestamp = false;
+    }
+
+    /**
+     * @param {bigint[][]} value 
+     */
+    put(value) {
+        if (value.length > 1) {
+            throw Error('WhisperTextStreamer only supports batch size of 1');
+        }
+        const tokens = value[0];
+
+        // Check if the token is a timestamp
+        if (tokens.length === 1) {
+            const offset = Number(tokens[0]) - this.timestamp_begin;
+            if (offset >= 0) {
+                const time = offset * this.time_precision;
+                if (this.waiting_for_timestamp) {
+                    this.on_chunk_end?.(time);
+                } else {
+                    this.on_chunk_start?.(time);
+                }
+                this.waiting_for_timestamp = !this.waiting_for_timestamp; // Toggle
+                value = [[]]; // Skip timestamp
+            }
+        }
+        return super.put(value);
+    }
+
+    end() {
+        super.end();
+        this.on_finalize?.();
+    }
+}