Update model.py

model.py

@@ -1,15 +1,14 @@
 import os
 from typing import Optional
 from transformers import Qwen3ForCausalLM, AutoTokenizer, AutoConfig
+from huggingface_hub import hf_hub_download
 import torch
 import torch.nn as nn
 
-from transformers.models.qwen3 import Qwen3Config
-
 
 # Define a custom model that wraps a causal LM and adds a regression head
 class CausalLMForRegression(nn.Module):
-    config_class = Qwen3Config
+    config_class = Qwen3ForCausalLM.config_class
     base_model_prefix = "model"
     
     def __init__(self, model_name):
@@ -88,35 +87,43 @@ class CausalLMForRegression(nn.Module):
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
-        config = kwargs.pop("config", None)        
-        if config is None:
-            from transformers import AutoConfig
-            config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
-        config.output_hidden_states = True
-
-        from transformers import Qwen3ForCausalLM
-        base_model = Qwen3ForCausalLM.from_pretrained(
+        cfg = kwargs.pop("config", None)
+        if cfg is None:
+            cfg = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        cfg.output_hidden_states = True
+
+        backbone = Qwen3ForCausalLM.from_pretrained(
             pretrained_model_name_or_path,
             *model_args,
-            config=config,
-            **kwargs 
+            config=cfg,
+            trust_remote_code=False,
+            **kwargs
         )
 
-        instance = cls.__new__(cls)
-        nn.Module.__init__(instance)
-        instance.model = base_model
-        instance.regression_head = nn.Linear(config.hidden_size, 1)
-        instance._keys_to_ignore_on_save = []
+        if os.path.isdir(pretrained_model_name_or_path):
+            head_path = os.path.join(pretrained_model_name_or_path,
+                                     "regression_head.bin")
+        else:
+            head_path = hf_hub_download(
+                repo_id=pretrained_model_name_or_path,
+                filename="regression_head.bin",
+                repo_type="model"
+            )
+
+        inst = cls.__new__(cls)
+        nn.Module.__init__(inst)
+        inst.model = backbone
+        inst.regression_head = nn.Linear(cfg.hidden_size, 1)
+        inst._keys_to_ignore_on_save = []
 
-        print(pretrained_model_name_or_path)
-        head_path = os.path.join(pretrained_model_name_or_path, "regression_head.bin")
         if os.path.exists(head_path):
-            instance.regression_head.load_state_dict(
+            inst.regression_head.load_state_dict(
                 torch.load(head_path, map_location="cpu")
             )
         else:
-            print("No regression head found – initialising randomly.")
-        return instance
+            print("'regression_head.bin' not found – initialising randomly.")
+
+        return inst
 
     @torch.no_grad()
     def generate(self, *args, **kwargs):