Ministral_3B_WebGPU_textonly

Running

Ministral_3B_WebGPU_textonly / src /model.js

Danmoreng

fix: Add image placeholder to prompt for multimodal model

3308917 11 months ago

3.05 kB

// src/model.js
import { AutoModelForImageTextToText, AutoProcessor, env, TextStreamer, RawImage } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers/dist/transformers.min.js';

// --- Configuration ---
const MODEL_ID = 'mistralai/Ministral-3-3B-Instruct-2512-ONNX';
const DEVICE = 'webgpu';

// --- State ---
let model;
let processor;

// --- Public API ---
/**
 * Loads the model and processor, and sets up the environment.
 * @param {function(object): void} progressCallback - A function to call with progress updates.
 */
export async function loadModel(progressCallback) {
    env.allowLocalModels = false;

    // The wasm files are bundled with the library, but we need to specify the path
    // to them when using a CDN.
    env.backends.onnx.wasm.wasmPaths = 'https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/';
    
    // Set device
    env.backends.onnx.device = DEVICE;

    try {
        console.log('Loading model...');
        model = await AutoModelForImageTextToText.from_pretrained(MODEL_ID, {
            progress_callback: progressCallback,
            device: DEVICE,
        });
        console.log('Model loaded.');

        console.log('Loading processor...');
        processor = await AutoProcessor.from_pretrained(MODEL_ID, {
            progress_callback: progressCallback,
        });
        console.log('Processor loaded.');

    } catch (error) {
        console.error('Error loading model:', error);
        throw error; // Propagate the error to the caller
    }
}

/**
 * Runs inference with the loaded model.
 * @param {string} inputText - The user's input text.
 * @param {function(string): void} streamCallback - A function to call with each generated token.
 */
export async function runInference(inputText, streamCallback) {
    if (!model || !processor) {
        throw new Error('Model not loaded. Please call loadModel() first.');
    }

    const messages = [
        {
            role: 'user',
            content: [
                { type: 'image' },
                { type: 'text', text: inputText }
            ],
        },
    ];
    const prompt = processor.apply_chat_template(messages, { add_generation_prompt: true });

    // Create a 1x1 black RawImage
    const blackPixel = new Uint8ClampedArray([0, 0, 0, 255]); // RGBA
    const dummyImage = new RawImage(blackPixel, 1, 1, 4);

    // Pass the dummy image to the processor
    const inputs = await processor(dummyImage, prompt);

    let fullOutput = '';
    const streamer = new TextStreamer(processor.tokenizer, {
        skip_prompt: true,
        skip_special_tokens: true,
        callback_function: (text) => {
            fullOutput += text;
            streamCallback(fullOutput);
        },
    });

    await model.generate({
        ...inputs,
        max_new_tokens: 512,
        streamer: streamer,
        do_sample: true,
        temperature: 0.7,
        top_p: 0.95,
        top_k: 50,
        eos_token_id: processor.tokenizer.eos_token_id,
        pad_token_id: processor.tokenizer.pad_token_id,
    });
}

	// src/model.js
	import { AutoModelForImageTextToText, AutoProcessor, env, TextStreamer, RawImage } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers/dist/transformers.min.js';

	// --- Configuration ---
	const MODEL_ID = 'mistralai/Ministral-3-3B-Instruct-2512-ONNX';
	const DEVICE = 'webgpu';

	// --- State ---
	let model;
	let processor;

	// --- Public API ---
	/**
	* Loads the model and processor, and sets up the environment.
	* @param {function(object): void} progressCallback - A function to call with progress updates.
	*/
	export async function loadModel(progressCallback) {
	env.allowLocalModels = false;

	// The wasm files are bundled with the library, but we need to specify the path
	// to them when using a CDN.
	env.backends.onnx.wasm.wasmPaths = 'https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/';

	// Set device
	env.backends.onnx.device = DEVICE;

	try {
	console.log('Loading model...');
	model = await AutoModelForImageTextToText.from_pretrained(MODEL_ID, {
	progress_callback: progressCallback,
	device: DEVICE,
	});
	console.log('Model loaded.');

	console.log('Loading processor...');
	processor = await AutoProcessor.from_pretrained(MODEL_ID, {
	progress_callback: progressCallback,
	});
	console.log('Processor loaded.');

	} catch (error) {
	console.error('Error loading model:', error);
	throw error; // Propagate the error to the caller
	}
	}

	/**
	* Runs inference with the loaded model.
	* @param {string} inputText - The user's input text.
	* @param {function(string): void} streamCallback - A function to call with each generated token.
	*/
	export async function runInference(inputText, streamCallback) {
	if (!model \|\| !processor) {
	throw new Error('Model not loaded. Please call loadModel() first.');
	}

	const messages = [
	{
	role: 'user',
	content: [
	{ type: 'image' },
	{ type: 'text', text: inputText }
	],
	},
	];
	const prompt = processor.apply_chat_template(messages, { add_generation_prompt: true });

	// Create a 1x1 black RawImage
	const blackPixel = new Uint8ClampedArray([0, 0, 0, 255]); // RGBA
	const dummyImage = new RawImage(blackPixel, 1, 1, 4);

	// Pass the dummy image to the processor
	const inputs = await processor(dummyImage, prompt);

	let fullOutput = '';
	const streamer = new TextStreamer(processor.tokenizer, {
	skip_prompt: true,
	skip_special_tokens: true,
	callback_function: (text) => {
	fullOutput += text;
	streamCallback(fullOutput);
	},
	});

	await model.generate({
	...inputs,
	max_new_tokens: 512,
	streamer: streamer,
	do_sample: true,
	temperature: 0.7,
	top_p: 0.95,
	top_k: 50,
	eos_token_id: processor.tokenizer.eos_token_id,
	pad_token_id: processor.tokenizer.pad_token_id,
	});
	}