Persist configured prompt/completion key

This commit is contained in:
Chime Ogbuji 2024-11-02 19:02:47 -04:00 committed by Awni Hannun
parent db9898d104
commit 40438b1371

View File

@ -95,7 +95,7 @@ def create_dataset(
if "messages" in sample: if "messages" in sample:
return ChatDataset(data, tokenizer) return ChatDataset(data, tokenizer)
elif prompt_feature in sample and completion_feature in sample: elif prompt_feature in sample and completion_feature in sample:
return CompletionsDataset(data, tokenizer) return CompletionsDataset(data, tokenizer, prompt_feature, completion_feature)
elif "text" in sample: elif "text" in sample:
return Dataset(data, tokenizer) return Dataset(data, tokenizer)
else: else: