Fixed app v2

app.py

@@ -1,19 +1,7 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from transformers import pipeline
 
-# 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)
+messages = [
+    {"role": "user", "content": "Who are you?"},
+]
+pipe = pipeline("text-generation", model="segolilylabs/Lily-Cybersecurity-7B-v0.2")
+pipe(messages)