The approach is a variant of : https://community.palantir.com/t/transforms-generator-providing-additional-parameter-for-every-transformation/581
High level:
- We have a file in Code repository, that contains a list of dictionary that configure how each model should be train (which data as input, etc.).
- This list is then reused in a single final transform that defines how the produced model should be “packed” together into a “routing model”.
In the below code example, you will find:
- The list of configuration for the different models
- The adapter for a single model (what defines how to read/write the model and call inference)
- The transform generator, that generates one transform per model to produce per the above configuration list
- The adapter for the router model - the one that “pack” all the models into one
- The singular transform that defines and publish the router model
Here is the full code example:
from sklearn.ensemble import HistGradientBoostingRegressor
from palantir_models.transforms import ModelInput, ModelOutput
from transforms.api import Input, Output, transform
import logging
import numpy as np
import palantir_models as pm
import copy
## Configure the models to be created
BASE_DATA = "ri.foundry.main.dataset.e209bfa3-b05b-4cfd-b9f4-4209bd18a940"
BASE_PATH = "/path/to/folder/"
MODELS_CONFIGS = [
{
"route": "A",
"model_path": f"{BASE_PATH}/model_for_a",
"model_input": ModelInput(f"{BASE_PATH}/model_for_a", use_sidecar=False),
},
{
"route": "B",
"model_path": f"{BASE_PATH}/model_for_b",
"model_input": ModelInput(f"{BASE_PATH}/model_for_b", use_sidecar=False),
},
{
"route": "C",
"model_path": f"{BASE_PATH}/model_for_c",
"model_input": ModelInput(f"{BASE_PATH}/model_for_c", use_sidecar=False),
},
]
# Router model config
ROUTER_MODEL_PATH = "/path/to/folder/router_model_abc"
# Training parameters
TARGET_COLUMN = "recommended_setting_pct"
MODEL_ROUTING_COL = "model_route"
MODEL_USED_COL = "model_used"
ESTIMATOR_PARAMS = {
"loss": "squared_error",
"learning_rate": 0.06,
"max_iter": 350,
"max_leaf_nodes": 31,
"min_samples_leaf": 15,
"l2_regularization": 0.5,
"early_stopping": True,
"validation_fraction": 0.15,
"n_iter_no_change": 30,
}
FEATURE_COLUMNS = [
"feature_1_c",
"feature_2_c",
"feature_3_c",
"feature_4_c",
]
INPUT_COLUMNS = [(MODEL_ROUTING_COL, str)] + [(column, float) for column in FEATURE_COLUMNS]
OUTPUT_COLUMNS = INPUT_COLUMNS + [(TARGET_COLUMN, float), (MODEL_USED_COL, str)]
## Utils functions
# Return a prediction frame, common for all models. So that we are sure the schema are matching
def prediction_frame(input_df, predictions, routes):
"""Return a model API output frame with bounded predictions and route labels."""
output_df = input_df[[MODEL_ROUTING_COL] + FEATURE_COLUMNS].copy()
output_df[TARGET_COLUMN] = np.clip(np.asarray(predictions, dtype=float), 0.0, 100.0)
output_df[MODEL_USED_COL] = np.asarray(routes, dtype=str)
return output_df
## Adapter for all models
"""Adapter used by the independently deployable A, B, and C specialists."""
class SpecializedModelAdapter(pm.ModelAdapter):
"""Serve one domain-specialized regression estimator."""
@pm.auto_serialize()
def __init__(self, estimator, specialist_profile):
self.estimator = estimator
self.specialist_profile = specialist_profile
@classmethod
def api(cls):
inputs = {"input_df": pm.Pandas(columns=INPUT_COLUMNS)}
outputs = {"output_df": pm.Pandas(columns=OUTPUT_COLUMNS)}
return inputs, outputs
def predict(self, input_df):
# Logic to perform the inference/prediction from the stored weights
predictions = self.estimator.predict(input_df[FEATURE_COLUMNS])
routes = np.full(len(input_df), self.specialist_profile, dtype=object)
return prediction_frame(input_df, predictions, routes)
## Definition of the models generation
DYNAMIC_MODEL_INPUTS = {}
for curr_model_config in MODELS_CONFIGS:
# Define the condition and the model to use, e.g. "A" => model A, "B" => Model B, etc.
DYNAMIC_MODEL_INPUTS[curr_model_config["route"]] = curr_model_config["model_input"]
def generate_model_transform(configuration):
# Note: The profile, or other configuration could be passed here as well
# profile = configuration["profile"]
# dynamic_dataset_inputs = configuration["datasets"]
# @configure(profile=spark_profile)
route = configuration["route"]
# Define the parameterized transform for one model training
@transform.using(
base_data=Input(BASE_DATA),
model_output=ModelOutput(configuration["model_path"]),
)
def my_transform(ctx, base_data, model_output):
logging.info(f"My Transform is running as Incremental : {ctx.is_incremental}")
## Step 1. Do the training of the model
base_frame = base_data.pandas()
# Optional: Filter for a subset of data relevant for the model here
# e.g. only relevant for A or B or ... here filtering for "training" data only
def rows_for_split(frame, split):
return frame.loc[frame["data_split"] == split].reset_index(drop=True)
training = rows_for_split(base_frame, "train")
estimator = HistGradientBoostingRegressor(random_state=100, **ESTIMATOR_PARAMS)
estimator.fit(training[FEATURE_COLUMNS], training[TARGET_COLUMN])
# If experiments are relevant
# experiment = create_specialist_experiment(
# model_output, profile, estimator, base_frame, specialization_frame
# )
# Step 2. Write the model on the output
model_output.publish(
model_adapter=SpecializedModelAdapter(estimator, route),
# experiment=experiment,
)
transform.using(
base_data=Input(BASE_DATA),
model_output=configuration["model_path"],
)
return my_transform
def create_model_transforms(configuration):
generated_transforms = []
# For each model to train, generate one transform
for model_to_train in MODELS_CONFIGS:
# Generate the one transform
curr_transform = generate_model_transform(model_to_train)
# Add it to the list of transforms to expose
generated_transforms.append(curr_transform)
# return the list of transforms generated
return generated_transforms
# Automatic Pipeline discovery inspects this module-level list of transforms
# and so will find this variable and the transforms it contains
TRANSFORMS = create_model_transforms(MODELS_CONFIGS)
## Adapter for wrapper model
class RouterModelAdapter(pm.ModelAdapter):
"""Route rows among fine-tuned adapters."""
@pm.auto_serialize()
def __init__(self, adapters_by_profile):
self.adapters_by_profile = adapters_by_profile
@classmethod
def api(cls):
inputs = {"input_df": pm.Pandas(columns=INPUT_COLUMNS)}
outputs = {"output_df": pm.Pandas(columns=OUTPUT_COLUMNS)}
return inputs, outputs
def predict(self, input_df):
normalized_profiles = input_df[MODEL_ROUTING_COL].astype(str).str.upper()
predictions = np.empty(len(input_df), dtype=float)
routes = np.empty(len(input_df), dtype=object)
for profile in normalized_profiles.unique():
route = profile if profile in self.adapters_by_profile else "BASE"
row_mask = normalized_profiles == profile
routed_output = self.adapters_by_profile[route].predict(input_df.loc[row_mask])
predictions[row_mask.to_numpy()] = routed_output[TARGET_COLUMN].to_numpy(dtype=float)
routes[row_mask.to_numpy()] = route
return prediction_frame(input_df, predictions, routes)
## Wrapper model
"""
Clone an immutable parent checkpoint before applying updates.
"""
def clone_adapter(previous_adapter, model_route):
if not hasattr(previous_adapter, "estimator"):
raise TypeError(
"Expected an in-process SpecializedModelAdapter; missing 'estimator' attribute"
)
return SpecializedModelAdapter(
estimator=copy.deepcopy(previous_adapter.estimator),
specialist_profile=model_route,
)
@transform.using(
base_data=Input(BASE_DATA),
routed_model=ModelOutput(ROUTER_MODEL_PATH),
# training_metrics=Output(
# f"{OUTPUT_FOLDER}/POC - Dynamic Incremental Router Metrics"
# ),
**DYNAMIC_MODEL_INPUTS, # List of all the models, sidecar disabled
)
def build_dynamic_model_router(
specialization_union,
base_model,
routed_model,
training_metrics,
**loaded_models,
):
"""Bundle an arbitrary configured number of specialist model inputs."""
# Step 1. create a dict of all the models to use and when
# Set the default model if no matching filtering condition
adapters_by_profile = {
"BASE": clone_adapter(base_model, "BASE")
}
# Add the other models
for block in MODELS_CONFIGS:
route = block["route"]
input_name = block["input_name"]
adapters_by_profile[route] = clone_adapter(
loaded_models[input_name], route
)
# Step 2. Create the router model from the dict of models
adapter = RouterModelAdapter(adapters_by_profile)
# Step 3. Keeping this focused on dynamic ModelInput registration and bundling.
# Production models may add experiments, tests, etc.
routed_model.publish(model_adapter=adapter)