fix: defensively rename VLM chat column to match model's forward() signature

This commit is contained in:
Roland Tannous 2026-02-17 23:49:13 +00:00
commit d2332622d1
3 changed files with 61 additions and 1 deletions

View file

@ -425,6 +425,7 @@ class UnslothTrainer:
dataset,
model_name=self.model_name,
tokenizer=self.tokenizer,
model=self.model,
is_vlm=self.is_vlm,
format_type=format_type,
dataset_name=dataset_source,
@ -447,6 +448,7 @@ class UnslothTrainer:
eval_dataset,
model_name=self.model_name,
tokenizer=self.tokenizer,
model=self.model,
is_vlm=self.is_vlm,
format_type=format_type,
dataset_name=dataset_source,

View file

@ -28,6 +28,8 @@ from .format_conversion import (
convert_alpaca_to_chatml,
convert_to_vlm_format,
convert_llava_to_vlm_format,
get_expected_chat_column,
rename_chat_column_in_list,
)
from .chat_templates import (
apply_chat_template_to_dataset,
@ -547,6 +549,7 @@ def format_and_template_dataset(
dataset,
model_name,
tokenizer,
model=None,
is_vlm = False,
format_type="auto",
# VLM-specific parameters
@ -735,12 +738,30 @@ def format_and_template_dataset(
dataset = [sample for sample in dataset]
warnings.append("Dataset already in standard VLM messages format")
# Defensive: rename chat column if model expects a different name
expected_col = get_expected_chat_column(model) if model is not None else None
# VLM data is a list of dicts — check what key the first item uses
current_col = "messages" # default from our converters
if isinstance(dataset, list) and len(dataset) > 0:
sample_keys = dataset[0].keys()
if "conversations" in sample_keys:
current_col = "conversations"
elif "messages" in sample_keys:
current_col = "messages"
if expected_col and expected_col != current_col:
warnings.append(
f"Model expects '{expected_col}' but dataset has '{current_col}' — renaming."
)
dataset = rename_chat_column_in_list(dataset, current_col, expected_col)
current_col = expected_col
# Return as list
return {
"dataset": dataset,
"detected_format": vlm_structure["format"],
"final_format": "vlm_messages",
"chat_column": "messages",
"chat_column": current_col,
"is_vlm": True,
"is_multimodal": multimodal_info["is_multimodal"],
"multimodal_info": multimodal_info,

View file

@ -8,6 +8,43 @@ This module contains functions for converting between dataset formats
from datasets import IterableDataset
def get_expected_chat_column(model):
"""
Inspect the model's forward() signature to determine if it expects
'messages' or 'conversations' as a column name.
Returns:
str or None: 'messages', 'conversations', or None if neither found.
"""
import inspect
try:
sig = inspect.signature(model.forward)
params = list(sig.parameters.keys())
if "messages" in params:
return "messages"
elif "conversations" in params:
return "conversations"
except (ValueError, TypeError):
pass
return None
def rename_chat_column_in_list(data, from_col, to_col):
"""
Rename a chat column key in a list of dicts (for VLM data).
"""
if from_col == to_col:
return data
renamed = []
for item in data:
new_item = {}
for k, v in item.items():
new_item[to_col if k == from_col else k] = v
renamed.append(new_item)
return renamed
def standardize_chat_format(
dataset,
tokenizer=None,