From 8fb72cbc0b94129141bae5fa4d84edd23b648af6 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 21 Jan 2025 15:39:30 -0500 Subject: [PATCH] use the extracted field_messages to parse the role fields (#2265) --- 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: