pgurazada1 commited on
Commit
0c99d9c
·
verified ·
1 Parent(s): 65d9bd9

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +37 -0
app.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import gradio as gr
3
+
4
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
5
+
6
+
7
+ tokenizer = AutoTokenizer.from_pretrained("ProtectAI/deberta-v3-base-prompt-injection")
8
+ model = AutoModelForSequenceClassification.from_pretrained("ProtectAI/deberta-v3-base-prompt-injection")
9
+
10
+ classifier = pipeline(
11
+ "text-classification",
12
+ model=model,
13
+ tokenizer=tokenizer,
14
+ truncation=True,
15
+ max_length=512,
16
+ device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
17
+ )
18
+
19
+ def predict(user_input: str):
20
+
21
+ return classifier(user_input)
22
+
23
+
24
+ textbox = gr.Textbox(placeholder="Enter user input presented for injection attack classification", lines=12)
25
+
26
+ interface = gr.Interface(
27
+ inputs=textbox, fn=predict, outputs="text",
28
+ title="Injection Attack Classifier",
29
+ description="This web API flags if the text presented as input to an LLM qualifies to be an injection attack",
30
+ allow_flagging="manual", flagging_options=["Useful", "Not Useful"]
31
+ )
32
+
33
+ with gr.Blocks() as demo:
34
+ interface.launch()
35
+
36
+ demo.queue(concurrency_count=4)
37
+ demo.launch()