ensure prompt key

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

src/llmcompressor/transformers/finetune/data/base.py

@@ -236,12 +236,12 @@ def filter_processor_args(self, dataset: DatasetType) -> DatasetType:
                 for key, param in signature.parameters.items()
                 if param.kind not in (Kind.VAR_POSITIONAL, Kind.VAR_KEYWORD)
             )
-        logger.debug(
-            f"Found processor args `{processor_kwargs}`. Removing all other columns"
-        )
+        logger.debug(f"Found processor args `{processor_kwargs}`")
 
         column_names = get_columns(dataset)
-        return dataset.remove_columns(list(set(column_names) - set(processor_kwargs)))
+        return dataset.remove_columns(
+            list(set(column_names) - set(processor_kwargs)) - set([self.PROMPT_KEY])
+        )
 
     def tokenize(self, data: LazyRow) -> Dict[str, Any]:
         # separate prompt