feat: Improve training settings and dataset handling

- Updated dataset to "pushshift-reddit-comments"
- Adjusted data field from "text" to "body"
- Increased tokenization max length to 256
- Enhanced training settings: more epochs, larger batch size, gradient accumulation, mixed precision
- Integrated scheduler with warmup steps
- Switched to using GPU if available

humanizer.py

@@ -1,6 +1,6 @@
 from datasets import load_dataset, DatasetDict
 from transformers import AutoTokenizer, T5ForConditionalGeneration, Seq2SeqTrainingArguments, Seq2SeqTrainer
-from transformers import EarlyStoppingCallback
+from transformers import EarlyStoppingCallback, get_linear_schedule_with_warmup
 from transformers.integrations import TensorBoardCallback
 import torch
 import logging
@@ -11,7 +11,7 @@ logger = logging.getLogger(__name__)
 
 # Load the dataset and take only 1000 samples
 logger.info("Loading dataset...")
-dataset = load_dataset("LucasChu/reddit_comments")
+dataset = load_dataset("fddemarco/pushshift-reddit-comments", split="train")
 dataset = dataset.shuffle(seed=42)
 logger.info("Dataset loaded, shuffled, and truncated to 10,000 samples.")
 
@@ -32,7 +32,7 @@ def generate_formal_text(text):
 
 # Prepare the dataset
 def prepare_data(example):
-    example["formal_text"] = generate_formal_text(example["text"])  # Changed from "comment" to "text"
+    example["formal_text"] = generate_formal_text(example["body"])  # Changed from "text" to "body"
     return example
 
 logger.info("Preparing dataset...")
@@ -44,8 +44,8 @@ model_name = "t5-base"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 def tokenize_function(examples):
-    model_inputs = tokenizer(examples["formal_text"], max_length=128, truncation=True, padding="max_length")
-    labels = tokenizer(examples["text"], max_length=128, truncation=True, padding="max_length")  # Changed from "comment" to "text"
+    model_inputs = tokenizer(examples["formal_text"], max_length=256, truncation=True, padding="max_length")
+    labels = tokenizer(examples["body"], max_length=256, truncation=True, padding="max_length")
     model_inputs["labels"] = labels["input_ids"]
     return model_inputs
 
@@ -63,19 +63,29 @@ model = T5ForConditionalGeneration.from_pretrained(model_name)
 
 training_args = Seq2SeqTrainingArguments(
     output_dir="./results",
-    num_train_epochs=1,
-    per_device_train_batch_size=16,
-    warmup_steps=100,
+    num_train_epochs=3,  # Increase number of epochs
+    per_device_train_batch_size=32,  # Increase batch size if memory allows
+    per_device_eval_batch_size=32,
+    warmup_steps=500,
     weight_decay=0.01,
     logging_dir="./logs",
     logging_steps=100,
-    evaluation_strategy="steps" if "test" in available_splits else "no",
-    eval_steps=500,
+    evaluation_strategy="steps",
+    eval_steps=1000,
     save_steps=1000,
-    use_cpu=True,
+    use_cpu=False,  # Use GPU if available
     load_best_model_at_end=True,
     metric_for_best_model="eval_loss",
-    greater_is_better=False
+    greater_is_better=False,
+    fp16=True,  # Enable mixed precision training if GPU supports it
+    gradient_accumulation_steps=2,  # Accumulate gradients to simulate larger batch sizes
+)
+
+optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
+scheduler = get_linear_schedule_with_warmup(
+    optimizer,
+    num_warmup_steps=500,
+    num_training_steps=len(tokenized_dataset["train"]) * training_args.num_train_epochs
 )
 
 trainer = Seq2SeqTrainer(
@@ -84,6 +94,7 @@ trainer = Seq2SeqTrainer(
     train_dataset=tokenized_dataset["train"],
     eval_dataset=tokenized_dataset.get("test"),
     tokenizer=tokenizer,
+    optimizers=(optimizer, scheduler),
     callbacks=[EarlyStoppingCallback(early_stopping_patience=3), TensorBoardCallback()]
 )
 logger.info("Model and trainer set up.")