Fix Unigram trainer prune loss using vocab size instead of per-piece alternatives count

[mayuriphad]

Bug

Fixes #2069.

UnigramTrainer::prune_sentence_pieces (tokenizers/src/models/unigram/trainer.rs) is a Rust port of SentencePiece's PruneSentencePieces. The loss computation for a candidate piece id estimates the new corpus-frequency sum if that piece were removed and its frequency reassigned to its resegmentation alternatives:

let logsum_alt = (sum + freq[id] * (alternatives.len() - 1) as f64).ln();

alternatives is Vec<Vec<usize>>, one inner vector per vocabulary piece, so alternatives.len() is just pieces.len() (the whole vocabulary size) — not the number of alternatives for id. The comment right above it even says alternatives[i].size(), matching SentencePiece's actual code (src/unigram_model_trainer.cc), so the [id] index was dropped when this was ported.

Since freq[id] * (alternatives.len() - 1) scales with total vocab size instead of the piece's own (usually 1-2) alternatives, this inflates the loss of high-frequency pieces by an amount tied to vocab size, which changes which pieces survive pruning. This was already reported in #1536 and re-filed with a fuller repro in #2069 (including an end-to-end comparison on a real corpus showing the trained vocabularies differ deterministically between the two formulas).

Fix

Use the per-piece alternative count instead of the whole-vocabulary size:

let logsum_alt = (sum + freq[id] * (alternatives[id].len() - 1) as f64).ln();

This branch only runs when alternatives[id] is non-empty, so the - 1 cannot underflow.

Test plan

Added test_prune_sentence_pieces_uses_per_piece_alternatives_len in tokenizers/src/models/unigram/trainer.rs, which calls prune_sentence_pieces directly with a small hand-constructed vocabulary containing two candidate pieces ("abcdef", freq 50, 6 alternatives; "xyz", freq 110, 3 alternatives) sized so only one of the two survives pruning. The relative frequencies/alternative counts are chosen (and verified by hand, see comments in the test) so the correct formula ranks "abcdef" above "xyz", while the buggy formula ranks them in the opposite order — this makes the test fail against the pre-fix code and pass against the fix.

I verified the fix by careful reading of prune_sentence_pieces plus the SentencePiece reference implementation, and by manually working out the loss values for both formulas for the test's inputs (shown in the test's comments). I was not able to build/run the Rust test suite in my current environment (no cargo/rustc available), so I was not able to execute cargo test -p tokenizers unigram::trainer locally — please run it in CI. Happy to iterate if the numbers don't line up exactly once it runs.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.