Update pentest_ai_streamlit.py

pentest_ai_streamlit.py

@@ -1,49 +1,36 @@
 import streamlit as st
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Load the model and tokenizer from Hugging Face
-MODEL_NAME = "Canstralian/pentest_ai"
-
-@st.cache_resource
+@st.cache(allow_output_mutation=True)
 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.")
+    model_path = "Canstralian/pentest_ai"
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        load_in_4bit=False,
+        load_in_8bit=True,
+        trust_remote_code=True,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    return model, tokenizer
+
+def generate_text(model, tokenizer, instruction):
+    tokens = tokenizer.encode(instruction, return_tensors='pt').to('cuda')
+    generated_tokens = model.generate(
+        tokens, 
+        max_length=1024, 
+        top_p=1.0, 
+        temperature=0.5, 
+        top_k=50
+    )
+    return tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
+
+model, tokenizer = load_model()
+
+st.title("Penetration Testing AI Assistant")
+instruction = st.text_area("Enter your question:")
+if st.button("Generate"):
+    response = generate_text(model, tokenizer, instruction)
+    st.write(response)