Update app.py

app.py

@@ -17,15 +17,15 @@ current_num = os.getenv("NUM")
 print(f"stage ${current_num}")
 
 api = HfApi(token=hf_token)
-models = f"dad1909/cybersentinal-2.0-{current_num}"
+# models = f"dad1909/cybersentinal-2.0-{current_num}"
 
-# model_base = "dad1909/cybersentinal-2.0"
+model_base = "unsloth/llama-3-8b-Instruct-bnb-4bit"
 
 print("Starting model and tokenizer loading...")
 
 # Load the model and tokenizer
 model, tokenizer = FastLanguageModel.from_pretrained(
-    model_name=models,
+    model_name=model_base,
     max_seq_length=max_seq_length,
     dtype=dtype,
     load_in_4bit=load_in_4bit,
@@ -119,14 +119,13 @@ trainer = SFTTrainer(
     dataset_num_proc=2,
     packing=False,
     args=TrainingArguments(
-        per_device_train_batch_size=5,
-        gradient_accumulation_steps=5,
+        per_device_train_batch_size=20,
+        gradient_accumulation_steps=20,
         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",
@@ -143,7 +142,7 @@ print("Training completed.")
 num = int(current_num)
 num += 1
 
-uploads_models = f"cybersentinal-2.0-{str(num)}"
+uploads_models = f"cybersentinal-3.0"
 
 print("Saving the trained model...")
 model.save_pretrained_merged("model", tokenizer, save_method="merged_16bit")