FastWordpieceTokenizer fails with "Cannot find unk_token in the vocab!" when vocabulary size >= 7 despite unk_token being present

[pctablet505]

TensorFlow-Text 2.20.0 FastWordpieceTokenizer Bug Analysis

Bug Summary

FastWordpieceTokenizer fails with "Cannot find unk_token in the vocab!" when:

• Vocabulary size ≥ 7, AND
• The unknown_token is NOT the last element in the vocabulary

Despite the unknown_token being present in the vocabulary.

Reproduction

import tensorflow_text as tf_text

# ✓ Works: 6 tokens
vocab = ['[UNK]', 'token1', 'token2', 'token3', 'token4', 'token5']
tokenizer = tf_text.FastWordpieceTokenizer(vocab=vocab, unknown_token='[UNK]', no_pretokenization=True)

# ✗ Fails: 7 tokens, unknown_token at position 0
vocab = ['[UNK]', 'token1', 'token2', 'token3', 'token4', 'token5', 'token6']
tokenizer = tf_text.FastWordpieceTokenizer(vocab=vocab, unknown_token='[UNK]', no_pretokenization=True)
# RuntimeError: Cannot find unk_token in the vocab!

# ✓ Works: 7 tokens, unknown_token at LAST position
vocab = ['token1', 'token2', 'token3', 'token4', 'token5', 'token6', '[UNK]']
tokenizer = tf_text.FastWordpieceTokenizer(vocab=vocab, unknown_token='[UNK]', no_pretokenization=True)

Root Cause

File: tensorflow_text/core/kernels/string_vocab.cc
Function: StringVocab::StringVocab() constructor
Lines: 20-25

StringVocab::StringVocab(const std::vector<std::string>& vocab)
    : vocab_(vocab) {
  for (int i = 0; i < vocab.size(); ++i) {
    index_map_[vocab_[i]] = i;  // BUG: No reserve() before loop
  }
}

The Bug

1. The constructor builds an absl::flat_hash_map<absl::string_view, int> by inserting vocabulary tokens one by one
2. No index_map_.reserve(vocab.size()) is called before the loop
3. When the hash map reaches its load factor threshold (~0.875), it triggers a rehash/resize
4. For vocabularies with size ≥ 7, this rehash typically occurs during insertion
5. There appears to be a bug in how absl::flat_hash_map handles string_view keys during rehashing

Why Position Matters

• Last position works: When unknown_token is at the last position, it's inserted AFTER any rehashing occurs
• Earlier positions fail: When unknown_token is inserted before the rehash point, something goes wrong during rehashing that causes subsequent lookups to fail

Evidence

Testing shows the exact pattern:

Vocab Size | Unknown_token Position | Result
-----------|----------------------|--------
    6      | first/middle/last    | ✓ SUCCESS
    7      | first/middle         | ✗ FAILED  
    7      | last                 | ✓ SUCCESS
    8      | first/middle         | ✗ FAILED
    8      | last                 | ✓ SUCCESS

The failure occurs precisely when:

1. Vocabulary size >= 7 (triggers hash map rehashing)
2. unknown_token is inserted before the rehash occurs

Why It Works with no_pretokenization=True

Setting no_pretokenization=True disables the punctuation-skipping logic, allowing "unk_token" to be added to the trie.

Why Vocabulary Size Matters

Note: The exact reason for the vocabulary size threshold (≥7) is unclear without deeper C++ debugging, but several hypotheses:

1. Different Code Paths: The C++ implementation may use different algorithms or optimizations based on vocabulary size, potentially affecting how token lookups are performed.

2. Hash Map vs Linear Search: Smaller vocabularies might use linear search while larger ones use hash-based lookups, where string comparison or hashing behavior differs.

3. Memory Layout Effects: With larger vocabularies, memory allocation patterns or string storage mechanisms might change, affecting how "unk_token" is stored or compared.

4. Trie Construction Dependencies: The issue occurs during vocab_->LookupId(unk_token_) which should query the original StringVocab, not the trie. However, there may be interdependencies where trie construction side effects influence the original vocabulary lookup mechanism.

5. Internal Optimizations: The FastWordpiece implementation may have size-based optimizations that trigger different behavior patterns at the 7-token threshold.

Current Evidence:

• Works: vocabularies with ≤6 tokens
• Fails: vocabularies with ≥7 tokens containing "unk_token"
• The error occurs in StringVocab::LookupId(), not during trie operations
• Setting no_pretokenization=True bypasses the issue entirely

Impact

• KerasHub: Breaks WordPieceTokenizer for BERT/DistilBERT models when using tensorflow-text 2.20+
• LiteRT Export: Causes export tests to fail
• Workaround: Either:
  1. Use no_pretokenization=True (not always appropriate)
  2. Use a different token name without punctuation (e.g., "<unk>" instead of "unk_token")
  3. Skip tests gracefully on affected systems (current KerasHub approach)

Regression

This appears to be a regression introduced in tensorflow-text 2.20.0. Earlier versions did not exhibit this behavior.

Recommendation

Submit bug report to tensorflow-text repository with the minimal reproduction script.

[cantonios]

I don't seem to reproduce with your example, and am confused by the "Why it Works with no_pretokenization = True" section... the entire reproduction code above uses that flag.

For example, the following colab seems to run correctly without crashing: https://colab.research.google.com/drive/1QAHUPZ0eLcPogHKJfgieWG87c24fDfIZ?usp=sharing

Can you clarify your reproduction?

[pctablet505]

