⚡️ Speed up function iob_to_biluo by 58%

[codeflash-ai]

📄 58% (0.58x) speedup for iob_to_biluo in spacy/training/iob_utils.py

⏱️ Runtime : 3.19 milliseconds → 2.02 milliseconds (best of 176 runs)

📝 Explanation and details

The optimized code achieves a 58% speedup by eliminating expensive list operations and reducing function call overhead.

Key Optimizations:

1. Eliminated costly pop(0) operations: The original code repeatedly called tags.pop(0), which is O(n) because it shifts all remaining elements. The optimization uses index-based iteration (i, j) with O(1) access instead.

2. Replaced separate helper functions with inline logic: Removed _consume_os() and _consume_ent() functions, eliminating generator overhead and function call costs. The profiler shows these functions consumed 43.9% and 54.5% of runtime respectively.

3. Optimized O-tag handling: Instead of yielding individual O-tags, the code identifies consecutive O-tags and extends the output with a single slice operation (out.extend(tags[start:i])).

4. Reduced list operations: Uses out.append() for single items and optimized list multiplication ([f"I-{label}"] * (length - 2)) instead of list comprehensions for repeated elements.

Performance Impact by Test Case:

• Large inputs benefit most: Single large entities show 90-96% speedup, leveraging the O(1) vs O(n) access pattern difference
• Multiple entities: 45-70% speedup from eliminated function calls and better memory access patterns
• Mixed O/entity sequences: 30-90% speedup from optimized O-tag batch processing
• Edge cases maintain similar speedup: Even small inputs see 15-60% improvements

The optimization is particularly effective for NLP pipelines where IOB tag sequences can be lengthy, making the quadratic behavior of the original implementation a significant bottleneck.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
parser/test_ner.py::test_issue2385	13.8μs	8.57μs	61.4%✅
training/test_training.py::test_iob_to_biluo	12.3μs	10.3μs	19.5%✅

🌀 Generated Regression Tests and Runtime

from typing import Iterable, Iterator, List

# imports
import pytest  # used for our unit tests
from spacy.training.iob_utils import iob_to_biluo

# unit tests

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

def test_basic_all_o():
    # Only 'O' tags, should remain unchanged
    codeflash_output = iob_to_biluo(['O', 'O', 'O']) # 2.01μs -> 1.51μs (32.9% faster)

def test_basic_single_entity():
    # Single 'B-ORG' should become 'U-ORG'
    codeflash_output = iob_to_biluo(['B-ORG']) # 2.77μs -> 1.86μs (48.5% faster)

def test_basic_simple_entity():
    # B-ORG followed by I-ORG should become B-ORG, L-ORG
    codeflash_output = iob_to_biluo(['B-ORG', 'I-ORG']) # 4.62μs -> 2.37μs (95.3% faster)

def test_basic_long_entity():
    # B-PER, I-PER, I-PER should become B-PER, I-PER, L-PER
    codeflash_output = iob_to_biluo(['B-PER', 'I-PER', 'I-PER']) # 4.82μs -> 3.13μs (53.8% faster)

def test_basic_multiple_entities():
    # Multiple entities and O tags
    codeflash_output = iob_to_biluo(['O', 'B-LOC', 'I-LOC', 'O', 'B-PER']) # 5.87μs -> 3.98μs (47.7% faster)

def test_basic_adjacent_entities():
    # Adjacent entities of different types
    codeflash_output = iob_to_biluo(['B-ORG', 'I-ORG', 'B-PER', 'I-PER', 'I-PER']) # 6.88μs -> 4.03μs (70.6% faster)

def test_basic_entity_at_end():
    # Entity at the end of the sequence
    codeflash_output = iob_to_biluo(['O', 'B-LOC', 'I-LOC']) # 3.90μs -> 2.61μs (49.5% faster)

def test_basic_entity_at_start():
    # Entity at the start of the sequence
    codeflash_output = iob_to_biluo(['B-LOC', 'I-LOC', 'O']) # 4.94μs -> 2.61μs (89.0% faster)

