⚡️ Speed up method RealmTokenizer._tokenize by 9%

[codeflash-ai]

📄 9% (0.09x) speedup for RealmTokenizer._tokenize in src/transformers/models/deprecated/realm/tokenization_realm.py

⏱️ Runtime : 43.8 milliseconds → 40.1 milliseconds (best of 106 runs)

📝 Explanation and details

The optimized code achieves a 9% speedup through several targeted micro-optimizations focused on reducing redundant operations and method call overhead:

Key Optimizations:

1. Reduced Method Call Overhead: The optimized version caches frequently accessed methods and attributes as local variables (e.g., wordpiece_tokenize = self.wordpiece_tokenizer.tokenize, all_special_tokens = self.all_special_tokens) to avoid repeated attribute lookups during loops.

2. Streamlined Set Operations: In BasicTokenizer.tokenize(), the code now conditionally creates the union of never_split sets only when needed, rather than always creating a new set, reducing unnecessary set operations.

3. Optimized String Operations: In WordpieceTokenizer.tokenize(), the code eliminates the intermediate list(token) conversion and directly works with string slicing (chars[start:end]), reducing memory allocations and improving string manipulation performance.

4. Inlined Utility Functions: Critical utility functions like load_vocab, whitespace_tokenize, and character classification helpers (_is_whitespace, _is_control, _is_punctuation) are now defined locally, eliminating import overhead and function call indirection.

5. Improved Loop Efficiency: The code moves invariant checks outside loops where possible and uses more efficient list comprehensions and generator expressions for better performance.

Impact on Workloads:

Based on the test results, the optimizations are particularly effective for:

• Large-scale tokenization: Tests with 500-1000 tokens show 5-16% improvements
• Mixed content processing: Tests combining known/unknown words benefit significantly (7-16% faster)
• Punctuation-heavy text: Large punctuation processing shows 14% improvement

The optimizations maintain identical functionality while providing consistent performance gains across various text processing scenarios, making this particularly valuable for high-throughput tokenization workloads.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 213 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import collections
import os
import tempfile

# imports
import pytest

from transformers.models.deprecated.realm.tokenization_realm import RealmTokenizer


def load_vocab(vocab_file):
    vocab = collections.OrderedDict()
    with open(vocab_file, "r", encoding="utf-8") as reader:
        tokens = reader.readlines()
    for index, token in enumerate(tokens):
        token = token.rstrip("\n")
        vocab[token] = index
    return vocab


# Minimal PreTrainedTokenizer stub
class PreTrainedTokenizer:
    def __init__(self, **kwargs):
        self.all_special_tokens = getattr(self, "all_special_tokens", [])


# --- Unit Tests ---


@pytest.fixture(scope="module")
def temp_vocab_file():
    # Create a temporary vocab file for tests
    vocab = [
        "[UNK]",
        "[CLS]",
        "[SEP]",
        "[PAD]",
        "[MASK]",
        "the",
        "quick",
        "brown",
        "fox",
        "jumps",
        "over",
        "lazy",
        "dog",
        "un",
        "##aff",
        "##able",
        "##ly",
        "affable",
        "hello",
        "world",
        "!",
        ".",
        ",",
        "##s",
        "##'",
        "##ed",
        "##ing",
        "##er",
        "##est",
        "a",
        "b",
        "c",
        "##c",
        "##b",
        "##a",
        "##dog",
        "##fox",
        "##o",
        "##g",
        "中国",
        "人",
        "##人",
        "的",
        "##的",
        "你",
        "好",
        "##好",
    ]
    with tempfile.NamedTemporaryFile(delete=False, mode="w", encoding="utf-8") as f:
        for token in vocab:
            f.write(token + "\n")
        fname = f.name
    yield fname
    os.remove(fname)


# --- Basic Test Cases ---


