Photo by Editor
# Introduction
Python decorators can be incredibly useful for projects that involve AI and machine learning system development. They excel at separating core logic like modeling and data pipelines from other boilerplate functions, like testing and validation, timing, logging, and so on.
This article describes five particularly useful Python decorators that, based on developer experience, have proven effective in making AI code cleaner.
The code examples below include a simple, basic approach based on standard Python libraries and best practices, e.g. functools.wraps. Their main purpose is to demonstrate the use of each specific decorator, so you don't have to worry about fixing the decorator logic in your AI coding project.
# 1. Concurrency Limiter
A very useful decorator when working around the (often annoying) limitations of the free class in the use of large third-party language models (LLMs). If you reach such limits as a result of sending many concurrent requests, this pattern presents a striking mechanism to make these calls more secure. With semaphores, the number of times an asynchronous function can run is limited:
import asyncio
from functools import wraps
def limit_concurrency(limit=5):
sem = asyncio.Semaphore(limit)
def decorator(func):
@wraps(func)
async def wrapper(*args, **kwargs):
async with sem:
return await func(*args, **kwargs)
return wrapper
return decorator
# Application
@limit_concurrency(5)
async def fetch_llm_batch(prompt):
return await async_api_client.generate(prompt)# 2. Built-in Machine Learning Cube
It is not surprising that in software as complex as that which controls machine learning programs, the norm print() statements are easily lost, especially when they are invested in production.
Using the following logging decorator, it is possible to “catch” executions and errors and format them into easily searchable JSON-formatted logs for quick remediation. The code example below can be used as a template for decorating, for example, a function that defines the training period for a model based on a neural network:
import logging, json, 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)}))
raise
return wrapper
# Application
@json_log
def train_epoch(model, training_data):
return model.fit(training_data)# 3. Injector feature
Add a decorator that is especially useful during the model deployment and targeting phases! Say you're moving your machine learning model from a notebook to a lightweight production environment, e.g. using the FastAPI endpoint. Manually verifying that the raw incoming data from the end users has the same transformation as the original training data can sometimes be a little painful. A feature injector helps ensure consistency in how features are generated from raw data, all under the hood.
This is very useful during the deployment and decision phase. When you move a model from a Jupyter notebook to a production environment, the big head ensures that the raw user input data gets the same transformation as your training data. This decorator ensures that those features are generated correctly under the hood before the data reaches your model.
The example below simplifies the process of adding a feature called 'is_weekend'based on whether the date column in the existing data frame contains a date associated with Saturday or Sunday:
from functools import wraps
def add_weekend_feature(func):
@wraps(func)
def wrapper(df, *args, **kwargs):
df = df.copy() # Prevents Pandas mutation warnings
df['is_weekend'] = df['date'].dt.dayofweek.isin([5, 6]).astype(int)
return func(df, *args, **kwargs)
return wrapper
# Application
@add_weekend_feature
def process_data(df):
# 'is_weekend' is guaranteed to exist here
return df.dropna()# 4. Deterministic Seed Setter
This one stands out in two specific phases of the AI/machine learning lifecycle: testing and hyperparameter tuning. These procedures usually include the use of random seeds as part of adjusting key parameters such as the learning rate of the model. Say you just adjusted its value, and suddenly, the accuracy of the model drops. In a case like this, you may need to know if the cause of this performance drop is a new hyperparameter setting or a random bad initialization of the weights. By locking the seed, we isolate the variables, thus making the results of tests like A/B more reliable.
import random, 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
# Application
@lock_seed(42)
def initialize_weights():
return np.random.randn(10, 10)# 5. Dev-Mode Fallback
A life-saving decorator, especially in the areas of site development and CI/CD testing. Say you're building an application layer on top of LLM — for example, a retrieval-augmented generation (RAG) system. If the decorator fails due to external factors, such as connection timeouts or API usage restrictions, instead of throwing an exception, the error is caught by this decorator and a predefined set of “pseudo test data” is returned.
Why save life? Because this method can ensure that your application does not stop completely if the external service fails temporarily.
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: # Catches timeouts and rate limits
return mock_data
return wrapper
return decorator
# Application
@fallback_mock(mock_data=[0.01, -0.05, 0.02])
def get_text_embeddings(text):
return external_api.embed(text)# Wrapping up
This article examines five effective Python decorators that will help make your AI and machine learning code clean in all different specific situations: from structured, easy-to-search hacking to controlled randomization in features like data sampling, testing, and more.
Iván Palomares Carrascosa is a leader, author, speaker, and consultant in AI, machine learning, deep learning and LLMs. He trains and guides others in using AI in the real world.