def test_basic_single_token_entity():
    # Single token entity with 'B-LOC'
    codeflash_output = iob_to_biluo(['O', 'B-LOC', 'O']) # 3.59μs -> 2.92μs (23.1% faster)

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

def test_edge_empty_input():
    # Empty input should return empty output
    codeflash_output = iob_to_biluo([]) # 510ns -> 663ns (23.1% slower)

def test_edge_single_o():
    # Single O tag
    codeflash_output = iob_to_biluo(['O']) # 1.63μs -> 1.27μs (27.9% faster)

def test_edge_single_invalid_b():
    # Single 'B-' (missing label) should raise ValueError
    with pytest.raises(ValueError):
        iob_to_biluo(['B-']) # 6.32μs -> 5.42μs (16.7% faster)

def test_edge_single_invalid_i():
    # Single 'I-' (missing label) should raise ValueError
    with pytest.raises(ValueError):
        iob_to_biluo(['I-']) # 5.00μs -> 4.36μs (14.6% faster)


def test_edge_iob_with_l_tag():
    # Sequence with 'L-' tag, which is not valid in IOB, should be treated as entity continuation
    # Here, 'B-ORG', 'L-ORG' should become B-ORG, L-ORG
    codeflash_output = iob_to_biluo(['B-ORG', 'L-ORG']) # 5.47μs -> 2.73μs (100% faster)



def test_edge_b_tag_followed_by_o():
    # 'B-ORG', 'O' should be U-ORG, O
    codeflash_output = iob_to_biluo(['B-ORG', 'O']) # 4.89μs -> 2.98μs (64.1% faster)

def test_edge_b_tag_followed_by_different_entity():
    # 'B-ORG', 'B-PER' should be U-ORG, U-PER
    codeflash_output = iob_to_biluo(['B-ORG', 'B-PER']) # 4.17μs -> 2.76μs (50.8% faster)


def test_edge_mixed_case_tags():
    # Lowercase 'o' should not be recognized, should raise
    with pytest.raises(ValueError):
        iob_to_biluo(['o', 'B-ORG']) # 7.32μs -> 6.12μs (19.6% faster)

def test_edge_non_string_tag():
    # Non-string tag should raise
    with pytest.raises(Exception):
        iob_to_biluo([None, 'B-ORG']) # 2.46μs -> 1.59μs (54.7% faster)

def test_edge_b_tag_with_dash_in_label():
    # Label with dash, e.g., 'B-ORG-UNIT', should be handled
    codeflash_output = iob_to_biluo(['B-ORG-UNIT', 'I-ORG-UNIT']) # 5.10μs -> 2.56μs (99.3% faster)

def test_edge_entity_with_l_tag_middle():
    # 'B-ORG', 'I-ORG', 'L-ORG' should become B-ORG, I-ORG, L-ORG
    codeflash_output = iob_to_biluo(['B-ORG', 'I-ORG', 'L-ORG']) # 5.12μs -> 3.35μs (53.0% faster)

def test_edge_entity_with_multiple_l_tags():
    # 'B-ORG', 'L-ORG', 'L-ORG' should treat as two entities
    codeflash_output = iob_to_biluo(['B-ORG', 'L-ORG', 'L-ORG']) # 4.70μs -> 2.97μs (58.2% faster)

def test_edge_entity_with_multiple_i_tags():
    # 'B-ORG', 'I-ORG', 'I-ORG', 'O' should become B-ORG, I-ORG, L-ORG, O
    codeflash_output = iob_to_biluo(['B-ORG', 'I-ORG', 'I-ORG', 'O']) # 5.70μs -> 3.42μs (66.9% faster)

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

def test_large_all_o_tags():
    # Large input of only 'O' tags
    tags = ['O'] * 1000
    codeflash_output = iob_to_biluo(tags) # 74.6μs -> 39.5μs (89.0% faster)

def test_large_single_entity():
    # Large single entity
    tags = ['B-ORG'] + ['I-ORG'] * 998 + ['O']
    expected = ['B-ORG'] + ['I-ORG'] * 997 + ['L-ORG', 'O']
    codeflash_output = iob_to_biluo(tags) # 147μs -> 75.1μs (96.2% faster)

