Update decoder.py

decoder.py

@@ -2,45 +2,43 @@ import torch.nn as nn
 import torch
 from transformers import Qwen2_5_VLForConditionalGeneration, AutoConfig, BitsAndBytesConfig
 
-
 class SketchDecoder(nn.Module):
-  """
-  Autoregressive generative model 
-  """
-  
-  def __init__(self,
-              **kwargs):
-    super().__init__()
-    self.vocab_size = 196042
-    self.bos_token_id = 151643
-    self.eos_token_id = 196041
-    self.pad_token_id = 151643
-
-    config = AutoConfig.from_pretrained(
-          "Qwen/Qwen2.5-VL-3B-Instruct",
-          #n_positions=8192,
-          vocab_size=self.vocab_size,
-          bos_token_id=self.bos_token_id,
-          eos_token_id=self.eos_token_id,
-          pad_token_id=self.pad_token_id)
-
-    quantization_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_compute_dtype=torch.float16,
-        bnb_4bit_quant_type="nf4"
-    )
-
-    self.transformer = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-        "Qwen/Qwen2.5-VL-3B-Instruct", 
-        config=config,
-        quantization_config=quantization_config,
-        torch_dtype=torch.bfloat16, attn_implementation="sdpa",
-        device_map ="auto",
-        ignore_mismatched_sizes=True
-    )
-
-    self.transformer.resize_token_embeddings(self.vocab_size)
+    """
+    Autoregressive generative model
+    """
+    def __init__(self, state_dict=None, **kwargs):
+        super().__init__()
+        self.vocab_size = 196042
+        self.bos_token_id = 151643
+        self.eos_token_id = 196041
+        self.pad_token_id = 151643
+        
+        config = AutoConfig.from_pretrained(
+            "Qwen/Qwen2.5-VL-3B-Instruct",
+            vocab_size=self.vocab_size,
+            bos_token_id=self.bos_token_id,
+            eos_token_id=self.eos_token_id,
+            pad_token_id=self.pad_token_id
+        )
+        
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_quant_type="nf4"
+        )
+        
+        self.transformer = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            "Qwen/Qwen2.5-VL-3B-Instruct",
+            config=config,
+            quantization_config=quantization_config,
+            torch_dtype=torch.bfloat16, 
+            attn_implementation="sdpa",
+            device_map="auto",
+            ignore_mismatched_sizes=True,
+            state_dict=state_dict  # <--- 关键修改：直接传入权重字典
+        )
+        
+        self.transformer.resize_token_embeddings(self.vocab_size)
 
-  def forward(self, *args, **kwargs):
-      raise NotImplementedError("Forward pass not included in open-source version")
-    
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError("Forward pass not included in open-source version")