weights type issue fixed

modeling_ablang2paired.py

@@ -69,34 +69,49 @@ class AbLang2PairedHFModel(PreTrainedModel):
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
-        # Check if we have custom weights
-        model_path = pretrained_model_name_or_path
-        custom_weights_path = os.path.join(model_path, "model.pt")
+        # Load config first
+        config = kwargs.get("config")
+        if config is None:
+            from transformers import AutoConfig
+            config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)
         
-        if os.path.exists(custom_weights_path):
-            # Load config
-            config = kwargs.get("config")
-            if config is None:
-                from transformers import AutoConfig
-                config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
-            
-            # Create model with only the config argument
-            model = cls(config)
-            
-            # Load custom weights
-            state_dict = torch.load(custom_weights_path, map_location="cpu", weights_only=True)
-            model.model.load_state_dict(state_dict)
-            
-            # Move model to appropriate device (GPU if available, otherwise CPU)
-            device = kwargs.get("device", None)
-            if device is None:
-                device = "cuda" if torch.cuda.is_available() else "cpu"
-            model = model.to(device)
+        # Create model with config
+        model = cls(config)
+        
+        # Try to load custom weights
+        try:
+            from transformers.utils import cached_file
+            custom_weights_path = cached_file(
+                pretrained_model_name_or_path,
+                "model.pt",
+                cache_dir=kwargs.get("cache_dir"),
+                force_download=kwargs.get("force_download", False),
+                resume_download=kwargs.get("resume_download", False),
+                proxies=kwargs.get("proxies"),
+                token=kwargs.get("token"),
+                revision=kwargs.get("revision"),
+                local_files_only=kwargs.get("local_files_only", False),
+            )
             
-            return model
-        else:
-            # Fall back to standard Hugging Face loading
-            return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+            if custom_weights_path is not None and os.path.exists(custom_weights_path):
+                # Load custom weights
+                state_dict = torch.load(custom_weights_path, map_location="cpu", weights_only=True)
+                model.model.load_state_dict(state_dict)
+                print(f"✅ Loaded custom weights from: {custom_weights_path}")
+            else:
+                print("⚠️ No custom weights found, using initialized model")
+                
+        except Exception as e:
+            print(f"⚠️ Could not load custom weights: {e}")
+            print("Using initialized model")
+        
+        # Move model to appropriate device (GPU if available, otherwise CPU)
+        device = kwargs.get("device", None)
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        model = model.to(device)
+        
+        return model
 
     def save_pretrained(self, save_directory, **kwargs):
         os.makedirs(save_directory, exist_ok=True)