Create pentest_ai_streamlit.py

pentest_ai_streamlit.py

@@ -0,0 +1,49 @@
+import streamlit as st
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+
+# Load the model and tokenizer from Hugging Face
+MODEL_NAME = "Canstralian/pentest_ai"
+
+@st.cache_resource
+def load_model():
+    # Load the tokenizer and model
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+# Load the pentest_ai model
+pentest_ai = load_model()
+
+# Streamlit interface setup
+st.title("Pentest AI Assistant")
+st.write("This tool allows you to ask penetration testing and cybersecurity-related queries, and it will generate AI-powered suggestions or commands.")
+
+# Text input for user's question
+user_input = st.text_area("Enter your question or command:")
+
+# Button to trigger generation
+if st.button("Generate Response"):
+    if user_input.strip() == "":
+        st.error("Please enter a valid input.")
+    else:
+        # Generate response using the model
+        with st.spinner("Generating response..."):
+            response = pentest_ai(user_input, max_length=150, num_return_sequences=1)[0]['generated_text']
+        
+        # Display the model's response
+        st.subheader("AI Response:")
+        st.write(response)
+
+# Add an example button to help users see a sample
+if st.button("Show Example"):
+    example_query = "How do I scan a network for open ports?"
+    with st.spinner("Generating response for example query..."):
+        example_response = pentest_ai(example_query, max_length=150, num_return_sequences=1)[0]['generated_text']
+    
+    st.subheader("Example Query:")
+    st.write(example_query)
+    st.subheader("AI Response:")
+    st.write(example_response)
+
+# Instructions for the user
+st.info("Note: This AI model provides general advice. Always ensure you're testing on systems you have permission to, and follow legal and ethical guidelines.")