Update app.py with optimized training settings

app.py

@@ -1,5 +1,4 @@
-from transformers import (GPT2Tokenizer, GPT2LMHeadModel, Trainer, 
-                          TrainingArguments, DataCollatorWithPadding)
+from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments
 from datasets import load_dataset
 
 # Load the text dataset from the specified file
@@ -9,54 +8,44 @@ dataset = load_dataset("text", data_files="training.txt")
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 tokenizer.pad_token = tokenizer.eos_token
 
-# Adjusted max_length for potentially reduced memory usage
-max_length = 256
-
+# Define a function to tokenize the dataset and prepare labels
 def tokenize_function(examples):
-    # Tokenize the text to input_ids, attention_mask, and ensure labels are set
+    # Tokenize the text to input_ids, attention_mask, with reduced max_length
     tokenized_inputs = tokenizer(
         examples["text"], 
         padding="max_length", 
         truncation=True, 
-        max_length=max_length,
+        max_length=256  # Reduced from 512 to 256
     )
     # Prepare labels: labels are the same as input_ids for language modeling
-    tokenized_inputs["labels"] = tokenized_inputs["input_ids"].detach().clone()
+    tokenized_inputs["labels"] = tokenized_inputs["input_ids"].copy()
     return tokenized_inputs
 
-# Apply tokenization to the entire dataset and remove 'text' column during mapping
-tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
-# No need to call remove_columns on tokenized_datasets again as 'text' is already removed
-
-# Use a DataCollator that dynamically pads the batches
-data_collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="pt")
+# Tokenize the entire dataset
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
+# Remove the 'text' column as it's no longer needed after tokenization
+tokenized_datasets = tokenized_datasets.remove_columns(["text"])
+# Set the format of the dataset to PyTorch tensors
+tokenized_datasets.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
 
 # Load the GPT-2 model
 model = GPT2LMHeadModel.from_pretrained("gpt2")
 
-# Define training arguments with optimized settings
+# Define training arguments with adjusted 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,  # Adjusted for gradient accumulation
+    num_train_epochs=2,  # Optionally reduced for quicker iteration
+    per_device_train_batch_size=2,  # Reduced from 4 to 2
+    gradient_accumulation_steps=8,  # Added to compensate for smaller batch size
     save_steps=10_000,
     save_total_limit=2,
-    no_cuda=False,  # Set based on your hardware
-    learning_rate=3e-5,  # Adjusted learning rate
-    weight_decay=0.01,
-    warmup_steps=100,
-    logging_dir='./logs',
-    logging_steps=100,
-    fp16=True,  # Enable fp16 for memory and speed improvement if your hardware supports it
 )
 
-# Initialize the Trainer
+# Initialize the Trainer with the training dataset including labels
 trainer = Trainer(
     model=model,
     args=training_args,
-    data_collator=data_collator,
     train_dataset=tokenized_datasets["train"],
 )