@huggingface/transformers 4.0.0-next.6 → 4.0.0-next.7

package/src/models/voxtral_realtime/modeling_voxtral_realtime.js

@@ -0,0 +1,239 @@
+import { PreTrainedModel } from '../modeling_utils.js';
+import { sessionRun } from '../session.js';
+import { getCacheShapes } from '../../configs.js';
+import { Tensor, ones } from '../../utils/tensor.js';
+import { DataTypeMap } from '../../utils/dtypes.js';
+import { pick } from '../../utils/core.js';
+import { DynamicCache } from '../../cache_utils.js';
+import { StoppingCriteria, StoppingCriteriaList } from '../../generation/stopping_criteria.js';
+
+// Causal conv padding constants
+const CONV1_LEFT_PAD = 2;
+const CONV2_LEFT_PAD = 1;
+
+/**
+ * WeakMap to hold encoder streaming states for each model instance during generation.
+ * This allows the state to be accessed and modified across the generation process
+ * without exposing it on the model instance itself.
+ * @private
+ * @type {WeakMap<VoxtralRealtimeForConditionalGeneration, Object>}
+ */
+const states = new WeakMap();
+
+/**
+ * Creates encoder streaming state for a VoxtralRealtime generation session.
+ * @param {VoxtralRealtimeForConditionalGeneration} model
+ * @param {Iterable<Tensor>|AsyncIterable<Tensor>} input_features
+ * @returns {Object} Encoder state object.
+ * @private
+ */
+function createEncoderState(model, input_features) {
+    const { text_config, audio_config } = /** @type {any} */ (model.config);
+    const encoder_session = model.sessions['audio_encoder'];
+
+    const { num_mel_bins, hidden_size: enc_hidden_size } = audio_config;
+    const PADDING_CACHE_CHANNELS = num_mel_bins + enc_hidden_size;
+
+    // Initialize encoder KV cache
+    const enc_kv_cache = new DynamicCache();
+    const enc_dtype = encoder_session?.config?.kv_cache_dtype ?? 'float32';
+    const enc_cls = enc_dtype === 'float16' ? DataTypeMap.float16 : DataTypeMap.float32;
+    const enc_shapes = getCacheShapes(audio_config, { batch_size: 1 });
+    for (const name in enc_shapes) {
+        const size = enc_shapes[name].reduce((a, b) => a * b, 1);
+        enc_kv_cache[name] = new Tensor(enc_dtype, new enc_cls(size), enc_shapes[name]);
+    }
+
+    const enc_padding_cache = new Tensor(enc_dtype, new enc_cls(PADDING_CACHE_CHANNELS * CONV1_LEFT_PAD), [
+        1,
+        PADDING_CACHE_CHANNELS,
+        CONV1_LEFT_PAD,
+    ]);
+
+    // Set up iterator from input_features
+    const chunks_iter = input_features[Symbol.asyncIterator]?.() ?? input_features[Symbol.iterator]?.();
+    if (!chunks_iter) {
+        throw new Error('input_features must be iterable or async iterable');
+    }
+
+    return {
+        encoder_session,
+        enc_kv_cache,
+        enc_padding_cache,
+        enc_past_seq_len: 0,
+        audio_embed_queue: [],
+        audio_embed_total_tokens: 0,
+        audio_queue_offset: 0,
+        audio_consumed: 0,
+        stream_exhausted: false,
+        chunks_iter,
+        text_hidden_size: text_config.hidden_size,
+    };
+}
+
+/**
+ * Encodes one audio chunk through the audio encoder.
+ * @param {Object} s Encoder state.
+ * @param {Tensor} chunk_features Mel spectrogram chunk [1, num_mel_bins, seq_len].
+ * @returns {Promise<Tensor>} Audio embeddings.
+ * @private
+ */
+async function encodeChunk(s, chunk_features) {
+    const audio_seq_len = chunk_features.dims[2];
+    const conv2_output_len = Math.floor((CONV2_LEFT_PAD + audio_seq_len - 3) / 2) + 1;
+
+    const position_ids = new Tensor(
+        'int64',
+        BigInt64Array.from({ length: conv2_output_len }, (_, i) => BigInt(s.enc_past_seq_len + i)),
+        [1, conv2_output_len],
+    );
+
+    const total_seq_len = s.enc_past_seq_len + conv2_output_len;
+    const attention_mask = ones([1, total_seq_len]);
+    const { audio_embeds, present_padding_cache, ...present_cache } = await sessionRun(s.encoder_session, {
+        input_features: chunk_features,
+        attention_mask,
+        position_ids,
+        past_padding_cache: s.enc_padding_cache,
+        ...s.enc_kv_cache,
+    });
+    // Dispose previous padding cache and update
+    if (s.enc_padding_cache.location === 'gpu-buffer') {
+        s.enc_padding_cache.dispose();
+    }
+    s.enc_padding_cache = present_padding_cache;
+
+    // Update encoder KV cache, disposing previous tensors
+    for (const name in present_cache) {
+        if (name.startsWith('present.')) {
+            const pastName = name.replace('present', 'past_key_values');
+            const prev = s.enc_kv_cache[pastName];
+            if (prev?.location === 'gpu-buffer') {
+                prev.dispose();
+            }
+            s.enc_kv_cache[pastName] = present_cache[name];
+        }
+    }
+    s.enc_past_seq_len = total_seq_len;
+    return audio_embeds;
+}
+
+/**
+ * Fills the audio embedding buffer until it has enough tokens.
+ * @param {Object} s Encoder state.
+ * @param {number} needed Total number of audio tokens needed.
+ * @private
+ */
+async function fillAudioBuffer(s, needed) {
+    while (s.audio_embed_total_tokens < needed && !s.stream_exhausted) {
+        const result = await s.chunks_iter.next();
+        if (result.done) {
+            s.stream_exhausted = true;
+            break;
+        }
+        const new_embeds = await encodeChunk(s, result.value);
+        s.audio_embed_queue.push({ data: new_embeds.data, tokens: new_embeds.dims[1] });
+        s.audio_embed_total_tokens += new_embeds.dims[1];
+    }
+}
+
+/**
+ * Adds audio embeddings to text embeddings from the queue.
+ * @param {Object} s Encoder state.
+ * @param {Tensor} inputs_embeds Text embeddings tensor (modified in-place).
+ * @param {number} current_len Number of tokens to consume.
+ * @private
+ */
+function addAudioEmbeddings(s, inputs_embeds, current_len) {
+    if (s.audio_embed_queue.length === 0) return;
+
+    const embed_data = inputs_embeds.data;
+    let embed_write_pos = 0;
+    let remaining = current_len;
+
+    while (remaining > 0 && s.audio_embed_queue.length > 0) {
+        const front = s.audio_embed_queue[0];
+        const available = front.tokens - s.audio_queue_offset;
+        const n = Math.min(remaining, available);
+
+        const src_offset = s.audio_queue_offset * s.text_hidden_size;
+        for (let i = 0; i < n * s.text_hidden_size; ++i) {
+            embed_data[embed_write_pos * s.text_hidden_size + i] += front.data[src_offset + i];
+        }
+
+        embed_write_pos += n;
+        remaining -= n;
+        s.audio_queue_offset += n;
+
+        if (s.audio_queue_offset >= front.tokens) {
+            s.audio_embed_queue.shift();
+            s.audio_queue_offset = 0;
+        }
+    }
+    s.audio_consumed += current_len - remaining;
+}
+
+/**
+ * Stopping criterion that triggers when the audio stream is exhausted
+ * and all buffered audio embeddings have been consumed.
+ * @private
+ */
+class AudioExhaustedCriteria extends StoppingCriteria {
+    constructor(enc_state) {
+        super();
+        this._s = enc_state;
+    }
+    _call(input_ids) {
+        const done = this._s.stream_exhausted && this._s.audio_embed_queue.length === 0;
+        return input_ids.map(() => done);
+    }
+}
+
+export class VoxtralRealtimePreTrainedModel extends PreTrainedModel {
+    forward_params = ['input_ids', 'attention_mask', 'position_ids', 'past_key_values'];
+}
+
+export class VoxtralRealtimeForConditionalGeneration extends VoxtralRealtimePreTrainedModel {
+    async forward({ input_ids, past_key_values, ...kwargs }) {
+        const current_len = input_ids.dims[1];
+
+        const enc = states.get(this);
+        if (enc) {
+            // Fill audio buffer and embed tokens with audio
+            await fillAudioBuffer(enc, enc.audio_consumed + current_len);
+        }
+
+        const { inputs_embeds } = await sessionRun(this.sessions['embed_tokens'], { input_ids });
+        if (enc) {
+            addAudioEmbeddings(enc, inputs_embeds, current_len);
+        }
+
+        const decoder_feeds = { inputs_embeds, ...kwargs };
+        this.addPastKeyValues(decoder_feeds, past_key_values);
+
+        const session = this.sessions['decoder_model_merged'];
+        const fixed = pick(decoder_feeds, session.inputNames);
+        return await sessionRun(session, fixed);
+    }
+
+    async generate({ input_features, stopping_criteria: userStoppingCriteria, ...kwargs }) {
+        if (!input_features) {
+            throw new Error('input_features (generator/iterable) must be provided');
+        }
+
+        const enc_state = createEncoderState(this, input_features);
+        states.set(this, enc_state);
+
+        const stopping_criteria = new StoppingCriteriaList();
+        stopping_criteria.push(new AudioExhaustedCriteria(enc_state));
+        if (userStoppingCriteria) stopping_criteria.extend(userStoppingCriteria);
+
+        try {
+            return await super.generate({ ...kwargs, stopping_criteria });
+        } finally {
+            // Cleanup encoder state
+            enc_state.enc_kv_cache.dispose();
+            states.delete(this);
+        }
+    }
+}

