fix: rename dataset column to 'text' for SFTTrainer compatibility

This commit is contained in:
Christian Medina
2026-07-02 16:39:21 -04:00
parent 36fe2c8d45
commit dd86a4df37

View File

@@ -140,13 +140,19 @@ def train(config_path):
text_column = config["dataset"][0].get("text_column", "text")
print(f"Using text column: {text_column}")
# Rename column to 'text' if needed (SFTTrainer expects 'text')
if text_column != "text":
print(f" Renaming '{text_column}' column to 'text'...")
dataset["train"] = dataset["train"].rename_column(text_column, "text")
if "test" in dataset and dataset["test"] is not None:
dataset["test"] = dataset["test"].rename_column(text_column, "text")
trainer = SFTTrainer(
model=model,
processing_class=tokenizer,
train_dataset=dataset["train"],
eval_dataset=dataset.get("test"),
args=training_args,
dataset_text_field=text_column,
)
# Train