From f9894018815f9c176f35715491ee2703a7dca6a3 Mon Sep 17 00:00:00 2001
From: Chime Ogbuji <chimezie@gmail.com>
Date: Fri, 6 Dec 2024 07:07:06 -0500
Subject: [PATCH] Default for hf_datasets configuration

---
 llms/mlx_lm/lora.py           | 1 +
 llms/mlx_lm/tuner/datasets.py | 3 ---
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/llms/mlx_lm/lora.py b/llms/mlx_lm/lora.py
index ea6ee973..87e3bb0c 100644
--- a/llms/mlx_lm/lora.py
+++ b/llms/mlx_lm/lora.py
@@ -61,6 +61,7 @@ CONFIG_DEFAULTS = {
     "config": None,
     "grad_checkpoint": False,
     "lr_schedule": None,
+    "hf_datasets": None,
     "lora_parameters": {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0},
 }
 
diff --git a/llms/mlx_lm/tuner/datasets.py b/llms/mlx_lm/tuner/datasets.py
index c8b39e31..ae6349c6 100644
--- a/llms/mlx_lm/tuner/datasets.py
+++ b/llms/mlx_lm/tuner/datasets.py
@@ -84,9 +84,6 @@ class CompletionsDataset:
     def __len__(self):
         return len(self._data)
 
-    def get_prompt_and_completion(self, idx: int):
-        return self._data[idx][self._prompt_key], self._data[idx][self._completion_key]
-
 
 class CompletionsDatasetCollection:
     def __init__(self, data: List[Union[ChatDataset, CompletionsDataset]]):