⚡️ Speed up method GPTSw3Tokenizer._tokenize by 23%

[codeflash-ai]

📄 23% (0.23x) speedup for GPTSw3Tokenizer._tokenize in src/transformers/models/gpt_sw3/tokenization_gpt_sw3.py

⏱️ Runtime : 6.21 milliseconds → 5.07 milliseconds (best of 11 runs)

📝 Explanation and details

The optimization replaces an inefficient character-by-character whitespace normalization with Python's built-in str.translate() method, delivering a 22% speedup.

Key optimization: In preprocess_text(), the original code used a list comprehension with set membership testing for every character:

text = "".join([char if char not in self.whitespaces else " " for char in text])

The optimized version precomputes a translation table during initialization and uses str.translate():

# In __init__:
self._whitespace_translation_table = str.maketrans({c: " " for c in self.whitespaces})

# In preprocess_text:
text = text.translate(self._whitespace_translation_table)

Why this is faster:

• str.translate() is implemented in C and operates at native speed with O(N) complexity
• The original approach had O(N×M) complexity where N is text length and M is whitespace set size (12 characters), requiring a set lookup for every character
• Translation tables eliminate the per-character conditional logic and string concatenation overhead

Performance impact: Line profiler shows whitespace normalization time dropped from 4.92ms to 216μs - a 95% reduction in that specific operation. The optimization is particularly effective for:

• Large texts with many whitespace characters (up to 44.8% faster for whitespace-heavy inputs)
• Repeated tokenization calls in production pipelines
• Long documents where character-level operations compound

This is a classic example of replacing Python loops with optimized C implementations, especially valuable in tokenization workflows that process large volumes of text.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import os

# Minimal GPTSw3Tokenizer implementation for testing
import tempfile

# imports
import pytest
import sentencepiece as spm

from transformers.models.gpt_sw3.tokenization_gpt_sw3 import GPTSw3Tokenizer


# Helper to create a temporary SentencePiece model for testing
def create_sp_model(sentences, vocab_size=30):
    with tempfile.TemporaryDirectory() as tmpdir:
        input_file = os.path.join(tmpdir, "input.txt")
        model_prefix = os.path.join(tmpdir, "spm")
        with open(input_file, "w", encoding="utf-8") as f:
            f.writelines(s + "\n" for s in sentences)
        spm.SentencePieceTrainer.Train(
            input=f"{input_file}",
            model_prefix=model_prefix,
            vocab_size=vocab_size,
            character_coverage=1.0,
            model_type="unigram",
        )
        model_file = model_prefix + ".model"
        # Copy to a persistent file for the test
        final_model = tempfile.NamedTemporaryFile(delete=False, suffix=".model")
        final_model.close()
        with open(model_file, "rb") as src, open(final_model.name, "wb") as dst:
            dst.write(src.read())
        return final_model.name


# Fixtures for reusable tokenizer/model
@pytest.fixture(scope="module")
def basic_sp_model():
    # Basic Swedish sentences with punctuation and unicode
    sentences = [
        "Svenska är kul!",
        "Hej världen.",
        "GPT är en stor språkmodell.",
        "Det här är ett test.",
        "   Mellanslag före och efter   ",
        "Emoji: 😊",
        "Accenter: café, naïve, résumé",
        "Tab\tnewline\ncarriage\rreturn",
        "Non-printing:\x00\x07text",
    ]
    model_file = create_sp_model(sentences, vocab_size=50)
    yield model_file
    os.unlink(model_file)


@pytest.fixture
def tokenizer(basic_sp_model):
    return GPTSw3Tokenizer(basic_sp_model)


# 1. Basic Test Cases
def test_basic_tokenization_simple(tokenizer):
    # Simple Swedish sentence
    codeflash_output = tokenizer._tokenize("Svenska är kul!")
    tokens = codeflash_output  # 24.3μs -> 23.8μs (2.01% faster)


def test_basic_tokenization_punctuation(tokenizer):
    # Sentence with punctuation
    codeflash_output = tokenizer._tokenize("Hej, världen!")
    tokens = codeflash_output  # 19.8μs -> 20.5μs (3.30% slower)


def test_basic_tokenization_unicode(tokenizer):
    # Sentence with Swedish characters
    codeflash_output = tokenizer._tokenize("Göteborg är en stad i Sverige.")
    tokens = codeflash_output  # 25.1μs -> 23.8μs (5.59% faster)


def test_basic_tokenization_whitespace(tokenizer):
    # Leading/trailing/multiple spaces
    codeflash_output = tokenizer._tokenize("   Mellanslag   ")
    tokens1 = codeflash_output  # 17.1μs -> 15.8μs (8.40% faster)
    codeflash_output = tokenizer._tokenize("Mellanslag")
    tokens2 = codeflash_output  # 7.20μs -> 5.99μs (20.2% faster)


def test_basic_tokenization_accents(tokenizer):
    # Accented characters
    codeflash_output = tokenizer._tokenize("café naïve résumé")
    tokens = codeflash_output  # 20.4μs -> 19.9μs (2.13% faster)


