fix(data): fix _pack_wrapped corrupting data on sliced input tables

[AuthRan]

What does this PR do?

Fixes #6669.

_pack_wrapped (the "wrapped" strategy of pack_dataset) computes
a fresh set of 0-based offsets sized to columns[0]'s own slice of
its child buffer, then reuses those offsets against every column's
raw, unsliced .values buffer:

offsets, values = columns[0].offsets, columns[0].values
values = values[offsets[0].as_py() : offsets[-1].as_py()]  # correctly sliced, but only used for num_elements
num_elements = len(values)
offsets = np.arange(0, num_elements, seq_length, ...)       # 0-based, relative to the *sliced* view
...
columns = [
    type(column).from_arrays(offsets..., column.values)      # `column.values` is the *raw, unsliced* buffer
    for column in columns
]

If the table handed to _pack_wrapped is itself a zero-copy slice of
a larger underlying buffer — which happens routinely, e.g. once
dataset.map(..., batched=True) moves past the first batch — each
column's real data starts at some non-zero offset into the raw
buffer. Pairing that raw buffer with 0-based offsets silently
duplicates earlier rows and drops the real ones, with no error or
warning:

from datasets import Dataset
from trl import pack_dataset

ds = Dataset.from_dict({"input_ids": [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]]})
packed = pack_dataset(ds, seq_length=4, strategy="wrapped", map_kwargs={"batch_size": 3})
print(packed.to_dict()["input_ids"])
# Actual:   [[1, 2, 3, 4], [5, 6], [1, 2, 3, 4], [5, 6]]   <- second batch duplicates the first
# Expected: [[1, 2, 3, 4], [5, 6], [7, 8, 9, 10], [11, 12]]

This is a regression from #5189. Only "wrapped" is affected;
"bfd"/"bfd_split" build fresh arrays via pc.take and don't share
this issue.

The fix

Slice each column's .values down to its own offset bounds before
pairing it with the new offsets, the same way it was already (only)
done for columns[0] to compute num_elements.

Before submitting

• Added a regression test reproducing the issue's repro (pack_dataset(..., strategy="wrapped", map_kwargs={"batch_size": 3})); confirmed it fails against the pre-fix code and passes against the fix.
• Manually verified the issue's larger repro (2500 rows) no longer loses any tokens.
• Ran ruff check / ruff format on the changed files.
• Existing TestPackDatasetWrapped / TestPackDatasetBfd tests still pass.

---

Note

Medium Risk
Fixes incorrect training data for wrapped packing under multi-batch map; scope is narrow but wrong packed tokens could have affected downstream LM training.

Overview
Fixes silent data corruption in pack_dataset when using the wrapped strategy on batched map input (e.g. map_kwargs={"batch_size": 3}). After the first batch, Hugging Face datasets often passes zero-copy slices of a larger Arrow buffer; _pack_wrapped built 0-based offsets for that slice but still attached them to each column’s full, unsliced .values buffer, which duplicated earlier rows and dropped the real ones with no error.

The change slices every column’s .values to [offsets[0]:offsets[-1]] before from_arrays, matching what was already done for columns[0] when computing num_elements. bfd / bfd_split are unchanged.

Adds test_with_multiple_map_batches to lock in correct packing across multiple map batches.

^{Reviewed by Cursor Bugbot for commit ff265d8. Bugbot is set up for automated code reviews on this repo. Configure here.}