Removed redundant input_length column from dataset. (#549)

- This is required because, whenever dataset padding is performed there will be an additional input being fed to the model. For causal lm models this did not gave any error as its forward methods has kwargs which caught unnecessary arguments. But for BERT kind of models this gave error.

Signed-off-by: meetkuma <[email protected]>

Author: quic-meetkuma

QEfficient/finetune/utils/dataset_utils.py

@@ -81,7 +81,7 @@ def padding_dataset(train_config, dataset, batch_size):
             # Hugging Face Dataset transformation
             dataset = dataset.map(lambda x: {"input_length": len(x["input_ids"])})
             dataset = dataset.sort("input_length")
-
+            dataset = dataset.remove_columns("input_length")
         else:
             dataset = sorted(dataset, key=lambda x: len(x["input_ids"]))