fix: rename dataset column to 'text' for SFTTrainer compatibility
This commit is contained in:
8
train.py
8
train.py
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user