do not remove prompt key

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
vllm-project · Dec 20, 2024 · 42f7892 · 42f7892
1 parent d27dad3
commit 42f7892
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/llmcompressor/transformers/finetune/data/base.py b/src/llmcompressor/transformers/finetune/data/base.py
@@ -234,7 +234,9 @@ def filter_tokenizer_args(self, dataset: DatasetType) -> DatasetType:
         )
 
         column_names = get_columns(dataset)
-        return dataset.remove_columns(list(set(column_names) - set(tokenizer_args)))
+        return dataset.remove_columns(
+            list(set(column_names) - set(tokenizer_args) - set([self.PROMPT_KEY]))
+        )
 
     def tokenize(self, data: LazyRow) -> Dict[str, Any]:
         # separate prompt