|
|
| import torch |
| from PIL import Image |
| from transformers import AutoModel, AutoTokenizer |
| import random |
| import math |
| import numpy as np |
|
|
| Image.MAX_IMAGE_PIXELS = 1000000000 |
|
|
| max_token = { |
| 'docVQA': 100, |
| 'textVQA': 100, |
| "docVQATest": 100 |
| } |
|
|
| class MiniCPM_V: |
|
|
| def __init__(self, model_path, ckpt, device=None)->None: |
| self.model_path = model_path |
| self.ckpt = ckpt |
| self.model = AutoModel.from_pretrained(self.model_path, trust_remote_code=True).eval() |
| if self.ckpt is not None: |
| self.ckpt = ckpt |
| self.state_dict = torch.load(self.ckpt, map_location=torch.device('cpu')) |
| self.model.load_state_dict(self.state_dict) |
| |
| self.model = self.model.to(dtype=torch.float16) |
| self.model.to(device) |
| |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True) |
| torch.cuda.empty_cache() |
|
|
| def generate(self, images, questions, datasetname): |
| image = Image.open(images[0]).convert('RGB') |
| try: |
| max_new_tokens = max_token[datasetname] |
| except: |
| max_new_tokens = 1024 |
| if (datasetname == 'docVQA') or (datasetname == "docVQATest") : |
| prompt = "Answer the question directly with single word." + "\n" + questions[0] |
| elif (datasetname == 'textVQA') : |
| prompt = "Answer the question directly with single word." + '\n'+ questions[0] |
| |
| msgs = [{'role': 'user', 'content': prompt}] |
| default_kwargs = dict( |
| max_new_tokens=max_new_tokens, |
| sampling=False, |
| num_beams=3 |
| ) |
| res = self.model.chat( |
| image=image, |
| msgs=msgs, |
| context=None, |
| tokenizer=self.tokenizer, |
| **default_kwargs |
| ) |
| |
| return [res] |
| |
| def generate_with_interleaved(self, images, questions, datasetname): |
| try: |
| max_new_tokens = max_token[datasetname] |
| except: |
| max_new_tokens = 1024 |
| |
| prompt = "Answer the question directly with single word." |
| |
| default_kwargs = dict( |
| max_new_tokens=max_new_tokens, |
| sampling=False, |
| num_beams=3 |
| ) |
| |
| content = [] |
| message = [ |
| {'type': 'text', 'value': prompt}, |
| {'type': 'image', 'value': images[0]}, |
| {'type': 'text', 'value': questions[0]} |
| ] |
| for x in message: |
| if x['type'] == 'text': |
| content.append(x['value']) |
| elif x['type'] == 'image': |
| image = Image.open(x['value']).convert('RGB') |
| content.append(image) |
| msgs = [{'role': 'user', 'content': content}] |
|
|
| res = self.model.chat( |
| image=None, |
| msgs=msgs, |
| context=None, |
| tokenizer=self.tokenizer, |
| **default_kwargs |
| ) |
|
|
| if isinstance(res, tuple) and len(res) > 0: |
| res = res[0] |
| print(f"Q: {content}, \nA: {res}") |
| return [res] |
|
|
|
|
| class MiniCPM_V_2_6: |
|
|
| def __init__(self, model_path, ckpt, device=None)->None: |
| seed = 0 |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| |
| self.model_path = model_path |
| self.ckpt = ckpt |
| self.model = AutoModel.from_pretrained(self.model_path, trust_remote_code=True).eval() |
| if self.ckpt is not None: |
| self.ckpt = ckpt |
| self.state_dict = torch.load(self.ckpt, map_location=torch.device('cpu')) |
| self.model.load_state_dict(self.state_dict) |
|
|
| self.model = self.model.to(dtype=torch.bfloat16) |
| self.model.to(device) |
| |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True) |
| torch.cuda.empty_cache() |
|
|
| def generate(self, images, questions, datasetname): |
| image = Image.open(images[0]).convert('RGB') |
| try: |
| max_new_tokens = max_token[datasetname] |
| except: |
| max_new_tokens = 1024 |
| if (datasetname == 'docVQA') or (datasetname == "docVQATest") : |
| prompt = "Answer the question directly with single word." + "\n" + questions[0] |
| elif (datasetname == 'textVQA') : |
| prompt = "Answer the question directly with single word." + '\n'+ questions[0] |
| |
| msgs = [{'role': 'user', 'content': prompt}] |
| default_kwargs = dict( |
| max_new_tokens=max_new_tokens, |
| sampling=False, |
| num_beams=3 |
| ) |
| res = self.model.chat( |
| image=image, |
| msgs=msgs, |
| context=None, |
| tokenizer=self.tokenizer, |
| **default_kwargs |
| ) |
| |
| return [res] |
| |
| def generate_with_interleaved(self, images, questions, datasetname): |
| try: |
| max_new_tokens = max_token[datasetname] |
| except: |
| max_new_tokens = 1024 |
| |
| prompt = "Answer the question directly with single word." |
| |
| default_kwargs = dict( |
| max_new_tokens=max_new_tokens, |
| sampling=False, |
| num_beams=3 |
| ) |
| |
| content = [] |
| message = [ |
| {'type': 'text', 'value': prompt}, |
| {'type': 'image', 'value': images[0]}, |
| {'type': 'text', 'value': questions[0]} |
| ] |
| for x in message: |
| if x['type'] == 'text': |
| content.append(x['value']) |
| elif x['type'] == 'image': |
| image = Image.open(x['value']).convert('RGB') |
| img_width, img_height = image.width, image.height |
| if (img_width * img_height) >= (1344 * 1344): |
| content.append(image) |
| else: |
| ratio = math.sqrt((1344 * 1344) / (img_width * img_height)) |
| max_img_width = int(img_width * ratio) |
| new_img_width = random.randint(img_width, max_img_width) |
| new_img_height = int(new_img_width / img_width * img_height) |
| resized_image = image.resize((new_img_width, new_img_height)) |
| content.append(resized_image) |
| msgs = [{'role': 'user', 'content': content}] |
|
|
| res = self.model.chat( |
| image=None, |
| msgs=msgs, |
| context=None, |
| tokenizer=self.tokenizer, |
| **default_kwargs |
| ) |
|
|
| if isinstance(res, tuple) and len(res) > 0: |
| res = res[0] |
| print(f"Q: {content}, \nA: {res}") |
| return [res] |
|
|
