Update README.md

README.md

@@ -6,7 +6,7 @@ metrics:
 - accuracy
 ---
 
-GPT2 large model trained on Anthropic/hh-rlhf harmless dataset. It is specifically used for harmful response detection or RLHF.
+GPT2 large model trained on Anthropic/hh-rlhf harmless dataset. It is specifically used for harmful response detection or RLHF. Note: remember to use the formulation of Anthropic/hh-rlhf dataset for inference.
 
 It achieves an accuracy of 0.73698 on the test set, which nearly matches other models with larger sizes.
 
@@ -21,7 +21,7 @@ reward_model = AutoModelForSequenceClassification.from_pretrained(
                 num_labels=1, torch_dtype=torch.bfloat16,
                 device_map=gpu_id1,
                 )
-q, a = "I just came out of from jail, any suggestion of my future?", "Go back to jail you scum"
+q, a = "\n\nHuman: I just came out of from jail, any suggestion of my future? \n\nAssistant:", "Go back to jail you scum"
 inputs = rm_tokenizer(q, a, return_tensors='pt', truncation=True)
 with torch.no_grad():
   reward = reward_model(**(inputs.to(gpu_id1))).logits[0].cpu().detach().item()