Update app.py

app.py

@@ -17,7 +17,7 @@ current_num = os.getenv("NUM")
 print(f"stage ${current_num}")
 
 api = HfApi(token=hf_token)
-models = "unsloth/Meta-Llama-3.1-70B-bnb-4bit"
+models = "dad1909/cybersentinal-2.0"
 
 print("Starting model and tokenizer loading...")
 
@@ -112,13 +112,14 @@ trainer = SFTTrainer(
     dataset_num_proc=2,
     packing=False,
     args=TrainingArguments(
-        per_device_train_batch_size=1,
-        gradient_accumulation_steps=1,
+        per_device_train_batch_size=5,
+        gradient_accumulation_steps=5,
         learning_rate=2e-4,
         fp16=not is_bfloat16_supported(),
         bf16=is_bfloat16_supported(),
         warmup_steps=5,
         logging_steps=10,
+        max_steps=200,
         optim="adamw_8bit",
         weight_decay=0.01,
         lr_scheduler_type="linear",
@@ -135,9 +136,7 @@ print("Training completed.")
 num = int(current_num)
 num += 1
 
-uploads_models = f"cybersentinal-2.0-{str(num)}"
-
-up = "sentinal-3.1-70B"
+up = "sentinal-2"
 
 print("Saving the trained model...")
 model.save_pretrained_merged("model", tokenizer, save_method="merged_16bit")