RedHatAI/Kimi-K3-NVFP4

1"""Kimi-K3 processor: wraps vision processor + tokenizer into a single interface.2 3Chat rendering (including XTML tool-result ordering) is handled by the4tokenizer's Python encoder; this processor adds multimodal media preprocessing.5"""6 7from transformers.feature_extraction_utils import BatchFeature8from transformers.processing_utils import ProcessorMixin9from transformers.utils import logging10 11from .media_utils import ensure_media_type12 13logger = logging.get_logger(__name__)14 15# ── KimiK3Processor ───────────────────────────────────────────────────16 17 18class KimiK3Processor(ProcessorMixin):19    r"""20    Constructs a KimiK3 processor which wraps a KimiK3 image processor21    and a tokenizer into a single processor.22 23    [`KimiK3Processor`] offers all the functionalities of24    [`KimiK3VisionProcessor`] and [`TikTokenTokenizer`].25 26    Args:27        image_processor ([`KimiK3VisionProcessor`], *optional*):28            The image processor is a required input.29        tokenizer ([`TikTokenTokenizer`], *optional*):30            The tokenizer is a required input.31        chat_template (`str`, *optional*): Kept for ProcessorMixin32            compatibility. Kimi K3 chat encoding is implemented in Python by33            the tokenizer.34    """35 36    attributes = ["image_processor", "tokenizer"]37    valid_kwargs = ["chat_template"]38    image_processor_class = "AutoImageProcessor"39    tokenizer_class = "AutoTokenizer"40 41    def __init__(42        self,43        image_processor=None,44        tokenizer=None,45        chat_template=None,46        **kwargs,47    ):48        super().__init__(image_processor,49                         tokenizer,50                         chat_template=chat_template)51        self.media_processor = image_processor52        self.image_placeholder = "<|kimi_image_placeholder|>"53 54    # ── Media preprocessing ────────────────────────────────────────────55 56    def update_raw_text(self, text: str, image_prompts: list[str]) -> str:57        # Replace image placeholders58        image_count = text.count(self.image_placeholder)59        if image_count > 0:60            assert image_count == len(image_prompts), (61                f"image placeholder count {image_count} != "62                f"image_prompts count {len(image_prompts)}")63            text_parts = text.split(self.image_placeholder)64            assert len(text_parts) == len(image_prompts) + 165            text = "".join([66                text_parts[i] + image_prompts[i]67                for i in range(len(image_prompts))68            ])69            text += text_parts[-1]70 71        return text72 73    def preprocess_medias(self,74                          medias: list[dict]) -> tuple[list[dict], list[str]]:75        """Process media items and generate corresponding prompts.76 77        Returns:78            A tuple of (updated_medias, image_prompts).79        """80        updated_medias = []81        image_prompts = []82        for media in medias:83            if media['type'] == 'image':84                updated_medias.append(media)85                img = ensure_media_type(86                    media,87                    transparent_bg_config=self.media_processor.88                    _transparent_bg_config,89                    transparent_bg_fill_stage=self.media_processor.90                    _transparent_bg_fill_stage,91                )['image']92                w, h = img.size93                image_prompts.append(94                    self.media_processor.make_image_prompt(w, h))95            else:96                raise ValueError(f"unsupported media type: {media['type']}")97        return updated_medias, image_prompts98 99    # ── Main entry points ──────────────────────────────────────────────100 101    def __call__(self,102                 messages: list[dict] = None,103                 medias: list[dict] = None,104                 text: str = None,105                 return_tensors: str = "pt",106                 **kwargs) -> BatchFeature:107        """108        Process multimodal inputs for Kimi-K3 model.109 110        Args:111            messages: List of message dicts with 'role' and 'content' fields.112                     If provided, medias and text will be extracted automatically.113            medias: Pre-extracted list of media dicts.114            text: Pre-formatted text string.115            return_tensors: Format of returned tensors. Default: 'pt'.116            **kwargs: Additional arguments passed to apply_chat_template.117 118        Returns:119            BatchFeature with fields: input_ids, attention_mask,120            pixel_values, grid_thws.121        """122        if messages is None and (medias is None or text is None):123            raise ValueError(124                "Provide either 'messages' or both 'medias' and 'text'")125 126        if medias is not None and text is not None:127            updated_medias, image_prompts = (self.preprocess_medias(medias))128            preprocessed = self.media_processor.preprocess(129                updated_medias, return_tensors=return_tensors)130            text = self.update_raw_text(text, image_prompts)131            text_inputs = self.tokenizer(text, return_tensors=return_tensors)132            return BatchFeature(data={**text_inputs, **preprocessed.data})133 134        if medias is None:135            medias = self._extract_medias_from_messages(messages)136        updated_medias, image_prompts = (self.preprocess_medias(medias))137        preprocessed = self.media_processor.preprocess(138            updated_medias, return_tensors=return_tensors)139 140        if text is None:141            text_inputs = self.tokenizer.apply_chat_template(142                messages,143                tokenize=True,144                return_tensors=return_tensors,145                return_dict=True,146                image_prompts=image_prompts,147                **kwargs)148            return BatchFeature(data={**text_inputs, **preprocessed.data})149 150        text = self.update_raw_text(text, image_prompts)151        text_inputs = self.tokenizer(text, return_tensors=return_tensors)152        return BatchFeature(data={**text_inputs, **preprocessed.data})153 154    @staticmethod155    def _extract_medias_from_messages(messages: list[dict]) -> list[dict]:156        """Extract media items from messages in a single pass."""157        medias = []158        for msg in messages:159            if msg['role'] != 'user' or not msg.get('content'):160                continue161 162            for content_part in msg['content']:163                if not isinstance(content_part, dict):164                    continue165 166                content_type = content_part.get('type')167                if content_type in ['image_url', 'image']:168                    image_data = content_part.get(content_type)169                    assert image_data is not None, f"image data is missing for content part: {content_part}"170                    medias.append({171                        'type': 'image',172                        'image': image_data,173                    })174        return medias175 176    def apply_chat_template(self, messages, **kwargs):177        return self.tokenizer.apply_chat_template(messages, **kwargs)178 179    def batch_decode(self, *args, **kwargs):180        return self.tokenizer.batch_decode(*args, **kwargs)181 182    def decode(self, *args, **kwargs):183        return self.tokenizer.decode(*args, **kwargs)184 185    @property186    def model_input_names(self):187        return ['input_ids', 'attention_mask', 'pixel_values', 'grid_thws']188