⚡️ Speed up method ClassifierTrainingParameters.serialize_model by 41%

[codeflash-ai]

📄 41% (0.41x) speedup for ClassifierTrainingParameters.serialize_model in src/mistralai/models/classifiertrainingparameters.py

⏱️ Runtime : 1.80 milliseconds → 1.28 milliseconds (best of 138 runs)

📝 Explanation and details

The optimized code achieves a 40% speedup through several key data structure and loop optimizations:

Key Performance Optimizations:

1. Set-based lookups instead of lists: Converting optional_fields and nullable_fields from lists to sets enables O(1) membership testing instead of O(n) linear search. This is critical since these lookups happen for every field in the serialization loop.

2. Reduced dictionary access overhead: The original code called serialized.get(k) followed by serialized.pop(k, None), performing two dictionary lookups. The optimized version uses a single serialized.pop(k, None) call, eliminating redundant dictionary access.

3. Cached expensive operations: Pre-computing fields_set = self.__pydantic_fields_set__ and model_fields = type(self).model_fields outside the loop avoids repeated attribute access during iteration.

4. Simplified set membership logic: Replaced the intersection-based check self.__pydantic_fields_set__.intersection({n}) with direct membership n in fields_set, which is more efficient for single-element lookups.

Performance Results by Test Case:

• Best gains (35-42% faster): Tests with explicit None values, many instances, and mixed field types benefit most from the set-based lookups
• Consistent improvements (20-30% faster): All test scenarios show meaningful speedup, indicating the optimizations help across different usage patterns
• Scalability: The 100-instance and 500-instance tests show 41-42% improvements, demonstrating that benefits compound with volume

The optimizations maintain identical behavior while significantly reducing computational overhead in the serialization hot path, making it ideal for applications that serialize many ClassifierTrainingParameters instances.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

from typing import Optional

# imports
import pytest
from mistralai.models.classifiertrainingparameters import \
    ClassifierTrainingParameters


# Simulate pydantic's BaseModel and UNSET/UNSET_SENTINEL for testing
class UNSET_SENTINEL_TYPE:
    pass
UNSET_SENTINEL = UNSET_SENTINEL_TYPE()
UNSET = UNSET_SENTINEL

class BaseModel:
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
        self.__pydantic_fields_set__ = set(kwargs.keys())
    # Simulate model_fields
    @property
    def model_fields(self):
        # For this class, we define the fields as per the original
        return {
            "training_steps": type("Field", (), {"alias": None})(),
            "learning_rate": type("Field", (), {"alias": None})(),
            "weight_decay": type("Field", (), {"alias": None})(),
            "warmup_fraction": type("Field", (), {"alias": None})(),
            "epochs": type("Field", (), {"alias": None})(),
            "seq_len": type("Field", (), {"alias": None})(),
        }
from mistralai.models.classifiertrainingparameters import \
    ClassifierTrainingParameters


# Helper handler function for serialization (simulates pydantic's .dict())
def simple_handler(obj):
    # Return a dict of all fields, including unset ones
    result = {}
    for k in obj.model_fields.keys():
        result[k] = getattr(obj, k)
    return result

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

# Basic Test Cases

def test_serialize_model_all_defaults():
    # All fields set to default (some UNSET, some default value)
    params = ClassifierTrainingParameters()
    codeflash_output = params.serialize_model(simple_handler); result = codeflash_output # 11.0μs -> 8.87μs (23.9% faster)
    # All other fields should not be present
    for field in ["training_steps", "weight_decay", "warmup_fraction", "epochs", "seq_len"]:
        pass

def test_serialize_model_all_set():
    # All fields set to non-default values
    params = ClassifierTrainingParameters(
        training_steps=100,
        learning_rate=0.01,
        weight_decay=0.05,
        warmup_fraction=0.1,
        epochs=5,
        seq_len=512,
    )
    codeflash_output = params.serialize_model(simple_handler); result = codeflash_output # 7.60μs -> 5.52μs (37.8% faster)

def test_serialize_model_some_set_some_unset():
    # Some fields set, some left as UNSET
    params = ClassifierTrainingParameters(
        training_steps=50,
        learning_rate=0.02,
        seq_len=256,
    )
    codeflash_output = params.serialize_model(simple_handler); result = codeflash_output # 8.48μs -> 6.67μs (27.2% faster)
    # Unset fields should not be present
    for field in ["weight_decay", "warmup_fraction", "epochs"]:
        pass

def test_serialize_model_none_values():
    # Fields explicitly set to None
    params = ClassifierTrainingParameters(
        training_steps=None,
        learning_rate=None,
        weight_decay=None,
        warmup_fraction=None,
        epochs=None,
        seq_len=None,
    )
    codeflash_output = params.serialize_model(simple_handler); result = codeflash_output # 7.30μs -> 5.16μs (41.6% faster)



def test_serialize_model_zero_and_negative_values():
    # Fields set to 0 and negative values
    params = ClassifierTrainingParameters(
        training_steps=0,
        learning_rate=-0.01,
        weight_decay=0,
        warmup_fraction=-0.5,
        epochs=-1,
        seq_len=0,
    )
    codeflash_output = params.serialize_model(simple_handler); result = codeflash_output # 9.41μs -> 6.73μs (39.8% faster)

def test_serialize_model_large_numbers():
    # Fields set to large numbers
    params = ClassifierTrainingParameters(
        training_steps=999999,
        learning_rate=1e6,
        weight_decay=1e8,
        warmup_fraction=1e10,
        epochs=1e12,
        seq_len=1024,
    )
    codeflash_output = params.serialize_model(simple_handler); result = codeflash_output # 6.97μs -> 5.22μs (33.4% faster)

def test_serialize_model_float_precision():
    # Fields set to floats with high precision
    params = ClassifierTrainingParameters(
        training_steps=1,
        learning_rate=0.123456789,
        weight_decay=0.987654321,
        warmup_fraction=1.000000001,
        epochs=2.999999999,
        seq_len=128,
    )
    codeflash_output = params.serialize_model(simple_handler); result = codeflash_output # 6.78μs -> 5.00μs (35.5% faster)


def test_serialize_model_many_instances():
    # Test serializing many instances in a loop (under 1000)
    instances = []
    for i in range(100):
        params = ClassifierTrainingParameters(
            training_steps=i,
            learning_rate=0.001 * i,
            weight_decay=0.002 * i,
            warmup_fraction=0.003 * i,
            epochs=i,
            seq_len=i + 100,
        )
        instances.append(params)
    # Serialize all and check correctness
    for i, params in enumerate(instances):
        codeflash_output = params.serialize_model(simple_handler); result = codeflash_output # 285μs -> 202μs (41.2% faster)

def test_serialize_model_large_data_structure():
    # Test with large values (but not exceeding 1000 elements)
    params = ClassifierTrainingParameters(
        training_steps=999,
        learning_rate=999.999,
        weight_decay=999.999,
        warmup_fraction=999.999,
        epochs=999.999,
        seq_len=999,
    )
    codeflash_output = params.serialize_model(simple_handler); result = codeflash_output # 6.29μs -> 4.56μs (38.0% faster)

def test_serialize_model_performance():
    # Performance test: serialize 500 instances and check time (not strict, just functional)
    import time
    instances = [
        ClassifierTrainingParameters(
            training_steps=i,
            learning_rate=0.001 * i,
            weight_decay=0.002 * i,
            warmup_fraction=0.003 * i,
            epochs=i,
            seq_len=i + 500,
        )
        for i in range(500)
    ]
    start = time.time()
    for params in instances:
        codeflash_output = params.serialize_model(simple_handler); result = codeflash_output # 1.38ms -> 969μs (42.0% 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.
#------------------------------------------------
from __future__ import annotations

from typing import Any, Optional

# imports
import pytest
from mistralai.models.classifiertrainingparameters import \
    ClassifierTrainingParameters


# Simulate UNSET and UNSET_SENTINEL as in the original code
class _UnsetType:
    def __repr__(self):
        return "UNSET"
UNSET = _UnsetType()
UNSET_SENTINEL = UNSET

# Simulate OptionalNullable (for the sake of the test, just use Optional)
OptionalNullable = Optional

# Simulate BaseModel and model_serializer
class BaseModel:
    # This will simulate pydantic's __pydantic_fields_set__
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)
        self.__pydantic_fields_set__ = set(kwargs.keys())
    @classmethod
    def model_fields(cls):
        # Return a dict of field name to a dummy object with .alias = None
        return {k: type("F", (), {"alias": None})() for k in cls.__annotations__}

