danielle2003 commited on
Commit
d6d26ef
·
1 Parent(s): d12fb3d
Files changed (1) hide show
  1. scripts/train.py +6 -3
scripts/train.py CHANGED
@@ -20,11 +20,14 @@ model = AutoModelForSequenceClassification.from_pretrained(model_name, num_label
20
  # Training arguments
21
  training_args = TrainingArguments(
22
  output_dir="./models",
 
 
 
 
23
  per_device_train_batch_size=8,
 
24
  num_train_epochs=3,
25
- evaluation_strategy="epoch",
26
- save_steps=1000,
27
- load_best_model_at_end=True,
28
  )
29
 
30
  # Trainer setup
 
20
  # Training arguments
21
  training_args = TrainingArguments(
22
  output_dir="./models",
23
+ evaluation_strategy="epoch", # Ensure this matches the save_strategy
24
+ save_strategy="epoch", # Change this to "epoch" to match evaluation_strategy
25
+ load_best_model_at_end=True, # Ensures best model is loaded
26
+ save_total_limit=2, # Keep only the last 2 models to save space
27
  per_device_train_batch_size=8,
28
+ per_device_eval_batch_size=8,
29
  num_train_epochs=3,
30
+ weight_decay=0.01,
 
 
31
  )
32
 
33
  # Trainer setup