Update custom_model.py

custom_model.py

@@ -1,8 +1,7 @@
 import torch
 import torch.nn as nn
-from transformers import PreTrainedModel, PreTrainedTokenizer, AutoTokenizer, AutoConfig
+from transformers import PreTrainedModel, AutoConfig
 
-# Define a custom model class
 class CustomModel(PreTrainedModel):
     config_class = AutoConfig  # Use AutoConfig to dynamically load the configuration class
 
@@ -14,11 +13,14 @@ class CustomModel(PreTrainedModel):
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
         try:
+            # Load the configuration
             config = cls.config_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+            # Initialize the model with the configuration
             model = cls(config)
-            # Load the weights from the pretrained model
-            model.load_state_dict(torch.load(pretrained_model_name_or_path))
+            # Load the model weights using the transformers library
+            state_dict = torch.load(f"{pretrained_model_name_or_path}/pytorch_model.bin", map_location="cpu")
+            model.load_state_dict(state_dict)
             return model
-        except FileNotFoundError:
-            print(f"Failed to load model from {pretrained_model_name_or_path}. Please check the model name and path.")
+        except Exception as e:
+            print(f"Failed to load model from {pretrained_model_name_or_path}. Error: {e}")
             return None