Upload 3 files

config.json

@@ -1,7 +1,7 @@
 {
-  "_name_or_path": "meta-llama/Llama-2-13b-hf",
+  "_name_or_path": "UltraRM-13b-32",
   "architectures": [
-    "LlamaForCausalLM"
+    "LlamaRewardModel"
   ],
   "bos_token_id": 1,
   "eos_token_id": 2,
@@ -19,9 +19,10 @@
   "pretraining_tp": 1,
   "rms_norm_eps": 1e-05,
   "rope_scaling": null,
+  "rope_theta": 10000.0,
   "tie_word_embeddings": false,
-  "torch_dtype": "float32",
-  "transformers_version": "4.31.0",
+  "torch_dtype": "float16",
+  "transformers_version": "4.33.2",
   "use_cache": true,
   "vocab_size": 32000
 }

modeling_llama_rm.py

@@ -0,0 +1,44 @@
+from transformers import PreTrainedModel, LlamaConfig, LlamaModel
+import torch.nn as nn
+import torch
+from typing import Optional, List
+
+class LlamaRewardModel(PreTrainedModel):
+    config_class = LlamaConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = LlamaModel(config)
+        self.regression_head = nn.Linear(self.config.hidden_size, 1, bias=False)
+
+    def forward( # args are the same as LlamaForCausalLM
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+
+        transformer_outputs = self.model.model(
+                                input_ids,
+                                attention_mask=attention_mask,
+                                position_ids=position_ids,
+                                past_key_values=past_key_values,
+                                inputs_embeds=inputs_embeds,                               
+                            )
+
+        hidden_states = transformer_outputs[0]
+        rewards = self.regression_head(hidden_states).squeeze(-1)
+        
+        ends = attention_mask.cumsum(dim=1).argmax(dim=1).view(-1,1)
+        rewards = torch.gather(rewards, 1, ends)
+        
+        return reward_models
+
+model = LlamaRewardModel.from_pretrained("UltraRM-13b-32").half()
+model.save_pretrained("UltraRM-13b")