app.py

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("ZySec-AI/SecurityLLM")
model = AutoModelForCausalLM.from_pretrained("ZySec-AI/SecurityLLM")

# Example prompt for generating a response related to security
security_prompt = "Analyze the following network log for potential security issues: 2024-11-08 12:30:00 SRC_IP=192.168.1.1 DEST_IP=10.0.0.5 PROTOCOL=TCP PACKET_SIZE=1500 SRC_PORT=443 DEST_PORT=80"

# Tokenize the prompt
inputs = tokenizer(security_prompt, return_tensors="pt")

# Generate a response from the model
output = model.generate(inputs['input_ids'], max_length=150, num_return_sequences=1, no_repeat_ngram_size=2)

# Decode and print the generated text
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
print("Generated Response:\n", generated_text)