Update app.py with optimized training settings

app.py

@@ -1,15 +1,9 @@
 from transformers import (GPT2Tokenizer, GPT2LMHeadModel, Trainer, 
-                          TrainingArguments, DataCollatorWithPadding, GradientAccumulationScheduler)
+                          TrainingArguments, DataCollatorWithPadding)
 from datasets import load_dataset
-import torch
 
-# Assuming your hardware supports it, enable gradient checkpointing
-model_config = {
-    "gradient_checkpointing": True,  # Enable gradient checkpointing
-}
-
-# Load the GPT-2 model with gradient checkpointing enabled
-model = GPT2LMHeadModel.from_pretrained("gpt2", **model_config)
+# Load the GPT-2 model
+model = GPT2LMHeadModel.from_pretrained("gpt2")
 
 # Initialize the GPT-2 tokenizer with a reduced max_length
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
@@ -34,13 +28,13 @@ tokenized_datasets.set_format(type="torch", columns=["input_ids", "attention_mas
 # Use a DataCollator that dynamically pads the batches
 data_collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="pt")
 
-# Define training arguments with a smaller batch size and optimized settings
+# Define training arguments with optimized settings
 training_args = TrainingArguments(
     output_dir="./output",
     overwrite_output_dir=True,
     num_train_epochs=3,
     per_device_train_batch_size=2,  # Decreased batch size
-    gradient_accumulation_steps=8,  # Adjust based on your new batch size to simulate larger batches
+    gradient_accumulation_steps=8,  # Adjusted for gradient accumulation
     save_steps=10_000,
     save_total_limit=2,
     no_cuda=False,
@@ -49,8 +43,7 @@ training_args = TrainingArguments(
     warmup_steps=100,
     logging_dir='./logs',
     logging_steps=100,
-    # Enable fp16 for memory and speed improvement if your hardware supports it
-    fp16=torch.cuda.is_available(),
+    fp16=True,  # Enable fp16 for memory and speed improvement if your hardware supports it
 )
 
 trainer = Trainer(