Update aifixcode_trainer.py

aifixcode_trainer.py

@@ -1,86 +1,177 @@
-### aifixcode_trainer.py
-
 """
-This script sets up a simple HuggingFace-based training + inference pipeline
-for bug-fixing AI using a CodeT5 model and supports continual training.
-You can upload this script to HuggingFace Space or Hub repo.
+This script sets up a HuggingFace-based training and inference pipeline
+for bug-fixing AI using a CodeT5 model. It is designed to be more
+robust and flexible than the original.
+
+Key improvements:
+- Uses argparse for configuration, making it easy to change settings
+  via the command line.
+- Adds checks to ensure data files exist.
+- Implements a compute_metrics function for better model evaluation.
+- Optimizes data preprocessing with dynamic padding.
+- Saves the best-performing model based on evaluation metrics.
+- Checks for GPU availability.
 """
 
+import os
+import argparse
+import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Trainer, TrainingArguments, DataCollatorForSeq2Seq
 from datasets import load_dataset, DatasetDict
-import torch
-import os
+from typing import Dict
+from evaluate import load
 
-# ========== CONFIG ==========
-MODEL_NAME = "Salesforce/codet5p-220m"
-MODEL_OUT_DIR = "./aifixcode-model"
-TRAIN_DATASET_PATH = "./data/train.json"
-VAL_DATASET_PATH = "./data/val.json"
+# ========== ARGUMENT PARSING ==========
+def parse_args():
+    """Parses command-line arguments for the training script."""
+    parser = argparse.ArgumentParser(description="Fine-tune a Seq2Seq model for code repair.")
+    parser.add_argument("--model_name", type=str, default="Salesforce/codet5p-220m",
+                        help="Pre-trained model name from HuggingFace.")
+    parser.add_argument("--output_dir", type=str, default="./aifixcode-model",
+                        help="Directory to save the trained model.")
+    parser.add_argument("--train_path", type=str, default="./data/train.json",
+                        help="Path to the training data JSON file.")
+    parser.add_argument("--val_path", type=str, default="./data/val.json",
+                        help="Path to the validation data JSON file.")
+    parser.add_argument("--epochs", type=int, default=3,
+                        help="Number of training epochs.")
+    parser.add_argument("--learning_rate", type=float, default=5e-5,
+                        help="Learning rate for the optimizer.")
+    parser.add_argument("--per_device_train_batch_size", type=int, default=4,
+                        help="Batch size per device for training.")
+    parser.add_argument("--per_device_eval_batch_size", type=int, default=4,
+                        help="Batch size per device for evaluation.")
+    parser.add_argument("--push_to_hub", action="store_true",
+                        help="Whether to push the model to the Hugging Face Hub.")
+    parser.add_argument("--hub_model_id", type=str, default="khulnasoft/aifixcode-model",
+                        help="Hugging Face Hub model ID to push to.")
+    return parser.parse_args()
 
