Update app.py

app.py

@@ -1,63 +1,75 @@
 
 
-
-
 import streamlit as st
-import tensorflow as tf
 import numpy as np
-import pickle
-
-# Set page title and header
-st.set_page_config(page_title="Prompt Injection Detection and Prevention")
-st.title("Prompt Injection Detection and Prevention")
-st.subheader("Classify prompts as malicious or valid and understand predictions using LIME.")
+import pandas as pd
+from tensorflow.keras.preprocessing.text import Tokenizer
+from tensorflow.keras.preprocessing.sequence import pad_sequences
+from tensorflow.keras.models import load_model
+from lime.lime_text import LimeTextExplainer
+import matplotlib.pyplot as plt
+import seaborn as sns
 
 # Load the trained model
 @st.cache_resource
-def load_model(model_path):
-    try:
-        return tf.keras.models.load_model(model_path)
-    except Exception as e:
-        st.error(f"Error loading model: {e}")
-        return None
-
-# Load the tokenizer
-@st.cache_resource
-def load_tokenizer(tokenizer_path):
-    try:
-        with open(tokenizer_path, "rb") as f:
-            return pickle.load(f)
-    except Exception as e:
-        st.error(f"Error loading tokenizer: {e}")
-        return None
-
-# Paths to your files (these should be present in your Hugging Face repository)
-MODEL_PATH = "model.h5"
-TOKENIZER_PATH = "tokenizer.pkl"
-
-# Load model and tokenizer
-model = load_model(MODEL_PATH)
-tokenizer = load_tokenizer(TOKENIZER_PATH)
-
-if model and tokenizer:
-    st.success("Model and tokenizer loaded successfully!")
-
-# User input for prompt classification
-st.write("## Classify a Prompt")
-user_input = st.text_area("Enter a prompt for classification:")
-if st.button("Classify"):
-    if user_input:
-        # Preprocess the user input
-        sequence = tokenizer.texts_to_sequences([user_input])
-        padded_sequence = tf.keras.preprocessing.sequence.pad_sequences(sequence, maxlen=50)
-
-        # Make prediction
-        prediction = model.predict(padded_sequence)
-        label = "Malicious" if prediction[0] > 0.5 else "Valid"
-        st.write(f"Prediction: **{label}** (Confidence: {prediction[0][0]:.2f})")
-    else:
-        st.error("Please enter a prompt for classification.")
-
-# Footer
-st.write("---")
-st.caption("Developed for detecting and preventing prompt injection attacks.")
+def load_trained_model():
+    model = load_model("deep_learning_model.h5")
+    return model
+
+model = load_trained_model()
+
+# Tokenizer setup
+tokenizer = Tokenizer(num_words=5000)
+max_length = 100
+
+# Load Data
+@st.cache_data
+def load_data():
+    data = pd.read_csv("train prompt.csv", sep=',', quoting=3, encoding='ISO-8859-1', on_bad_lines='skip', engine='python')
+    data['label'] = data['label'].replace({'valid': 0, 'malicious': 1})
+    return data
+
+data = load_data()
+tokenizer.fit_on_texts(data['input'].values)
+
+# Preprocessing functions
+def preprocess_prompt(prompt, tokenizer, max_length):
+    sequence = tokenizer.texts_to_sequences([prompt])
+    padded_sequence = pad_sequences(sequence, maxlen=max_length)
+    return padded_sequence
+
+def detect_prompt(prompt, model, tokenizer, max_length):
+    processed_prompt = preprocess_prompt(prompt, tokenizer, max_length)
+    prediction = model.predict(processed_prompt)[0][0]
+    class_label = 'Malicious' if prediction >= 0.5 else 'Valid'
+    confidence_score = prediction * 100 if prediction >= 0.5 else (1 - prediction) * 100
+    return class_label, confidence_score
+
+# Load model
+model = load_trained_model()
+
+# Streamlit app
+st.title("Prompt Injection Attack Detection")
+st.write("This application detects malicious prompts to prevent injection attacks.")
+
+prompt = st.text_input("Enter a prompt to analyze:")
+
+if prompt:
+    class_label, confidence = detect_prompt(prompt, model, tokenizer, max_length)
+    st.write(f"### Prediction: {class_label}")
+    st.write(f"Confidence: {confidence:.2f}%")
+
+    # LIME explanation
+    st.write("Generating LIME Explanation...")
+    explainer = LimeTextExplainer(class_names=["Valid", "Malicious"])
+
+    def predict_fn(prompts):
+        sequences = tokenizer.texts_to_sequences(prompts)
+        padded_sequences = pad_sequences(sequences, maxlen=max_length)
+        predictions = model.predict(padded_sequences)
+        return np.hstack([1 - predictions, predictions])
+
+    explanation = explainer.explain_instance(prompt, predict_fn, num_features=10)
+    fig = explanation.as_pyplot_figure()
+    st.pyplot(fig)
+