app.py

import gradio as gr
from transformers import pipeline, set_seed
from transformers import AutoTokenizer, AutoModelForSequenceClassification

set_seed(42)
num_labels=2
id2label = {0:'benign',1:'phishing'}
label2id = {'benign':0,'phishing':1}
checkpoint = 'bgspaditya/distilbert-phish'
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)

def predict(url):
    url_classifier = pipeline(task='text-classification', model=model, tokenizer=tokenizer)
    result = url_classifier(url)
    return {'label': result[0]['label'], 'score': result[0]['score']}

gradio_app = gr.Interface(
    predict,
    inputs=gr.Textbox(label="Enter URL"),
    outputs=gr.Label(label="Result"),
    title="Phishing URL Detection",
)

if __name__ == "__main__":
    gradio_app.launch()