elftsdmr commited on
Commit
ac54b7e
1 Parent(s): 56e3fda

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +30 -1
app.py CHANGED
@@ -1,3 +1,32 @@
1
  import gradio as gr
 
 
2
 
3
- gr.Interface.load("models/elftsdmr/malware-url-detect").launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import gradio as gr
2
+ from transformers import BertTokenizer, BertForSequenceClassification
3
+ from transformers import pipeline
4
 
5
+ urlbert = BertForSequenceClassification.from_pretrained('elftsdmr/malware-url-detect',num_labels=2)
6
+ tokenizer = BertTokenizer.from_pretrained('elftsdmr/malware-url-detect')
7
+
8
+ url_classifier = pipeline("text-classification", model=urlbert, tokenizer=tokenizer)
9
+
10
+
11
+
12
+
13
+
14
+ def text_to_url(text):
15
+ return url_classifier(text)[0]["label"]
16
+
17
+
18
+ demo = gr.Blocks(theme=gr.themes.Soft())
19
+
20
+ with demo:
21
+
22
+ text = gr.Textbox()
23
+ label = gr.Label()
24
+
25
+
26
+ b2 = gr.Button("Classify URL")
27
+
28
+
29
+ b2.click(text_to_url, inputs=text, outputs=label)
30
+
31
+ if __name__ == "__main__":
32
+ demo.launch(share=True)