malurl-detect / app.py
bgspaditya's picture
Create app.py
263bd4f verified
import gradio as gr
from transformers import pipeline, set_seed
from transformers import AutoTokenizer, AutoModelForSequenceClassification
# Set seed and define model parameters
set_seed(42)
checkpoint = 'bgspaditya/malurl-electra-10e'
id2label = {0:'benign',1:'defacement',2:'malware',3:'phishing'}
label2id = {'benign':0,'defacement':1,'malware':2,'phishing':3}
num_labels=4
# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(checkpoint, use_fast=True, force_download=True)
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=num_labels, id2label=id2label, label2id=label2id, force_download=True)
# Define predict function
def predict(url):
url_classifier = pipeline(task='text-classification', model=model, tokenizer=tokenizer)
result = url_classifier(url)
predicted_label = result[0]['label']
return predicted_label
# Define Gradio interface
gradio_app = gr.Interface(
predict,
inputs=gr.Textbox(label="Enter URL"),
outputs=gr.Label(label="Result"),
title="Phishing URL Detection",
)
# Launch the Gradio interface
if __name__ == "__main__":
gradio_app.launch()