tfarhan10 commited on
Commit
af2fb62
·
verified ·
1 Parent(s): bff12f9

Upload app.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. app.py +29 -0
app.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import pipeline
3
+
4
+ # Load the jailbreak classifcation model
5
+ task = "text-classification"
6
+ pretrained_model = "jackhhao/jailbreak-classifier"
7
+ classifier = pipeline(task, model=pretrained_model)
8
+
9
+ #Define the prediction function
10
+ def classify_text(text):
11
+ result = classifier(text)
12
+ label = result[0]['label']
13
+ score = result[0]['score']
14
+ return f"Label: {label}, Confidence: {score:.4f}"
15
+
16
+ #Create the interface
17
+ iface = gr.Interface(
18
+ fn = classify_text,
19
+ inputs = gr.Textbox(lines=5, label="Enter a text"),
20
+ outputs = "text",
21
+ title = "Jailbreak Classification",
22
+ description = "A simple interface to classify text as jailbreak or not jailbreak"
23
+
24
+ )
25
+
26
+ # Launch the app
27
+ if __name__ == "__main__":
28
+ iface.launch()
29
+