huolongguo10 commited on
Commit
d2214de
·
1 Parent(s): 4dd94f1

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +17 -0
README.md CHANGED
@@ -16,4 +16,21 @@ tags:
16
  ```
17
  LABEL_0: secure
18
  LABEL_1: insecure(可能包含xss payload)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19
  ```
 
16
  ```
17
  LABEL_0: secure
18
  LABEL_1: insecure(可能包含xss payload)
19
+ ```
20
+
21
+ ## 使用
22
+ ```python
23
+ import transformers
24
+ from transformers import BertTokenizer, DataCollatorWithPadding
25
+ from transformers import AutoModelForSequenceClassification
26
+ tokenizer = BertTokenizer.from_pretrained('huolongguo10/check_sec')
27
+ model = AutoModelForSequenceClassification.from_pretrained('huolongguo10/check_sec', num_labels=2)
28
+ import torch
29
+ def check(text):
30
+ inputs = tokenizer(text, return_tensors="pt")
31
+ with torch.no_grad():
32
+ logits = model(**inputs).logits
33
+ predicted_class_id = logits.argmax().item()
34
+ print(f'{logits.argmax().item()}:{text}')
35
+ return 'secure' if predicted_class_id==0 else 'insecure'
36
  ```