Update modeling_deepseekocr.py

modeling_deepseekocr.py

@@ -1,28 +1,31 @@
-from .modeling_deepseekv2 import DeepseekV2Model, DeepseekV2ForCausalLM
-from .configuration_deepseek_v2 import DeepseekV2Config
-from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+import os
+import math
+import re
+from tqdm import tqdm
+from abc import ABC
 from typing import List, Optional, Tuple, Union
-from transformers.cache_utils import Cache
-import requests
+
+from addict import Dict
 from PIL import Image, ImageOps, ImageDraw, ImageFont
-from io import BytesIO
+import numpy as np
+
 import torch
 import torch.nn as nn
 from torch.nn import CrossEntropyLoss
 from torchvision import transforms
-from torchvision.transforms.functional import InterpolationMode
-import os
-from .deepencoder import build_sam_vit_b, build_clip_l, MlpProjector
-from addict import Dict
+
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers import DeepseekV2Model, DeepseekV2ForCausalLM
+from transformers import DeepseekV2Config
+from transformers.models.deepseek_v2.modeling_deepseek_v2 import (
+    DeepseekV2Attention, DeepseekV2MLP, DeepseekV2MoE, DeepseekV2RMSNorm, DeepseekV2DecoderLayer)
+from transformers.models.llama.modeling_llama import LlamaAttention, LlamaRotaryEmbedding
 from transformers import TextStreamer
+from .deepencoder import build_sam_vit_b, build_clip_l, MlpProjector
 from .conversation import get_conv_template
-from abc import ABC
-import math
-import re
-from tqdm import tqdm
-import numpy as np
-import time
 
+torch_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
 
 def load_image(image_path):
 
@@ -348,6 +351,23 @@ class NoEOSTextStreamer(TextStreamer):
         print(text, flush=True, end="")
 
 
+def decoder_layer_init(self, config: DeepseekV2Config, layer_idx: int):
+    nn.Module.__init__(self)
+    self.hidden_size = config.hidden_size
+
+    if config.use_mla:
+        self.self_attn = DeepseekV2Attention(config=config, layer_idx=layer_idx)
+    else:
+        config.head_dim = config.hidden_size // config.num_attention_heads
+        self.self_attn = LlamaAttention(config, layer_idx)
+    self.mlp = DeepseekV2MoE(config) if layer_idx >= config.first_k_dense_replace else DeepseekV2MLP(config)
+
+    self.input_layernorm = DeepseekV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+    self.post_attention_layernorm = DeepseekV2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+
+DeepseekV2DecoderLayer.__init__ = decoder_layer_init
+
 class DeepseekOCRConfig(DeepseekV2Config):
     model_type = "DeepseekOCR"
 
@@ -366,8 +386,7 @@ class DeepseekOCRModel(DeepseekV2Model):
         self.image_newline = nn.Parameter(torch.randn(n_embed) * embed_std)
         self.view_seperator = nn.Parameter(torch.randn(n_embed) * embed_std)
 
