Update model & tokenizer load snippet

README.md

@@ -41,6 +41,66 @@ We recommend using the latest version of HF Transformers, or any `transformers>=
 Below we provide a code snippet demonstrating how to load the tokenizer and model and score a candidate instruction. We strongly recommend to format the instruction input as shown to maintain consistency with the format of the data used during training of MDCureRM. As the model outputs values normalized to the 0-1 range, we scale outputted scores up to the 1-5 range for more interpretable results. Relative weighting of fine-grained rewards may be configured as desired to obtain the final score; we reproduce the weights used in our implementation in `reward_weights` below.
 
 ```python
+from transformers import AutoTokenizer, AutoModel, LlamaConfig, PreTrainedModel, LlamaForSequenceClassification
+import torch.nn as nn
+import torch
+
+# Login to HF to access LLAMA model
+from huggingface_hub import login
+login("") # HF token
+
+class RewardModelConfig(LlamaConfig):
+    model_type = "RewardModel"
+
+    def __init__(self, reward_dim=None, base_model_name=None, **kwargs):
+        super().__init__(**kwargs)
+        
+        self.reward_dim = reward_dim
+        self.base_model_name = base_model_name
+
+class RewardModel(PreTrainedModel):
+    config_class = RewardModelConfig
+
+    def create_base_model(self):
+        
+        # use sequence classification model for consistency with https://huggingface.co/sfairXC/FsfairX-LLaMA3-RM-v0.1 
+        BACKBONE_MODEL =  LlamaForSequenceClassification.from_pretrained( 
+            self.config.base_model_name,
+            config=LlamaConfig.from_pretrained(self.config.base_model_name),
+        )
+        BACKBONE_MODEL.config.pad_token_id = BACKBONE_MODEL.config.eos_token_id
+        BACKBONE_MODEL.config.output_hidden_states = True
+
+        for param in BACKBONE_MODEL.parameters():
+            param.requires_grad = False
+
+        return BACKBONE_MODEL
+
+    def __init__(self, config):
+        super(RewardModel, self).__init__(config)
+        
+        # use .base_model to remove lm_head
+        self.BASE_MODEL = self.create_base_model().base_model 
+
+        # regression head for reward prediction
+        self.regression_head = nn.Linear(self.BASE_MODEL.config.hidden_size, config.reward_dim)
+        
+    def forward(self, input_ids, attention_mask=None, rewards=None, **kwargs):
+
+        # forward pass through the base model
+        outputs = self.BASE_MODEL(input_ids, attention_mask=attention_mask, **kwargs)
+        
+        hidden_states = outputs.hidden_states[-1]
+
+        # access hidden state corresponding to the last token in each sequence across the batch
+        last_token_hidden_state = hidden_states[:, -1, :] 
+        reward_predictions = self.regression_head(last_token_hidden_state)
+
+        return reward_predictions
+
+    def prepare_inputs_for_generation(self, *args, **kwargs):
+        return self.BASE_MODEL.prepare_inputs_for_generation(*args, **kwargs)
+
 model = AutoModel.from_pretrained("yale-nlp/MDCureRM").to(torch.device("cuda"))
 tokenizer = AutoTokenizer.from_pretrained("yale-nlp/MDCureRM", use_fast=True)
 tokenizer.pad_token = tokenizer.eos_token