Update README.md

README.md

@@ -36,16 +36,17 @@ It can play a role in the safe RLHF algorithm, helping the Beaver model become m
 - **Reward Model:** <https://huggingface.co/PKU-Alignment/beaver-7b-v1.0-reward>
 - **Cost Model:** <https://huggingface.co/PKU-Alignment/beaver-7b-v1.0-cost>
 - **Dataset Paper:** <https://arxiv.org/abs/2307.04657>
-- **Paper:** *Coming soon...*
+- **Paper:** <https://arxiv.org/abs/2310.12773>
 
 ## How to Use the Reward Model
 
 ```python
+import torch
 from transformers import AutoTokenizer
 from safe_rlhf.models import AutoModelForScore
 
-model = AutoModelForScore.from_pretrained('PKU-Alignment/beaver-7b-v1.0-reward', device_map='auto')
-tokenizer = AutoTokenizer.from_pretrained('PKU-Alignment/beaver-7b-v1.0-reward', use_fast=False)
+model = AutoModelForScore.from_pretrained('PKU-Alignment/beaver-7b-v1.0-reward', torch_dtype=torch.bfloat16, device_map='auto')
+tokenizer = AutoTokenizer.from_pretrained('PKU-Alignment/beaver-7b-v1.0-reward')
 
 input = 'BEGINNING OF CONVERSATION: USER: hello ASSISTANT:Hello! How can I help you today?'
 
@@ -54,34 +55,45 @@ output = model(**input_ids)
 print(output)
 
 # ScoreModelOutput(
-#     scores=tensor([[[-19.6476],
-#         [-20.2238],
-#         [-21.4228],
-#         [-19.2506],
-#         [-20.2728],
-#         [-23.8799],
-#         [-22.6898],
-#         [-21.5825],
-#         [-21.0855],
-#         [-20.2068],
-#         [-23.8296],
-#         [-21.4940],
-#         [-21.9484],
-#         [-13.1220],
-#         [ -6.4499],
-#         [ -8.1982],
-#         [ -7.2492],
-#         [ -9.3377],
-#         [-13.5010],
-#         [-10.4932],
-#         [ -9.7837],
-#         [ -6.4540],
-#         [ -6.0084],
-#         [ -5.8093],
-#         [ -6.6134],
-#         [ -5.8995],
-#         [ -9.1505],
-#         [-11.3254]]], grad_fn=<ToCopyBackward0>),
-#     end_scores=tensor([[-11.3254]], grad_fn=<ToCopyBackward0>)
+#     scores=tensor([[[-19.7500],
+#          [-19.3750],
+#          [-20.1250],
+#          [-18.0000],
+#          [-20.0000],
+#          [-23.8750],
+#          [-23.5000],
+#          [-22.0000],
+#          [-21.0000],
+#          [-20.1250],
+#          [-23.7500],
+#          [-21.6250],
+#          [-21.7500],
+#          [-12.9375],
+#          [ -6.4375],
+#          [ -8.1250],
+#          [ -7.3438],
+#          [ -9.1875],
+#          [-13.6250],
+#          [-10.5625],
+#          [ -9.9375],
+#          [ -6.4375],
+#          [ -6.0938],
+#          [ -5.8438],
+#          [ -6.6562],
+#          [ -5.9688],
+#          [ -9.1875],
+#          [-11.4375]]], grad_fn=<ToCopyBackward0>),
+#     end_scores=tensor([[-11.4375]], grad_fn=<ToCopyBackward0>),
+#     last_hidden_state=tensor([[[ 0.7461, -0.6055, -0.4980,  ...,  0.1670,  0.7812, -0.3242],
+#          [ 0.7383, -0.5391, -0.1836,  ..., -0.1396,  0.5273, -0.2256],
+#          [ 0.6836, -0.7031, -0.3730,  ...,  0.2100,  0.5000, -0.6328],
+#          ...,
+#          [-1.7969,  1.0234,  1.0234,  ..., -0.8047,  0.2500, -0.8398],
+#          [ 2.0469, -1.3203,  0.8984,  ..., -0.7734, -1.4141, -1.6797],
+#          [ 4.3438, -0.6953,  0.9648,  ..., -0.1787,  0.6680, -3.0000]]],
+#        dtype=torch.bfloat16, grad_fn=<ToCopyBackward0>),
+#     end_last_hidden_state=tensor([[ 4.3438, -0.6953,  0.9648,  ..., -0.1787,  0.6680, -3.0000]],
+#        dtype=torch.bfloat16, grad_fn=<ToCopyBackward0>),
+#     end_index=tensor([27])
 # )
 ```