Update app.py

app.py

@@ -4,6 +4,7 @@ import pickle
 import numpy as np
 import requests
 from ProGPT import Conversation
+from sklearn.preprocessing import LabelEncoder
 
 # Load saved components
 with open('preprocessing_params.pkl', 'rb') as f:
@@ -20,12 +21,12 @@ with open('html_tokenizer.pkl', 'rb') as f:
 # Load the model with custom loss
 @tf.keras.utils.register_keras_serializable()
 class EWCLoss(tf.keras.losses.Loss):
-    def __init__(self, model, fisher_information, importance=1.0, reduction='auto', name=None):
+    def __init__(self, model=None, fisher_information=None, importance=1.0, reduction='auto', name=None):
         super(EWCLoss, self).__init__(reduction=reduction, name=name)
         self.model = model
         self.fisher_information = fisher_information
         self.importance = importance
-        self.prev_weights = [layer.numpy() for layer in model.trainable_weights]
+        self.prev_weights = [layer.numpy() for layer in model.trainable_weights] if model else None
 
     def call(self, y_true, y_pred):
         standard_loss = tf.keras.losses.binary_crossentropy(y_true, y_pred)
@@ -45,17 +46,17 @@ class EWCLoss(tf.keras.losses.Loss):
     
     @classmethod
     def from_config(cls, config):
-        # Load fisher information from external file
         with open('fisher_information.pkl', 'rb') as f:
             fisher_information = pickle.load(f)
         return cls(model=None, fisher_information=fisher_information, **config)
 
-# Load the model
-model = tf.keras.models.load_model('new_phishing_detection_model.keras',
-                                   custom_objects={'EWCLoss': EWCLoss})
+# Load the model first without the custom loss
+model = tf.keras.models.load_model('new_phishing_detection_model.keras', compile=False)
 
-# Recompile the model
+# Reconstruct the EWC loss
 ewc_loss = EWCLoss(model=model, fisher_information=fisher_information, importance=1000)
+
+# Compile the model with EWC loss and metrics
 model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0005),
               loss=ewc_loss,
               metrics=['accuracy', tf.keras.metrics.Precision(), tf.keras.metrics.Recall()])
@@ -74,9 +75,9 @@ def preprocess_input(input_text, tokenizer, max_length):
 def get_prediction(input_text, input_type):
     is_url = input_type == "URL"
     if is_url:
-        input_data = preprocess_input(input_text, url_tokenizer, preprocessing_params['max_url_length'])
+        input_data = preprocess_input(input_text, url_tokenizer, preprocessing_params['max_new_url_length'])
     else:
-        input_data = preprocess_input(input_text, html_tokenizer, preprocessing_params['max_html_length'])
+        input_data = preprocess_input(input_text, html_tokenizer, preprocessing_params['max_new_html_length'])
     
     prediction = model.predict([input_data, input_data])[0][0]
     return prediction
@@ -106,12 +107,25 @@ def chatbot_response(user_input):
     response = chatbot.prompt(user_input)
     return response
 
+def interface(input_text, input_type):
+    result = phishing_detection(input_text, input_type)
+    latest_sites = latest_phishing_sites()
+    chatbot_res = chatbot_response(input_text)
+    return result, latest_sites, chatbot_res
+
 iface = gr.Interface(
-    fn=phishing_detection,
-    inputs=[gr.inputs.Textbox(lines=5, placeholder="Enter URL or HTML code"), gr.inputs.Radio(["URL", "HTML"], type="value", label="Input Type")],
-    outputs="text",
+    fn=interface,
+    inputs=[
+        gr.inputs.Textbox(lines=5, placeholder="Enter URL or HTML code"), 
+        gr.inputs.Radio(["URL", "HTML"], type="value", label="Input Type")
+    ],
+    outputs=[
+        "text",
+        gr.outputs.Textbox(label="Latest Phishing Sites"),
+        gr.outputs.Textbox(label="Chatbot Response")
+    ],
     title="Phishing Detection with Enhanced EWC Model",
-    description="Check if a URL or HTML is Phishing",
+    description="Check if a URL or HTML is Phishing. Latest phishing sites from PhishTank and a chatbot assistant for phishing issues.",
     theme="default"
 )