Spaces:
Sleeping
Sleeping
# 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") |