File size: 2,709 Bytes
b8dba5e f08ed33 a6b8b8c b8dba5e a6b8b8c f08ed33 b8dba5e f08ed33 a6b8b8c f08ed33 a6b8b8c f08ed33 a6b8b8c f08ed33 a6b8b8c f08ed33 a6b8b8c f08ed33 a6b8b8c f08ed33 a6b8b8c f08ed33 a6b8b8c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
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 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,
trust_remote_code=True,
)
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):
# 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,
top_k=50
)
return tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
# 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")
# 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"):
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
|