From 7bf9741831894fee0408beb8d01b5beab08b0e47 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Thu, 16 Jan 2025 08:36:00 -0500 Subject: [PATCH] use the extracted field_messages to parse the role fields --- scripts/chat_datasets.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/chat_datasets.py b/scripts/chat_datasets.py index 5eb5bde1e..6210b1138 100644 --- 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: