Update app.py

app.py

@@ -17,15 +17,14 @@ current_num = os.getenv("NUM")
 print(f"stage ${current_num}")
 
 api = HfApi(token=hf_token)
-models = f"dad1909/cybersentinal-2.0-{current_num}"
-
-# model_base = "dad1909/cybersentinal-2.0"
+# models = f"dad1909/cybersentinal-2.0-{current_num}"
+model_base = "unsloth/Meta-Llama-3.1-70B-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,
@@ -113,14 +112,14 @@ 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=2,
+        gradient_accumulation_steps=2,
         learning_rate=2e-4,
         fp16=not is_bfloat16_supported(),
         bf16=is_bfloat16_supported(),
         warmup_steps=5,
         logging_steps=10,
-        max_steps=200,
+        max_steps=100,
         optim="adamw_8bit",
         weight_decay=0.01,
         lr_scheduler_type="linear",
@@ -139,13 +138,15 @@ num += 1
 
 uploads_models = f"cybersentinal-2.0-{str(num)}"
 
+up = "sentinal-3.1-70B"
+
 print("Saving the trained model...")
 model.save_pretrained_merged("model", tokenizer, save_method="merged_16bit")
 print("Model saved successfully.")
 
 print("Pushing the model to the hub...")
 model.push_to_hub_merged(
-    uploads_models,
+    up,
     tokenizer,
     save_method="merged_16bit",
     token=hf_token