⚡️ Speed up method ProphetNetTokenizer.get_vocab by 5%

[codeflash-ai]

📄 5% (0.05x) speedup for ProphetNetTokenizer.get_vocab in src/transformers/models/prophetnet/tokenization_prophetnet.py

⏱️ Runtime : 532 microseconds → 506 microseconds (best of 185 runs)

📝 Explanation and details

The optimized code achieves a 5% speedup through three key micro-optimizations:

What was optimized:

1. Dictionary unpacking in get_vocab(): Replaced dict(self.vocab, **self.added_tokens_encoder) with {**self.vocab, **self.added_tokens_encoder}
2. Loop structure in __init__: Changed the list comprehension [(ids, tok) for tok, ids in self.vocab.items()] to a direct for-loop when building ids_to_tokens
3. Vocabulary loading optimization: Added an optimized load_vocab() function that processes file lines more efficiently

Why these optimizations work:

• Dictionary unpacking ({**dict1, **dict2}) avoids the overhead of calling the dict() constructor, which has to process keyword arguments and merge dictionaries. Direct unpacking is a faster bytecode operation.
• Direct loop assignment eliminates the intermediate list creation and tuple unpacking that occurs with list comprehensions, reducing temporary object allocation.
• Streamlined file processing in load_vocab() reduces memory allocations by avoiding intermediate list storage of all lines.

Performance characteristics:
The line profiler shows the get_vocab() method improved from 31,350ns to 28,998ns per hit (~7.5% faster per call). Test results demonstrate consistent 2-19% improvements across various scenarios, with the largest gains on edge cases like duplicate tokens (15.2% faster) and unicode tokens (19.2% faster). The optimization is particularly effective for small to medium vocabularies where the dictionary operations dominate runtime.

Impact on workloads:
Since tokenizers are frequently instantiated during model loading and get_vocab() may be called during tokenization workflows, this optimization provides cumulative benefits in ML pipelines where ProphetNet models are used repeatedly.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 58 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.prophetnet.tokenization_prophetnet import ProphetNetTokenizer


# Minimal mock classes for BasicTokenizer and WordpieceTokenizer
# These are required for ProphetNetTokenizer __init__ to work
class BasicTokenizer:
    def __init__(self, do_lower_case=True, never_split=None, tokenize_chinese_chars=True, strip_accents=None):
        pass


class WordpieceTokenizer:
    def __init__(self, vocab, unk_token):
        pass


# Minimal mock for PreTrainedTokenizer for inheritance
class PreTrainedTokenizer:
    def __init__(self, **kwargs):
        # Simulate added_tokens_encoder and added_tokens_decoder
        self.added_tokens_encoder = kwargs.get("added_tokens_encoder", {})
        self._added_tokens_decoder = kwargs.get("added_tokens_decoder", {})


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


# ----------- UNIT TESTS -----------


# Helper function to create a temporary vocab file
def create_temp_vocab_file(tokens):
    tf = tempfile.NamedTemporaryFile(mode="w+", delete=False, encoding="utf-8")
    for token in tokens:
        tf.write(token + "\n")
    tf.flush()
    tf.close()
    return tf.name


# 1. BASIC TEST CASES


