mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-03 08:14:29 +08:00
113 lines
4.2 KiB
Python
113 lines
4.2 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
# Adapted from tensorrt_llm/quantization/image_processing.py
|
|
"""Utility classes for image processing."""
|
|
|
|
import torch
|
|
|
|
|
|
class BaseImageProcessor:
|
|
"""Base class for image processors."""
|
|
|
|
def __init__(self, tokenizer, device="auto"):
|
|
"""Constructor."""
|
|
self.tokenizer = tokenizer
|
|
self.device = device
|
|
|
|
def __call__(self, **kwargs):
|
|
"""Call the tokenizer."""
|
|
return self.tokenizer(**kwargs)
|
|
|
|
def preprocess_function(self, examples):
|
|
"""Preprocess function."""
|
|
raise NotImplementedError("Each image processor must implement its own preprocess method")
|
|
|
|
def collate_function(self, examples):
|
|
"""Collate function to process images during data loading."""
|
|
raise NotImplementedError("Each image processor must implement its own collate method")
|
|
|
|
|
|
# A light Encapsulation for Huggingface MllamaImageProcessor
|
|
|
|
|
|
class MllamaImageProcessor(BaseImageProcessor):
|
|
"""Image processor for Mllama."""
|
|
|
|
def preprocess_function(self, examples):
|
|
"""Preprocess function."""
|
|
# Prepare prompts in a generic chat format
|
|
question = examples.get("question", "Describe this image.")
|
|
|
|
if examples["image"] is not None:
|
|
if self.tokenizer.chat_template is not None:
|
|
prompt = self.tokenizer.apply_chat_template(
|
|
[
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "image"}, {"type": "text", "text": question}],
|
|
}
|
|
],
|
|
add_generation_prompt=True,
|
|
)
|
|
else:
|
|
prompt = f"<|image|><|begin_of_text|>{question}"
|
|
|
|
# Process images using the processor's image processor
|
|
values = self.tokenizer(text=prompt, images=examples["image"], return_tensors="pt").to(
|
|
self.device
|
|
)
|
|
else:
|
|
if self.tokenizer.chat_template is not None:
|
|
prompt = self.tokenizer.apply_chat_template(
|
|
[
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "text", "text": question}],
|
|
}
|
|
],
|
|
add_generation_prompt=True,
|
|
)
|
|
else:
|
|
prompt = question
|
|
|
|
values = self.tokenizer(text=prompt, images=None, return_tensors="pt").to(self.device)
|
|
|
|
values["pixel_values"] = None
|
|
values["aspect_ratio_ids"] = None
|
|
values["aspect_ratio_mask"] = None
|
|
values["cross_attention_mask"] = None
|
|
|
|
return values
|
|
|
|
def collate_function(self, batch):
|
|
"""Collate function to process images during data loading."""
|
|
batch[0]["input_ids"] = torch.LongTensor(batch[0]["input_ids"]).to(self.device)
|
|
batch[0]["attention_mask"] = torch.LongTensor(batch[0]["attention_mask"]).to(self.device)
|
|
|
|
if batch[0]["pixel_values"] is not None:
|
|
batch[0]["pixel_values"] = torch.Tensor(batch[0]["pixel_values"]).to(self.device)
|
|
batch[0]["aspect_ratio_ids"] = torch.LongTensor(batch[0]["aspect_ratio_ids"]).to(
|
|
self.device
|
|
)
|
|
batch[0]["aspect_ratio_mask"] = torch.LongTensor(batch[0]["aspect_ratio_mask"]).to(
|
|
self.device
|
|
)
|
|
batch[0]["cross_attention_mask"] = torch.LongTensor(
|
|
batch[0]["cross_attention_mask"]
|
|
).to(self.device)
|
|
|
|
return batch[0]
|