def test_large_multiple_entities():
    # Multiple entities interleaved with O tags
    tags = []
    expected = []
    for i in range(100):
        tags.extend(['O', 'B-PER', 'I-PER', 'O', 'B-LOC', 'I-LOC', 'I-LOC'])
        expected.extend(['O', 'B-PER', 'L-PER', 'O', 'B-LOC', 'I-LOC', 'L-LOC'])
    codeflash_output = iob_to_biluo(tags) # 214μs -> 141μs (51.8% faster)

def test_large_varied_entities():
    # Alternating single and multi-token entities
    tags = []
    expected = []
    for i in range(250):
        tags.extend(['O', 'B-ORG', 'O', 'B-PER', 'I-PER', 'I-PER'])
        expected.extend(['O', 'U-ORG', 'O', 'B-PER', 'I-PER', 'L-PER'])
    codeflash_output = iob_to_biluo(tags) # 465μs -> 319μs (45.7% faster)

def test_large_entities_with_l_tags():
    # Entities with L- tags in large input
    tags = []
    expected = []
    for i in range(200):
        tags.extend(['B-LOC', 'I-LOC', 'L-LOC', 'O'])
        expected.extend(['B-LOC', 'I-LOC', 'L-LOC', 'O'])
    codeflash_output = iob_to_biluo(tags) # 229μs -> 153μs (49.4% faster)


def test_large_entity_with_dash_label():
    # Large entity with dash in label
    tags = ['B-ORG-UNIT'] + ['I-ORG-UNIT'] * 998 + ['O']
    expected = ['B-ORG-UNIT'] + ['I-ORG-UNIT'] * 997 + ['L-ORG-UNIT', 'O']
    codeflash_output = iob_to_biluo(tags) # 145μs -> 76.0μs (91.8% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from typing import Iterable, Iterator, List

# imports
import pytest  # used for our unit tests
from spacy.training.iob_utils import iob_to_biluo

# unit tests

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

def test_single_o_tag():
    # Single 'O' should remain 'O'
    codeflash_output = iob_to_biluo(["O"]) # 1.81μs -> 1.34μs (34.9% faster)

def test_single_b_tag():
    # Single 'B-ORG' should become 'U-ORG'
    codeflash_output = iob_to_biluo(["B-ORG"]) # 2.97μs -> 1.91μs (55.7% faster)

def test_single_entity_two_tokens():
    # 'B-PER', 'I-PER' should become 'B-PER', 'L-PER'
    codeflash_output = iob_to_biluo(["B-PER", "I-PER"]) # 4.62μs -> 2.31μs (101% faster)

def test_single_entity_three_tokens():
    # 'B-LOC', 'I-LOC', 'I-LOC' should become 'B-LOC', 'I-LOC', 'L-LOC'
    codeflash_output = iob_to_biluo(["B-LOC", "I-LOC", "I-LOC"]) # 4.87μs -> 3.11μs (56.5% faster)

def test_mixed_o_and_entity():
    # Mix of O and entity
    codeflash_output = iob_to_biluo(["O", "B-PER", "I-PER", "O"]) # 5.15μs -> 3.30μs (56.3% faster)

def test_multiple_entities():
    # Two entities separated by O
    codeflash_output = iob_to_biluo(["B-PER", "I-PER", "O", "B-LOC", "I-LOC"]) # 6.17μs -> 3.67μs (68.4% faster)

def test_multiple_single_token_entities():
    # Multiple single-token entities
    codeflash_output = iob_to_biluo(["B-PER", "O", "B-LOC", "O", "B-ORG"]) # 5.36μs -> 3.94μs (36.1% faster)

def test_entity_at_start_and_end():
    # Entity at start and end
    codeflash_output = iob_to_biluo(["B-ORG", "O", "O", "B-PER"]) # 4.13μs -> 3.18μs (30.0% faster)

