use the extracted field_messages to parse the role fields

2025-01-16 08:36:00 -05:00
1 changed files with 1 additions and 1 deletions
--- a/scripts/chat_datasets.py
+++ b/scripts/chat_datasets.py
@@ -30,7 +30,7 @@ def parse_dataset(dataset=None, split="train"):
        )
    ds_cfg["field_messages"] = field_messages

-    message_fields = features["conversations"][0].keys()
+    message_fields = features[field_messages][0].keys()
    message_field_role = None
    for key in ["from", "role"]:
        if key in message_fields: