Update model.py

model.py

@@ -1,6 +1,6 @@
 import os
 from typing import Optional
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import Qwen3ForCausalLM, AutoTokenizer
 import torch
 import torch.nn as nn
 
@@ -88,34 +88,26 @@ class CausalLMForRegression(nn.Module):
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
-
+        # make sure hidden states are returned
         kwargs.setdefault("output_hidden_states", True)
 
-        base_model = AutoModelForCausalLM.from_pretrained(
-            pretrained_model_name_or_path,
-            *model_args,
-            **kwargs
+        base_model = Qwen3ForCausalLM.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.regression_head = nn.Linear(base_model.config.hidden_size, 1)
 
-        # Load the regression head separately
-        head_path = os.path.join(
-            pretrained_model_name_or_path, "regression_head.bin"
-        )
+        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)
+            instance.regression_head.load_state_dict(
+                torch.load(head_path, map_location="cpu")
+            )
         else:
             print("No regression head found – initialising randomly.")
+        instance._keys_to_ignore_on_save = []
         return instance
 
     @torch.no_grad()