Spaces:
Running
Running
from trl import SFTTrainer | |
from transformers import TrainingArguments | |
def train_model(model, tokenizer, train_dataset, dataset_text_field, max_seq_length, dataset_num_proc, packing, training_args): | |
trainer = SFTTrainer( | |
model=model, | |
tokenizer=tokenizer, | |
train_dataset=train_dataset, | |
dataset_text_field=dataset_text_field, | |
max_seq_length=max_seq_length, | |
dataset_num_proc=dataset_num_proc, | |
packing=packing, | |
args=TrainingArguments(**training_args), | |
) | |
# Train the model | |
train_results = trainer.train() | |
# Optionally, you can return more specific training information if necessary | |
return train_results | |