Update app.py

app.py

@@ -1,88 +1,59 @@
 import streamlit as st
 import tensorflow as tf
-from tensorflow.keras.preprocessing.text import Tokenizer
-from tensorflow.keras.preprocessing.sequence import pad_sequences
 import numpy as np
-from lime.lime_text import LimeTextExplainer
-import matplotlib.pyplot as plt
+import pickle
 
-# Streamlit Title
+# Set page title and header
+st.set_page_config(page_title="Prompt Injection Detection and Prevention")
 st.title("Prompt Injection Detection and Prevention")
-st.write("Classify prompts as malicious or valid and understand predictions using LIME.")
+st.subheader("Classify prompts as malicious or valid and understand predictions using LIME.")
 
-# Cache Model Loading
+# Load the trained model
 @st.cache_resource
-def load_model(filepath):
-    return tf.keras.models.load_model(filepath)
+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
 
-# Tokenizer Setup
+# Load the tokenizer
 @st.cache_resource
-def setup_tokenizer():
-    tokenizer = Tokenizer(num_words=5000)
-    # Predefined vocabulary for demonstration purposes; replace with your actual tokenizer setup.
-    tokenizer.fit_on_texts(["example prompt", "malicious attack", "valid input prompt"])
-    return tokenizer
-
-# Preprocessing Function
-def preprocess_prompt(prompt, tokenizer, max_length=100):
-    sequence = tokenizer.texts_to_sequences([prompt])
-    return pad_sequences(sequence, maxlen=max_length)
-
-# Prediction Function
-def detect_prompt(prompt, tokenizer, model):
-    processed_prompt = preprocess_prompt(prompt, tokenizer)
-    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
-
-# LIME Explanation
-def lime_explain(prompt, model, tokenizer, max_length=100):
-    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
-    )
-    return explanation
-
-# Load Model Section
-st.subheader("Load Your Trained Model")
-model_path = st.text_input("Enter the path to your trained model (.h5):")
-model = None
-tokenizer = None
-
-if model_path:
+def load_tokenizer(tokenizer_path):
     try:
-        model = load_model(model_path)
-        tokenizer = setup_tokenizer()
-        st.success("Model Loaded Successfully!")
-        
-        # User Prompt Input
-        st.subheader("Classify Your Prompt")
-        user_prompt = st.text_input("Enter a prompt to classify:")
-
-        if user_prompt:
-            class_label, confidence_score = detect_prompt(user_prompt, tokenizer, model)
-            st.write(f"Predicted Class: **{class_label}**")
-            st.write(f"Confidence Score: **{confidence_score:.2f}%**")
-
-            # LIME Explanation
-            st.subheader("LIME Explanation")
-            explanation = lime_explain(user_prompt, model, tokenizer)
-            explanation_as_html = explanation.as_html()
-            st.components.v1.html(explanation_as_html, height=500)
-
+        with open(tokenizer_path, "rb") as f:
+            return pickle.load(f)
     except Exception as e:
-        st.error(f"Error Loading Model: {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.write("Developed for detecting and preventing prompt injection attacks.")
+st.caption("Developed for detecting and preventing prompt injection attacks.")