def model_serializer(mode=None):
    # Decorator that does nothing for our purposes
    def wrapper(fn):
        return fn
    return wrapper
from mistralai.models.classifiertrainingparameters import \
    ClassifierTrainingParameters


# Helper handler function
def default_handler(obj):
    # Simulate serialization: just return the __dict__ minus private keys
    d = {}
    for k in obj.__class__.__annotations__:
        v = getattr(obj, k, UNSET)
        d[k] = v
    return d.copy()

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

# 1. Basic Test Cases

def test_serialize_model_all_defaults():
    """All fields default: only learning_rate should be present with default value."""
    obj = ClassifierTrainingParameters()
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 9.79μs -> 8.20μs (19.4% faster)
    # All other fields should be absent (since they're UNSET)
    for k in ['training_steps', 'weight_decay', 'warmup_fraction', 'epochs', 'seq_len']:
        pass

def test_serialize_model_set_some_fields():
    """Set some fields, check only those and learning_rate are present."""
    obj = ClassifierTrainingParameters(training_steps=100, seq_len=512)
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 8.02μs -> 6.45μs (24.2% faster)
    # Unset fields should not be present
    for k in ['weight_decay', 'warmup_fraction', 'epochs']:
        pass

def test_serialize_model_set_all_fields():
    """Set all fields, all should be present in the output."""
    obj = ClassifierTrainingParameters(
        training_steps=100,
        learning_rate=0.01,
        weight_decay=0.1,
        warmup_fraction=0.2,
        epochs=10,
        seq_len=256,
    )
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 6.54μs -> 5.07μs (29.0% faster)

def test_serialize_model_set_nullable_to_none():
    """Set nullable fields to None, should be present with value None."""
    obj = ClassifierTrainingParameters(
        training_steps=None,
        weight_decay=None,
        warmup_fraction=None,
        epochs=None,
        seq_len=None,
    )
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 7.44μs -> 5.51μs (35.0% faster)
    # All nullable fields set to None should be present
    for k in ['training_steps', 'weight_decay', 'warmup_fraction', 'epochs', 'seq_len']:
        pass

def test_serialize_model_learning_rate_none():
    """learning_rate is not nullable, setting it to None should be present as None."""
    obj = ClassifierTrainingParameters(learning_rate=None)
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 8.40μs -> 6.65μs (26.4% faster)
    # Other fields should not be present
    for k in ['training_steps', 'weight_decay', 'warmup_fraction', 'epochs', 'seq_len']:
        pass

# 2. Edge Test Cases


def test_serialize_model_set_to_zero_and_falsey():
    """Set fields to 0, which should be present in output."""
    obj = ClassifierTrainingParameters(
        training_steps=0,
        learning_rate=0.0,
        weight_decay=0.0,
        warmup_fraction=0.0,
        epochs=0,
        seq_len=0,
    )
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 9.64μs -> 7.05μs (36.6% faster)

def test_serialize_model_missing_handler_fields():
    """Simulate handler returning incomplete dict."""
    def handler_missing(obj):
        # Only returns learning_rate
        return {"learning_rate": getattr(obj, "learning_rate", UNSET)}
    obj = ClassifierTrainingParameters(training_steps=5)
    codeflash_output = obj.serialize_model(handler_missing); result = codeflash_output # 6.40μs -> 4.51μs (41.7% faster)

def test_serialize_model_handler_returns_extra_fields():
    """Handler returns extra fields, they should be ignored."""
    def handler_extra(obj):
        d = default_handler(obj)
        d['extra_field'] = 123
        return d
    obj = ClassifierTrainingParameters(training_steps=1)
    codeflash_output = obj.serialize_model(handler_extra); result = codeflash_output # 9.58μs -> 7.88μs (21.5% faster)

def test_serialize_model_fields_set_vs_default():
    """Field set to default value but explicitly set should be present."""
    obj = ClassifierTrainingParameters(learning_rate=0.0001)
    codeflash_output = obj.serialize_model(default_handler); result = codeflash_output # 8.70μs -> 6.82μs (27.7% faster)

To edit these changes git checkout codeflash/optimize-ClassifierTrainingParameters.serialize_model-mh4hhxu8 and push.