PanoVLD-500M / processing_panovlm.py

PanoVLD-500M

1ba3d87 1 day ago

6.6 kB

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""Standalone HF processor for PanoVLM: image processor + tokenizer.
Expands each ``<|image|>`` placeholder into ``image_seq_length`` repeated image
tokens wrapped by ``<|begin_of_image|>``/``<|end_of_image|>``, then tokenizes.
``image_seq_length`` is fixed for a given image resolution and is stored in the
processor config at upload time.
``apply_chat_template`` also accepts the list-structured message content used by
newer HF multimodal callers, and normalizes all message content to the string
form the model expects.
"""

from __future__ import annotations

from transformers.feature_extraction_utils import BatchFeature
from transformers.processing_utils import ProcessorMixin


class PanoVLMProcessor(ProcessorMixin):
    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "AutoImageProcessor"
    tokenizer_class = "AutoTokenizer"

    # Hardcoded PanoVLM image placeholder tokens (not configurable / not serialized).
    image_token = "<|image|>"
    boi_token = "<|begin_of_image|>"
    eoi_token = "<|end_of_image|>"

    def __init__(
        self,
        image_processor=None,
        tokenizer=None,
        image_seq_length=None,
        chat_template=None,
        **kwargs,
    ):
        # Set extra scalar attrs before super().__init__ (mirrors LlavaProcessor pattern).
        # ProcessorMixin.__init__ in transformers 5.x only accepts modality args
        # (image_processor, tokenizer) plus chat_template — extra scalars like
        # image_seq_length must be stored before the super() call.
        self.image_seq_length = image_seq_length
        super().__init__(image_processor, tokenizer, chat_template=chat_template)

    def _expand_image_tokens(self, text: str) -> str:
        """Replace each ``<|image|>`` with BOI + N image tokens + EOI."""
        if self.image_seq_length is None:
            raise ValueError(
                "image_seq_length is not set on this processor. Pass "
                "image_seq_length=... when constructing, or load a processor saved "
                "with it set."
            )
        block = (
            self.boi_token + (self.image_token * self.image_seq_length) + self.eoi_token
        )
        return text.replace(self.image_token, block)

    def _prompt_text(self, text: str) -> str:
        """The text to tokenize: image tokens expanded, trailing space removed.
        rstrip before tokenizing: chat templates end with a trailing space (e.g.
        "ASSISTANT: ") that SentencePiece tokenizes as a standalone U+2581 token. That
        token is not expected at this position (the space belongs to the first
        generated token), so it would break the first-token distribution.
        """
        return self._expand_image_tokens(text).rstrip()

    @classmethod
    def _normalize_content(cls, content):
        """Render message content to the exact string form the model expects.
        PanoVLM's chat template concatenates
        ``message['content']`` directly, so content must be a plain string. Callers
        that follow the newer HF multimodal convention (e.g. lmms-eval's generic
        ``huggingface`` wrapper) pass a list of ``{"type": ...}`` parts instead;
        image/video parts become ``<|image|>`` placeholders (expanded later in
        ``__call__``) and text parts are concatenated as-is.
        The resulting string is then split on ``<|image|>``, each text segment
        stripped, empty ones dropped, and everything joined with no separator:
        the model expects ``"<|image|>What is ..."``, so a separator here would
        put tokens after an image that the model does not expect there.
        """
        if not isinstance(content, str):
            pieces = []
            for part in content:
                kind = part.get("type")
                if kind in ("image", "video"):
                    pieces.append(cls.image_token)
                elif kind == "text":
                    pieces.append(part.get("text", ""))
                else:
                    raise ValueError(
                        f"PanoVLM chat content part of type {kind!r} is not "
                        "supported; expected 'text', 'image' or 'video'."
                    )
            content = "".join(pieces)
        normalized = cls.image_token.join(
            chunk.strip() for chunk in content.split(cls.image_token)
        )
        # Keep whitespace-only content untouched.
        return normalized or content

    def apply_chat_template(self, conversation, *args, **kwargs):
        """Normalize message content to the expected string form, then render.
        See :meth:`_normalize_content`. Accepts a single conversation or a batch of
        conversations, matching ``ProcessorMixin.apply_chat_template``.
        """
        is_batched = isinstance(conversation, (list, tuple)) and isinstance(
            conversation[0], (list, tuple)
        )
        conversations = conversation if is_batched else [conversation]
        # Rebuild the messages instead of mutating: the caller (and, for lmms-eval,
        # the code that extracts the images from the same message objects) still
        # needs the original content parts.
        normalized = [
            [
                {**message, "content": self._normalize_content(message["content"])}
                for message in messages
            ]
            for messages in conversations
        ]
        return super().apply_chat_template(
            normalized if is_batched else normalized[0], *args, **kwargs
        )

    def __call__(self, text=None, images=None, return_tensors=None, **kwargs):
        if text is None:
            raise ValueError("PanoVLMProcessor requires `text`.")
        if isinstance(text, str):
            text = [text]
        text = [self._prompt_text(t) for t in text]

        data = {}
        if images is not None:
            image_inputs = self.image_processor(images, return_tensors=return_tensors)
            data.update(image_inputs)
        text_inputs = self.tokenizer(text, return_tensors=return_tensors, **kwargs)
        data.update(text_inputs)
        return BatchFeature(data=data, tensor_type=return_tensors)

    def batch_decode(self, *args, **kwargs):
        return self.tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        return self.tokenizer.decode(*args, **kwargs)

	# Copyright (c) Meta Platforms, Inc. and affiliates.
	# All rights reserved.
	#
	# This source code is licensed under the BSD-style license found in the
	# LICENSE file in the root directory of this source tree.

	"""Standalone HF processor for PanoVLM: image processor + tokenizer.

	Expands each ``<\|image\|>`` placeholder into ``image_seq_length`` repeated image
	tokens wrapped by ``<\|begin_of_image\|>``/``<\|end_of_image\|>``, then tokenizes.
	``image_seq_length`` is fixed for a given image resolution and is stored in the
	processor config at upload time.

	``apply_chat_template`` also accepts the list-structured message content used by
	newer HF multimodal callers, and normalizes all message content to the string
	form the model expects.
	"""

	from __future__ import annotations

	from transformers.feature_extraction_utils import BatchFeature
	from transformers.processing_utils import ProcessorMixin


	class PanoVLMProcessor(ProcessorMixin):
	attributes = ["image_processor", "tokenizer"]
	image_processor_class = "AutoImageProcessor"
	tokenizer_class = "AutoTokenizer"

	# Hardcoded PanoVLM image placeholder tokens (not configurable / not serialized).
	image_token = "<\|image\|>"
	boi_token = "<\|begin_of_image\|>"
	eoi_token = "<\|end_of_image\|>"

	def __init__(
	self,
	image_processor=None,
	tokenizer=None,
	image_seq_length=None,
	chat_template=None,
	**kwargs,
	):
	# Set extra scalar attrs before super().__init__ (mirrors LlavaProcessor pattern).
	# ProcessorMixin.__init__ in transformers 5.x only accepts modality args
	# (image_processor, tokenizer) plus chat_template — extra scalars like
	# image_seq_length must be stored before the super() call.
	self.image_seq_length = image_seq_length
	super().__init__(image_processor, tokenizer, chat_template=chat_template)

	def _expand_image_tokens(self, text: str) -> str:
	"""Replace each ``<\|image\|>`` with BOI + N image tokens + EOI."""
	if self.image_seq_length is None:
	raise ValueError(
	"image_seq_length is not set on this processor. Pass "
	"image_seq_length=... when constructing, or load a processor saved "
	"with it set."
	)
	block = (
	self.boi_token + (self.image_token * self.image_seq_length) + self.eoi_token
	)
	return text.replace(self.image_token, block)

	def _prompt_text(self, text: str) -> str:
	"""The text to tokenize: image tokens expanded, trailing space removed.

	rstrip before tokenizing: chat templates end with a trailing space (e.g.
	"ASSISTANT: ") that SentencePiece tokenizes as a standalone U+2581 token. That
	token is not expected at this position (the space belongs to the first
	generated token), so it would break the first-token distribution.
	"""
	return self._expand_image_tokens(text).rstrip()

	@classmethod
	def _normalize_content(cls, content):
	"""Render message content to the exact string form the model expects.

	PanoVLM's chat template concatenates
	``message['content']`` directly, so content must be a plain string. Callers
	that follow the newer HF multimodal convention (e.g. lmms-eval's generic
	``huggingface`` wrapper) pass a list of ``{"type": ...}`` parts instead;
	image/video parts become ``<\|image\|>`` placeholders (expanded later in
	``__call__``) and text parts are concatenated as-is.

	The resulting string is then split on ``<\|image\|>``, each text segment
	stripped, empty ones dropped, and everything joined with no separator:
	the model expects ``"<\|image\|>What is ..."``, so a separator here would
	put tokens after an image that the model does not expect there.
	"""
	if not isinstance(content, str):
	pieces = []
	for part in content:
	kind = part.get("type")
	if kind in ("image", "video"):
	pieces.append(cls.image_token)
	elif kind == "text":
	pieces.append(part.get("text", ""))
	else:
	raise ValueError(
	f"PanoVLM chat content part of type {kind!r} is not "
	"supported; expected 'text', 'image' or 'video'."
	)
	content = "".join(pieces)
	normalized = cls.image_token.join(
	chunk.strip() for chunk in content.split(cls.image_token)
	)
	# Keep whitespace-only content untouched.
	return normalized or content

	def apply_chat_template(self, conversation, args, *kwargs):
	"""Normalize message content to the expected string form, then render.

	See :meth:`_normalize_content`. Accepts a single conversation or a batch of
	conversations, matching ``ProcessorMixin.apply_chat_template``.
	"""
	is_batched = isinstance(conversation, (list, tuple)) and isinstance(
	conversation[0], (list, tuple)
	)
	conversations = conversation if is_batched else [conversation]
	# Rebuild the messages instead of mutating: the caller (and, for lmms-eval,
	# the code that extracts the images from the same message objects) still
	# needs the original content parts.
	normalized = [
	[
	{**message, "content": self._normalize_content(message["content"])}
	for message in messages
	]
	for messages in conversations
	]
	return super().apply_chat_template(
	normalized if is_batched else normalized[0], args, *kwargs
	)

	def __call__(self, text=None, images=None, return_tensors=None, **kwargs):
	if text is None:
	raise ValueError("PanoVLMProcessor requires `text`.")
	if isinstance(text, str):
	text = [text]
	text = [self._prompt_text(t) for t in text]

	data = {}
	if images is not None:
	image_inputs = self.image_processor(images, return_tensors=return_tensors)
	data.update(image_inputs)
	text_inputs = self.tokenizer(text, return_tensors=return_tensors, **kwargs)
	data.update(text_inputs)
	return BatchFeature(data=data, tensor_type=return_tensors)

	def batch_decode(self, args, *kwargs):
	return self.tokenizer.batch_decode(args, *kwargs)

	def decode(self, args, *kwargs):
	return self.tokenizer.decode(args, *kwargs)