Upload 2 files

mm_utils.py

@@ -21,15 +21,6 @@ DEFAULT_IM_START_TOKEN = "<im_start>"
 DEFAULT_IM_END_TOKEN = "<im_end>"
 
 
-def disable_torch_init():
-    """
-    Disable the redundant torch default initialization to accelerate model creation.
-    """
-    import torch
-    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
-    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
-
-
 def load_image_from_base64(image):
     return Image.open(BytesIO(base64.b64decode(image)))
 
@@ -63,28 +54,6 @@ def process_images(images, image_processor, model_cfg):
     return new_images
 
 
-def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):
-    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('<image>')]
-
-    def insert_separator(X, sep):
-        return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]
-
-    input_ids = []
-    offset = 0
-    if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:
-        offset = 1
-        input_ids.append(prompt_chunks[0][0])
-
-    for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
-        input_ids.extend(x[offset:])
-
-    if return_tensors is not None:
-        if return_tensors == 'pt':
-            return torch.tensor(input_ids, dtype=torch.long)
-        raise ValueError(f'Unsupported tensor type: {return_tensors}')
-    return input_ids
-
-
 def get_model_name_from_path(model_path):
     model_path = model_path.strip("/")
     model_paths = model_path.split("/")

mmalaya_arch.py

@@ -4,6 +4,7 @@ import torch
 import torch.nn as nn
 from transformers import Blip2Model, Blip2Processor, Blip2Config
 from .mm_utils import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
+from .mm_utils import conv_templates
 
 
 class BLIP2VisionTower(nn.Module):
@@ -265,6 +266,48 @@ class MMAlayaMetaForCausalLM(ABC):
 
         return None, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels
 
-    def initialize_vision_tokenizer(self, tokenizer):
+    def initialize_tokenizer(self, tokenizer):
         tokenizer.add_tokens([DEFAULT_IMAGE_TOKEN], special_tokens=True)
         self.resize_token_embeddings(len(tokenizer))
+
+    def prepare_for_inference(
+        self, 
+        prompt, 
+        tokenizer, 
+        image,
+        image_token_index=IMAGE_TOKEN_INDEX, 
+        return_tensors=None
+        ):
+        # 加载对话模板
+        conv = conv_templates["mmalaya_llama"].copy()
+        inp = DEFAULT_IMAGE_TOKEN + '\n' + prompt
+        conv.append_message(conv.roles[0], inp)
+        conv.append_message(conv.roles[1], None)
+        prompt = conv.get_prompt()
+
+        prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('<image>')]
+
+        def insert_separator(X, sep):
+            return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]
+
+        input_ids = []
+        offset = 0
+        if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:
+            offset = 1
+            input_ids.append(prompt_chunks[0][0])
+
+        for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
+            input_ids.extend(x[offset:])
+
+        if return_tensors is not None:
+            if return_tensors == 'pt':
+                return torch.tensor(input_ids, dtype=torch.long)
+            raise ValueError(f'Unsupported tensor type: {return_tensors}')
+        
+        # 加载generate stop条件
+        stopping_criteria = KeywordsStoppingCriteria([conv.sep2], tokenizer, input_ids)
+        # 加载图像
+        image_processor = model.get_vision_tower().image_processor
+        image_tensor = image_processor(image, return_tensors='pt')['pixel_values'].half().cuda()
+
+        return input_ids, image_tensor, stopping_criteria