Fixed app v2

app.py

@@ -1,3 +1,19 @@
-import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-gr.load("models/WhiteRabbitNeo/Llama-3-WhiteRabbitNeo-8B-v2.0").launch()
+# 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)