Flare77/HuLuLLM
013
1"""Processor class for HuluMed with 3D support."""2 3import copy4import importlib.util5import os6import os.path as osp7import warnings8from collections import defaultdict9from typing import Any, List, Union, Dict, Optional, Tuple, TypedDict10 11import cv212import ffmpeg13import imageio14import json15import numpy as np16import torch17import transformers18from decord import VideoReader, cpu19from PIL import Image20from transformers.feature_extraction_utils import BatchFeature21from transformers.image_utils import ImageInput22from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack23from transformers.tokenization_utils_base import PreTokenizedInput, TextInput24 25try:26 import nibabel as nib27 NIBABEL_AVAILABLE = True28except ImportError:29 NIBABEL_AVAILABLE = False30 warnings.warn("nibabel is not installed. 3D medical imaging support will be limited. Install with: pip install nibabel")31 32try:33 from . import image_processing_hulumed34 from .image_processing_hulumed import (35 is_valid_image, is_valid_video,36 )37except ModuleNotFoundError:38 spec = importlib.util.spec_from_file_location(39 "image_processing_hulumed",40 osp.join(osp.dirname(__file__), "image_processing_hulumed.py"),41 )42 image_processing_hulumed = importlib.util.module_from_spec(spec)43 spec.loader.exec_module(image_processing_hulumed)44 is_valid_image = getattr(image_processing_hulumed, "is_valid_image")45 is_valid_video = getattr(image_processing_hulumed, "is_valid_video")46 47DEFAULT_IMAGE_TOKEN = "<image>"48IGNORE_INDEX = -10049 50Conversation = List[Dict[str, Any]]51SingleImage = Union[Image.Image, np.ndarray, torch.Tensor]52SingleVideo = Union[List[SingleImage], np.ndarray, torch.Tensor]53BatchedImage = List[Union[SingleImage, SingleVideo]]54BatchedNamedImage = List[Tuple[str, Union[SingleImage, SingleVideo]]]55 56 57def _custom_import(class_name: str):58 try:59 attribute_class = getattr(transformers, class_name)60 except AttributeError:61 attribute_class = getattr(image_processing_hulumed, class_name)62 return attribute_class63 64 65def is_named_image(image) -> bool:66 return isinstance(image, (list, tuple)) and \67 len(image) == 2 and \68 isinstance(image[0], str) and \69 image[0] in ["image", "video", "3d"] and \70 (is_valid_image(image[1]) or is_valid_video(image[1]))71 72 73def make_batched_images(images) -> Tuple[List[str], List[ImageInput]]:74 if isinstance(images, (list, tuple)) and all(is_named_image(image) for image in images):75 modals = [image[0] if image[0] != "3d" else "video" for image in images]76 data = [image[1] for image in images]77 return modals, data78 elif isinstance(images, (list, tuple)) and all(is_valid_image(image) or is_valid_video(image) for image in images):79 batch = []80 for image in images:81 if is_valid_video(image):82 batch.append(("video", image))83 elif is_valid_image(image):84 batch.append(("image", image))85 else:86 raise ValueError(f"Could not make batched images from {images}")87 return [x[0] for x in batch], [x[1] for x in batch]88 elif is_named_image(images):89 modal = images[0] if images[0] != "3d" else "video"90 return [modal], [images[1]]91 elif is_valid_video(images):92 return ["video"], [images]93 elif is_valid_image(images):94 return ["image"], [images]95 96 raise ValueError(f"Could not make batched images from {images}")97 98 99def frame_sample(duration, mode='uniform', num_frames=None, vid_fps=None, fps=None):100 if mode == 'uniform':101 assert num_frames is not None, "Number of frames must be provided for uniform sampling."102 if duration <= num_frames:103 return np.arange(duration).astype(int)104 return np.linspace(0, duration-1, num_frames, dtype=int)105 elif mode == 'fps':106 assert vid_fps is not None, "FPS must be provided for FPS sampling."107 assert fps is not None, "FPS must be provided for FPS sampling."108 segment_len = min(vid_fps // fps, duration)109 return np.arange(segment_len // 2, duration, segment_len, dtype=int)110 else:111 raise ValueError(f'Unsupported frame sampling mode: {mode}')112 113 114def load_video_from_ids(video_path, s=None, e=None, fps=None, max_frames=128, temporal_factor=1):115 if s is not None and e is not None:116 s = s if s >= 0. else 0.117 e = e if e >= 0. else 0.118 if s > e:119 s, e = e, s120 elif s == e:121 e = s + 1122 123 if os.path.isdir(video_path):124 frame_files = sorted(os.listdir(video_path))125 vid_fps = 3126 num_frames_of_video = len(frame_files)127 elif video_path.endswith('.gif'):128 gif_reader = imageio.get_reader(video_path)129 vid_fps = 25130 num_frames_of_video = len(gif_reader)131 else:132 vreader = VideoReader(video_path, ctx=cpu(0), num_threads=2)133 vid_fps = vreader.get_avg_fps()134 num_frames_of_video = len(vreader)135 136 f_start = 0 if s is None else max(int(s * vid_fps) - 1, 0)137 f_end = num_frames_of_video - 1 if e is None else min(int(e * vid_fps) - 1, num_frames_of_video - 1)138 frame_indices = list(range(f_start, f_end + 1))139 140 duration = len(frame_indices)141 if fps is not None and duration / vid_fps < max_frames:142 sampled_frame_indices = [frame_indices[i] for i in frame_sample(duration, mode='fps', vid_fps=vid_fps, fps=fps)]143 else:144 sampled_frame_indices = [frame_indices[i] for i in frame_sample(duration, mode='uniform', num_frames=max_frames)]145 146 if os.path.isdir(video_path):147 frames = np.array([cv2.cvtColor(cv2.imread(os.path.join(video_path, frame_files[frame_idx])), cv2.COLOR_BGR2RGB) for frame_idx in sampled_frame_indices])148 elif video_path.endswith('.gif'):149 frames = np.array([cv2.cvtColor(frame, cv2.COLOR_RGBA2RGB) for idx, frame in enumerate(gif_reader) if idx in sampled_frame_indices])150 else:151 frames = vreader.get_batch(sampled_frame_indices).asnumpy()152 153 frames = frames.transpose(0, 3, 1, 2)154 timestamps = [x / vid_fps for x in sampled_frame_indices]155 156 if temporal_factor > 1:157 pad_length = temporal_factor - len(frames) % temporal_factor158 frames = np.concatenate([frames, frames[-1:].repeat(pad_length, axis=0)])159 [timestamps.append(timestamps[-1] + 1 / fps) for _ in range(pad_length)]160 161 frames = [frame for frame in frames]162 163 return frames, timestamps164 165 166class ChatTemplateKwargs(TypedDict, total=False):167 chat_template: Optional[str]168 add_system_prompt: Optional[bool]169 add_generation_prompt: Optional[bool]170 171 172class HulumedProcessorKwargs(ProcessingKwargs, ChatTemplateKwargs, total=False):173 chat_template_kwargs: ChatTemplateKwargs = {174 **ChatTemplateKwargs.__annotations__,175 }176 177 _defaults = {178 "text_kwargs": {179 "padding": False,180 },181 "images_kwargs": {182 183 },184 "chat_template_kwargs": {185 "chat_template": None,186 "add_system_prompt": False,187 "add_generation_prompt": False,188 },189 }190 191 192class HulumedProcessor(ProcessorMixin):193 attributes = ["image_processor", "tokenizer"]194 image_processor_class = "HulumedImageProcessor"195 tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")196 valid_kwargs = ["chat_template", "image_merge_size", "video_merge_size", "fps", "max_frames"]197 198 def __init__(199 self,200 image_processor=None,201 tokenizer=None,202 chat_template: str = None,203 image_merge_size: int = 1,204 video_merge_size: int = 2,205 fps: Optional[int] = 1,206 max_frames: Optional[int] = 128,207 ):208 self.image_processor = image_processor209 self.tokenizer = tokenizer210 if chat_template is None:211 chat_template = self.tokenizer.chat_template212 self.chat_template = chat_template213 214 self.image_merge_size = image_merge_size215 self.video_merge_size = video_merge_size216 self.fps = fps217 self.max_frames = max_frames218 219 self.generation_prompt = self._infer_generation_prompt()220 self.generation_prompt_ids = self.tokenizer.encode(self.generation_prompt, return_tensors="pt")221 self.generation_prompt_length = len(self.generation_prompt_ids[0])222 self.image_token_id = self.tokenizer.convert_tokens_to_ids(DEFAULT_IMAGE_TOKEN)223 self.eos_token_id = self.tokenizer.eos_token_id224 225 @classmethod226 def _get_arguments_from_pretrained(cls, pretrained_model_name_or_path, **kwargs):227 args = []228 for attribute_name in cls.attributes:229 class_name = getattr(cls, f"{attribute_name}_class")230 if isinstance(class_name, tuple):231 classes = tuple(_custom_import(n) if n is not None else None for n in class_name)232 use_fast = kwargs.get("use_fast", True)233 if use_fast and classes[1] is not None:234 attribute_class = classes[1]235 else:236 attribute_class = classes[0]237 else:238 attribute_class = _custom_import(class_name)239 240 args.append(attribute_class.from_pretrained(pretrained_model_name_or_path, **kwargs))241 return args242 243 def get_generation_prompt(self):244 return self.generation_prompt245 246 def get_generation_prompt_ids(self):247 return self.generation_prompt_ids248 249 def _infer_generation_prompt(self):250 pseudo_message = [{"role": "user", "content": ""}]251 instruction = self.apply_chat_template(pseudo_message, tokenize=False, add_generation_prompt=True)252 conversation = self.apply_chat_template(pseudo_message, tokenize=False, add_generation_prompt=False)253 return instruction.replace(conversation, "")254 255 def _get_downsampled_grid_sizes(self, image_inputs: Dict[str, Any]):256 grid_sizes = []257 for grid_size, merge_size in zip(image_inputs.get("grid_sizes", []), image_inputs.get("merge_sizes", [])):258 if not torch.all(grid_size[1:] % merge_size == 0):259 warnings.warn(f"Grid size {grid_size} is not divisible by merge size. Some undesired errors may occur.")260 if grid_size[0] == 1:261 grid_sizes.append(grid_size[1:] / merge_size)262 elif grid_size[0] > 1:263 grid_sizes.extend([grid_size[1:] / merge_size] * grid_size[0])264 return grid_sizes265 266 def _get_visual_seq_len(self, grid_size: torch.Tensor):267 num_tokens = int(grid_size.prod().item())268 return num_tokens269 270 def load_images(self, image_path: Union[str, List[str], Image.Image, List[Image.Image]]):271 if isinstance(image_path, str) and os.path.isfile(image_path):272 images = [Image.open(image_path).convert('RGB')]273 elif isinstance(image_path, str) and os.path.isdir(image_path):274 images = [Image.open(os.path.join(image_path, f)).convert('RGB') for f in sorted(os.listdir(image_path))]275 elif isinstance(image_path, list) and isinstance(image_path[0], str):276 images = [Image.open(f).convert('RGB') for f in image_path]277 elif isinstance(image_path, list) and isinstance(image_path[0], Image.Image):278 images = [np.array(x) for x in image_path]279 elif isinstance(image_path, Image.Image):280 images = [np.array(image_path)]281 else:282 raise ValueError(f"Unsupported image path type: {type(image_path)}")283 return images284 285 def load_nii(286 self,287 nii_path: str,288 num_slices: Optional[int] = None,289 axis: int = 2,290 window_center: Optional[float] = None,291 window_width: Optional[float] = None,292 normalize: bool = True,293 ):294 if not NIBABEL_AVAILABLE:295 raise ImportError("nibabel is required for NIfTI support. Install with: pip install nibabel")296 297 if not os.path.exists(nii_path):298 raise FileNotFoundError(f"NIfTI file not found: {nii_path}")299 300 nii_img = nib.load(nii_path)301 volume = nii_img.get_fdata()302 303 if axis == 0:304 slices = [volume[i, :, :] for i in range(volume.shape[0])]305 elif axis == 1:306 slices = [volume[:, i, :] for i in range(volume.shape[1])]307 elif axis == 2:308 slices = [volume[:, :, i] for i in range(volume.shape[2])]309 else:310 raise ValueError(f"Invalid axis: {axis}. Must be 0, 1, or 2.")311 312 if num_slices is not None and num_slices < len(slices):313 indices = np.linspace(0, len(slices) - 1, num_slices, dtype=int)314 slices = [slices[i] for i in indices]315 316 processed_slices = []317 for slice_2d in slices:318 if window_center is not None and window_width is not None:319 lower = window_center - window_width / 2320 upper = window_center + window_width / 2321 slice_2d = np.clip(slice_2d, lower, upper)322 323 if normalize:324 slice_min = slice_2d.min()325 slice_max = slice_2d.max()326 if slice_max > slice_min:327 slice_2d = (slice_2d - slice_min) / (slice_max - slice_min) * 255.0328 else:329 slice_2d = np.zeros_like(slice_2d)330 331 slice_2d = slice_2d.astype(np.uint8)332 slice_rgb = np.stack([slice_2d] * 3, axis=0)333 334 processed_slices.append(slice_rgb)335 336 return processed_slices337 338 def load_video(339 self,340 video_path: str,341 start_time: Optional[float] = None,342 end_time: Optional[float] = None,343 fps: Optional[float] = None,344 max_frames: Optional[float] = None,345 size: Optional[int] = None,346 size_divisible: int = 1,347 precise_time: bool = False,348 verbose: bool = False,349 temporal_factor: int = 1350 ):351 fps = self.fps if fps is None else fps352 max_frames = self.max_frames if max_frames is None else max_frames353 354 if start_time is not None and end_time is not None and end_time - start_time < 1:355 return load_video_from_ids(video_path, start_time, end_time, fps=fps, max_frames=max_frames)356 if os.path.isdir(video_path):357 return load_video_from_ids(video_path, start_time, end_time, fps=fps, max_frames=max_frames)358 if video_path.endswith('.gif'):359 return load_video_from_ids(video_path, start_time, end_time, fps=fps, max_frames=max_frames)360 361 probe = ffmpeg.probe(video_path)362 duration = float(probe['format']['duration'])363 video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)364 w, h = int(video_stream['width']), int(video_stream['height'])365 366 kwargs, input_kwargs, output_kwargs = {}, {}, {}367 do_trim = start_time is not None or end_time is not None368 if start_time is not None:369 new_start_time = max(float(video_stream['start_time']), start_time)370 duration -= new_start_time - start_time371 start_time = new_start_time372 else:373 start_time = float(video_stream['start_time'])374 if end_time is not None:375 duration = min(duration, end_time - start_time)376 if do_trim:377 kwargs = {'ss': start_time, 't': duration}378 if precise_time:379 output_kwargs.update(kwargs)380 else:381 input_kwargs.update(kwargs)382 383 if size is not None:384 scale_factor = size / min(w, h)385 new_w, new_h = round(w * scale_factor), round(h * scale_factor)386 else:387 new_w, new_h = w, h388 new_w = new_w // size_divisible * size_divisible389 new_h = new_h // size_divisible * size_divisible390 391 stream = ffmpeg.input(video_path, **input_kwargs)392 if fps is not None:393 stream = ffmpeg.filter(stream, "fps", fps=fps, round="down")394 if new_w != w or new_h != h:395 stream = ffmpeg.filter(stream, 'scale', new_w, new_h)396 stream = ffmpeg.output(stream, "pipe:", format="rawvideo", pix_fmt="rgb24", **output_kwargs)397 out, _ = ffmpeg.run(stream, capture_stdout=True, quiet=not verbose)398 399 frames = np.frombuffer(out, np.uint8).reshape([-1, new_h, new_w, 3]).transpose([0, 3, 1, 2])400 401 if fps is not None:402 timestamps = np.arange(start_time, start_time + duration + 1 / fps, 1 / fps)[:len(frames)]403 else:404 timestamps = np.linspace(start_time, start_time + duration, len(frames))405 406 if max_frames is not None and len(frames) > max_frames:407 indices = np.linspace(0, len(frames) - 1, max_frames, dtype=int)408 frames = frames[indices]409 timestamps = timestamps[indices]410 411 if temporal_factor > 1:412 pad_length = temporal_factor - len(frames) % temporal_factor413 frames = np.concatenate([frames, frames[-1:].repeat(pad_length, axis=0)])414 timestamps = np.concatenate([timestamps, timestamps[-1:].repeat(pad_length) + np.arange(1, pad_length + 1) / fps])415 416 frames = [frame for frame in frames]417 timestamps = [timestamp for timestamp in timestamps]418 419 return frames, timestamps420 421 def _load_multimodal_data(self, conversation: Conversation):422 multimodal_info = defaultdict(list)423 new_conversation = []424 for message in conversation:425 new_message = {"role": message["role"]}426 if not isinstance(message["content"], (list, tuple)):427 new_message["content"] = message["content"]428 new_conversation.append(new_message)429 continue430 431 new_contents = []432 for content in message["content"]:433 if not isinstance(content, dict):434 new_contents.append(content)435 continue436 assert "type" in content, "Content must have 'type' field."437 438 if content["type"] in ["image", "video", "3d"] and content["type"] in content and isinstance(content[content["type"]], dict):439 load_args = content[content["type"]]440 data_id = json.dumps({k: v for k, v in load_args.items() if k not in ["start_time", "end_time"]})441 new_content = copy.deepcopy(content)442 multimodal_info[data_id].append(new_content)443 new_contents.append(new_content)444 else:445 new_contents.append(content)446 447 new_message["content"] = new_contents448 new_conversation.append(new_message)449 450 for data_id, contents in multimodal_info.items():451 data_type = contents[0]["type"]452 453 if data_type == "image":454 image = self.load_images(contents[0][data_type]["image_path"])[0]455 for content in contents:456 content["image"] = [image.copy()]457 458 elif data_type == "3d":459 load_args = contents[0]["3d"]460 nii_path = load_args["image_path"]461 num_slices = load_args.get("nii_num_slices", None)462 axis = load_args.get("nii_axis", 2)463 window_center = load_args.get("window_center", None)464 window_width = load_args.get("window_width", None)465 466 slices = self.load_nii(467 nii_path=nii_path,468 num_slices=num_slices,469 axis=axis,470 window_center=window_center,471 window_width=window_width,472 )473 474 for content in contents:475 content["type"] = "video"476 content["video"] = slices477 content["num_frames"] = len(slices)478 content.pop("3d", None)479 480 elif data_type == "video":481 start_times = [content["video"].get("start_time", 0.) for content in contents]482 end_times = [content["video"].get("end_time", float("inf")) for content in contents]483 484 load_args = contents[0][data_type]485 start_time, end_time = min(start_times), max(end_times)486 if start_time > 0:487 load_args["start_time"] = start_time488 if end_time < float("inf"):489 load_args["end_time"] = end_time490 images, timestamps = self.load_video(**load_args)491 492 for content, start_time, end_time in zip(contents, start_times, end_times):493 cur_images, cur_timestamps = [], []494 for image, timestamp in zip(images, timestamps):495 if start_time <= timestamp <= end_time:496 cur_images.append(image.copy())497 cur_timestamps.append(timestamp)498 499 content[data_type] = cur_images500 content["num_frames"] = len(cur_images)501 content["timestamps"] = cur_timestamps502 503 return new_conversation504 505 def _gather_multimodal_data(self, conversation: Conversation):506 images = []507 for message in conversation:508 if not isinstance(message["content"], (list, tuple)):509 continue510 for content in message["content"]:511 if not isinstance(content, dict):512 continue513 if content["type"] == "video":514 video = content["video"]515 assert is_valid_video(video), f"Invalid video data: {video}."516 images.append(("video", video))517 elif content["type"] == "image":518 image = content["image"]519 images.append(("image", image))520 images = images if len(images) > 0 else None521 return images522 523 def _process_conversation_with_label(524 self,525 conversation: Conversation,526 image_inputs: Dict[str, Any],527 **kwargs,528 ):529 assert kwargs.pop("return_tensors", "pt") == "pt", "Only PyTorch tensors are supported when return_labels=True."530 assert "add_generation_prompt" not in kwargs, "'add_generation_prompt' argument is not supported when return_labels=True."531 532 output_kwargs = self._merge_kwargs(533 HulumedProcessorKwargs,534 tokenizer_init_kwargs=self.tokenizer.init_kwargs,535 **kwargs,536 )537 output_kwargs["chat_template_kwargs"].pop("add_generation_prompt")538 539 grid_sizes = self._get_downsampled_grid_sizes(image_inputs)540 text_inputs = {"input_ids": [], "labels": []}541 sample_types_list = []542 image_idx = 0543 544 for message_idx, message in enumerate(conversation):545 prompt = self.apply_chat_template(546 [message],547 tokenize=False,548 add_generation_prompt=False,549 **output_kwargs["chat_template_kwargs"],550 )551 prompt_chunks = prompt.split(DEFAULT_IMAGE_TOKEN)552 prompt = []553 for chunk_idx in range(len(prompt_chunks) - 1):554 prompt.append(prompt_chunks[chunk_idx])555 num_tokens = self._get_visual_seq_len(grid_sizes[image_idx])556 prompt.append(DEFAULT_IMAGE_TOKEN * num_tokens)557 image_idx += 1558 prompt.append(prompt_chunks[-1])559 prompt = "".join(prompt)560 561 input_ids = self.tokenizer.encode(prompt, return_tensors="pt", **output_kwargs["text_kwargs"])[0]562 text_inputs["input_ids"].append(input_ids)563 564 targets = torch.full_like(input_ids, IGNORE_INDEX)565 sample_types = torch.full_like(input_ids, IGNORE_INDEX)566 if message["role"] == "assistant":567 targets[self.generation_prompt_length:-1] = input_ids[self.generation_prompt_length:-1].clone()568 elif message["role"] == "stream":569 diff = torch.diff((input_ids == self.image_token_id).float())570 image_end_indices = torch.nonzero(diff < 0)[:, 0]571 targets[image_end_indices + 1] = input_ids[image_end_indices + 1]572 sample_types = targets.clone()573 sample_types[torch.logical_and(sample_types > 0, sample_types != self.eos_token_id)] = 0574 targets[-2] = input_ids[-2]575 576 if message_idx > 0 and conversation[message_idx - 1]["role"] == "stream":577 targets[0] = input_ids[0]578 sample_types[0] = input_ids[0]579 580 text_inputs["labels"].append(targets)581 sample_types_list.append(sample_types)582 583 text_inputs = {k: torch.cat(v) for k, v in text_inputs.items()}584 sample_types = torch.cat(sample_types_list)585 types, counts = torch.unique(sample_types[sample_types > -1], return_counts=True)586 587 if len(types) > 0:588 target_num_samples = counts.amin()589 for type_id, type_count in zip(types, counts):590 if type_count > target_num_samples:591 indices = torch.nonzero(sample_types == type_id)[:, 0]592 random_selector = torch.randperm(indices.size(0))[:-target_num_samples]593 text_inputs["labels"][indices[random_selector]] = IGNORE_INDEX594 595 assert len(grid_sizes) == image_idx, "Number of images does not match the number of image tokens in the text."596 597 return text_inputs598 599 def _process_conversation_without_label(600 self,601 conversation: Conversation,602 image_inputs: Dict[str, Any],603 **kwargs,604 ):605 output_kwargs = self._merge_kwargs(606 HulumedProcessorKwargs,607 tokenizer_init_kwargs=self.tokenizer.init_kwargs,608 **kwargs,609 )610 prompt = self.apply_chat_template(611 conversation,612 tokenize=False,613 **output_kwargs["chat_template_kwargs"],614 )615 return self.process_text(prompt, image_inputs, **output_kwargs["text_kwargs"])616 617 def _process_conversation(618 self,619 conversation: Conversation,620 images: Optional[Union[BatchedImage, BatchedNamedImage]] = None,621 return_labels: bool = False,622 **kwargs: Unpack[HulumedProcessorKwargs],623 ) -> BatchFeature:624 assert isinstance(conversation, list), "Conversation must be a list of messages."625 626 if images is None:627 conversation = self._load_multimodal_data(conversation)628 images = self._gather_multimodal_data(conversation)629 630 if not images: 631 images = None632 elif isinstance(images, (list, tuple)):633 images = [img for img in images if img and (not isinstance(img, (list, tuple)) or len(img) > 0)]634 if not images:635 images = None636 output_kwargs = self._merge_kwargs(637 HulumedProcessorKwargs,638 tokenizer_init_kwargs=self.tokenizer.init_kwargs,639 **kwargs,640 )641 642 if images is not None:643 if "merge_size" not in output_kwargs["images_kwargs"]:644 has_video_or_3d = any(645 content.get("type") in ["video", "3d"] or "video" in content or "3d" in content646 for message in conversation647 if isinstance(message.get("content"), list)648 for content in message["content"]649 if isinstance(content, dict)650 )651 652 output_kwargs["images_kwargs"]["merge_size"] = 2 if has_video_or_3d else 1653 654 image_inputs = self.process_images(images, **output_kwargs["images_kwargs"])655 else:656 image_inputs = {}657 658 if return_labels:659 text_inputs = self._process_conversation_with_label(conversation, image_inputs, **kwargs)660 else:661 text_inputs = self._process_conversation_without_label(conversation, image_inputs, **kwargs)662 663 return BatchFeature(data={**text_inputs, **image_inputs})664 665 def _process_plain(666 self,667 text: Union[TextInput, PreTokenizedInput] = None,668 images: Optional[Union[BatchedImage, BatchedNamedImage]] = None,669 return_labels: bool = False,670 **kwargs: Unpack[HulumedProcessorKwargs],671 ) -> BatchFeature:672 if text is None:673 raise ValueError("You must provide 'text' or 'conversation'.")674 if return_labels:675 raise ValueError("return_labels is not supported for plain text processing.")676 677 output_kwargs = self._merge_kwargs(678 HulumedProcessorKwargs,679 tokenizer_init_kwargs=self.tokenizer.init_kwargs,680 **kwargs,681 )682 683 if images is not None:684 image_inputs = self.process_images(images, **output_kwargs["images_kwargs"])685 else:686 image_inputs = {}687 688 text_inputs = self.process_text(text, image_inputs, **output_kwargs["text_kwargs"])689 690 return BatchFeature(data={**text_inputs, **image_inputs})691 692 def process_images(self, images: Union[BatchedImage, BatchedNamedImage], **kwargs):693 modals, images = make_batched_images(images)694 695 if "merge_size" not in kwargs:696 kwargs["merge_size"] = [697 self.video_merge_size if modal == "video" else self.image_merge_size698 for modal in modals699 ]700 701 image_inputs = self.image_processor(images=images, **kwargs)702 image_inputs["modals"] = modals703 return image_inputs704 705 def process_text(706 self,707 text: TextInput,708 image_inputs: Dict[str, Any],709 **kwargs,710 ):711 grid_sizes = self._get_downsampled_grid_sizes(image_inputs)712 713 kwargs.pop("padding", None)714 kwargs.pop("padding_side", None)715 716 if len(grid_sizes) > 0:717 image_idx = 0718 while DEFAULT_IMAGE_TOKEN in text:719 num_tokens = self._get_visual_seq_len(grid_sizes[image_idx])720 text = text.replace(DEFAULT_IMAGE_TOKEN, "<placeholder>" * num_tokens, 1)721 image_idx += 1722 text = text.replace("<placeholder>", DEFAULT_IMAGE_TOKEN)723 724 assert len(grid_sizes) == image_idx, "Number of images does not match the number of image tokens in the text."725 726 text_inputs = self.tokenizer(text, **kwargs)727 return text_inputs728 729 def __call__(730 self,731 text: Optional[TextInput] = None,732 conversation: Optional[Conversation] = None,733 images: Optional[Union[BatchedImage, BatchedNamedImage]] = None,734 return_labels: bool = False,735 **kwargs: Unpack[HulumedProcessorKwargs],736 ) -> BatchFeature:737 if conversation is not None:738 if text is not None:739 raise ValueError("You cannot provide both 'conversation' and 'text'.")740 return self._process_conversation(conversation, images, return_labels, **kwargs)741 return self._process_plain(text, images, return_labels, **kwargs)742 743 def batch_decode(self, *args, skip_special_tokens=True, use_think=False, **kwargs):744 outputs = self.tokenizer.batch_decode(*args, skip_special_tokens=skip_special_tokens, **kwargs)745 746 if not use_think:747 outputs = [self._remove_think_tags(output) for output in outputs]748 749 return outputs750 751 def decode(self, *args, skip_special_tokens=True, use_think=False, **kwargs):752 output = self.tokenizer.decode(*args, skip_special_tokens=skip_special_tokens, **kwargs)753 754 if not use_think:755 output = self._remove_think_tags(output)756 757 return output758 759 def _remove_think_tags(self, text: str) -> str:760 import re761 pattern = r'<think>.*?</think>'762 cleaned = re.sub(pattern, '', text, flags=re.DOTALL)763 cleaned = re.sub(r'\n\s*\n', '\n\n', cleaned)764 cleaned = cleaned.strip()765 return cleaned766 767 def apply_chat_template(768 self,769 conversation: Conversation,770 chat_template: Optional[str] = None,771 tokenize: bool = False,772 add_system_prompt: bool = False,773 add_generation_prompt: bool = False,774 image_token: Optional[str] = DEFAULT_IMAGE_TOKEN,775 **kwargs,776 ) -> str:777 if chat_template is None:778 if self.chat_template is not None:779 chat_template = self.chat_template780 else:781 raise ValueError(782 "No chat template is set for this processor. Please either set the `chat_template` attribute, "783 "or provide a chat template as an argument."784 )785 return self.tokenizer.apply_chat_template(786 conversation,787 chat_template=chat_template,788 tokenize=tokenize,789 add_system_prompt=add_system_prompt,790 add_generation_prompt=add_generation_prompt,791 image_token=image_token,792 **kwargs793 )794 795 @property796 def model_input_names(self):797 tokenizer_input_names = self.tokenizer.model_input_names798 image_processor_input_names = self.image_processor.model_input_names799 return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names)) + ["modals"]800 801 def _merge_kwargs(802 self,803 ModelProcessorKwargs: ProcessingKwargs,804 tokenizer_init_kwargs: Optional[Dict] = None,805 **kwargs,806 ) -> Dict[str, Dict]:807 output_kwargs = {808 "text_kwargs": {},809 "images_kwargs": {},810 "audio_kwargs": {},811 "videos_kwargs": {},812 "chat_template_kwargs": {},813 "common_kwargs": {},814 }815 816 default_kwargs = {817 "text_kwargs": {},818 "images_kwargs": {},819 "audio_kwargs": {},820 "videos_kwargs": {},821 "chat_template_kwargs": {},822 "common_kwargs": {},823 }824 825 used_keys = set()826 827 for modality in default_kwargs:828 default_kwargs[modality] = ModelProcessorKwargs._defaults.get(modality, {}).copy()829 for modality_key in ModelProcessorKwargs.__annotations__[modality].__annotations__.keys():830 if modality_key in tokenizer_init_kwargs:831 value = (832 getattr(self.tokenizer, modality_key)833 if hasattr(self.tokenizer, modality_key)834 else tokenizer_init_kwargs[modality_key]835 )836 default_kwargs[modality][modality_key] = value837 838 output_kwargs.update(default_kwargs)839 840 non_modality_kwargs = set(kwargs) - set(output_kwargs)841 for modality in output_kwargs:842 for modality_key in ModelProcessorKwargs.__annotations__[modality].__annotations__.keys():843 if modality in kwargs:844 kwarg_value = kwargs[modality].pop(modality_key, "__empty__")845 if kwarg_value != "__empty__" and modality_key in non_modality_kwargs:846 raise ValueError(847 f"Keyword argument {modality_key} was passed twice: "848 f"in a dictionary for {modality} and as a **kwarg."849 )850 elif modality_key in kwargs:851 kwarg_value = kwargs.get(modality_key, "__empty__")852 else:853 kwarg_value = "__empty__"854 if kwarg_value != "__empty__":855 output_kwargs[modality][modality_key] = kwarg_value856 used_keys.add(modality_key)857 858 if any(key in default_kwargs for key in kwargs):859 for modality, subdict in kwargs.items():860 if modality in default_kwargs:861 for subkey, subvalue in subdict.items():862 if subkey not in used_keys:863 output_kwargs[modality][subkey] = subvalue864 used_keys.add(subkey)865 else:866 for key in kwargs:867 if key not in used_keys:868 output_kwargs["common_kwargs"][key] = kwargs[key]869 870 for modality in output_kwargs:871 output_kwargs[modality].update(output_kwargs["common_kwargs"])872 873 return output_kwargs