fix: pass dataset_text_field to SFTTrainer from config
This commit is contained in:
5
train.py
5
train.py
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user