mirror of
https://github.com/OpenBMB/MiniCPM-V.git
synced 2026-02-05 10:19:18 +08:00
86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
import torch
|
|
from PIL import Image
|
|
from transformers import AutoModel, AutoTokenizer
|
|
|
|
from .base import BaseModel
|
|
from ..smp import *
|
|
from ..utils import DATASET_TYPE
|
|
|
|
|
|
class MiniCPM_V(BaseModel):
|
|
|
|
INSTALL_REQ = False
|
|
INTERLEAVE = False
|
|
|
|
def __init__(self, model_path='openbmb/MiniCPM-V', **kwargs):
|
|
assert model_path is not None
|
|
self.model_path = model_path
|
|
print(f'load from {self.model_path}')
|
|
self.model = AutoModel.from_pretrained(self.model_path, trust_remote_code=True)
|
|
self.model = self.model.to(dtype=torch.bfloat16)
|
|
self.model.eval().cuda()
|
|
self.kwargs = kwargs
|
|
self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
|
|
torch.cuda.empty_cache()
|
|
self.num_beams = 1 if self.model_path == 'openbmb/MiniCPM-V' else 3
|
|
|
|
def use_custom_prompt(self, dataset):
|
|
assert dataset is not None
|
|
if listinstr(['MMMU'], dataset):
|
|
return True
|
|
return False
|
|
|
|
def build_prompt(self, line, dataset=None):
|
|
assert dataset is None or isinstance(dataset, str)
|
|
assert self.use_custom_prompt(dataset)
|
|
tgt_path = self.dump_image(line, dataset)
|
|
|
|
question = line['question']
|
|
options = {
|
|
cand: line[cand]
|
|
for cand in string.ascii_uppercase
|
|
if cand in line and not pd.isna(line[cand])
|
|
}
|
|
options_prompt = 'Options:\n'
|
|
for key, item in options.items():
|
|
options_prompt += f'{key}. {item}\n'
|
|
hint = line['hint'] if ('hint' in line and not pd.isna(line['hint'])) else None
|
|
prompt = ''
|
|
if hint is not None:
|
|
prompt += f'Hint: {hint}\n'
|
|
prompt += f'{question}\n'
|
|
if len(options):
|
|
prompt += options_prompt
|
|
prompt = 'Study the image carefully and pick the option associated with the correct answer. \
|
|
Focus solely on selecting the option and avoid including any other content.\n' + prompt
|
|
message = [dict(type='text', value=prompt)]
|
|
message.extend([dict(type='image', value=p) for p in tgt_path])
|
|
|
|
return message
|
|
|
|
def generate_inner(self, message, dataset=None):
|
|
prompt, image_path = self.message_to_promptimg(message)
|
|
image = Image.open(image_path).convert('RGB')
|
|
msgs = [{'role': 'user', 'content': prompt}]
|
|
if DATASET_TYPE(dataset) == 'multi-choice':
|
|
max_new_tokens = 20
|
|
elif DATASET_TYPE(dataset) == 'Y/N':
|
|
max_new_tokens = 100
|
|
else:
|
|
max_new_tokens = 1024
|
|
|
|
default_kwargs = dict(
|
|
max_new_tokens=max_new_tokens,
|
|
sampling=False,
|
|
num_beams=self.num_beams
|
|
)
|
|
default_kwargs.update(self.kwargs)
|
|
res, _, _ = self.model.chat(
|
|
image=image,
|
|
msgs=msgs,
|
|
context=None,
|
|
tokenizer=self.tokenizer,
|
|
**default_kwargs
|
|
)
|
|
return res
|