File size: 2,814 Bytes
38b27cd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """
Fine-tune CodeT5 on a buggy-code -> fixed-code dataset.
Dataset: CodeXGLUE "code-refinement" (Java, small subset) — pairs of
(buggy_function, fixed_function). This is the standard benchmark for this
exact task, so it's a strong, defensible choice for a project report.
Swap `DATASET_NAME` for "google/code_x_glue_cc_code_refinement" (small) or
use the "Bugs2Fix" dataset if you want Python-only pairs instead.
Run:
pip install -r requirements.txt
python train.py
"""
from datasets import load_dataset
from transformers import (
AutoTokenizer,
AutoModelForSeq2SeqLM,
Seq2SeqTrainer,
Seq2SeqTrainingArguments,
DataCollatorForSeq2Seq,
)
MODEL_NAME = "Salesforce/codet5-base"
DATASET_NAME = "google/code_x_glue_cc_code_refinement"
DATASET_CONFIG = "small" # or "medium"
OUTPUT_DIR = "./checkpoints/codet5-bugfix-finetuned"
MAX_INPUT_LEN = 256
MAX_TARGET_LEN = 256
def main():
print("Loading dataset...")
raw_dataset = load_dataset(DATASET_NAME, DATASET_CONFIG)
print("Loading tokenizer & model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
def preprocess(examples):
inputs = ["fix bug: " + code for code in examples["buggy"]]
targets = examples["fixed"]
model_inputs = tokenizer(
inputs, max_length=MAX_INPUT_LEN, truncation=True, padding="max_length"
)
labels = tokenizer(
targets, max_length=MAX_TARGET_LEN, truncation=True, padding="max_length"
)
model_inputs["labels"] = labels["input_ids"]
return model_inputs
print("Tokenizing...")
tokenized_dataset = raw_dataset.map(
preprocess, batched=True, remove_columns=raw_dataset["train"].column_names
)
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
training_args = Seq2SeqTrainingArguments(
output_dir=OUTPUT_DIR,
eval_strategy="epoch",
save_strategy="epoch",
learning_rate=3e-5,
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
num_train_epochs=5,
weight_decay=0.01,
predict_with_generate=True,
fp16=True, # requires GPU
logging_steps=50,
save_total_limit=2,
load_best_model_at_end=True,
)
trainer = Seq2SeqTrainer(
model=model,
args=training_args,
train_dataset=tokenized_dataset["train"],
eval_dataset=tokenized_dataset["validation"],
data_collator=data_collator,
tokenizer=tokenizer,
)
print("Starting fine-tuning...")
trainer.train()
print(f"Saving final model to {OUTPUT_DIR}")
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
if __name__ == "__main__":
main()
|