Update app.py

app.py

@@ -1,47 +1,53 @@
 import gradio as gr
-import pandas as pd
 import joblib
+import pandas as pd
 
-# Load your trained model
+# Load trained model
 model = joblib.load("cybersecurity_model.joblib")
 
-# Exact features the model was trained on
-FEATURES = [
-    'A_frequency', 'NS_frequency', 'CNAME_frequency', 'SOA_frequency', 'NULL_frequency',
-    'PTR_frequency', 'HINFO_frequency', 'MX_frequency', 'TXT_frequency', 'AAAA_frequency',
-    'SRV_frequency', 'OPT_frequency', 'rr_type', 'rr_count', 'rr_name_entropy',
-    'rr_name_length', 'distinct_ns', 'distinct_ip', 'unique_country', 'unique_asn',
-    'distinct_domains', 'reverse_dns', 'a_records', 'unique_ttl', 'ttl_mean',
-    'ttl_variance', 'FQDN_count', 'subdomain_length', 'upper', 'lower',
-    'numeric', 'entropy', 'special', 'labels', 'labels_max', 'labels_average',
-    'longest_word', 'sld', 'len', 'subdomain'
+# List of features your model expects (update as per your actual features)
+FEATURE_NAMES = [
+    'A_frequency','NS_frequency','CNAME_frequency','SOA_frequency','NULL_frequency',
+    'PTR_frequency','HINFO_frequency','MX_frequency','TXT_frequency','AAAA_frequency',
+    'SRV_frequency','OPT_frequency','rr_type','rr_count','rr_name_entropy','rr_name_length',
+    'distinct_ns','distinct_ip','unique_country','unique_asn','distinct_domains','reverse_dns',
+    'a_records','unique_ttl','ttl_mean','ttl_variance','FQDN_count','subdomain_length','upper',
+    'lower','numeric','entropy','special','labels','labels_max','labels_average','longest_word',
+    'sld','len','subdomain'
 ]
 
+# Inference function
 def predict_from_csv(file):
     try:
         df = pd.read_csv(file)
-
-        # Ensure only expected features are used
-        data = df[FEATURES]
-
-        # Make predictions
-        predictions = model.predict(data)
-
-        df['Prediction'] = predictions
-
-        return df[['Prediction']]
-
+        
+        # Ensure required columns are present
+        missing_cols = set(FEATURE_NAMES) - set(df.columns)
+        if missing_cols:
+            return f"❌ Missing columns in CSV: {', '.join(missing_cols)}"
+
+        # Select and predict
+        X = df[FEATURE_NAMES]
+        preds = model.predict(X)
+        
+        # Format output nicely
+        result = pd.DataFrame({"Prediction": preds})
+        return result
     except Exception as e:
-        return f"❌ Error: {str(e)}"
+        return f"❌ Error reading file: {str(e)}"
 
-# Gradio Interface
-iface = gr.Interface(
+# Gradio UI
+interface = gr.Interface(
     fn=predict_from_csv,
     inputs=gr.File(label="📁 Upload CSV File with Network Traffic Features"),
-    outputs="dataframe",
+    outputs=gr.Dataframe(label="📊 Predictions"),
     title="🚨 Cybersecurity Attack Detector",
-    description="AI-powered model to detect attacks from DNS/network traffic data. Upload a CSV file with preprocessed features to get predictions."
+    description=(
+        "AI-powered model to detect attacks from DNS/network traffic data. "
+        "Upload a CSV file with preprocessed features to get predictions."
+    ),
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    interface.launch()