train.py

# 4. train.py
# Fine-tuning logic
from transformers import Trainer, TrainingArguments
from model_utils import BugClassifier

def fine_tune_model(dataset):
    model = BugClassifier().model
    training_args = TrainingArguments(
        output_dir="./results",
        evaluation_strategy="epoch",
        save_strategy="epoch",
        per_device_train_batch_size=8,
        num_train_epochs=3,
        logging_dir="./logs",
        logging_steps=10,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
    )

    trainer.train()
    model.save_pretrained("./fine_tuned_model")