def test_entity_with_no_i_tags():
    # Only B tags, each should become U
    codeflash_output = iob_to_biluo(["B-LOC", "B-PER", "B-ORG"]) # 4.57μs -> 3.12μs (46.3% faster)

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

def test_empty_input():
    # Empty input should return empty output
    codeflash_output = iob_to_biluo([]) # 510ns -> 652ns (21.8% slower)

def test_invalid_tag_raises():
    # 'B' without a label should raise ValueError
    with pytest.raises(ValueError):
        iob_to_biluo(["B"]) # 6.53μs -> 5.72μs (14.1% faster)

def test_invalid_i_tag_raises():
    # 'I' without a label should raise ValueError
    with pytest.raises(ValueError):
        iob_to_biluo(["I"]) # 5.10μs -> 4.45μs (14.6% faster)

def test_invalid_b_tag_in_middle():
    # 'B' without label in middle should raise ValueError
    with pytest.raises(ValueError):
        iob_to_biluo(["O", "B", "O"]) # 5.17μs -> 5.14μs (0.681% faster)

def test_invalid_i_tag_in_middle():
    # 'I' without label in middle should raise ValueError
    with pytest.raises(ValueError):
        iob_to_biluo(["O", "I", "O"]) # 5.17μs -> 4.87μs (6.29% faster)


def test_i_tag_without_b_before():
    # 'I-PER' at the start should be interpreted as a single entity
    # Since the function does not check for this error, it will treat 'I-PER' as 'U-PER'
    # But according to the implementation, it will treat 'I-PER' as 'U-PER'
    codeflash_output = iob_to_biluo(["I-PER"]) # 3.42μs -> 2.15μs (59.4% faster)

def test_b_tag_followed_by_o():
    # 'B-ORG' followed by 'O' should become 'U-ORG', 'O'
    codeflash_output = iob_to_biluo(["B-ORG", "O"]) # 4.07μs -> 2.67μs (52.2% faster)

def test_entity_with_l_tag():
    # 'B-LOC', 'L-LOC' should become 'B-LOC', 'L-LOC'
    codeflash_output = iob_to_biluo(["B-LOC", "L-LOC"]) # 4.75μs -> 2.35μs (102% faster)

def test_entity_with_i_and_l_tags():
    # 'B-LOC', 'I-LOC', 'L-LOC' should become 'B-LOC', 'I-LOC', 'L-LOC'
    codeflash_output = iob_to_biluo(["B-LOC", "I-LOC", "L-LOC"]) # 4.95μs -> 3.17μs (56.2% faster)

def test_entity_with_multiple_l_tags():
    # 'B-LOC', 'L-LOC', 'L-LOC' should consume both L-LOC as part of entity
    codeflash_output = iob_to_biluo(["B-LOC", "L-LOC", "L-LOC"]) # 4.62μs -> 2.90μs (59.4% faster)

def test_entity_with_mixed_i_and_l_tags():
    # 'B-ORG', 'I-ORG', 'L-ORG', 'I-ORG' should consume till first L-ORG, then treat 'I-ORG' as new entity
    codeflash_output = iob_to_biluo(["B-ORG", "I-ORG", "L-ORG", "I-ORG"]) # 4.90μs -> 3.04μs (61.1% faster)

def test_entity_with_label_with_dash():
    # 'B-FOO-BAR', 'I-FOO-BAR' should become 'B-FOO-BAR', 'L-FOO-BAR'
    codeflash_output = iob_to_biluo(["B-FOO-BAR", "I-FOO-BAR"]) # 4.20μs -> 2.24μs (87.1% faster)

def test_multiple_consecutive_o_tags():
    # Multiple consecutive O tags
    codeflash_output = iob_to_biluo(["O", "O", "O"]) # 2.04μs -> 1.57μs (29.9% faster)

def test_entity_with_only_l_tag():
    # 'L-PER' at the start should be interpreted as 'U-PER'
    codeflash_output = iob_to_biluo(["L-PER"]) # 2.70μs -> 1.84μs (46.6% faster)

