From 40438b137143441793e7dc33a2e7b350ad6c9cec Mon Sep 17 00:00:00 2001 From: Chime Ogbuji Date: Sat, 2 Nov 2024 19:02:47 -0400 Subject: [PATCH] Persist configured prompt/completion key --- llms/mlx_lm/tuner/datasets.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llms/mlx_lm/tuner/datasets.py b/llms/mlx_lm/tuner/datasets.py index 692bdb5c..81e5d293 100644 --- a/llms/mlx_lm/tuner/datasets.py +++ b/llms/mlx_lm/tuner/datasets.py @@ -95,7 +95,7 @@ def create_dataset( if "messages" in sample: return ChatDataset(data, tokenizer) 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: return Dataset(data, tokenizer) else: