@huggingface/transformers/src/models/paligemma/processing_paligemma.js

State-of-the-art Machine Learning for the web. Run 🤗 Transformers directly in your browser, with no need for a server!

import { Processor } from '../../processing_utils.js';
import { AutoImageProcessor } from '../auto/image_processing_auto.js';
import { AutoTokenizer } from '../auto/tokenization_auto.js';
import { logger } from '../../utils/logger.js';

const IMAGE_TOKEN = '<image>';

function build_string_from_input(prompt, bos_token, image_seq_len, image_token, num_images) {
    return `${image_token.repeat(image_seq_len * num_images)}${bos_token}${prompt}\n`;
}

export class PaliGemmaProcessor extends Processor {
    static tokenizer_class = AutoTokenizer;
    static image_processor_class = AutoImageProcessor;
    static uses_processor_config = false;

    /**
     * @typedef {import('../../utils/image.js').RawImage} RawImage
     */

    // `images` is required, `text` is optional
    async _call(/** @type {RawImage|RawImage[]} */ images, text = null, kwargs = {}) {
        if (!text) {
            logger.warn(
                'You are using PaliGemma without a text prefix. It will perform as a picture-captioning model.',
            );
            text = '';
        }

        if (!Array.isArray(images)) {
            images = [images];
        }

        if (!Array.isArray(text)) {
            text = [text];
        }

        const bos_token = this.tokenizer.bos_token;
        // @ts-expect-error TS2339
        const image_seq_length = this.image_processor.config.image_seq_length;
        let input_strings;
        if (text.some((t) => t.includes(IMAGE_TOKEN))) {
            input_strings = text.map((sample) => {
                const expanded_sample = sample.replaceAll(IMAGE_TOKEN, IMAGE_TOKEN.repeat(image_seq_length));
                const bos_rfind_index = expanded_sample.lastIndexOf(IMAGE_TOKEN);
                const bos_index = bos_rfind_index === -1 ? 0 : bos_rfind_index + IMAGE_TOKEN.length;
                return expanded_sample.slice(0, bos_index) + bos_token + expanded_sample.slice(bos_index) + '\n';
            });
        } else {
            logger.warn(
                'You are passing both `text` and `images` to `PaliGemmaProcessor`. The processor expects special ' +
                    'image tokens in the text, as many tokens as there are images per each text. It is recommended to ' +
                    'add `<image>` tokens in the very beginning of your text. For this call, we will infer how many images ' +
                    'each text has and add special tokens.',
            );

            input_strings = text.map((sample) =>
                build_string_from_input(sample, bos_token, image_seq_length, IMAGE_TOKEN, images.length),
            );
        }

        const text_inputs = this.tokenizer(input_strings, kwargs);
        const image_inputs = await this.image_processor(images, kwargs);

        return {
            ...image_inputs,
            ...text_inputs,
        };
    }
}