|
import math |
|
from typing import List, Optional |
|
import json |
|
import torch |
|
import torchvision |
|
|
|
from threading import Thread |
|
from copy import deepcopy |
|
from PIL import Image |
|
from transformers import AutoProcessor, Qwen2PreTrainedModel, Qwen2ForCausalLM, TextIteratorStreamer |
|
from .processing_mplugowl3 import mPLUGOwl3Processor |
|
from .image_processing_mplugowl3 import mPLUGOwl3ImageProcessor |
|
from .configuration_mplugowl3 import mPLUGOwl3Config |
|
|
|
from transformers.models.siglip.modeling_siglip import SiglipVisionTransformer |
|
from .x_sdpa import ScaleDotProductAttention |
|
from .modeling_hyper_qwen2 import HyperQwen2ForCausalLM |
|
from torch import nn |
|
|
|
|
|
class mPLUGOwl3PreTrainedModel(Qwen2PreTrainedModel): |
|
config_class = mPLUGOwl3Config |
|
|
|
|
|
class mPLUGOwl3Model(mPLUGOwl3PreTrainedModel): |
|
def __init__(self, config): |
|
super().__init__(config) |
|
self.language_model = HyperQwen2ForCausalLM(config) |
|
self.vision_model = self.init_vision_module() |
|
self.vision_dim = self.vision_model.embed_dim |
|
self.embed_dim = self.language_model.config.hidden_size |
|
self.vision2text_model = nn.Linear(self.vision_dim, self.embed_dim) |
|
self.processor = None |
|
|
|
self.terminators = ['<|im_end|>', '<|endoftext|>'] |
|
|
|
def init_vision_module(self): |
|
|
|
self.config.vision_config._attn_implementation = self.config.vision_config._attn_implementation |
|
model = SiglipVisionTransformer(self.config.vision_config) |
|
|
|
setattr(model, 'embed_dim', model.embeddings.embed_dim) |
|
setattr(model, 'patch_size', model.embeddings.patch_size) |
|
return model |
|
|
|
|
|
def get_input_embeddings(self): |
|
return self.language_model.get_input_embeddings() |
|
|
|
def set_input_embeddings(self, value): |
|
self.language_model.embed_tokens = value |
|
|
|
def get_output_embeddings(self): |
|
return self.language_model.lm_head |
|
|
|
def set_output_embeddings(self, new_embeddings): |
|
self.language_model.lm_head = new_embeddings |
|
|
|
def set_decoder(self, decoder): |
|
self.language_model = decoder |
|
|
|
def get_decoder(self): |
|
return self.language_model |
|
|
|
def forward_image(self, pixel_values): |
|
if pixel_values is None: |
|
return None |
|
dtype = self.language_model.model.embed_tokens.weight.dtype |
|
with torch.inference_mode(): |
|
image_embeds = self.vision_model(pixel_values.to(dtype), output_hidden_states=True).hidden_states[-2] |
|
|
|
if self.vision2text_model is not None: |
|
image_embeds = self.vision2text_model(image_embeds) |
|
else: |
|
pass |
|
|
|
return image_embeds |
|
|
|
def forward(self, pixel_values=None, **kwargs): |
|
image_embeds = self.forward_image(pixel_values) |
|
|
|
return self.language_model( |
|
image_embeds=image_embeds, |
|
**kwargs |
|
) |
|
|
|
def _decode(self, input_ids, image_embeds, media_offset, tokenizer, attention_mask, decode_text=False, **kwargs): |
|
terminators = [tokenizer.convert_tokens_to_ids(i) for i in self.terminators] |
|
output = self.language_model.generate( |
|
input_ids=input_ids, |
|
image_embeds=image_embeds, |
|
media_offset=media_offset, |
|
pad_token_id=0, |
|
eos_token_id=terminators, |
|
attention_mask=attention_mask, |
|
**kwargs |
|
) |
|
|
|
output = output[:,input_ids.shape[1]:] |
|
if decode_text: |
|
return self._decode_text(output, tokenizer) |
|
return output |
|
|
|
def _decode_stream(self, input_ids, image_embeds, media_offset, tokenizer, **kwargs): |
|
terminators = [tokenizer.convert_tokens_to_ids(i) for i in self.terminators] |
|
streamer = TextIteratorStreamer(tokenizer=tokenizer) |
|
generation_kwargs = { |
|
'input_ids': input_ids, |
|
'image_embeds': image_embeds, |
|
'media_offset': media_offset, |
|
'pad_token_id': 0, |
|
'eos_token_id': terminators, |
|
'streamer': streamer |
|
} |
|
generation_kwargs.update(kwargs) |
|
|
|
thread = Thread(target=self.language_model.generate, kwargs=generation_kwargs) |
|
thread.start() |
|
|
|
return streamer |
|
|
|
def _decode_text(self, result_ids, tokenizer): |
|
terminators = [tokenizer.convert_tokens_to_ids(i) for i in self.terminators] |
|
result_text = [] |
|
for result in result_ids: |
|
result = result[result != 0] |
|
if result[-1] in terminators: |
|
result = result[:-1] |
|
result_text.append(tokenizer.decode(result).strip()) |
|
return result_text |
|
|
|
def init_processor(self, tokenizer): |
|
ip = mPLUGOwl3ImageProcessor(image_size=384) |
|
self.processor = mPLUGOwl3Processor(image_processor=ip, tokenizer=tokenizer) |
|
processor = self.processor |
|
return processor |
|
|
|
def generate( |
|
self, |
|
input_ids=None, |
|
pixel_values=None, |
|
media_offset=None, |
|
attention_mask=None, |
|
tokenizer=None, |
|
return_vision_hidden_states=False, |
|
stream=False, |
|
decode_text=False, |
|
**kwargs |
|
): |
|
assert input_ids is not None |
|
|
|
with torch.inference_mode(): |
|
image_embeds = self.forward_image(pixel_values) |
|
|
|
if stream: |
|
result = self._decode_stream(input_ids=input_ids, image_embeds=image_embeds, media_offset=media_offset, tokenizer=tokenizer, **kwargs) |
|
else: |
|
result = self._decode(input_ids=input_ids, image_embeds=image_embeds, media_offset=media_offset, tokenizer=tokenizer, attention_mask=attention_mask, decode_text=decode_text, **kwargs) |
|
|
|
if return_vision_hidden_states: |
|
return result, image_embeds |
|
|
|
return result |
|
|
|
def chat( |
|
self, |
|
images, |
|
videos, |
|
msgs, |
|
tokenizer, |
|
processor=None, |
|
vision_hidden_states=None, |
|
max_new_tokens=2048, |
|
min_new_tokens=0, |
|
sampling=True, |
|
max_inp_length=8192, |
|
system_prompt='', |
|
stream=False, |
|
max_slice_nums=None, |
|
use_image_id=None, |
|
**kwargs |
|
): |
|
print(msgs) |
|
if processor is None: |
|
if self.processor is None: |
|
self.processor = AutoProcessor.from_pretrained(self.config._name_or_path, trust_remote_code=True) |
|
processor = self.processor |
|
|
|
|
|
inputs = processor( |
|
prompts_lists, |
|
input_images_lists, |
|
max_slice_nums=max_slice_nums, |
|
use_image_id=use_image_id, |
|
return_tensors="pt", |
|
max_length=max_inp_length |
|
).to(self.device) |
|
|
|
if sampling: |
|
generation_config = { |
|
"top_p": 0.8, |
|
"top_k": 100, |
|
"temperature": 0.7, |
|
"do_sample": True, |
|
"repetition_penalty": 1.05 |
|
} |
|
else: |
|
generation_config = { |
|
"num_beams": 3, |
|
"repetition_penalty": 1.2, |
|
} |
|
|
|
if min_new_tokens > 0: |
|
generation_config['min_new_tokens'] = min_new_tokens |
|
|
|
generation_config.update( |
|
(k, kwargs[k]) for k in generation_config.keys() & kwargs.keys() |
|
) |
|
|
|
inputs.pop("image_sizes") |
|
with torch.inference_mode(): |
|
res = self.generate( |
|
**inputs, |
|
tokenizer=tokenizer, |
|
max_new_tokens=max_new_tokens, |
|
vision_hidden_states=vision_hidden_states, |
|
stream=stream, |
|
decode_text=True, |
|
**generation_config |
|
) |
|
|
|
if stream: |
|
def stream_gen(): |
|
for text in res: |
|
for term in self.terminators: |
|
text = text.replace(term, '') |
|
yield text |
|
return stream_gen() |
|
|
|
else: |
|
if batched: |
|
answer = res |
|
else: |
|
answer = res[0] |
|
return answer |
|
|
|
|
