File size: 2,702 Bytes
22e0e62 |
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 |
### 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.
"""
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Trainer, TrainingArguments, DataCollatorForSeq2Seq
from datasets import load_dataset, DatasetDict
import torch
import os
# ========== CONFIG ==========
MODEL_NAME = "Salesforce/codet5p-220m"
MODEL_OUT_DIR = "./aifixcode-model"
TRAIN_DATASET_PATH = "./data/train.json"
VAL_DATASET_PATH = "./data/val.json"
# ========== 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):
dataset = DatasetDict({
"train": load_dataset("json", data_files=train_path)["train"],
"validation": load_dataset("json", data_files=val_path)["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"]
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
)
# ========== TRAIN ==========
print("Starting training...")
trainer.train()
# ========== SAVE FINAL MODEL ==========
print("Saving model...")
trainer.save_model(MODEL_OUT_DIR)
tokenizer.save_pretrained(MODEL_OUT_DIR)
print("Training complete and model saved!")
|