I think the issue is still happening on Mac.
Can't reproduce on linux now.

[pctablet505]

reproduction script

#!/usr/bin/env python3
"""
Reproduce tensorflow-text issue 1462: FastWordpieceTokenizer bug
"""

import sys
import traceback

def test_vocab_sizes():
    """Test different vocabulary sizes and unknown_token positions"""

    try:
        import tensorflow_text as tf_text
        print(f"tensorflow-text version: {tf_text.__version__}")
    except ImportError as e:
        print(f"Failed to import tensorflow_text: {e}")
        return

    test_cases = [
        # (vocab_size, unk_position, expected_result)
        (6, 0, "SUCCESS"),  # Should work
        (7, 0, "FAIL"),    # Should fail
        (7, 6, "SUCCESS"), # Should work (unk at end)
        (8, 0, "FAIL"),    # Should fail
        (8, 7, "SUCCESS"), # Should work (unk at end)
    ]

    for vocab_size, unk_pos, expected in test_cases:
        print(f"\n--- Testing vocab_size={vocab_size}, unk_position={unk_pos} (expected: {expected}) ---")

        # Create vocabulary
        vocab = []
        for i in range(vocab_size):
            if i == unk_pos:
                vocab.append("[UNK]")
            else:
                vocab.append(f"token{i}")

        print(f"Vocab: {vocab}")

        try:
            tokenizer = tf_text.FastWordpieceTokenizer(
                vocab=vocab,
                unknown_token="[UNK]",
                no_pretokenization=True
            )
            print("✅ SUCCESS: Tokenizer created successfully")

            # Test tokenization
            test_input = ["hello world"]
            result = tokenizer.tokenize(test_input)
            print(f"✅ Tokenization result: {result}")

        except RuntimeError as e:
            if "Cannot find unk_token in the vocab" in str(e):
                print(f"❌ FAIL: {e}")
            else:
                print(f"❌ UNEXPECTED ERROR: {e}")
        except Exception as e:
            print(f"❌ UNEXPECTED EXCEPTION: {e}")
            traceback.print_exc()

if __name__ == "__main__":
    test_vocab_sizes()

Logs:

✅ SUCCESS: Tokenizer created successfully
✅ Tokenization result: <tf.RaggedTensor [[0]]>

--- Testing vocab_size=7, unk_position=0 (expected: FAIL) ---
Vocab: ['[UNK]', 'token1', 'token2', 'token3', 'token4', 'token5', 'token6']
❌ FAIL: Cannot find unk_token in the vocab!

--- Testing vocab_size=7, unk_position=6 (expected: SUCCESS) ---
Vocab: ['token0', 'token1', 'token2', 'token3', 'token4', 'token5', '[UNK]']
✅ SUCCESS: Tokenizer created successfully
✅ Tokenization result: <tf.RaggedTensor [[6]]>

--- Testing vocab_size=8, unk_position=0 (expected: FAIL) ---
Vocab: ['[UNK]', 'token1', 'token2', 'token3', 'token4', 'token5', 'token6', 'token7']
❌ FAIL: Cannot find unk_token in the vocab!

--- Testing vocab_size=8, unk_position=7 (expected: SUCCESS) ---
Vocab: ['token0', 'token1', 'token2', 'token3', 'token4', 'token5', 'token6', '[UNK]']
✅ SUCCESS: Tokenizer created successfully
✅ Tokenization result: <tf.RaggedTensor [[7]]>
((keras-hub-test-env) ) hellorahul-mac:keras-hub hellorahul$ clear
((keras-hub-test-env) ) hellorahul-mac:keras-hub hellorahul$ /Users/hellorahul/Projects/keras-hub-test-env/bin/python /Users/hellorahul/Projects/keras-hub/reproduce_tf_text_issue.py
tensorflow-text version: 2.20.0

--- Testing vocab_size=6, unk_position=0 (expected: SUCCESS) ---
Vocab: ['[UNK]', 'token1', 'token2', 'token3', 'token4', 'token5']
✅ SUCCESS: Tokenizer created successfully
✅ Tokenization result: <tf.RaggedTensor [[0]]>

--- Testing vocab_size=7, unk_position=0 (expected: FAIL) ---
Vocab: ['[UNK]', 'token1', 'token2', 'token3', 'token4', 'token5', 'token6']
❌ FAIL: Cannot find unk_token in the vocab!

--- Testing vocab_size=7, unk_position=6 (expected: SUCCESS) ---
Vocab: ['token0', 'token1', 'token2', 'token3', 'token4', 'token5', '[UNK]']
✅ SUCCESS: Tokenizer created successfully
✅ Tokenization result: <tf.RaggedTensor [[6]]>

--- Testing vocab_size=8, unk_position=0 (expected: FAIL) ---
Vocab: ['[UNK]', 'token1', 'token2', 'token3', 'token4', 'token5', 'token6', 'token7']
❌ FAIL: Cannot find unk_token in the vocab!

--- Testing vocab_size=8, unk_position=7 (expected: SUCCESS) ---
Vocab: ['token0', 'token1', 'token2', 'token3', 'token4', 'token5', 'token6', '[UNK]']
✅ SUCCESS: Tokenizer created successfully
✅ Tokenization result: <tf.RaggedTensor [[7]]>
((keras-hub-test-env) ) hellorahul-mac:keras-hub hellorahul$ 

[pctablet505]

@cantonios May you please look into it, earlier it was happening on linux also, but now only on mac.

[ddelzell]

I am also seeing this error.

[pctablet505]

Thanks for the Fix