Update app.py

app.py

@@ -2,19 +2,25 @@ import gradio as gr
 from transformers import pipeline, set_seed
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
+# Set seed and define model parameters
 set_seed(42)
-num_labels=2
-id2label = {0:'benign',1:'phishing'}
-label2id = {'benign':0,'phishing':1}
+num_labels = 2
+id2label = {0: 'benign', 1: 'phishing'}
+label2id = {'benign': 0, 'phishing': 1}
 checkpoint = 'bgspaditya/distilbert-phish'
+
+# 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)
-    return {'label': result[0]['label'], 'score': result[0]['score']}
+    predicted_label = result[0]['label']
+    return predicted_label
 
+# Define Gradio interface
 gradio_app = gr.Interface(
     predict,
     inputs=gr.Textbox(label="Enter URL"),
@@ -22,5 +28,6 @@ gradio_app = gr.Interface(
     title="Phishing URL Detection",
 )
 
+# Launch the Gradio interface
 if __name__ == "__main__":
     gradio_app.launch()