Update pentest_ai_streamlit.py

pentest_ai_streamlit.py

@@ -1,13 +1,15 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import json
 
+# Load the model and tokenizer
 @st.cache(allow_output_mutation=True)
 def load_model():
     model_path = "Canstralian/pentest_ai"
     model = AutoModelForCausalLM.from_pretrained(
         model_path,
-        torch_dtype=torch.float16,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,  # Use float16 if CUDA is available
         device_map="auto",
         load_in_4bit=False,
         load_in_8bit=True,
@@ -16,21 +18,55 @@ def load_model():
     tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
     return model, tokenizer
 
+# Function to generate text from the model
 def generate_text(model, tokenizer, instruction):
-    tokens = tokenizer.encode(instruction, return_tensors='pt').to('cuda')
+    # Check if CUDA is available and send tensors to the appropriate device
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    tokens = tokenizer.encode(instruction, return_tensors='pt').to(device)
     generated_tokens = model.generate(
-        tokens, 
-        max_length=1024, 
-        top_p=1.0, 
-        temperature=0.5, 
+        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()
+# Load the JSON data (simulated here for simplicity)
+@st.cache(allow_output_mutation=True)
+def load_json_data():
+    json_data = [
+        {"name": "Raja Clarke", "email": "[email protected]", "country": "Chile", "company": "Urna Nunc Consulting"},
+        {"name": "Melissa Hobbs", "email": "[email protected]", "country": "France", "company": "Gravida Mauris Limited"},
+        {"name": "John Doe", "email": "[email protected]", "country": "USA", "company": "Example Corp"},
+        {"name": "Jane Smith", "email": "[email protected]", "country": "Canada", "company": "Innovative Solutions Inc"}
+    ]
+    return json_data
 
+# Streamlit UI
 st.title("Penetration Testing AI Assistant")
-instruction = st.text_area("Enter your question:")
+
+# Load model and tokenizer
+model, tokenizer = load_model()
+
+# Generate some text based on user input
+instruction = st.text_area("Enter your question for the AI assistant:")
 if st.button("Generate"):
-    response = generate_text(model, tokenizer, instruction)
-    st.write(response)
+    if instruction:
+        response = generate_text(model, tokenizer, instruction)
+        st.subheader("Generated Response:")
+        st.write(response)
+    else:
+        st.warning("Please enter a question to generate a response.")
+
+# Displaying user data from JSON
+st.subheader("User Data (from JSON)")
+user_data = load_json_data()
+
+# Display user details in a readable format
+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("---")  # Separator