Update app.py

app.py

@@ -5,6 +5,7 @@ from trl import SFTTrainer
 from transformers import TrainingArguments
 from datasets import load_dataset
 import gradio as gr
+import json
 from huggingface_hub import HfApi
 
 max_seq_length = 4096
@@ -16,7 +17,7 @@ print("Starting model and tokenizer loading...")
 
 # Load the model and tokenizer
 model, tokenizer = FastLanguageModel.from_pretrained(
-    model_name="unsloth/llama-3-8b-Instruct-bnb-4bit",
+    model_name="dad1909/CyberSentinel-initial",
     max_seq_length=max_seq_length,
     dtype=dtype,
     load_in_4bit=load_in_4bit,
@@ -117,7 +118,7 @@ trainer = SFTTrainer(
         bf16=is_bfloat16_supported(),
         warmup_steps=5,
         logging_steps=10,
-        max_steps=50,
+        max_steps=200,
         optim="adamw_8bit",
         weight_decay=0.01,
         lr_scheduler_type="linear",
@@ -137,7 +138,7 @@ print("Model saved successfully.")
 
 print("Pushing the model to the hub...")
 model.push_to_hub_merged(
-    "CyberSentinel-initial",
+    "CyberSentinel-1",
     tokenizer,
     save_method="merged_16bit",
     token=hf_token