Update README.md

README.md

@@ -41,7 +41,7 @@ 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
+from transformers import AutoTokenizer, AutoModel, AutoConfig, LlamaConfig, PreTrainedModel, LlamaForSequenceClassification
 import torch.nn as nn
 import torch
 
@@ -101,6 +101,9 @@ class RewardModel(PreTrainedModel):
     def prepare_inputs_for_generation(self, *args, **kwargs):
         return self.BASE_MODEL.prepare_inputs_for_generation(*args, **kwargs)
 
+AutoConfig.register("RewardModel", RewardModelConfig)
+AutoModel.register(RewardModelConfig, RewardModel)
+
 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