[Bug Report] tokenize_and_concatenate doesn't tokenize correctly.

[BorisTheBrave]

Describe the bug

tokenize_and_concatenate slices strings into 20 chunks by character before tokenizing. This can cut a token in two, leading to token pairs that would never normally occur.

Code example
I don't have an example, I found this while debugging a larger project.

In a debugger in that method after tokenizing and dropping padding tokens I observed the following:

> chunks[2][-10:], chunks[3][:10]
('t on the M', 'ilitary Ne')
> tokenizer.decode([4460])
' Mil'
> tokenizer.decode([337])
' M'
> tokenizer.decode([346])
'il'
> np.where((tokens[:-1] == 337) & (tokens[1:] == 346))[0]
array([79848])
> tokens[79848:79848+2]
array([337, 346]) #  SHOULD NEVER OCCUR

I think it's obvious from code inspection what the problem is.

Checklist

• I have checked that there is no similar issue in the repo (required)

[brainsnog]

I'd like to work on this. The fix I'm planning is to change the chunking strategy in tokenize_and_concatenate from character-based splitting to word-boundary splitting, so no token is ever cut across chunk boundaries. I'll open a draft PR once I have a working fix with a regression test.

[brainsnog]

I've opened PR #1201 with a fix for this. The approach splits chunks at whitespace boundaries to prevent words being cut across chunk boundaries

[BorisTheBrave]

as it happens, I was making some notes on this issue myself. I think there might be some pathological cases for splitting on whitespace (like, what if there isn't any?)

My notes:

transformerlens forms the full set like this:

full_text = tokenizer.eos_token.join(text)

We can use this to efficiently find eos_boundaries:

eos_positions = cumsum(map(len, text))
eos_near_partition = bisect(len(full_text)*i/20, eos_positions)

These can be safely sliced with no feat of token problems
...
actually... perhaps there's an even easier way?

tokenized_dataset = dataset.map(
    tokenize_function,
    batched=True,
    num_proc=(num_proc if not streaming else None),
    remove_columns=[column_name],
)

am i right in saying this chunks the dataset in to num_proc separate calls to tokenize?
What's stopping us using numproc=num_proc*20, and turning the chunking down to 1.

Sure, the chunks won't be perfectly even, but that's not super important

The map() function supports working with batches of examples. Operate on batches by setting batched=True. The default batch size is 1000, but you can adjust it with the batch_size parameter.
yeah, just the batch size to 1000/20?

[brainsnog]

Thanks for looking at this, really helpful notes. You’re right that pure whitespace splitting has a pathological case if a chunk contains no spaces at all. I can add a fallback to handle that.

The EOS boundary approach is interesting and arguably more principled,splitting at natural document boundaries rather than mid-document. I’m happy to go in that direction if you think it’s cleaner.

On the num_proc idea I want to make sure I understand correctly. Are you suggesting removing the manual chunking entirely and relying on dataset.map with a smaller batch_size to distribute the tokenization work? That would simplify the function considerably. Do you see any downside in terms of tokenization quality or performance?

Happy to revise the approach based on your preference. Just let me know which direction you’d like me to take it.

[BorisTheBrave]

To be clear, I'm not a project maintainer, just a user who has been hit by this issue. I think your questions are better directed at someone who has the ability to approve PRs.

[brainsnog]

I've pushed a second commit addressing the edge case raised in the issue. The whitespace lookahead is now bounded to chunk_length // 10 characters — if no whitespace is found within that window, the cut falls back to the original character boundary rather than consuming the rest of the string. Both the original regression test and a new pathological no-whitespace test pass.

[BorisTheBrave]

Right, so this doesn't actually fix the issue then, it just makes it happen in a smaller selection of cases.

[jlarson4]

@BorisTheBrave You're right that this fix did not 100% cover all possible cases of this bug, I'll reopen it for the time being. I have an EOS related fix in #1273 that should close the gap further.