kimi_k3_processor.py

"""Kimi-K3 processor: wraps vision processor + tokenizer into a single interface.

Chat rendering (including XTML tool-result ordering) is handled by the
tokenizer's Python encoder; this processor adds multimodal media preprocessing.
"""

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

from .media_utils import ensure_media_type

logger = logging.get_logger(__name__)

# ── KimiK3Processor ───────────────────────────────────────────────────


class KimiK3Processor(ProcessorMixin):
    r"""
    Constructs a KimiK3 processor which wraps a KimiK3 image processor
    and a tokenizer into a single processor.

    [`KimiK3Processor`] offers all the functionalities of
    [`KimiK3VisionProcessor`] and [`TikTokenTokenizer`].

    Args:
        image_processor ([`KimiK3VisionProcessor`], *optional*):
            The image processor is a required input.
        tokenizer ([`TikTokenTokenizer`], *optional*):
            The tokenizer is a required input.
        chat_template (`str`, *optional*): Kept for ProcessorMixin
            compatibility. Kimi K3 chat encoding is implemented in Python by
            the tokenizer.
    """

    attributes = ["image_processor", "tokenizer"]
    valid_kwargs = ["chat_template"]
    image_processor_class = "AutoImageProcessor"
    tokenizer_class = "AutoTokenizer"

    def __init__(
        self,
        image_processor=None,
        tokenizer=None,
        chat_template=None,
        **kwargs,
    ):
        super().__init__(image_processor, tokenizer, chat_template=chat_template)
        self.media_processor = image_processor
        self.image_placeholder = "<|kimi_image_placeholder|>"

    # ── Media preprocessing ────────────────────────────────────────────

    def update_raw_text(self, text: str, image_prompts: list[str]) -> str:
        # Replace image placeholders
        image_count = text.count(self.image_placeholder)
        if image_count > 0:
            assert image_count == len(image_prompts), (
                f"image placeholder count {image_count} != "
                f"image_prompts count {len(image_prompts)}"
            )
            text_parts = text.split(self.image_placeholder)
            assert len(text_parts) == len(image_prompts) + 1
            text = "".join(
                [text_parts[i] + image_prompts[i] for i in range(len(image_prompts))]
            )
            text += text_parts[-1]

        return text

    def preprocess_medias(self, medias: list[dict]) -> tuple[list[dict], list[str]]:
        """Process media items and generate corresponding prompts.

        Returns:
            A tuple of (updated_medias, image_prompts).
        """
        updated_medias = []
        image_prompts = []
        for media in medias:
            if media["type"] == "image":
                updated_medias.append(media)
                img = ensure_media_type(
                    media,
                    transparent_bg_config=self.media_processor._transparent_bg_config,
                    transparent_bg_fill_stage=self.media_processor._transparent_bg_fill_stage,
                )["image"]
                w, h = img.size
                image_prompts.append(self.media_processor.make_image_prompt(w, h))
            else:
                raise ValueError(f"unsupported media type: {media['type']}")
        return updated_medias, image_prompts

    # ── Main entry points ──────────────────────────────────────────────

    def __call__(
        self,
        messages: list[dict] = None,
        medias: list[dict] = None,
        text: str = None,
        return_tensors: str = "pt",
        **kwargs,
    ) -> BatchFeature:
        """
        Process multimodal inputs for Kimi-K3 model.

        Args:
            messages: List of message dicts with 'role' and 'content' fields.
                     If provided, medias and text will be extracted automatically.
            medias: Pre-extracted list of media dicts.
            text: Pre-formatted text string.
            return_tensors: Format of returned tensors. Default: 'pt'.
            **kwargs: Additional arguments passed to apply_chat_template.

        Returns:
            BatchFeature with fields: input_ids, attention_mask,
            pixel_values, grid_thws.
        """
        if messages is None and (medias is None or text is None):
            raise ValueError("Provide either 'messages' or both 'medias' and 'text'")

        if medias is not None and text is not None:
            updated_medias, image_prompts = self.preprocess_medias(medias)
            preprocessed = self.media_processor.preprocess(
                updated_medias, return_tensors=return_tensors
            )
            text = self.update_raw_text(text, image_prompts)
            text_inputs = self.tokenizer(text, return_tensors=return_tensors)
            return BatchFeature(data={**text_inputs, **preprocessed.data})

        if medias is None:
            medias = self._extract_medias_from_messages(messages)
        updated_medias, image_prompts = self.preprocess_medias(medias)
        preprocessed = self.media_processor.preprocess(
            updated_medias, return_tensors=return_tensors
        )

        if text is None:
            text_inputs = self.tokenizer.apply_chat_template(
                messages,
                tokenize=True,
                return_tensors=return_tensors,
                return_dict=True,
                image_prompts=image_prompts,
                **kwargs,
            )
            return BatchFeature(data={**text_inputs, **preprocessed.data})

        text = self.update_raw_text(text, image_prompts)
        text_inputs = self.tokenizer(text, return_tensors=return_tensors)
        return BatchFeature(data={**text_inputs, **preprocessed.data})

    @staticmethod
    def _extract_medias_from_messages(messages: list[dict]) -> list[dict]:
        """Extract media items from messages in a single pass."""
        medias = []
        for msg in messages:
            if msg["role"] != "user" or not msg.get("content"):
                continue

            for content_part in msg["content"]:
                if not isinstance(content_part, dict):
                    continue

                content_type = content_part.get("type")
                if content_type in ["image_url", "image"]:
                    image_data = content_part.get(content_type)
                    assert (
                        image_data is not None
                    ), f"image data is missing for content part: {content_part}"
                    medias.append(
                        {
                            "type": "image",
                            "image": image_data,
                        }
                    )
        return medias

    def apply_chat_template(self, messages, **kwargs):
        return self.tokenizer.apply_chat_template(messages, **kwargs)

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

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

    @property
    def model_input_names(self):
        return ["input_ids", "attention_mask", "pixel_values", "grid_thws"]