fixing unit tests

neuralmagic · Apr 9, 2024 · f47d219 · f47d219
1 parent 71c997e
commit f47d219
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/sparseml/transformers/finetune/runner.py b/src/sparseml/transformers/finetune/runner.py
@@ -119,9 +119,10 @@ def _get_split_name(inp_str):
                 tokenizer=tokenizer,
             )
 
-            if "input_ids" in self._data_args.dataset.column_names:
+            dataset = self._data_args.dataset
+            if hasattr(dataset, "column_names") and "input_ids" in dataset.column_names:
                 # dataset is already tokenized
-                tokenized_datasets[split_name] = self._data_args.dataset
+                tokenized_datasets[split_name] = dataset
             else:
                 # dataset needs to be tokenized
                 raw_dataset = dataset_manager.get_raw_dataset()