def test_entity_with_only_l_tag_in_middle():
    # 'O', 'L-PER', 'O' should treat 'L-PER' as 'U-PER'
    codeflash_output = iob_to_biluo(["O", "L-PER", "O"]) # 3.84μs -> 3.07μs (25.3% faster)

def test_entity_with_only_l_tag_and_i_tag():
    # 'L-PER', 'I-PER' should treat each as single entity
    codeflash_output = iob_to_biluo(["L-PER", "I-PER"]) # 4.28μs -> 2.20μs (94.1% faster)

def test_entity_with_only_l_tag_and_b_tag():
    # 'L-PER', 'B-PER' should treat each as single entity
    codeflash_output = iob_to_biluo(["L-PER", "B-PER"]) # 3.86μs -> 2.68μs (44.1% faster)

def test_entity_with_multiple_labels():
    # 'B-PER', 'I-LOC', 'I-ORG' should treat each as single entity
    codeflash_output = iob_to_biluo(["B-PER", "I-LOC", "I-ORG"]) # 4.67μs -> 3.23μs (44.7% faster)

def test_entity_with_nonstandard_case():
    # Lowercase tags should not be recognized as valid, but function does not check case
    # So 'b-per', 'i-per' should be treated as single entities
    codeflash_output = iob_to_biluo(["b-per", "i-per"]) # 3.75μs -> 2.46μs (52.4% faster)

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

def test_large_all_o_tags():
    # Large input with all 'O'
    tags = ["O"] * 1000
    codeflash_output = iob_to_biluo(tags) # 74.5μs -> 38.9μs (91.2% faster)

def test_large_single_entity():
    # Large entity: 'B-LOC' + 998 'I-LOC'
    tags = ["B-LOC"] + ["I-LOC"] * 998
    expected = ["B-LOC"] + ["I-LOC"] * 998 + ["L-LOC"]
    # The function as implemented will consume all 'I-LOC' and treat as entity of length 999
    codeflash_output = iob_to_biluo(tags) # 143μs -> 74.0μs (94.5% faster)

def test_large_mixed_entities_and_o():
    # Mix of entities and O tags
    tags = []
    expected = []
    # 100 entities of length 5, separated by 5 O tags
    for i in range(100):
        tags += ["B-ORG"] + ["I-ORG"] * 3 + ["I-ORG"]
        tags += ["O"] * 5
        expected += ["B-ORG", "I-ORG", "I-ORG", "L-ORG"]
        expected += ["O"] * 5
    codeflash_output = iob_to_biluo(tags) # 188μs -> 116μs (62.1% faster)

def test_large_many_single_token_entities():
    # 500 single-token entities, alternating labels
    tags = []
    expected = []
    for i in range(500):
        label = "PER" if i % 2 == 0 else "LOC"
        tags.append(f"B-{label}")
        expected.append(f"U-{label}")
    codeflash_output = iob_to_biluo(tags) # 265μs -> 164μs (61.6% faster)

def test_large_all_different_entities():
    # 1000 entities, each with a unique label
    tags = [f"B-LABEL{i}" for i in range(1000)]
    expected = [f"U-LABEL{i}" for i in range(1000)]
    codeflash_output = iob_to_biluo(tags) # 562μs -> 351μs (59.9% faster)

def test_large_alternating_o_and_entity():
    # Alternating O and single-token entity, 500 times
    tags = []
    expected = []
    for i in range(500):
        tags.append("O")
        expected.append("O")
        tags.append("B-PER")
        expected.append("U-PER")
    codeflash_output = iob_to_biluo(tags) # 300μs -> 229μs (30.9% faster)

def test_large_entity_with_l_tags():
    # Large entity: 'B-LOC' + 998 'L-LOC'
    tags = ["B-LOC"] + ["L-LOC"] * 998
    # All 'L-LOC' are treated as part of the entity
    expected = ["B-LOC"] + ["L-LOC"] * 998
    codeflash_output = iob_to_biluo(tags) # 138μs -> 72.9μs (90.4% 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-iob_to_biluo-mhlix7cr and push.