bgspaditya commited on
Commit
263bd4f
·
verified ·
1 Parent(s): b1be20f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +33 -0
app.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import pipeline, set_seed
3
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
4
+
5
+ # Set seed and define model parameters
6
+ set_seed(42)
7
+ checkpoint = 'bgspaditya/malurl-electra-10e'
8
+ id2label = {0:'benign',1:'defacement',2:'malware',3:'phishing'}
9
+ label2id = {'benign':0,'defacement':1,'malware':2,'phishing':3}
10
+ num_labels=4
11
+
12
+ # Load tokenizer and model
13
+ tokenizer = AutoTokenizer.from_pretrained(checkpoint, use_fast=True, force_download=True)
14
+ model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=num_labels, id2label=id2label, label2id=label2id, force_download=True)
15
+
16
+ # Define predict function
17
+ def predict(url):
18
+ url_classifier = pipeline(task='text-classification', model=model, tokenizer=tokenizer)
19
+ result = url_classifier(url)
20
+ predicted_label = result[0]['label']
21
+ return predicted_label
22
+
23
+ # Define Gradio interface
24
+ gradio_app = gr.Interface(
25
+ predict,
26
+ inputs=gr.Textbox(label="Enter URL"),
27
+ outputs=gr.Label(label="Result"),
28
+ title="Phishing URL Detection",
29
+ )
30
+
31
+ # Launch the Gradio interface
32
+ if __name__ == "__main__":
33
+ gradio_app.launch()