Component Callbacks#

Component callbacks are user-defined functions that adhere to specific signatures depending on the component.

Prediction Callbacks#

Prediction callbacks are defined as functions that take three arguments:

These callbacks are invoked in the unionml.model.Model.predict() method after predictions have been generated for a particular batch of features.


Currently, callbacks apply to all contexts in which predictions occur including: batch prediction via Flyte, online predictions in the FastAPI app or BentoML, and serverless predictions via AWS Lambda or BentoML.

from typing import List

import pandas as pd
from sklearn.linear_model import LogisticRegression
from unionml import Model, Dataset

dataset = Dataset()
model = Model(dataset=dataset, init=LogisticRegression)

def reader() -> pd.DataFrame: ...

def trainer(
    model_object: LogisticRegression,
    features: pd.DataFrame,
    target: pd.DataFrame,
) -> LogisticRegression:

def prediction_callback(
    model_object: LogisticRegression,
    features: pd.DataFrame,
    predictions: List[float],
    ...  # do something here

def predictor(model_object: LogisticRegression, features: pd.DataFrame) -> List[float]:
    return [float(x) for x in model_object.predict(features)]

Callbacks can be used for any purpose, but the primary use case is to log predictions using some third-party package or service for the purposes of model monitoring.