@@ -650,7 +650,7 @@ class TextEngine(BaseEngine):
}
hf_dataset = HFDataset.from_list(raw_data)
- tokenized = hf_dataset.map(tokenize_fn, batched=True, remove_columns=hf_dataset.column_names)
+ tokenized = hf_dataset.map(tokenize_fn, batched=True)
return tokenized