package/src/models/voxtral_realtime/processing_voxtral_realtime.js

@@ -0,0 +1,113 @@
+import { AutoFeatureExtractor } from '../auto/feature_extraction_auto.js';
+import { AutoTokenizer } from '../auto/tokenization_auto.js';
+import { Processor } from '../../processing_utils.js';
+import { Tensor } from '../../utils/tensor.js';
+import { validate_audio_inputs } from '../../feature_extraction_utils.js';
+
+// Voxtral Realtime audio config constants (from mistral_common AudioConfig)
+const NUM_LEFT_PAD_TOKENS = 32;
+const NUM_DELAY_TOKENS = 6;
+const AUDIO_LENGTH_PER_TOK = 8;
+const OFFLINE_STREAMING_BUFFER_TOKENS = 10;
+
+/** Token ID for [STREAMING_PAD] in the Voxtral tokenizer. */
+const STREAMING_PAD_TOKEN_ID = 32;
+
+export class VoxtralRealtimeProcessor extends Processor {
+    static tokenizer_class = AutoTokenizer;
+    static feature_extractor_class = AutoFeatureExtractor;
+    static uses_processor_config = false;
+
+    /** Number of mel frames in the first audio chunk. */
+    get num_mel_frames_first_audio_chunk() {
+        return (NUM_DELAY_TOKENS + 1) * AUDIO_LENGTH_PER_TOK;
+    }
+
+    /** Number of raw audio samples in the first audio chunk. */
+    get num_samples_first_audio_chunk() {
+        const { hop_length, n_fft } = this.feature_extractor.config;
+        return (this.num_mel_frames_first_audio_chunk - 1) * hop_length + Math.floor(n_fft / 2);
+    }
+
+    /** Number of raw audio samples per subsequent audio chunk. */
+    get num_samples_per_audio_chunk() {
+        const { hop_length, n_fft } = this.feature_extractor.config;
+        return AUDIO_LENGTH_PER_TOK * hop_length + n_fft;
+    }
+
+    /** Number of right-pad tokens for non-streaming mode. */
+    get num_right_pad_tokens() {
+        return NUM_DELAY_TOKENS + 1 + OFFLINE_STREAMING_BUFFER_TOKENS;
+    }
+
+    /** Number of mel frames per text token. */
+    get audio_length_per_tok() {
+        return AUDIO_LENGTH_PER_TOK;
+    }
+
+    /** Number of raw audio samples per token. */
+    get raw_audio_length_per_tok() {
+        return AUDIO_LENGTH_PER_TOK * this.feature_extractor.config.hop_length;
+    }
+
+    /**
+     * Process audio input for VoxtralRealtime.
+     *
+     * In streaming mode with `is_first_audio_chunk=true`, the audio is left-padded
+     * with silence and mel features are extracted with `center=true`.
+     * Returns `{ input_ids, input_features }`.
+     *
+     * In streaming mode with `is_first_audio_chunk=false`, the audio chunk is
+     * processed with `center=false` and only `{ input_features }` is returned.
+     *
+     * In non-streaming mode, the audio is right-padded to ensure the model
+     * transcribes the full audio, then processed with `center=true`.
+     * Returns `{ input_features }`.
+     *
+     * @param {Float32Array|Float64Array} audio The audio waveform.
+     * @param {Object} [options]
+     * @param {boolean} [options.is_streaming=false] Whether processing in streaming mode.
+     * @param {boolean} [options.is_first_audio_chunk=true] Whether this is the first audio chunk.
+     * @returns {Promise<Object>}
+     */
+    async _call(audio, { is_streaming = false, is_first_audio_chunk = true } = {}) {
+        validate_audio_inputs(audio, 'VoxtralRealtimeProcessor');
+
+        if (!is_streaming && !is_first_audio_chunk) {
+            throw new Error('In non-streaming mode (`is_streaming=false`), `is_first_audio_chunk` must be `true`.');
+        }
+
+        if (is_first_audio_chunk) {
+            if (is_streaming) {
+                // Streaming first chunk: left-pad audio with silence, extract mel with center=true, build input_ids
+                const num_left_pad_samples = NUM_LEFT_PAD_TOKENS * this.raw_audio_length_per_tok;
+                const padded_audio = new Float32Array(num_left_pad_samples + audio.length);
+                padded_audio.set(audio, num_left_pad_samples);
+
+                const audio_encoding = await this.feature_extractor(padded_audio, { center: true });
+
+                // Build input_ids: BOS + (num_left_pad_tokens + num_delay_tokens) * [STREAMING_PAD]
+                const num_pad_tokens = NUM_LEFT_PAD_TOKENS + NUM_DELAY_TOKENS;
+                const num_input_tokens = 1 + num_pad_tokens;
+                const input_ids_data = new BigInt64Array(num_input_tokens).fill(BigInt(STREAMING_PAD_TOKEN_ID));
+                input_ids_data[0] = 1n; // BOS
+                const input_ids = new Tensor('int64', input_ids_data, [1, num_input_tokens]);
+
+                return {
+                    input_ids,
+                    ...audio_encoding,
+                };
+            } else {
+                // Non-streaming: right-pad audio to ensure full transcription, extract mel with center=true
+                const right_pad_samples = this.num_right_pad_tokens * this.raw_audio_length_per_tok;
+                const padded_audio = new Float32Array(audio.length + right_pad_samples);
+                padded_audio.set(audio);
+
+                return await this.feature_extractor(padded_audio, { center: true });
+            }
+        } else {
+            // Subsequent streaming chunks: extract mel with center=false
+            return await this.feature_extractor(audio, { center: false });
+        }
+    }
+}

