Update model.py

model.py

@@ -82,29 +82,35 @@ class CausalLMForRegression(nn.Module):
         return tokenizer
 
     @classmethod
-    def from_pretrained(cls, output_dir):
-        from_local = os.path.exists(output_dir)
-        loading_kwargs = {"use_safetensors": False} if from_local else {}
-        
-        model = AutoModelForCausalLM.from_pretrained(output_dir, **loading_kwargs)
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+
+        kwargs.setdefault("output_hidden_states", True)
 
-        # Explicitly enable `output_hidden_states` after loading
-        model.config.output_hidden_states = True  
+        base_model = AutoModelForCausalLM.from_pretrained(
+            pretrained_model_name_or_path,
+            *model_args,
+            **kwargs
+        )
 
         # Create an uninitialized instance of CausalLMForRegression
         instance = cls.__new__(cls)
         nn.Module.__init__(instance)
+        
+        instance.model = base_model
+        instance.regression_head = nn.Linear(
+            base_model.config.hidden_size, 1
+        )
         instance._keys_to_ignore_on_save = []
-        instance.model = model
 
         # Load the regression head separately
-        instance.regression_head = nn.Linear(model.config.hidden_size, 1)
-        try:
-            regression_head_path = os.path.join(output_dir, "regression_head.bin")
-            state = torch.load(regression_head_path, map_location="cpu")
+        head_path = os.path.join(
+            pretrained_model_name_or_path, "regression_head.bin"
+        )
+        if os.path.exists(head_path):
+            state = torch.load(head_path, map_location="cpu")
             instance.regression_head.load_state_dict(state)
-        except FileNotFoundError:
-            print(f"No regression head found. Initializing with random weights!")
+        else:
+            print("No regression head found – initialising randomly.")
         return instance
 
     @torch.no_grad()