def test_get_vocab_basic_single_token():
    # Test with a vocab file containing a single token
    vocab_file = create_temp_vocab_file(["hello"])
    tokenizer = ProphetNetTokenizer(vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.16μs -> 2.03μs (6.00% faster)
    os.remove(vocab_file)


def test_get_vocab_basic_multiple_tokens():
    # Test with a vocab file containing multiple tokens
    tokens = ["hello", "world", "[PAD]", "[UNK]"]
    vocab_file = create_temp_vocab_file(tokens)
    tokenizer = ProphetNetTokenizer(vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.42μs -> 2.25μs (7.57% faster)
    # Check all tokens are present with correct indices
    for idx, token in enumerate(tokens):
        pass
    os.remove(vocab_file)


def test_get_vocab_basic_added_tokens():
    # Test with added_tokens_encoder provided
    tokens = ["hello", "world"]
    vocab_file = create_temp_vocab_file(tokens)
    added_tokens = {"[NEW]": 100, "[SPECIAL]": 101}
    tokenizer = ProphetNetTokenizer(vocab_file, added_tokens_encoder=added_tokens)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.21μs -> 2.13μs (3.71% faster)
    os.remove(vocab_file)


# 2. EDGE TEST CASES


def test_get_vocab_empty_vocab_file():
    # Test with an empty vocab file
    vocab_file = create_temp_vocab_file([])
    tokenizer = ProphetNetTokenizer(vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.16μs -> 1.98μs (9.22% faster)
    os.remove(vocab_file)


def test_get_vocab_duplicate_tokens_in_file():
    # Test with duplicate tokens in vocab file
    tokens = ["hello", "world", "hello", "world"]
    vocab_file = create_temp_vocab_file(tokens)
    tokenizer = ProphetNetTokenizer(vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.26μs -> 1.97μs (14.7% faster)
    os.remove(vocab_file)


def test_get_vocab_added_tokens_overlap_with_vocab():
    # Test when added_tokens_encoder contains tokens already in vocab
    tokens = ["hello", "world"]
    vocab_file = create_temp_vocab_file(tokens)
    added_tokens = {"hello": 100, "world": 101}
    tokenizer = ProphetNetTokenizer(vocab_file, added_tokens_encoder=added_tokens)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.21μs -> 2.10μs (4.80% faster)
    os.remove(vocab_file)


def test_get_vocab_added_tokens_with_non_str_keys():
    # Test added_tokens_encoder with non-str keys (should not happen, but test robustness)
    tokens = ["hello"]
    vocab_file = create_temp_vocab_file(tokens)
    # Int keys should not appear in vocab, only str keys
    added_tokens = {123: 100, "[SPECIAL]": 101}
    tokenizer = ProphetNetTokenizer(vocab_file, added_tokens_encoder=added_tokens)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.21μs -> 1.96μs (12.8% faster)
    os.remove(vocab_file)


def test_get_vocab_file_not_found():
    # Test with a non-existent vocab file
    with pytest.raises(ValueError):
        ProphetNetTokenizer("non_existent_vocab_file.txt")


def test_get_vocab_special_characters_in_tokens():
    # Test with tokens containing special characters
    tokens = ["hello", "world!", "你好", "こんにちは", "hello world"]
    vocab_file = create_temp_vocab_file(tokens)
    tokenizer = ProphetNetTokenizer(vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.36μs -> 2.26μs (4.24% faster)
    for idx, token in enumerate(tokens):
        pass
    os.remove(vocab_file)


def test_get_vocab_tokens_with_leading_trailing_spaces():
    # Test tokens with leading/trailing whitespace
    tokens = ["  hello", "world  ", "  spaced  "]
    vocab_file = create_temp_vocab_file(tokens)
    tokenizer = ProphetNetTokenizer(vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.40μs -> 2.03μs (18.0% faster)
    for idx, token in enumerate(tokens):
        pass
    os.remove(vocab_file)


def test_get_vocab_newline_and_empty_strings():
    # Test tokens with empty string and newline characters
    tokens = ["hello", "", "world", "\n"]
    vocab_file = create_temp_vocab_file(tokens)
    tokenizer = ProphetNetTokenizer(vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.09μs -> 2.13μs (1.97% slower)
    os.remove(vocab_file)


# 3. LARGE SCALE TEST CASES


def test_get_vocab_large_vocab_file():
    # Test with a large vocab file (1000 tokens)
    tokens = [f"token_{i}" for i in range(1000)]
    vocab_file = create_temp_vocab_file(tokens)
    tokenizer = ProphetNetTokenizer(vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 43.5μs -> 43.9μs (0.925% slower)
    for idx, token in enumerate(tokens):
        pass
    os.remove(vocab_file)


def test_get_vocab_large_vocab_file_with_added_tokens():
    # Large vocab file + added tokens
    tokens = [f"token_{i}" for i in range(990)]
    added_tokens = {f"added_{i}": 1000 + i for i in range(10)}
    vocab_file = create_temp_vocab_file(tokens)
    tokenizer = ProphetNetTokenizer(vocab_file, added_tokens_encoder=added_tokens)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 41.8μs -> 40.4μs (3.52% faster)
    for idx, token in enumerate(tokens):
        pass
    for i in range(10):
        pass
    os.remove(vocab_file)


def test_get_vocab_large_added_tokens_overlap():
    # Large vocab file, large overlap with added tokens
    tokens = [f"token_{i}" for i in range(950)] + [f"overlap_{i}" for i in range(50)]
    added_tokens = {f"overlap_{i}": 2000 + i for i in range(50)}
    vocab_file = create_temp_vocab_file(tokens)
    tokenizer = ProphetNetTokenizer(vocab_file, added_tokens_encoder=added_tokens)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 42.9μs -> 40.6μs (5.82% faster)
    # Overlap tokens should have added token values
    for i in range(50):
        pass
    os.remove(vocab_file)


def test_get_vocab_performance_large():
    # Performance test: ensure it runs quickly for large vocab
    tokens = [f"token_{i}" for i in range(999)]
    vocab_file = create_temp_vocab_file(tokens)
    tokenizer = ProphetNetTokenizer(vocab_file)
    import time

    start = time.time()
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 42.9μs -> 41.8μs (2.48% faster)
    end = time.time()
    os.remove(vocab_file)


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

# imports
import pytest

from transformers.models.prophetnet.tokenization_prophetnet import ProphetNetTokenizer


# function to test (ProphetNetTokenizer.get_vocab) is assumed to be defined above


# Helper function to create a vocab file for testing
def create_vocab_file(tokens, tmp_path):
    vocab_file = tmp_path / "prophetnet.tokenizer"
    with open(vocab_file, "w", encoding="utf-8") as f:
        f.writelines(token + "\n" for token in tokens)
    return str(vocab_file)


# ------------------ BASIC TEST CASES ------------------


def test_get_vocab_basic(tmp_path):
    # Test with a simple vocab file
    tokens = ["hello", "world", "[PAD]", "[UNK]"]
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.39μs -> 2.29μs (4.36% faster)
    # Check that all tokens are present and mapped to correct indices
    for idx, token in enumerate(tokens):
        pass


def test_get_vocab_with_added_tokens(tmp_path):
    # Test with added tokens after initialization
    tokens = ["a", "b", "c"]
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    # Add a new token, should appear in get_vocab
    added_token = "<NEW>"
    tokenizer.add_tokens([added_token])
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.39μs -> 2.27μs (5.60% faster)
    # All original tokens still present
    for idx, token in enumerate(tokens):
        pass


def test_get_vocab_empty_vocab_file(tmp_path):
    # Test with an empty vocab file
    tokens = []
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.12μs -> 2.07μs (2.22% faster)


def test_get_vocab_special_tokens(tmp_path):
    # Test that special tokens are present in vocab
    tokens = ["[PAD]", "[UNK]", "[SEP]", "[MASK]", "[X_SEP]"]
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.46μs -> 2.24μs (9.83% faster)
    for token in tokens:
        pass


# ------------------ EDGE TEST CASES ------------------


def test_get_vocab_duplicate_tokens(tmp_path):
    # Vocab file with duplicate tokens (should only keep first occurrence)
    tokens = ["foo", "bar", "foo", "baz"]
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.37μs -> 2.06μs (15.2% faster)


def test_get_vocab_token_with_newline_and_spaces(tmp_path):
    # Tokens with spaces and newlines
    tokens = ["foo ", " bar", "\tbaz\n", "qux"]
    # Clean tokens for what should be in vocab (rstrip('\n'))
    expected_tokens = ["foo ", " bar", "\tbaz", "qux"]
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.52μs -> 2.23μs (13.4% faster)
    for idx, token in enumerate(expected_tokens):
        pass


def test_get_vocab_unicode_tokens(tmp_path):
    # Unicode tokens
    tokens = ["你好", "こんにちは", "안녕하세요", "😀"]
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.49μs -> 2.09μs (19.2% faster)
    for idx, token in enumerate(tokens):
        pass


def test_get_vocab_token_with_empty_string(tmp_path):
    # Vocab file with empty string as a token
    tokens = ["foo", "", "bar"]
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.37μs -> 2.13μs (10.9% faster)


def test_get_vocab_added_token_collision(tmp_path):
    # Added token collides with existing token
    tokens = ["foo", "bar"]
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    tokenizer.add_tokens(["bar"])  # 'bar' already exists
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.11μs -> 2.01μs (4.67% faster)


def test_get_vocab_nonexistent_vocab_file(tmp_path):
    # Test error when vocab file does not exist
    non_existent_file = tmp_path / "does_not_exist.txt"
    with pytest.raises(ValueError):
        ProphetNetTokenizer(vocab_file=str(non_existent_file))


def test_get_vocab_added_token_with_special_characters(tmp_path):
    # Added token contains special characters
    tokens = ["foo"]
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    special_token = "t@k#n$%^&*()!"
    tokenizer.add_tokens([special_token])
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 2.22μs -> 2.09μs (5.98% faster)


# ------------------ LARGE SCALE TEST CASES ------------------


def test_get_vocab_large_vocab(tmp_path):
    # Large vocab file (up to 1000 tokens)
    tokens = [f"token_{i}" for i in range(1000)]
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 42.6μs -> 41.8μs (1.71% faster)


def test_get_vocab_large_added_tokens(tmp_path):
    # Add a large number of tokens after initialization
    tokens = ["base"]
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    added_tokens = [f"added_{i}" for i in range(999)]
    tokenizer.add_tokens(added_tokens)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 137μs -> 123μs (11.5% faster)
    # Check that all added tokens are present and mapped to correct indices
    for i, token in enumerate(added_tokens):
        pass


def test_get_vocab_large_vocab_and_added_tokens(tmp_path):
    # Large vocab file and large added tokens
    tokens = [f"token_{i}" for i in range(500)]
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    added_tokens = [f"extra_{i}" for i in range(499)]
    tokenizer.add_tokens(added_tokens)
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 90.4μs -> 88.0μs (2.73% faster)


def test_get_vocab_performance_large(tmp_path):
    # Performance: ensure get_vocab returns quickly for large vocab
    tokens = [f"token_{i}" for i in range(1000)]
    vocab_file = create_vocab_file(tokens, tmp_path)
    tokenizer = ProphetNetTokenizer(vocab_file=vocab_file)
    import time

    start = time.time()
    codeflash_output = tokenizer.get_vocab()
    vocab = codeflash_output  # 44.3μs -> 43.9μs (0.957% faster)
    duration = time.time() - start


# 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-ProphetNetTokenizer.get_vocab-misjj66c and push.