Update README.md

README.md

@@ -71,25 +71,28 @@ tokenized_dataset = dataset.map(tokenize_function, batched=True)
 
 # Training Arguments
 training_args = TrainingArguments(
-    output_dir="./output/TCLM-beta/",
-    num_train_epochs=1,
-    per_device_train_batch_size=2,
-    gradient_accumulation_steps=1,
-    evaluation_strategy="epoch",
-    save_strategy="epoch",
-    learning_rate=1e-4,
-    weight_decay=0.0,
-    lr_scheduler_type="cosine",
-    warmup_ratio=0.0,
-    logging_dir="./logs",
-    logging_steps=10,
-    fp16=True,
-    save_total_limit=1,
-    load_best_model_at_end=True,
-    metric_for_best_model="loss",
-    greater_is_better=False
+    output_dir="./output/TCLM-beta/",  # Directory to save model checkpoints
+    num_train_epochs=3,  # Increase epochs for better fine-tuning results
+    per_device_train_batch_size=4,  # Adjust based on GPU memory, increase if possible
+    gradient_accumulation_steps=4,  # Accumulate gradients to simulate a larger batch size
+    evaluation_strategy="steps",  # Evaluate more frequently for detailed tracking
+    eval_steps=500,  # Evaluate every 500 steps to track progress without over-evaluating
+    save_strategy="steps",  # Save checkpoints during training
+    save_steps=500,  # Save model every 500 steps
+    save_total_limit=2,  # Limit to the two best models to save disk space
+    learning_rate=5e-5,  # Lower learning rate for fine-tuning
+    weight_decay=0.01,  # Slight weight decay to prevent overfitting
+    lr_scheduler_type="cosine",  # Cosine schedule for smoother learning rate decay
+    warmup_ratio=0.06,  # Warmup to stabilize initial training
+    logging_dir="./logs",  # Directory to save training logs
+    logging_steps=50,  # Log progress every 50 steps for better monitoring
+    fp16=True,  # Enable mixed precision for faster training with less memory
+    load_best_model_at_end=True,  # Load the best model at the end based on evaluation metric
+    metric_for_best_model="eval_loss",  # Use evaluation loss to determine the best model
+    greater_is_better=False,  # Lower loss is better
 )
 
+
 # Trainer Setup
 trainer = Trainer(
     model=model,