|
import gradio as gr |
|
import transformers |
|
from transformers import BertTokenizer, DataCollatorWithPadding |
|
from transformers import AutoModelForSequenceClassification |
|
tokenizer = BertTokenizer.from_pretrained('huolongguo10/check_sec') |
|
model = AutoModelForSequenceClassification.from_pretrained('huolongguo10/check_sec', num_labels=2) |
|
import torch |
|
def check(text): |
|
with torch.no_grad(): |
|
logits = model(**inputs).logits |
|
predicted_class_id = logits.argmax().item() |
|
return 'secure' if predicted_class_id==0 else 'insecure' |
|
with gr.Blocks() as demo: |
|
text = gr.Textbox(label="Text") |
|
output = gr.Textbox(label="Output Box") |
|
org = gr.Textbox(label="By normal check") |
|
greet_btn = gr.Button("Check!") |
|
greet_btn.click(fn=check, inputs=text, outputs=output, api_name="check") |
|
gr.Markdown('''# check_sec |
|
检查web参数安全性，目前只支持xss |
|
## 类型 |
|
``` |
|
LABEL_0: secure |
|
LABEL_1: insecure(可能包含xss payload) |
|
``` |
|
''') |
|
|
|
demo.launch() |
