Update app.py with optimized training settings

app.py

@@ -2,32 +2,40 @@ from transformers import (GPT2Tokenizer, GPT2LMHeadModel, Trainer,
                           TrainingArguments, DataCollatorWithPadding)
 from datasets import load_dataset
 
-# Load the GPT-2 model
-model = GPT2LMHeadModel.from_pretrained("gpt2")
+# Load the text dataset from the specified file
+dataset = load_dataset("text", data_files="training.txt")
 
-# Initialize the GPT-2 tokenizer with a reduced max_length
+# Initialize the GPT-2 tokenizer
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
 tokenizer.pad_token = tokenizer.eos_token
-max_length = 256  # Reduced max_length
+
+# Adjusted max_length for potentially reduced memory usage
+max_length = 256
 
 def tokenize_function(examples):
-    return tokenizer(
+    # Tokenize the text to input_ids, attention_mask, and ensure labels are set
+    tokenized_inputs = tokenizer(
         examples["text"], 
         padding="max_length", 
         truncation=True, 
         max_length=max_length,
         return_tensors="pt"
     )
+    # Prepare labels: labels are the same as input_ids for language modeling
+    tokenized_inputs["labels"] = tokenized_inputs["input_ids"].detach().clone()
+    return tokenized_inputs
 
-# Load and preprocess the dataset
-dataset = load_dataset("text", data_files="training.txt")
-tokenized_datasets = dataset.map(tokenize_function, batched=True)
+# Apply tokenization to the entire dataset
+tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
 tokenized_datasets = tokenized_datasets.remove_columns(["text"])
 tokenized_datasets.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
 
 # Use a DataCollator that dynamically pads the batches
 data_collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="pt")
 
+# Load the GPT-2 model
+model = GPT2LMHeadModel.from_pretrained("gpt2")
+
 # Define training arguments with optimized settings
 training_args = TrainingArguments(
     output_dir="./output",
@@ -37,7 +45,7 @@ training_args = TrainingArguments(
     gradient_accumulation_steps=8,  # Adjusted for gradient accumulation
     save_steps=10_000,
     save_total_limit=2,
-    no_cuda=False,
+    no_cuda=False,  # Set based on your hardware
     learning_rate=3e-5,  # Adjusted learning rate
     weight_decay=0.01,
     warmup_steps=100,
@@ -46,6 +54,7 @@ training_args = TrainingArguments(
     fp16=True,  # Enable fp16 for memory and speed improvement if your hardware supports it
 )
 
+# Initialize the Trainer with the training dataset including labels and data collator for dynamic padding
 trainer = Trainer(
     model=model,
     args=training_args,
@@ -53,7 +62,9 @@ trainer = Trainer(
     train_dataset=tokenized_datasets["train"],
 )
 
+# Start the training process
 trainer.train()
 
+# Save the fine-tuned model and tokenizer
 model.save_pretrained("fine_tuned_gpt2_model")
 tokenizer.save_pretrained("fine_tuned_gpt2_model")