Create aifixcode_trainer.py

aifixcode_trainer.py

@@ -0,0 +1,86 @@
+### 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!")