Update README.md

README.md

@@ -37,6 +37,8 @@ import torch.nn.functional as F
 tokenizer = AutoTokenizer.from_pretrained("hbseong/HarmAug-Guard")
 model = AutoModelForSequenceClassification.from_pretrained("hbseong/HarmAug-Guard")
 
+# If response is not given, the model will predict the unsafe score of the prompt.
+# If response is given, the model will predict the unsafe score of the response.
 def predict(prompt, response=None):
     if response == None:
         inputs = tokenizer(prompt, return_tensors="pt")