-
-
+        self.rotary_emb = LlamaRotaryEmbedding(config=config)
     
     def forward(
         self,
@@ -387,12 +406,11 @@ class DeepseekOCRModel(DeepseekV2Model):
 
 
 
-
         if inputs_embeds is None:
             # inputs_embeds = self.embed_tokens(input_ids)
             inputs_embeds = self.get_input_embeddings()(input_ids)
 
-
+        inputs_embeds = inputs_embeds.clone()
 
         sam_model = getattr(self, 'sam_model', None)
         # sam_model = self.sam_model
@@ -475,10 +493,6 @@ class DeepseekOCRModel(DeepseekV2Model):
                         global_features_2 = vision_model(image_ori, global_features_1) 
                         global_features = torch.cat((global_features_2[:, 1:], global_features_1.flatten(2).permute(0, 2, 1)), dim=-1) 
                         global_features = self.projector(global_features)
-                        print('=====================')
-                        print('BASE: ', global_features.shape)
-                        print('NO PATCHES')
-                        print('=====================')
                         _, hw, n_dim = global_features.shape
                         h = w = int(hw ** 0.5)
 
@@ -496,17 +510,17 @@ class DeepseekOCRModel(DeepseekV2Model):
                     images_in_this_batch.append(global_local_features)
                 
 
-                # print(inputs_embeds.shape)
-
                 if images_in_this_batch:
                     images_in_this_batch = torch.cat(images_in_this_batch, dim=0)
-                    # exit()
-
-                    inputs_embeds[idx].masked_scatter_(images_seq_mask[idx].unsqueeze(-1).cuda(), images_in_this_batch)
+                    images_in_this_batch = images_in_this_batch.to(
+                        device=inputs_embeds.device, dtype=inputs_embeds.dtype
+                    )
+                    mask = images_seq_mask[idx].unsqueeze(-1).to(inputs_embeds.device)   # bool [T, 1]
+                    updated_row = inputs_embeds[idx].masked_scatter(mask, images_in_this_batch)
+                    inputs_embeds[idx] = updated_row
 
                 idx += 1
             
-
         return super(DeepseekOCRModel, self).forward(
             input_ids=None, attention_mask=attention_mask, past_key_values=past_key_values,
             inputs_embeds=inputs_embeds, use_cache=use_cache, position_ids = position_ids,
@@ -528,8 +542,6 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
 
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
 
-        # self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
         # Initialize weights and apply final processing
         self.post_init()
 
@@ -578,10 +590,6 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
             
         )
 
-
-        
-        # print(transformer_outputs)
-
         hidden_states = outputs[0]
         logits = self.lm_head(hidden_states)
         logits = logits.float()
@@ -622,8 +630,8 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
         if past_key_values is not None:
             if isinstance(past_key_values, Cache):
                 cache_length = past_key_values.get_seq_length()
-                past_length = past_key_values.seen_tokens
-                max_cache_length = past_key_values.get_max_length()
+                past_length = past_key_values.get_seq_length()
+                max_cache_length = None
             else:
                 cache_length = past_length = past_key_values[0][0].shape[2]
                 max_cache_length = None
@@ -799,9 +807,9 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
 
 
                 
-                images_list.append(image_transform(global_view).to(torch.bfloat16))
+                images_list.append(image_transform(global_view).to(torch_dtype))
 
-                # global_view_tensor = image_transform(global_view).to(torch.bfloat16)
+                # global_view_tensor = image_transform(global_view).to(torch_dtype)
 
                 width_crop_num, height_crop_num = crop_ratio
 
@@ -812,7 +820,7 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
                     """process the local views"""
                     
                     for i in range(len(images_crop_raw)):
-                        images_crop_list.append(image_transform(images_crop_raw[i]).to(torch.bfloat16))
+                        images_crop_list.append(image_transform(images_crop_raw[i]).to(torch_dtype))
                 
                 if image_size == 640:
                     valid_img_tokens += len(images_crop_list) * 100
@@ -846,7 +854,7 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
                 # else:
                 global_view = ImageOps.pad(image, (image_size, image_size),
                                         color=tuple(int(x * 255) for x in image_transform.mean))
-                images_list.append(image_transform(global_view).to(torch.bfloat16))
+                images_list.append(image_transform(global_view).to(torch_dtype))
 
                 if base_size == 1024:
                     valid_img_tokens += int(256 * ratio)
@@ -888,9 +896,6 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
 
         input_ids = torch.LongTensor(tokenized_str)
 
-
-        
-
         images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool)
 
 
@@ -911,7 +916,7 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
 
         if not eval_mode:
             streamer = NoEOSTextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
-            with torch.autocast("cuda", dtype=torch.bfloat16):
+            with torch.autocast("cuda", dtype=torch_dtype):
                 with torch.no_grad():
                     output_ids = self.generate(
                         input_ids.unsqueeze(0).cuda(),
@@ -929,7 +934,7 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
                         )
 
         else:
-            with torch.autocast("cuda", dtype=torch.bfloat16):
+            with torch.autocast("cuda", dtype=torch_dtype):
                 with torch.no_grad():
                     output_ids = self.generate(
                         input_ids.unsqueeze(0).cuda(),
@@ -1034,4 +1039,4 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
                 plt.savefig(f'{output_path}/geo.jpg')
                 plt.close()
 
-            result.save(f"{output_path}/result_with_boxes.jpg")
+            result.save(f"{output_path}/result_with_boxes.jpg")