Python: 5 Essential Decorators to Optimize Your AI Code

The Utility of Python Decorators in AI Projects

In the world of developing artificial intelligence and machine learning systems, Python decorators play a crucial role. They allow for the separation of essential logic, such as modeling and data pipeline management, from repetitive tasks like testing, validation, timing, and logging. By decoupling these concerns, developers can write clearer and more maintainable code.

This article explores five Python decorators that have proven particularly effective in improving code cleanliness in AI projects. The provided code examples rely on Python's standard libraries and best practices, such as using functools.wraps. The goal is to demonstrate how each decorator can be utilized, enabling you to easily adapt them to your AI coding projects.

Concurrency Limiter: Controlling Asynchronous Execution

When using third-party large language models (LLMs), developers often encounter limits imposed by free tiers. These limits can be quickly reached if too many asynchronous requests are sent simultaneously. The concurrency limiter decorator introduces a mechanism to safeguard these calls using semaphores, which restrict the number of simultaneous executions of an asynchronous function.

from functools import wraps
import asyncio

def limit_concurrency(limit=5):
    sem = asyncio.Semaphore(limit)
    
    def decorator(func):
        async def wrapper(*args, **kwargs):
            async with sem:
                return await func(*args, **kwargs)
        return wrapper
    return decorator

@limit_concurrency(5)
async def fetch_llm_batch([prompt](/glossaire/prompt)):
    return await async_api_client.generate(prompt)

Structured Logger for Machine Learning

In the complex environments of machine learning, simple print() statements can easily get lost, especially in production. The structured logging decorator captures executions and errors, formatting them into structured JSON logs. These logs are not only easy to search but also essential for quick debugging. The following example shows how to decorate a function that defines a training epoch in a neural network model.

import logging
import json
import time
from functools import wraps

def json_log(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        try:
            res = func(*args, **kwargs)
            logging.info(json.dumps({"step": func.__name__, "status": "success", "time": time.time() - start}))
            return res
        except Exception as e:
            logging.error(json.dumps({"step": func.__name__, "error": str(e)}))
    return wrapper

@json_log
def train_epoch(model, training_data):
    return model.fit(training_data)

Feature Injector: Ensuring Data Consistency

When deploying and inferring models, it is crucial to ensure that raw data from end users undergoes the same transformations as those used during training. The feature injector decorator guarantees this consistency by automating the generation of features from raw data. The example below shows how to automatically add a 'is_weekend' feature to an existing dataframe.

from functools import wraps

def add_weekend_feature(func):
    @wraps(func)
    def wrapper(df, *args, **kwargs):
        df = df.copy()  # Avoids Pandas mutation warnings
        df['is_weekend'] = df['date'].dt.dayofweek.isin([5, 6]).astype(int)
        return func(df, *args, **kwargs)
    return wrapper

@add_weekend_feature
def process_data(df):
    # 'is_weekend' is guaranteed to exist here
    return df.dropna()

Deterministic Seed Setter: Ensuring Reliable Experiments

In the phases of experimentation and hyperparameter tuning, using a random seed is common. However, to isolate variables and make test results more reliable, it is often necessary to lock this seed. This allows for determining whether performance variations are due to changes in hyperparameters or random weight initialization. The deterministic seed decorator facilitates this task.

import random
import numpy as np
from functools import wraps

def lock_seed(seed=42):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            random.seed(seed)
            np.random.seed(seed)
            return func(*args, **kwargs)
        return wrapper
    return decorator

@lock_seed()
def initialize_weights():
    return np.random.randn(10, 10)

Development Mode Fallback: Ensuring Continuity

In local development environments and CI/CD testing, it is common to encounter errors due to external factors, such as connection timeouts or API usage limits. The development mode fallback decorator intercepts these errors and returns a predefined set of "mock test data." This allows the application to continue functioning even if an external service temporarily fails.

from functools import wraps

def fallback_mock(mock_data):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            except Exception:  # Captures timeouts and rate limits
                return mock_data
        return wrapper
    return decorator

@fallback_mock(mock_data=[0.01, -0.05, 0.02])
def get_text_embeddings(text):
    return external_api.embed(text)

In conclusion, these five Python decorators offer effective solutions for making AI and machine learning code cleaner and more maintainable. They cover various aspects, from structured logging to random seed management, while ensuring operational continuity in the event of external service failures.