# Copyright (c) ModelScope Contributors. All rights reserved. import numpy as np import torch from PIL import Image from typing import Any, Dict, List, Literal, Optional, Tuple from ..base import Template from ..constant import MLLMTemplateType from ..register import TemplateMeta, register_template from ..template_inputs import StdTemplateInputs from ..utils import Context, findall from .utils import ChatmlTemplateMeta class MolmoTemplate(Template): placeholder_tokens = [''] def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, inputs: StdTemplateInputs) -> List[Context]: return [] def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: encoded = super()._encode(inputs) # image images_inputs = self.processor.process(images=inputs.images or None, text='') images_input_ids = images_inputs.pop('input_ids').tolist() user_token = self._tokenize(' User') assert len(user_token) == 1 idx = findall(images_input_ids, user_token[0]) assert len(idx) == 1 labels = encoded['labels'] encoded['input_ids'] = images_input_ids[:idx[0]] + encoded['input_ids'] if labels: encoded['labels'] = [-100] * idx[0] + labels if 'images' in images_inputs: images_inputs['images'] = images_inputs['images'].to(self.model_info.torch_dtype) encoded.update(images_inputs) return encoded def generate(self, model, **kwargs): kwargs.pop('attention_mask', None) generation_config = kwargs.pop('generation_config') batch = { k: kwargs.pop(k, None) for k in ['input_ids', 'attention_mask', 'images', 'image_input_idx', 'image_masks'] } return model.generate_from_batch(batch, generation_config, **kwargs) def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: res = super()._data_collator(batch, padding_to=padding_to) # prepare batchfy inputs keys = ['images', 'image_input_idx', 'image_masks'] images_res = self.fetch_inputs(batch, keys) for key in keys: val = images_res.get(key) if val: images_res[key] = torch.stack(val) res.update(images_res) return res register_template( TemplateMeta( MLLMTemplateType.molmo, prefix=[], prompt=[' User: {{QUERY}} Assistant:'], chat_sep=None, suffix=['<|endoftext|>'], template_cls=MolmoTemplate, )) class Molmo2Template(Template): """Molmo2 template for image and video understanding. Uses ChatML format with BOS auto-insertion. Media placeholders (<|image|>, <|video|>) are expanded via _extend_tokens. Video loading/sampling is delegated entirely to processor.video_processor. """ use_model = True placeholder_tokens = [ '<|image|>', '<|video|>', '', '', '', ] def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, inputs: StdTemplateInputs) -> List[Context]: if media_type == 'image': return ['<|image|>'] elif media_type == 'video': return ['<|video|>'] else: raise ValueError(f'Unsupported media_type: {media_type}') def _prepare_mm_inputs(self, inputs: StdTemplateInputs) -> Tuple[Dict[str, Any], List[List[int]], List[List[int]]]: media_inputs: Dict[str, Any] = {} image_expansions: List[List[int]] = [] video_expansions: List[List[int]] = [] tokenizer = self.tokenizer if inputs.images: image_inputs = self.processor.image_processor(images=inputs.images, return_tensors='pt') for image_grid in image_inputs['image_grids']: image_tokens = self.processor.get_image_tokens(image_grid.cpu().numpy()) image_expansions.append(tokenizer.encode(''.join(image_tokens), add_special_tokens=False)) media_inputs.update(image_inputs) if inputs.videos: if len(inputs.videos) != 1: raise ValueError('Molmo2 currently only supports single-video samples.') video_inputs = self.processor.video_processor( videos=inputs.videos, return_tensors='pt', return_metadata=True, ) video_metadata = video_inputs.pop('video_metadata') for video_grid, metadata in zip(video_inputs['video_grids'], video_metadata): video_string = self.processor.get_video_string( video_grid.cpu().numpy(), np.asarray(metadata.timestamps, dtype=np.float32), ) video_expansions.append(tokenizer.encode(video_string, add_special_tokens=False)) media_inputs.update(video_inputs) return media_inputs, image_expansions, video_expansions def _build_token_type_ids(self, input_ids: List[int]) -> List[int]: image_token_ids = {int(token_id) for token_id in getattr(self.processor, 'image_token_ids', [])} return [1 if token_id in image_token_ids else 0 for token_id in input_ids] def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: encoded = super()._encode(inputs) media_inputs, image_expansions, video_expansions = self._prepare_mm_inputs(inputs) input_ids = encoded['input_ids'] labels = encoded['labels'] loss_scale = encoded.get('loss_scale') # Expand image placeholders image_placeholder = self.tokenizer.convert_tokens_to_ids('<|image|>') idx_list = findall(input_ids, image_placeholder) if idx_list: input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, lambda i: image_expansions[i]) # Expand video placeholders video_placeholder = self.tokenizer.convert_tokens_to_ids('<|video|>') idx_list = findall(input_ids, video_placeholder) if idx_list: input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, lambda i: video_expansions[i]) encoded['input_ids'] = input_ids encoded['labels'] = labels encoded['loss_scale'] = loss_scale encoded['token_type_ids'] = self._build_token_type_ids(input_ids) encoded.update(media_inputs) return encoded def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]: res = super()._data_collator_mm_data(batch) for key in ['image_grids', 'video_grids', 'image_token_pooling', 'video_token_pooling', 'image_num_crops']: value = self.concat_tensor(batch, key, 0) if value is not None: res[key] = value return res register_template(ChatmlTemplateMeta( MLLMTemplateType.molmo2, template_cls=Molmo2Template, ))