def test_basic_tokenization_emoji(tokenizer):
    codeflash_output = tokenizer._tokenize("Emoji: 😊")
    tokens = codeflash_output  # 18.2μs -> 17.2μs (5.63% faster)


# 2. Edge Test Cases
def test_edge_empty_string(tokenizer):
    # Empty input should yield empty list
    codeflash_output = tokenizer._tokenize("")
    tokens = codeflash_output  # 8.71μs -> 8.31μs (4.78% faster)


def test_edge_only_whitespace(tokenizer):
    # Only whitespace input should yield empty list
    codeflash_output = tokenizer._tokenize("     ")
    tokens = codeflash_output  # 10.2μs -> 9.30μs (9.33% faster)


def test_edge_non_printing_characters(tokenizer):
    # Non-printing characters should be removed
    codeflash_output = tokenizer._tokenize("\x00\x01\x02test\x07\x08")
    tokens = codeflash_output  # 15.5μs -> 15.0μs (3.29% faster)


def test_edge_tab_newline_carriage(tokenizer):
    # Tabs, newlines, carriage returns should be normalized
    codeflash_output = tokenizer._tokenize("Tab\tnewline\ncarriage\rreturn")
    tokens = codeflash_output  # 21.6μs -> 20.4μs (6.06% faster)
    # Should tokenize as if whitespace normalized
    codeflash_output = tokenizer._tokenize("Tab newline carriage return")
    tokens_ref = codeflash_output  # 11.1μs -> 9.78μs (13.8% faster)


def test_edge_multiple_whitespace_types(tokenizer):
    # Various unicode whitespaces
    text = "foo bar baz qux quux　corge grault garply waldofred�plugh"
    # All should be normalized to spaces
    codeflash_output = tokenizer._tokenize(text)
    tokens = codeflash_output  # 34.4μs -> 31.4μs (9.44% faster)
    codeflash_output = tokenizer._tokenize("foo bar baz qux quux corge grault garply waldo fred plugh")
    tokens_ref = codeflash_output  # 21.3μs -> 18.7μs (13.5% faster)


def test_edge_nfc_normalization(tokenizer):
    # NFC normalization: composed vs decomposed unicode
    text1 = "Svenska är kul!"
    # 'ä' as composed
    text2 = "Svenska a\u0308r kul!"  # 'a' + combining diaeresis
    codeflash_output = tokenizer._tokenize(text1)
    tokens1 = codeflash_output  # 18.8μs -> 17.7μs (6.65% faster)
    codeflash_output = tokenizer._tokenize(text2)
    tokens2 = codeflash_output  # 11.8μs -> 10.1μs (17.0% faster)


def test_edge_long_token(tokenizer):
    # Very long word
    long_word = "a" * 100
    codeflash_output = tokenizer._tokenize(long_word)
    tokens = codeflash_output  # 39.6μs -> 35.0μs (13.2% faster)


def test_edge_special_characters(tokenizer):
    # Special characters and symbols
    text = "!@#$%^&*()_+-=[]{}|;':\",.<>/?`~"
    codeflash_output = tokenizer._tokenize(text)
    tokens = codeflash_output  # 21.2μs -> 19.2μs (10.5% faster)


def test_edge_mixed_language(tokenizer):
    # Mixed Swedish and English
    text = "Hej world! GPT är cool."
    codeflash_output = tokenizer._tokenize(text)
    tokens = codeflash_output  # 21.4μs -> 20.6μs (3.80% faster)


# 3. Large Scale Test Cases
def test_large_scale_long_sentence(tokenizer):
    # Very long sentence (1000 tokens)
    sentence = " ".join(["Svenska"] * 1000)
    codeflash_output = tokenizer._tokenize(sentence)
    tokens = codeflash_output  # 1.79ms -> 1.41ms (27.7% faster)


def test_large_scale_many_sentences(tokenizer):
    # 500 different sentences
    sentences = [f"Test sentence {i}." for i in range(500)]
    text = " ".join(sentences)
    codeflash_output = tokenizer._tokenize(text)
    tokens = codeflash_output  # 2.07ms -> 1.67ms (24.0% faster)


def test_large_scale_repeated_tokens(tokenizer):
    # Repeated tokens should not cause issues
    text = "abc " * 999
    codeflash_output = tokenizer._tokenize(text)
    tokens = codeflash_output  # 944μs -> 726μs (29.9% faster)


def test_large_scale_unicode(tokenizer):
    # Large input with unicode
    text = ("Göteborg är bäst! 😊 " * 200).strip()
    codeflash_output = tokenizer._tokenize(text)
    tokens = codeflash_output  # 973μs -> 882μs (10.4% faster)


def test_large_scale_edge_case(tokenizer):
    # Large input with only whitespace and non-printing chars
    text = (" " * 500) + ("\x00" * 500)
    codeflash_output = tokenizer._tokenize(text)
    tokens = codeflash_output  # 65.4μs -> 45.2μs (44.8% faster)


# 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-GPTSw3Tokenizer._tokenize-mi9ysxtd and push.