From 92e42c439d119c366d9e51df0061189e76550a41 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Tue, 23 Jun 2026 16:43:16 +0800 Subject: [PATCH 1/2] [data] fix KeyError on invalid role tag in ranking dataset converters The OpenAI and ShareGPT dataset converters log "Invalid role tag" and set broken_data=True for a pairwise example whose chosen/rejected role is not an accepted tag. But the `if broken_data:` skip guard is evaluated at the top of the if/elif chain, so setting it inside the ranking branch is too late: the code then indexes tag_mapping[invalid_role] and raises KeyError, crashing align_dataset instead of skipping the example. Skip such examples (prompt/response = []) as the guard intends; add a converter test for the ranking invalid-role path. Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- src/llamafactory/data/converter.py | 52 +++++++++++++++--------------- tests/data/test_converter.py | 20 ++++++++++++ 2 files changed, 46 insertions(+), 26 deletions(-) diff --git a/src/llamafactory/data/converter.py b/src/llamafactory/data/converter.py index ad49deded4..d257490852 100644 --- a/src/llamafactory/data/converter.py +++ b/src/llamafactory/data/converter.py @@ -198,19 +198,19 @@ def __call__(self, example: dict[str, Any]) -> dict[str, Any]: or rejected[self.dataset_attr.role_tag] not in accept_tags[-1] ): logger.warning_rank0(f"Invalid role tag in {[chosen, rejected]}.") - broken_data = True - - prompt = aligned_messages - response = [ - { - "role": tag_mapping[chosen[self.dataset_attr.role_tag]], - "content": chosen[self.dataset_attr.content_tag], - }, - { - "role": tag_mapping[rejected[self.dataset_attr.role_tag]], - "content": rejected[self.dataset_attr.content_tag], - }, - ] + prompt, response = [], [] + else: + prompt = aligned_messages + response = [ + { + "role": tag_mapping[chosen[self.dataset_attr.role_tag]], + "content": chosen[self.dataset_attr.content_tag], + }, + { + "role": tag_mapping[rejected[self.dataset_attr.role_tag]], + "content": rejected[self.dataset_attr.content_tag], + }, + ] else: # normal example prompt = aligned_messages[:-1] response = aligned_messages[-1:] @@ -319,19 +319,19 @@ def __call__(self, example: dict[str, Any]) -> dict[str, Any]: or rejected[self.dataset_attr.role_tag] not in accept_tags[-1] ): logger.warning_rank0(f"Invalid role tag in {[chosen, rejected]}.") - broken_data = True - - prompt = aligned_messages - response = [ - { - "role": tag_mapping[chosen[self.dataset_attr.role_tag]], - "content": chosen[self.dataset_attr.content_tag], - }, - { - "role": tag_mapping[rejected[self.dataset_attr.role_tag]], - "content": rejected[self.dataset_attr.content_tag], - }, - ] + prompt, response = [], [] + else: + prompt = aligned_messages + response = [ + { + "role": tag_mapping[chosen[self.dataset_attr.role_tag]], + "content": chosen[self.dataset_attr.content_tag], + }, + { + "role": tag_mapping[rejected[self.dataset_attr.role_tag]], + "content": rejected[self.dataset_attr.content_tag], + }, + ] else: # normal example prompt = aligned_messages[:-1] response = aligned_messages[-1:] diff --git a/tests/data/test_converter.py b/tests/data/test_converter.py index 6b411aed53..fb5355bc38 100644 --- a/tests/data/test_converter.py +++ b/tests/data/test_converter.py @@ -62,3 +62,23 @@ def test_sharegpt_converter(): "_videos": None, "_audios": None, } + + +@pytest.mark.runs_on(["cpu", "mps"]) +def test_sharegpt_converter_ranking_skips_invalid_role(): + # A pairwise (ranking) example whose chosen/rejected carries a role tag not + # in the accepted set must be skipped, not crash with KeyError on tag_mapping. + dataset_attr = DatasetAttr("hf_hub", "llamafactory/tiny-supervised-dataset") + dataset_attr.ranking = True + dataset_attr.chosen = "chosen" + dataset_attr.rejected = "rejected" + data_args = DataArguments() + example = { + "conversations": [{"from": "human", "value": "Solve the math problem.\n3 + 4"}], + "chosen": {"from": "not_a_role", "value": "The answer is 7."}, + "rejected": {"from": "gpt", "value": "The answer is 8."}, + } + dataset_converter = get_dataset_converter("sharegpt", dataset_attr, data_args) + result = dataset_converter(example) + assert result["_prompt"] == [] + assert result["_response"] == [] From e15ce17f4d04aeb9ca7bc916053419929c9f8db3 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Fri, 26 Jun 2026 08:00:16 +0800 Subject: [PATCH 2/2] [data] also guard OpenAI converter main loop against unknown role tags MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ranking-branch fix in the previous commit (`tag_mapping[chosen/rejected]`) covered pairwise examples, but the OpenAI converter's main message loop had the same gap: `tag_mapping[role]` on line 280 (now 284) is reached before the validation loop at line 288 (now 292) has a chance to flag the bad role, so a dataset entry with any unrecognised role tag in its message list raises `KeyError` instead of being gracefully skipped. Add an explicit `role not in tag_mapping` guard (matching the pattern already used in the Sharegpt converter) that logs a warning and sets `broken_data`, then extend the test suite with a case that drove the crash to confirm RED → GREEN. Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> --- src/llamafactory/data/converter.py | 5 +++++ tests/data/test_converter.py | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/src/llamafactory/data/converter.py b/src/llamafactory/data/converter.py index d257490852..cdd7308549 100644 --- a/src/llamafactory/data/converter.py +++ b/src/llamafactory/data/converter.py @@ -275,6 +275,11 @@ def __call__(self, example: dict[str, Any]) -> dict[str, Any]: ) tool_responses = [] + if role not in tag_mapping: + logger.warning_rank0(f"Invalid role tag in {messages}.") + broken_data = True + break + aligned_messages.append( { "role": tag_mapping[role], diff --git a/tests/data/test_converter.py b/tests/data/test_converter.py index fb5355bc38..ffb357efd0 100644 --- a/tests/data/test_converter.py +++ b/tests/data/test_converter.py @@ -82,3 +82,21 @@ def test_sharegpt_converter_ranking_skips_invalid_role(): result = dataset_converter(example) assert result["_prompt"] == [] assert result["_response"] == [] + + +@pytest.mark.runs_on(["cpu", "mps"]) +def test_openai_converter_main_loop_skips_invalid_role(): + # An openai-format message list containing a role tag not in tag_mapping must + # be skipped (prompt=[], response=[]), not crash with KeyError. + dataset_attr = DatasetAttr("hf_hub", "llamafactory/tiny-supervised-dataset") + data_args = DataArguments() + example = { + "conversations": [ + {"from": "human", "value": "Solve the math problem.\n3 + 4"}, + {"from": "not_a_role", "value": "The answer is 7."}, + ] + } + dataset_converter = get_dataset_converter("openai", dataset_attr, data_args) + result = dataset_converter(example) + assert result["_prompt"] == [] + assert result["_response"] == []