def test_tokenize_simple_sentence(temp_vocab_file):
    # Simple sentence, all tokens in vocab
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "the quick brown fox jumps over the lazy dog"
    expected = ["the", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog"]
    codeflash_output = tokenizer._tokenize(text)  # 68.5μs -> 66.8μs (2.50% faster)


def test_tokenize_with_punctuation(temp_vocab_file):
    # Sentence with punctuation
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "hello, world!"
    # "hello" and "world" are in vocab, "," and "!" are in vocab
    expected = ["hello", ",", "world", "!"]
    codeflash_output = tokenizer._tokenize(text)  # 33.2μs -> 32.8μs (1.09% faster)


def test_tokenize_wordpiece_split(temp_vocab_file):
    # Word not in vocab but split into wordpieces
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "unaffable"
    # "un" + "##aff" + "##able"
    expected = ["un", "##aff", "##able"]
    codeflash_output = tokenizer._tokenize(text)  # 28.5μs -> 27.6μs (3.08% faster)


def test_tokenize_mixed_case(temp_vocab_file):
    # Case insensitivity
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "The Quick BROWN"
    expected = ["the", "quick", "brown"]
    codeflash_output = tokenizer._tokenize(text)  # 34.4μs -> 34.0μs (0.967% faster)


def test_tokenize_unknown_word(temp_vocab_file):
    # Word not in vocab and cannot be split into known wordpieces
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "xyzzy"
    expected = ["[UNK]"]
    codeflash_output = tokenizer._tokenize(text)  # 20.8μs -> 21.5μs (3.29% slower)


def test_tokenize_with_special_tokens(temp_vocab_file):
    # Input contains special tokens, which should not be split
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "[CLS] hello world [SEP]"
    expected = ["[CLS]", "hello", "world", "[SEP]"]
    codeflash_output = tokenizer._tokenize(text)  # 37.6μs -> 37.0μs (1.63% faster)


# --- Edge Test Cases ---


def test_tokenize_empty_string(temp_vocab_file):
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = ""
    expected = []
    codeflash_output = tokenizer._tokenize(text)  # 6.89μs -> 7.74μs (11.0% slower)


def test_tokenize_whitespace_only(temp_vocab_file):
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "    \t   \n"
    expected = []
    codeflash_output = tokenizer._tokenize(text)  # 13.4μs -> 14.8μs (9.65% slower)


def test_tokenize_long_word_over_max_length(temp_vocab_file):
    # Word longer than max_input_chars_per_word should yield [UNK]
    tokenizer = RealmTokenizer(temp_vocab_file)
    long_word = "a" * 101
    expected = ["[UNK]"]
    codeflash_output = tokenizer._tokenize(long_word)  # 102μs -> 104μs (1.64% slower)


def test_tokenize_multiple_spaces(temp_vocab_file):
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "the    quick\tbrown\nfox"
    expected = ["the", "quick", "brown", "fox"]
    codeflash_output = tokenizer._tokenize(text)  # 40.6μs -> 40.3μs (0.647% faster)


def test_tokenize_with_accents(temp_vocab_file):
    # Should strip accents when lowercasing
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "Café"
    # "cafe" not in vocab, so should yield [UNK]
    expected = ["[UNK]"]
    codeflash_output = tokenizer._tokenize(text)  # 25.0μs -> 25.1μs (0.036% slower)


def test_tokenize_with_chinese_characters(temp_vocab_file):
    # Chinese characters are surrounded by whitespace and tokenized
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "你好中国人"
    # Should split: 你, 好, 中, 国, 人
    # Only "你", "好", "中国", "人" are in vocab, so "你", "好", "中国", "人"
    # But since tokenizer splits chars, and "中国" is in vocab, but "中" and "国" are not
    # The tokenizer will split to ["你", "好", "中", "国", "人"], but only "你", "好", "中国", "人" in vocab
    # In practice, BasicTokenizer will split each CJK char, so ["你", "好", "中", "国", "人"]
    # But only "你", "好", "中国", "人" in vocab, so "你"->"你", "好"->"好", "中"->"[UNK]", "国"->"[UNK]", "人"->"人"
    expected = ["你", "好", "[UNK]", "[UNK]", "人"]
    codeflash_output = tokenizer._tokenize(text)  # 30.4μs -> 29.6μs (2.68% faster)


def test_tokenize_never_split_token(temp_vocab_file):
    # Provide a custom never_split token
    tokenizer = RealmTokenizer(temp_vocab_file, never_split=["<SPECIAL>"])
    text = "the <SPECIAL> fox"
    expected = ["the", "<SPECIAL>", "fox"]
    codeflash_output = tokenizer._tokenize(text)  # 32.8μs -> 32.1μs (2.25% faster)


def test_tokenize_strip_accents_false(temp_vocab_file):
    # If strip_accents is False, accents are not stripped
    tokenizer = RealmTokenizer(temp_vocab_file, strip_accents=False)
    text = "Café"
    # Lowercased: "café", which is not in vocab, so [UNK]
    expected = ["[UNK]"]
    codeflash_output = tokenizer._tokenize(text)  # 22.9μs -> 23.4μs (1.98% slower)


def test_tokenize_with_only_punctuation(temp_vocab_file):
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "!!!"
    # Each "!" is in vocab, so should be three "!"
    expected = ["!", "!", "!"]
    codeflash_output = tokenizer._tokenize(text)  # 20.3μs -> 20.2μs (0.390% faster)


def test_tokenize_with_apostrophe_and_suffix(temp_vocab_file):
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "dog's"
    # "dog" in vocab, "##'s" in vocab, so should be ["dog", "##'s"], but "##'s" is not in vocab, but "##'" is
    # So "dog", "##'", "s" ("s" in vocab)
    expected = ["dog", "##'", "s"]
    codeflash_output = tokenizer._tokenize(text)  # 24.0μs -> 23.8μs (0.824% faster)


def test_tokenize_with_repeated_unknowns(temp_vocab_file):
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "foo bar baz"
    # None of these in vocab, so each is [UNK]
    expected = ["[UNK]", "[UNK]", "[UNK]"]
    codeflash_output = tokenizer._tokenize(text)  # 34.1μs -> 33.5μs (1.99% faster)


def test_tokenize_with_mixed_known_and_unknown(temp_vocab_file):
    tokenizer = RealmTokenizer(temp_vocab_file)
    text = "hello foo world"
    expected = ["hello", "[UNK]", "world"]
    codeflash_output = tokenizer._tokenize(text)  # 36.2μs -> 34.4μs (5.16% faster)


# --- Large Scale Test Cases ---


def test_tokenize_large_input(temp_vocab_file):
    tokenizer = RealmTokenizer(temp_vocab_file)
    # Repeat a known sentence 100 times
    text = "the quick brown fox jumps over the lazy dog " * 100
    expected = ["the", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog"] * 100
    codeflash_output = tokenizer._tokenize(text.strip())  # 4.78ms -> 4.51ms (5.89% faster)


def test_tokenize_large_mixed_input(temp_vocab_file):
    tokenizer = RealmTokenizer(temp_vocab_file)
    # Mix known and unknown words, 500 tokens
    known = "the quick brown fox"
    unknown = "foo bar baz"
    text = ((known + " " + unknown + " ") * 50).strip()
    expected = ["the", "quick", "brown", "fox", "[UNK]", "[UNK]", "[UNK]"] * 50
    codeflash_output = tokenizer._tokenize(text)  # 1.81ms -> 1.68ms (7.56% faster)


def test_tokenize_long_wordpiece_chain(temp_vocab_file):
    tokenizer = RealmTokenizer(temp_vocab_file)
    # "unaffably" -> "un", "##aff", "##able", "##ly" (all in vocab)
    text = "unaffably"
    expected = ["un", "##aff", "##able", "##ly"]
    codeflash_output = tokenizer._tokenize(text)  # 30.5μs -> 28.9μs (5.50% faster)


def test_tokenize_max_batch(temp_vocab_file):
    tokenizer = RealmTokenizer(temp_vocab_file)
    # 1000 tokens, alternating known/unknown
    text = ("the foo " * 500).strip()
    expected = ["the", "[UNK]"] * 500
    codeflash_output = tokenizer._tokenize(text)  # 4.57ms -> 4.17ms (9.48% faster)


def test_tokenize_large_punctuation(temp_vocab_file):
    tokenizer = RealmTokenizer(temp_vocab_file)
    # 500 exclamation marks
    text = "!" * 500
    expected = ["!"] * 500
    codeflash_output = tokenizer._tokenize(text)  # 667μs -> 585μs (14.0% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import collections
import os
import tempfile

# imports
import pytest

from transformers.models.deprecated.realm.tokenization_realm import RealmTokenizer


def load_vocab(vocab_file):
    vocab = collections.OrderedDict()
    with open(vocab_file, "r", encoding="utf-8") as reader:
        tokens = reader.readlines()
    for index, token in enumerate(tokens):
        token = token.rstrip("\n")
        vocab[token] = index
    return vocab


# --- Fixtures and helpers ---


@pytest.fixture(scope="module")
def vocab_file():
    # Create a temporary vocab file for tests
    vocab = [
        "[UNK]",
        "[SEP]",
        "[PAD]",
        "[CLS]",
        "[MASK]",
        "hello",
        "world",
        "!",
        "un",
        "##aff",
        "##able",
        "the",
        "quick",
        "brown",
        "fox",
        "##es",
        "jumps",
        "over",
        "lazy",
        "dog",
        "##s",
        "##ly",
        "##ing",
        "a",
        "##b",
        "##c",
        "##d",
        "##e",
        "##f",
        "##g",
        "##h",
        "##i",
        "##j",
        "##k",
        "##l",
        "##m",
        "##n",
        "##o",
        "##p",
        "##q",
        "##r",
        "##s",
        "##t",
        "##u",
        "##v",
        "##w",
        "##x",
        "##y",
        "##z",
        "##!",
        "##.",
        "##,",
        "##?",
        "##-",
        "##_",
        "##'",
        '##"',
        "##(",
        "##)",
        "##:",
        "##;",
        "##/",
        "###",
        "##$",
        "##%",
        "##&",
        "##*",
        "##@",
        "##^",
        "##~",
        "##`",
        "##[",
        "##]",
        "##{",
        "##}",
        "##<",
        "##>",
        "##|",
        "##\\",
        "##=",
        "##+",
        "##0",
        "##1",
        "##2",
        "##3",
        "##4",
        "##5",
        "##6",
        "##7",
        "##8",
        "##9",
        "unaffable",
        "##affable",
        "##able",
        "中国",
        "的",
        "人",
        "##国",
        "##的",
        "##人",
        "##中",
        "##国",
        "##中",
        "##国",
        "##人",
        "##中国",
        "##中国人",
    ]
    with tempfile.NamedTemporaryFile(mode="w+", delete=False, encoding="utf-8") as f:
        for token in vocab:
            f.write(token + "\n")
        f.flush()
        yield f.name
    os.remove(f.name)


@pytest.fixture
def tokenizer(vocab_file):
    return RealmTokenizer(vocab_file=vocab_file)


# --- Basic Test Cases ---


def test_basic_single_word(tokenizer):
    # Basic single word in vocab
    codeflash_output = tokenizer._tokenize("hello")
    tokens = codeflash_output  # 25.3μs -> 25.3μs (0.020% slower)


def test_basic_sentence(tokenizer):
    # Sentence with words in vocab
    codeflash_output = tokenizer._tokenize("hello world!")
    tokens = codeflash_output  # 34.9μs -> 34.5μs (1.19% faster)


def test_basic_wordpiece(tokenizer):
    # Word that should be split into wordpieces
    codeflash_output = tokenizer._tokenize("unaffable")
    tokens = codeflash_output  # 26.9μs -> 27.6μs (2.82% slower)


def test_basic_mixed_case(tokenizer):
    # Should lowercase by default
    codeflash_output = tokenizer._tokenize("HELLO world")
    tokens = codeflash_output  # 31.3μs -> 31.5μs (0.410% slower)


def test_basic_special_tokens(tokenizer):
    # Special tokens should not be split
    codeflash_output = tokenizer._tokenize("[CLS] hello [SEP]")
    tokens = codeflash_output  # 32.0μs -> 32.2μs (0.460% slower)


def test_basic_punctuation(tokenizer):
    # Punctuation splitting
    codeflash_output = tokenizer._tokenize("hello,world!")
    tokens = codeflash_output  # 32.9μs -> 33.0μs (0.454% slower)


def test_basic_wordpiece_not_in_vocab(tokenizer):
    # Should return [UNK] for unknown word
    codeflash_output = tokenizer._tokenize("unknownword")
    tokens = codeflash_output  # 41.4μs -> 37.0μs (11.8% faster)


# --- Edge Test Cases ---


def test_edge_empty_string(tokenizer):
    # Empty string should return empty list
    codeflash_output = tokenizer._tokenize("")
    tokens = codeflash_output  # 8.75μs -> 9.31μs (5.97% slower)


def test_edge_whitespace_only(tokenizer):
    # Only whitespace should return empty list
    codeflash_output = tokenizer._tokenize("   \t  \n")
    tokens = codeflash_output  # 14.4μs -> 15.6μs (7.19% slower)


def test_edge_long_word(tokenizer):
    # Word longer than max_input_chars_per_word (default 100)
    long_word = "a" * 101
    codeflash_output = tokenizer._tokenize(long_word)
    tokens = codeflash_output  # 104μs -> 105μs (1.34% slower)


def test_edge_control_characters(tokenizer):
    # Control characters should be removed
    codeflash_output = tokenizer._tokenize("hello\u0000world")
    tokens = codeflash_output  # 33.5μs -> 32.4μs (3.26% faster)


def test_edge_accents(tokenizer):
    # Accented characters should be stripped if do_lower_case
    codeflash_output = tokenizer._tokenize("Café")
    tokens = codeflash_output  # 25.4μs -> 26.1μs (2.66% slower)


def test_edge_chinese_characters(tokenizer):
    # Chinese characters should be tokenized separately
    codeflash_output = tokenizer._tokenize("中国的人")
    tokens = codeflash_output  # 30.7μs -> 29.6μs (3.77% faster)


def test_edge_mixed_chinese_english(tokenizer):
    codeflash_output = tokenizer._tokenize("hello中国world")
    tokens = codeflash_output  # 40.8μs -> 39.6μs (3.11% faster)


def test_edge_multiple_punctuation(tokenizer):
    codeflash_output = tokenizer._tokenize("hello!!!")
    tokens = codeflash_output  # 28.5μs -> 28.1μs (1.59% faster)


def test_edge_never_split(tokenizer, vocab_file):
    # Test never_split argument
    t = RealmTokenizer(vocab_file=vocab_file, never_split=["foo"])
    codeflash_output = t._tokenize("foo bar")
    tokens = codeflash_output  # 23.5μs -> 24.3μs (3.24% slower)


def test_edge_strip_accents_false(tokenizer, vocab_file):
    # Test strip_accents=False disables accent stripping
    t = RealmTokenizer(vocab_file=vocab_file, strip_accents=False)
    codeflash_output = t._tokenize("Café")
    tokens = codeflash_output  # 20.9μs -> 22.3μs (6.20% slower)


def test_edge_do_basic_tokenize_false(tokenizer, vocab_file):
    # Only wordpiece tokenizer should be used
    t = RealmTokenizer(vocab_file=vocab_file, do_basic_tokenize=False)
    codeflash_output = t._tokenize("hello world!")
    tokens = codeflash_output  # 5.91μs -> 4.46μs (32.4% faster)


def test_edge_tokenize_chinese_chars_false(tokenizer, vocab_file):
    # Should not split Chinese chars if tokenize_chinese_chars=False
    t = RealmTokenizer(vocab_file=vocab_file, tokenize_chinese_chars=False)
    codeflash_output = t._tokenize("中国的人")
    tokens = codeflash_output  # 23.8μs -> 23.8μs (0.370% faster)


def test_edge_wordpiece_partial_match(tokenizer, vocab_file):
    # 'unaffable' is in vocab, should match whole word, not split
    t = RealmTokenizer(vocab_file=vocab_file)
    codeflash_output = t._tokenize("unaffable")
    tokens = codeflash_output  # 25.8μs -> 26.3μs (1.77% slower)


def test_edge_wordpiece_greedy(tokenizer, vocab_file):
    # 'unaffable' should split into ['un', '##aff', '##able'] if full word not in vocab
    # Remove 'unaffable' from vocab for this test
    with tempfile.NamedTemporaryFile(mode="w+", delete=False, encoding="utf-8") as f:
        vocab = [tok for tok in load_vocab(vocab_file) if tok != "unaffable"]
        for token in vocab:
            f.write(token + "\n")
        f.flush()
        t = RealmTokenizer(vocab_file=f.name)
        codeflash_output = t._tokenize("unaffable")
        tokens = codeflash_output  # 28.6μs -> 28.5μs (0.661% faster)
    os.remove(f.name)


# --- Large Scale Test Cases ---


def test_large_scale_many_words(tokenizer):
    # Tokenize a sentence with many words (<=1000 tokens)
    sentence = " ".join(["hello"] * 500 + ["world"] * 499)
    codeflash_output = tokenizer._tokenize(sentence)
    tokens = codeflash_output  # 6.25ms -> 5.94ms (5.13% faster)


def test_large_scale_long_text(tokenizer):
    # Tokenize a long paragraph with punctuation
    text = ("The quick brown fox jumps over the lazy dog! " * 40).strip()
    codeflash_output = tokenizer._tokenize(text)
    tokens = codeflash_output  # 1.94ms -> 1.82ms (7.00% faster)
    # Each sentence: ['the', 'quick', 'brown', 'fox', 'jumps', 'over', 'the', 'lazy', 'dog', '!']
    expected = ["the", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog", "!"] * 40


def test_large_scale_all_vocab(tokenizer, vocab_file):
    # Tokenize all vocab tokens in one string
    vocab = list(load_vocab(vocab_file).keys())
    text = " ".join(vocab)
    codeflash_output = tokenizer._tokenize(text)
    tokens = codeflash_output  # 580μs -> 522μs (11.1% faster)


def test_large_scale_wordpiece_split(tokenizer, vocab_file):
    # Tokenize 1000 unknown words, should return [UNK] * 1000
    text = " ".join(["notinvocab"] * 1000)
    codeflash_output = tokenizer._tokenize(text)
    tokens = codeflash_output  # 11.9ms -> 11.0ms (8.03% faster)


def test_large_scale_mixed(tokenizer, vocab_file):
    # Mix known and unknown words, punctuation, Chinese chars
    text = "hello " * 250 + "中国 " * 250 + "unknownword " * 250 + "world! " * 249
    codeflash_output = tokenizer._tokenize(text.strip())
    tokens = codeflash_output  # 10.1ms -> 8.65ms (16.4% faster)
    expected = ["hello"] * 250 + ["中国"] * 250 + ["[UNK]"] * 250 + ["world", "!"] * 249


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-RealmTokenizer._tokenize-mi9vp76r and push.