-# ========== LOAD MODEL + TOKENIZER ==========
-print("Loading model and tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
-
-# ========== LOAD DATASET ==========
-print("Loading dataset...")
-def load_json_dataset(train_path, val_path):
+# ========== DATA LOADING ==========
+def load_json_dataset(train_path: str, val_path: str) -> DatasetDict:
+    """Loads and returns a dataset dictionary from JSON files."""
+    if not os.path.exists(train_path) or not os.path.exists(val_path):
+        raise FileNotFoundError(f"One or both data files not found: {train_path}, {val_path}")
+    
+    print("Loading dataset...")
     dataset = DatasetDict({
-        "train": load_dataset("json", data_files=train_path)["train"],
-        "validation": load_dataset("json", data_files=val_path)["train"]
+        "train": load_dataset("json", data_files=train_path, split="train"),
+        "validation": load_dataset("json", data_files=val_path, split="train")
     })
     return dataset
 
-dataset = load_json_dataset(TRAIN_DATASET_PATH, VAL_DATASET_PATH)
-
-# ========== PREPROCESS ==========
-print("Tokenizing dataset...")
-def preprocess(example):
-    input_code = example["input"]
-    target_code = example["output"]
-    model_inputs = tokenizer(input_code, truncation=True, padding="max_length", max_length=512)
-    labels = tokenizer(target_code, truncation=True, padding="max_length", max_length=512)
-    model_inputs["labels"] = labels["input_ids"]
+# ========== DATA PREPROCESSING ==========
+def preprocess_function(examples: Dict[str, list], tokenizer) -> Dict[str, list]:
+    """Tokenizes a batch of input and target code.
+    
+    This function uses dynamic padding by default, which is more
+    memory-efficient than padding all sequences to a fixed max length.
+    """
+    inputs = [ex for ex in examples["input"]]
+    targets = [ex for ex in examples["output"]]
+    
+    model_inputs = tokenizer(inputs, text_target=targets, max_length=512, truncation=True)
     return model_inputs
 
-encoded_dataset = dataset.map(preprocess, batched=True)
-
-# ========== TRAINING SETUP ==========
-print("Setting up trainer...")
-training_args = TrainingArguments(
-    output_dir=MODEL_OUT_DIR,
-    evaluation_strategy="epoch",
-    save_strategy="epoch",
-    learning_rate=5e-5,
-    per_device_train_batch_size=4,
-    per_device_eval_batch_size=4,
-    num_train_epochs=3,
-    weight_decay=0.01,
-    logging_dir="./logs",
-    logging_strategy="epoch",
-    push_to_hub=True,
-    hub_model_id="khulnasoft/aifixcode-model",
-    hub_strategy="every_save"
-)
-
-data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
-
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=encoded_dataset["train"],
-    eval_dataset=encoded_dataset["validation"],
-    tokenizer=tokenizer,
-    data_collator=data_collator
-)
+# ========== METRIC CALCULATION ==========
+def compute_metrics(eval_pred):
+    """Computes BLEU and Rouge metrics for model evaluation."""
+    bleu_metric = load("bleu")
+    rouge_metric = load("rouge")
+    
+    predictions, labels = eval_pred
+    
+    # Replace -100 in labels as we can't decode them
+    labels = [[item if item != -100 else tokenizer.pad_token_id for item in row] for row in labels]
+    
+    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
+    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+    
+    # Compute BLEU score
+    bleu_result = bleu_metric.compute(predictions=decoded_preds, references=decoded_labels)
+    
+    # Compute ROUGE score
+    rouge_result = rouge_metric.compute(predictions=decoded_preds, references=decoded_labels)
+    
+    return {
+        "bleu": bleu_result["bleu"],
+        "rouge1": rouge_result["rouge1"],
+        "rouge2": rouge_result["rouge2"],
+        "rougeL": rouge_result["rougeL"],
+    }
 
-# ========== TRAIN ==========
-print("Starting training...")
-trainer.train()
+# ========== MAIN EXECUTION BLOCK ==========
+def main():
+    """Main function to set up and run the training pipeline."""
+    args = parse_args()
+    
+    # Check for GPU availability
+    if not torch.cuda.is_available():
+        print("Warning: A GPU is not available. Training will be very slow on CPU.")
 
-# ========== SAVE FINAL MODEL ==========
-print("Saving model...")
-trainer.save_model(MODEL_OUT_DIR)
-tokenizer.save_pretrained(MODEL_OUT_DIR)
+    # Load model and tokenizer
+    print(f"Loading model '{args.model_name}' and tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name)
+    
+    # Load and preprocess dataset
+    try:
+        dataset = load_json_dataset(args.train_path, args.val_path)
+    except FileNotFoundError as e:
+        print(e)
+        return
+        
+    print("Tokenizing dataset...")
+    tokenized_dataset = dataset.map(
+        lambda examples: preprocess_function(examples, tokenizer),
+        batched=True,
+        remove_columns=dataset["train"].column_names
+    )
+    
+    # Training arguments setup
+    print("Setting up trainer...")
+    training_args = TrainingArguments(
+        output_dir=os.path.join(args.output_dir, "checkpoints"),
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        learning_rate=args.learning_rate,
+        per_device_train_batch_size=args.per_device_train_batch_size,
+        per_device_eval_batch_size=args.per_device_eval_batch_size,
+        num_train_epochs=args.epochs,
+        weight_decay=0.01,
+        logging_dir=os.path.join(args.output_dir, "logs"),
+        logging_strategy="epoch",
+        push_to_hub=args.push_to_hub,
+        hub_model_id=args.hub_model_id if args.push_to_hub else None,
+        hub_strategy="every_save",
+        load_best_model_at_end=True,  # Saves the best model
+        metric_for_best_model="rougeL", # Specify the metric to use for saving the best model
+        greater_is_better=True,
+        report_to="tensorboard"
+    )
+    
+    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)
+    
+    # Initialize and train the trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset["train"],
+        eval_dataset=tokenized_dataset["validation"],
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics
+    )
+    
+    print("Starting training...")
+    trainer.train()
+    
+    # Save final model
+    print("Saving final model...")
+    final_model_dir = os.path.join(args.output_dir, "final")
+    trainer.save_model(final_model_dir)
+    tokenizer.save_pretrained(final_model_dir)
+    print("Training complete and model saved!")
 
-print("Training complete and model saved!")
+if __name__ == "__main__":
+    main()