Upload kimi_k25_processor.py with huggingface_hub

kimi_k25_processor.py

@@ -0,0 +1,165 @@
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.processing_utils import ProcessorMixin
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class KimiK25Processor(ProcessorMixin):
+    r"""
+    Constructs a KimiK25 processor which wraps a KimiK25 image processor and a tokenizer into a single processor.
+
+    [`KimiK25Processor`] offers all the functionalities of [`KimiK25ImageProcessor`] and [`TikTokenTokenizer`]. See the
+    [`~KimiK25Processor.__call__`] and [`~KimiK25Processor.decode`] for more information.
+
+    Args:
+        image_processor ([`KimiK25ImageProcessor`], *optional*):
+            The image processor is a required input.
+        tokenizer ([`TikTokenTokenizer`], *optional*):
+            The tokenizer is a required input.
+        chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
+            in a chat into a tokenizable string.
+    """
+
+    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
+        # A special temporal placeholder to be replaced by actual video placeholders
+        self.video_placeholder = "<|kimi_k25_video_placeholder|>"
+
+    def update_raw_text(self, text: str, video_prompts: list[str]) -> str:
+        # replace video prompt in text with video chunk prompts
+        video_count = text.count(self.video_placeholder)
+        if video_count == 0:
+            return text
+        assert video_count == len(video_prompts)
+        text_parts = text.split(self.video_placeholder)
+        assert len(text_parts) == len(video_prompts) + 1
+        text = "".join([
+            text_parts[i] + video_prompts[i] for i in range(len(video_prompts))
+        ])
+        text += text_parts[-1]
+        return text
+
+    def preprocess_medias(self, medias: list[dict]) -> list[dict]:
+        updated_medias = []
+        video_prompts = []
+        for media in medias:
+            if media['type'] == 'image':
+                updated_medias.append(media)
+            elif media['type'] == 'video':
+                video_chunks = self.media_processor.split_video_chunks(
+                    media['video'])
+                updated_medias.extend(video_chunks)
+                video_prompts.append("".join(
+                    [vc['prompt'] for vc in video_chunks]))
+            else:
+                raise ValueError(f"unsupported media type: {media['type']}")
+        return updated_medias, video_prompts
+
+    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-K2.5 model.
+
+        This processor accepts ordered messages and extracts both media and text in a single pass.
+        text will be automatically updated if video input detected in messages
+
+        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. If None, extracted from messages.
+            text: Pre-formatted text string. If None, generated via apply_chat_template.
+            return_tensors: Format of returned tensors ('pt', 'np', 'tf'). Default: 'pt'.
+            **kwargs: Additional arguments passed to tokenizer.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, video_prompts = self.preprocess_medias(medias)
+            preprocessed = self.media_processor.preprocess(
+                updated_medias, return_tensors=return_tensors)
+            text = self.update_raw_text(text, video_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, video_prompts = self.preprocess_medias(medias)
+        preprocessed = self.media_processor.preprocess(
+            updated_medias, return_tensors=return_tensors)
+
+        # Generate text if not provided
+        if text is None:
+            text = self.tokenizer.apply_chat_template(messages, **kwargs)
+
+        text = self.update_raw_text(text, video_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.
+        
+        This is an optimized version that processes messages only once.
+        Kept as internal method since external callers should use __call__.
+        """
+        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 ['video_url', 'video']:
+                    medias.append({
+                        'type': 'video',
+                        'video': content_part['video_url']['url'],
+                        'first_frame_timestamp': 0.0
+                    })
+                elif content_type in ['image_url', 'image']:
+                    medias.append({
+                        'type': 'image',
+                        'image': content_part['image_url'],
+                    })
+        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']