diff --git a/src/llamafactory/data/converter.py b/src/llamafactory/data/converter.py
index ad49deded4..79dc5ece1c 100644
--- a/src/llamafactory/data/converter.py
+++ b/src/llamafactory/data/converter.py
@@ -141,9 +141,6 @@ def __call__(self, example: dict[str, Any]) -> dict[str, Any]:
self.dataset_attr.function_tag: Role.FUNCTION.value,
self.dataset_attr.system_tag: Role.SYSTEM.value,
}
- odd_tags = (self.dataset_attr.user_tag, self.dataset_attr.observation_tag)
- even_tags = (self.dataset_attr.assistant_tag, self.dataset_attr.function_tag)
- accept_tags = (odd_tags, even_tags)
messages = example[self.dataset_attr.messages]
if (
self.dataset_attr.system_tag
@@ -156,20 +153,46 @@ def __call__(self, example: dict[str, Any]) -> dict[str, Any]:
system = example[self.dataset_attr.system] if self.dataset_attr.system else ""
aligned_messages = []
+ tool_responses = []
broken_data = False
- for turn_idx, message in enumerate(messages):
- if message[self.dataset_attr.role_tag] not in accept_tags[turn_idx % 2]:
+ for message in messages:
+ role = message[self.dataset_attr.role_tag]
+ content = message[self.dataset_attr.content_tag]
+
+ if role not in tag_mapping:
logger.warning_rank0(f"Invalid role tag in {messages}.")
broken_data = True
break
+ if role == self.dataset_attr.observation_tag:
+ tool_responses.append(content)
+ continue
+ elif len(tool_responses) > 0:
+ observation_content = "\n\n\n".join(tool_responses)
+ aligned_messages.append(
+ {
+ "role": Role.OBSERVATION.value,
+ "content": observation_content,
+ }
+ )
+ tool_responses = []
+
aligned_messages.append(
{
- "role": tag_mapping[message[self.dataset_attr.role_tag]],
- "content": message[self.dataset_attr.content_tag],
+ "role": tag_mapping[role],
+ "content": content,
}
)
+ odd_tags = (Role.USER.value, Role.OBSERVATION.value)
+ even_tags = (Role.ASSISTANT.value, Role.FUNCTION.value)
+ accept_tags = (odd_tags, even_tags)
+ for turn_idx, message in enumerate(aligned_messages):
+ if message["role"] not in accept_tags[turn_idx % 2]:
+ logger.warning_rank0(f"Invalid role tag in {messages}.")
+ broken_data = True
+ break
+
if (not self.dataset_attr.ranking and len(aligned_messages) % 2 != 0) or (
self.dataset_attr.ranking and len(aligned_messages) % 2 == 0
):
@@ -193,9 +216,10 @@ def __call__(self, example: dict[str, Any]) -> dict[str, Any]:
): # pairwise example
chosen = example[self.dataset_attr.chosen]
rejected = example[self.dataset_attr.rejected]
+ response_tags = (self.dataset_attr.assistant_tag, self.dataset_attr.function_tag)
if (
- chosen[self.dataset_attr.role_tag] not in accept_tags[-1]
- or rejected[self.dataset_attr.role_tag] not in accept_tags[-1]
+ chosen[self.dataset_attr.role_tag] not in response_tags
+ or rejected[self.dataset_attr.role_tag] not in response_tags
):
logger.warning_rank0(f"Invalid role tag in {[chosen, rejected]}.")
broken_data = True
diff --git a/tests/data/test_converter.py b/tests/data/test_converter.py
index 6b411aed53..ae0cc8de22 100644
--- a/tests/data/test_converter.py
+++ b/tests/data/test_converter.py
@@ -62,3 +62,35 @@ def test_sharegpt_converter():
"_videos": None,
"_audios": None,
}
+
+
+@pytest.mark.runs_on(["cpu", "mps"])
+def test_sharegpt_converter_merges_consecutive_observations():
+ dataset_attr = DatasetAttr("hf_hub", "llamafactory/tiny-supervised-dataset")
+ data_args = DataArguments()
+ example = {
+ "conversations": [
+ {"from": "human", "value": "Check both symbols."},
+ {"from": "function_call", "value": "lookup_a\nlookup_b"},
+ {"from": "observation", "value": '{"symbol": "A", "price": 10}'},
+ {"from": "observation", "value": '{"symbol": "B", "price": 20}'},
+ {"from": "gpt", "value": "A is 10 and B is 20."},
+ ]
+ }
+ dataset_converter = get_dataset_converter("sharegpt", dataset_attr, data_args)
+ assert dataset_converter(example) == {
+ "_prompt": [
+ {"role": Role.USER.value, "content": "Check both symbols."},
+ {"role": Role.FUNCTION.value, "content": "lookup_a\nlookup_b"},
+ {
+ "role": Role.OBSERVATION.value,
+ "content": '{"symbol": "A", "price": 10}\n\n\n{"symbol": "B", "price": 20}',
+ },
+ ],
+ "_response": [{"role": Role.ASSISTANT.value, "content": "A is 10 and B is 20."}],
+ "_system": "",
+ "_tools": "",
+ "_images": None,
+ "_videos": None,
+ "_audios": None,
+ }