Update app.py

app.py

@@ -1,7 +1,6 @@
-import streamlit as st
+import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-import json
 import logging
 import re
 
@@ -13,7 +12,6 @@ logging.basicConfig(
 )
 
 # Model and tokenizer loading function with caching
-@st.cache_resource
 def load_model():
     """
     Loads and caches the pre-trained language model and tokenizer.
@@ -36,7 +34,6 @@ def load_model():
         return model, tokenizer
     except Exception as e:
         logging.error(f"Error loading model: {e}")
-        st.error("Failed to load model. Please check the logs.")
         return None, None
 
 def sanitize_input(text):
@@ -83,57 +80,35 @@ def generate_text(model, tokenizer, instruction):
         logging.error(f"Error generating text: {e}")
         return "Error in text generation."
 
-@st.cache_data
-def load_json_data():
+# Gradio Interface Function
+def gradio_interface(instruction):
     """
-    Loads JSON data, simulating the loading process with a sample list.
-    Returns:
-        list: A list of dictionaries with sample user data.
+    Interface function for Gradio to interact with the model and generate text.
     """
-    try:
-        json_data = [
-            {"name": "Raja Clarke", "email": "consectetuer@yahoo.edu", "country": "Chile", "company": "Urna Nunc Consulting"},
-            {"name": "Melissa Hobbs", "email": "massa.non@hotmail.couk", "country": "France", "company": "Gravida Mauris Limited"},
-            {"name": "John Doe", "email": "john.doe@example.com", "country": "USA", "company": "Example Corp"},
-            {"name": "Jane Smith", "email": "jane.smith@example.org", "country": "Canada", "company": "Innovative Solutions Inc"}
-        ]
-        logging.info("User JSON data loaded successfully.")
-        return json_data
-    except Exception as e:
-        logging.error(f"Error loading JSON data: {e}")
-        return []
-
-# Streamlit App
-st.title("Penetration Testing AI Assistant")
+    # Load the model and tokenizer
+    model, tokenizer = load_model()
 
-# Load the model and tokenizer
-model, tokenizer = load_model()
+    if not model or not tokenizer:
+        return "Failed to load model or tokenizer. Please check your configuration."
 
-if not model or not tokenizer:
-    st.error("Failed to load model or tokenizer. Please check your configuration.")
-
-# User instruction input
-instruction = st.text_input("Enter an instruction for the model:")
-
-# Generate text button
-if instruction:
+    # Generate the text
     try:
         generated_text = generate_text(model, tokenizer, instruction)
-        st.subheader("Generated Text:")
-        st.write(generated_text)
+        return generated_text
     except ValueError as ve:
-        st.error(f"Invalid input: {ve}")
+        return f"Invalid input: {ve}"
     except Exception as e:
         logging.error(f"Error during text generation: {e}")
-        st.error("An error occurred. Please try again.")
+        return "An error occurred. Please try again."
 
-# Display JSON user data
-st.subheader("User Data (from JSON)")
-user_data = load_json_data()
+# Create Gradio Interface
+iface = gr.Interface(
+    fn=gradio_interface,
+    inputs=gr.Textbox(label="Enter an instruction for the model:", placeholder="Type your instruction here..."),
+    outputs=gr.Textbox(label="Generated Text:"),
+    title="Penetration Testing AI Assistant",
+    description="This tool allows you to interact with a pre-trained AI model for penetration testing assistance. Enter an instruction to generate a response.",
+)
 
-for user in user_data:
-    st.write(f"**Name:** {user['name']}")
-    st.write(f"**Email:** {user['email']}")
-    st.write(f"**Country:** {user['country']}")
-    st.write(f"**Company:** {user['company']}")
-    st.write("---")
+# Launch the Gradio interface
+iface.launch()