cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Christian Medina
2026-07-02 16:39:21 -04:00
parent 36fe2c8d45
commit dd86a4df37
1 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
train.py
View File

@@ -140,13 +140,19 @@ def train(config_path):
text_column = config["dataset"][0].get("text_column", "text") text_column = config["dataset"][0].get("text_column", "text")
print(f"Using text column: {text_column}") 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( trainer = SFTTrainer(
model=model, model=model,
processing_class=tokenizer, processing_class=tokenizer,
train_dataset=dataset["train"], train_dataset=dataset["train"],
eval_dataset=dataset.get("test"), eval_dataset=dataset.get("test"),
args=training_args, args=training_args,
dataset_text_field=text_column,
) )
# Train # Train