Add pruning functionality to tokenizer creation to remove zero frequency tokens

[chandanms]

Description

This PR implements tokenizer pruning functionality to address high cosine similarity issues between token embeddings. The solution removes unused tokens from the trained tokenizer vocabulary and reassigns sequential token IDs to create a more compact and efficient tokenizer.

• Added prune_tokenizer() function that identifies and removes tokens with zero frequency in the dataset
• Preserves special tokens ([UNK], [EOS]) and their functionality, original tokenizer settings (normalizers, pre-tokenizers, post-processors, decoders) and serialized IDs

Other minor updates in the PR are

• changing the default names to more updated ones (ex: lennart-finke/SimpleStories to SimpleStories/SimpleStories)
• streamlining the tokenizer creating steps like removing the dataset splitting (it adds no value in case of tokenizer), making the script easily executable with different datasets etc

Related Issue

Fixes the token embedding cosine similarity issue where ~2% of zero-frequency tokens from WordPieceTrainer were collapsing into the same embedding direction as [UNK], causing cosine similarities >0.9 and nearly identical embeddings.

Motivation and Context

The original issue was that trained model's embedding layer had cosine similarities >0.9. Investigation revealed that WordPieceTrainer was generating tokens with zero frequency in the actual dataset, which during model training would collapse to have similar embeddings as the [UNK] token. This resulted in high cosine similarities between token embeddings.

The pruning approach removes these problematic tokens and creates a cleaner, more efficient tokenizer that only contains tokens actually present in the training data.

How Has This Been Tested?

A test code has been written in test_tokenizer.py which demonstrates the wordpiece tokenizer issue of creating artifacts when vocab_size is too high and how pruning removes the zero frequency tokens

Does this PR introduce a breaking change?

No. The pruning step can be optionally disabled by not calling prune_tokenizer() in the main workflow.

[danbraunai]

Nice, feel free to merge after addressing my comments.

I commented on this once, but I think all default arguments should be removed that aren't actually used in this codebase. I spotted a few of them.

I changed the base of this PR to a dev branch. I don't think we should change the main branch, at least not right now, as people will expect that to match the paper.

[danbraunai]

Not great having the file be called 35M but the model name 1.25M.

[chandanms]

I think its better to change the filename to train_config.yaml instead

[danbraunai]

Yep that sounds fine. Do you want to do that and then merge?

[danbraunai]

@chandanms changes look great! Still one comment about the 35M_config.yaml naming. Feel free to merge after handling that.

[chandanms]

@danbraunai Made the changes to file. Btw, not sure if I am missing something but I don't think I have the merge PR rights!