fix: pass dataset_text_field to SFTTrainer from config

This commit is contained in:
Christian Medina
2026-07-02 16:24:34 -04:00
parent 728073ef6a
commit 36fe2c8d45

View File

@@ -136,12 +136,17 @@ def train(config_path):
# SFT Trainer # SFT Trainer
from trl import SFTTrainer from trl import SFTTrainer
# Get text column from config (default to 'text' if not specified)
text_column = config["dataset"][0].get("text_column", "text")
print(f"Using text column: {text_column}")
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