From 33630d6b29ed0184d47b67c58f5e3e9d6de14273 Mon Sep 17 00:00:00 2001
From: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
Date: Thu, 25 Jun 2026 02:56:55 +0800
Subject: [PATCH] [data] fix KeyError on unmapped role tag in openai converter
 main loop

The OpenAIDatasetConverter builds aligned_messages with tag_mapping[role]
but never checks role against tag_mapping first, so a single message whose
role tag is not in the mapping raises KeyError and aborts the whole dataset
conversion. The later validity loop runs on aligned_messages, which is too
late to catch it. The SharegptDatasetConverter main loop already guards this
case and skips the example, so this just brings the openai path in line.

Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
---
 src/llamafactory/data/converter.py |  5 +++++
 tests/data/test_converter.py       | 20 ++++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/src/llamafactory/data/converter.py b/src/llamafactory/data/converter.py
index ad49deded4..061cf7c397 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 6b411aed53..0ba2b320df 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_openai_converter_skips_invalid_role():
+    # A message that carries a role tag not present in tag_mapping must be
+    # treated as an abnormal example and skipped, not crash with a KeyError when
+    # the converter builds aligned_messages. The sharegpt converter already
+    # guards its build loop; the openai converter only validated afterwards.
+    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"] == []