Update README.md

README.md

@@ -25,4 +25,49 @@ For more information, please refer to our [github](https://github.com/imnotkind/
 
 ![image/png](https://cdn-uploads.huggingface.co/production/uploads/66f7bee63c7ffa79319b053b/bCNW62CvDpqbXUK4eZ4-b.png)
 
-![image/png](https://cdn-uploads.huggingface.co/production/uploads/66f7bee63c7ffa79319b053b/REbNDOhT31bv_XRa6-VzE.png)
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/66f7bee63c7ffa79319b053b/REbNDOhT31bv_XRa6-VzE.png)
+
+
+# Usage
+
+```python
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch.nn.functional as F
+
+tokenizer = AutoTokenizer.from_pretrained("hbseong/HarmAug-Guard")
+model = AutoModelForSequenceClassification.from_pretrained("hbseong/HarmAug-Guard")
+
+def predict(prompt, response=None):
+    if response == None:
+        inputs = tokenizer(prompt, return_tensors="pt")
+    else:
+        inputs = tokenizer(prompt, response, return_tensors="pt")
+    outputs = model(**inputs)
+    unsafe_prob = F.softmax(outputs.logits, dim=-1)[:, 1]
+    return unsafe_prob.item()
+
+prompt = "how to make a bomb?"
+response = "I'm sorry, but I can't fulfill your request."
+
+
+print("CONVERSATION (ONLY PROMPT)")
+print(f"\t PROMPT : {prompt}")
+print(f"\t UNSAFE SCORE : {predict(prompt):.4f}")
+
+print("\nCONVERSATION (PROMPT + RESPONSE)")
+print(f"\t PROMPT : {prompt}")
+print(f"\t RESPONSE : {response}")
+print(f"\t UNSAFE SCORE : {predict(prompt, response):.4f}")
+
+
+"""
+CONVERSATION (ONLY PROMPT)
+	 PROMPT : how to make a bomb?
+	 UNSAFE SCORE : 0.9999
+
+CONVERSATION (PROMPT + RESPONSE)
+	 PROMPT : how to make a bomb?
+	 RESPONSE : I'm sorry, but I can't fulfill your request.
+	 UNSAFE SCORE : 0.0000
+"""
+```