edit for inference

model.py

@@ -11,9 +11,14 @@ class ClassificationOutput(ModelOutput):
     hidden_states: Optional[Tuple[torch.FloatTensor]] = None
 
 class MoralEmotionVLClassifier(nn.Module):
-    def __init__(self, model_id, num_labels=1, device="auto", label_names=None):
+    def __init__(self, model_id_or_save_dir, num_labels=1, device="auto", max_memory=None, label_names=None, train=True):
         super().__init__()
-        
+
+        self.device = device
+        self.max_memory = max_memory
+        self.model_id_or_save_dir = model_id_or_save_dir
+
+        # Bits and bytes config for model quantization
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True, 
             bnb_4bit_use_double_quant=True, 
@@ -21,23 +26,37 @@ class MoralEmotionVLClassifier(nn.Module):
             bnb_4bit_compute_dtype=torch.float16
         )
         
-        self.base_model = AutoModelForVision2Seq.from_pretrained(
-            model_id,
-            device_map='auto' if device == 'auto' else {"": device},
-            torch_dtype=torch.float16,
-            quantization_config=bnb_config
-        )
+        # Load base model (vision-to-text)
+        if device == 'auto':
+            self.base_model = AutoModelForVision2Seq.from_pretrained(
+                self.model_id_or_save_dir,
+                device_map=self.device,
+                torch_dtype=torch.float16,
+                quantization_config=bnb_config if train else None,
+                ignore_mismatched_sizes=not train,
+            )
+        
+        else:
+            self.base_model = AutoModelForVision2Seq.from_pretrained(
+                self.model_id_or_save_dir,
+                device_map={"": device},
+                torch_dtype=torch.float16,
+                quantization_config=bnb_config if train else None,
+                max_memory=self.max_memory,
+                ignore_mismatched_sizes=not train,
+            )
 
         self.config = self.base_model.config
         self.config.num_labels = num_labels
         self.gradient_checkpointing_enable = self.base_model.gradient_checkpointing_enable
         
+        # Modify the final classification head (lm_head)
         original_lm_head = self.base_model.lm_head
         hidden_size = original_lm_head.in_features
         head_device = original_lm_head.weight.device
         head_dtype = original_lm_head.weight.dtype
 
-        # change to classification head
+        # Change to classification head for the number of labels required
         self.base_model.lm_head = nn.Linear(
             hidden_size, 
             num_labels,
@@ -45,16 +64,37 @@ class MoralEmotionVLClassifier(nn.Module):
             dtype=head_dtype
         )
 
+        if not train:
+            try:
+                from safetensors import safe_open
+                import os
+                
+                safetensors_path = os.path.join(model_id_or_save_dir, "model.safetensors")
+                if os.path.exists(safetensors_path):
+                    with safe_open(safetensors_path, framework="pt") as f:
+                        lm_head_weight = f.get_tensor("lm_head.weight")
+                        lm_head_bias = f.get_tensor("lm_head.bias") if "lm_head.bias" in f.keys() else None
+                        
+                        target_device = self.base_model.lm_head.weight.device
+                        self.base_model.lm_head.weight.data = lm_head_weight.to(target_device)
+                        if lm_head_bias is not None:
+                            self.base_model.lm_head.bias.data = lm_head_bias.to(target_device)
+                        print('\nload the custom layer weights successed!\n')
+            except Exception as e:
+                print(f"Warning: Could not load lm_head weights: {e}")
+
+
         # label mapping
         self.num_labels = num_labels
         self.label_names = label_names if label_names is not None else []
         self.label2id = {label: i for i, label in enumerate(self.label_names)}
         self.id2label = {i: label for i, label in enumerate(self.label_names)}
-    
+
     def forward(self, **kwargs):
+        # Forward pass through the model
         outputs = self.base_model(**kwargs)        
         logits = outputs.logits        
-        classification_logits = logits[:, -1, :]
+        classification_logits = logits[:, -1, :]  # Assuming we want to use the last token's logits
         
         return ClassificationOutput(
             logits=classification_logits,