Update model.py

model.py

@@ -11,12 +11,12 @@ class ClassificationOutput(ModelOutput):
     hidden_states: Optional[Tuple[torch.FloatTensor]] = None
 
 class MoralEmotionVLClassifier(nn.Module):
-    def __init__(self, model_id_or_save_dir, num_labels=1, device="auto", max_memory=None, label_names=None, train=True):
+    def __init__(self, model_id, num_labels=1, device="auto", max_memory=None, label_names=None):
         super().__init__()
 
         self.device = device
         self.max_memory = max_memory
-        self.model_id_or_save_dir = model_id_or_save_dir
+        self.model_id = model_id
 
         # Bits and bytes config for model quantization
         bnb_config = BitsAndBytesConfig(
@@ -27,24 +27,13 @@ class MoralEmotionVLClassifier(nn.Module):
         )
         
         # 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.base_model = AutoModelForVision2Seq.from_pretrained(
+            self.model_id,
+            device_map="auto" if device == "auto" else {"": device},
+            torch_dtype=torch.float16,
+            quantization_config=bnb_config,
+            max_memory=self.max_memory if device == "auto" else None
+        )
 
         self.config = self.base_model.config
         self.config.num_labels = num_labels
@@ -64,26 +53,6 @@ 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 []
@@ -91,12 +60,11 @@ class MoralEmotionVLClassifier(nn.Module):
         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, :]  # Assuming we want to use the last token's logits
+        classification_logits = logits[:, -1, :]
         
         return ClassificationOutput(
             logits=classification_logits,
             hidden_states=outputs.hidden_states if hasattr(outputs, 'hidden_states') else None
-        )
+        )