Update app.py

app.py

@@ -8,7 +8,7 @@ from nltk.stem import WordNetLemmatizer
 from tensorflow.keras.preprocessing.text import Tokenizer
 from tensorflow.keras.preprocessing.sequence import pad_sequences
 from tensorflow import keras
-import pandas as pd
+from sklearn.preprocessing import LabelEncoder
 
 # Ensure necessary NLTK resources are downloaded
 nltk.download('punkt')
@@ -58,6 +58,10 @@ with open('url_tokenizer.pkl', 'rb') as file:
 with open('html_tokenizer.pkl', 'rb') as file:
     html_tokenizer = pickle.load(file)
 
+# Load the label encoder
+with open('label_encoder.pkl', 'rb') as file:
+    label_encoder = pickle.load(file)
+
 # Define the prediction function
 def predict_phishing(url, html):
     cleaned_url = preprocess_url(url)
@@ -72,10 +76,10 @@ def predict_phishing(url, html):
     new_predictions_prob = model.predict([new_url_padded, new_html_padded])
     new_predictions = (new_predictions_prob > 0.5).astype(int)
     
-    predicted_category = "Spam" if new_predictions[0][0] == 1 else "Legitimate"
+    predicted_category = label_encoder.inverse_transform(new_predictions)[0]
     predicted_probability = f"{new_predictions_prob[0][0]:.4f}"
     
-    return predicted_category, predicted_probability
+    return predicted_category.capitalize(), predicted_probability
 
 # Create Gradio Interface
 interface = gr.Interface(