package/src/models/whisper/feature_extraction_whisper.js

@@ -1,7 +1,6 @@
 import { FeatureExtractor, validate_audio_inputs } from '../../feature_extraction_utils.js';
 import { Tensor } from '../../utils/tensor.js';
 import { mel_filter_bank, spectrogram, window_function } from '../../utils/audio.js';
-import { max } from '../../utils/maths.js';
 import { logger } from '../../utils/logger.js';
 
 export class WhisperFeatureExtractor extends FeatureExtractor {
@@ -28,7 +27,7 @@ export class WhisperFeatureExtractor extends FeatureExtractor {
      * @returns {Promise<Tensor>} An object containing the log-Mel spectrogram data as a Float32Array and its dimensions as an array of numbers.
      */
     async _extract_fbank_features(waveform) {
-        const features = await spectrogram(
+        return await spectrogram(
             waveform,
             this.window, // window
             this.config.n_fft, // frame_length
@@ -36,7 +35,7 @@ export class WhisperFeatureExtractor extends FeatureExtractor {
             {
                 power: 2.0,
                 mel_filters: this.config.mel_filters,
-                log_mel: 'log10',
+                log_mel: 'log10_max_norm',
 
                 // Custom
                 max_num_frames: Math.min(
@@ -45,15 +44,6 @@ export class WhisperFeatureExtractor extends FeatureExtractor {
                 ),
             },
         );
-
-        const data = features.data;
-        const maxValue = max(/** @type {Float32Array} */ (data))[0];
-
-        for (let i = 0; i < data.length; ++i) {
-            data[i] = (Math.max(data[i], maxValue - 8.0) + 4.0) / 4.0;
-        }
-
-        return features;
     }
 
     /**

package/src/transformers.js

@@ -53,6 +53,8 @@ export * from './utils/tensor.js';
 export { softmax, log_softmax, dot, cos_sim } from './utils/maths.js';
 export { random } from './utils/random.js';
 
+export { DynamicCache } from './cache_utils.js';
+
 // Cache and file management
 export { ModelRegistry } from './utils/model_registry/ModelRegistry.js';
 

package/src/utils/audio.js

@@ -452,8 +452,10 @@ function power_to_db(spectrogram, reference = 1.0, min_value = 1e-10, db_range =
  * If supplied, applies this filter bank to create a mel spectrogram.
  * @param {number} [options.mel_floor=1e-10] Minimum value of mel frequency banks.
  * @param {string} [options.log_mel=null] How to convert the spectrogram to log scale. Possible options are:
- * `null` (don't convert), `"log"` (take the natural logarithm) `"log10"` (take the base-10 logarithm), `"dB"` (convert to decibels).
+ * `null` (don't convert), `"log"` (take the natural logarithm), `"log10"` (take the base-10 logarithm), `"dB"` (convert to decibels),
+ * `"log10_max_norm"` (take `log10`, then apply `(max(x, maxVal - 8) + 4) / 4` normalization, where `maxVal` is computed from data or given by `max_log_mel`).
  * Can only be used when `power` is not `null`.
+ * @param {number} [options.max_log_mel=null] When `log_mel` is `"log10_max_norm"`, use this fixed value as the max instead of computing from data.
  * @param {number} [options.reference=1.0] Sets the input spectrogram value that corresponds to 0 dB. For example, use `max(spectrogram)[0]` to set
  * the loudest part to 0 dB. Must be greater than zero.
  * @param {number} [options.min_value=1e-10] The spectrogram will be clipped to this minimum value before conversion to decibels, to avoid taking `log(0)`.
@@ -486,6 +488,7 @@ export async function spectrogram(
         mel_filters = null,
         mel_floor = 1e-10,
         log_mel = null,
+        max_log_mel = null,
         reference = 1.0,
         min_value = 1e-10,
         db_range = null,
@@ -669,6 +672,17 @@ export async function spectrogram(
                     mel_spec_data[i] = Math.log10(mel_spec_data[i]);
                 }
                 break;
+            case 'log10_max_norm': {
+                for (let i = 0; i < o; ++i) {
+                    mel_spec_data[i] = Math.log10(mel_spec_data[i]);
+                }
+                const logMax = max_log_mel ?? max(mel_spec_data)[0];
+                const threshold = logMax - 8.0;
+                for (let i = 0; i < o; ++i) {
+                    mel_spec_data[i] = (Math.max(mel_spec_data[i], threshold) + 4.0) / 4.0;
+                }
+                break;
+            }
             case 'dB':
                 if (power === 1.0) {
                     amplitude_to_db(mel_spec_data, reference, min_value, db_range);
@@ -679,7 +693,9 @@ export async function spectrogram(
                 }
                 break;
             default:
-                throw new Error(`log_mel must be one of null, 'log', 'log10' or 'dB'. Got '${log_mel}'`);
+                throw new Error(
+                    `log_mel must be one of null, 'log', 'log10', 'log10_max_norm', or 'dB'. Got '${log_mel}'`,
+                );
         }
     }