Update model

__init__.py

@@ -0,0 +1,4 @@
+from .configuration_reward_model import RewardModelConfig
+from .modeling_reward_model import RewardModel
+
+__all__ = ["RewardModelConfig", "RewardModel"]

configuration_reward_model.py

@@ -0,0 +1,8 @@
+from transformers import PretrainedConfig
+
+class RewardModelConfig(PretrainedConfig):
+    model_type = "reward_model"
+    
+    def __init__(self, attributes=None, **kwargs):
+        self.attributes = attributes or []
+        super().__init__(**kwargs)

modeling_reward_model.py

@@ -0,0 +1,29 @@
+import torch
+import torch.nn as nn
+from transformers import AutoConfig, PreTrainedModel
+from .configuration_reward_model import RewardModelConfig
+
+# Register config
+AutoConfig.register("reward_model", RewardModelConfig)
+
+class RewardModel(PreTrainedModel):
+    config_class = RewardModelConfig
+    
+    def __init__(self, config, transformer, regression_weights):
+        super().__init__(config)
+        self.transformer = transformer
+        self.regression_heads = nn.Linear(
+            regression_weights.shape[1],
+            regression_weights.shape[0],
+            bias=False
+        )
+        self.regression_heads.weight.data = torch.from_numpy(regression_weights).float()
+
+    def forward(self, input_ids, attention_mask=None):
+        outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)
+        last_hidden = outputs.last_hidden_state[:, -1, :]
+        return self.regression_heads(last_hidden)
+
+# Register model
+from transformers import AutoModel
+AutoModel.register(RewardModelConfig, RewardModel)