Fix BFD formatting

Author: mariosasko

tests/test_data_utils.py

@@ -1111,7 +1111,10 @@ def test_with_dataset(self):
             "seq_lengths": [[4], [3, 1]],
         }
         dataset = pack_dataset(dataset, seq_length, strategy="bfd")
+        expected_format = dataset.format
         assert dataset.to_dict() == expected_output
+        assert "seq_lengths" in expected_format["columns"]
+        expected_format["columns"].remove("seq_lengths")
         assert format == dataset.format
 
     def test_with_iterable_dataset(self):

trl/data_utils.py

@@ -894,6 +894,8 @@ def pack_dataset(
         )
     else:  # PackingStrategy.WRAPPED
         dataset = dataset.map(_pack_wrapped, batched=True, fn_kwargs={"seq_length": seq_length}, **map_kwargs)
+    if strategy in {PackingStrategy.BFD, PackingStrategy.BFD_SPLIT} and "columns" in format:
+        format["columns"].append("seq_lengths")
     dataset = dataset.with_format(**format)
     return dataset