diff --git a/train.py b/train.py
index 10761ed..dd127ae 100644
--- a/train.py
+++ b/train.py
@@ -136,12 +136,17 @@ def train(config_path):
     # SFT Trainer
     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(
         model=model,
         processing_class=tokenizer,
         train_dataset=dataset["train"],
         eval_dataset=dataset.get("test"),
         args=training_args,
+        dataset_text_field=text_column,
     )
     
     # Train