265 lines
10 KiB
Python
265 lines
10 KiB
Python
import importlib
|
|
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
import cloudpickle
|
|
from packaging.version import Version
|
|
|
|
from mlflow.exceptions import INVALID_PARAMETER_VALUE, MlflowException
|
|
from mlflow.recipes.artifacts import DataframeArtifact, TransformerArtifact
|
|
from mlflow.recipes.cards import BaseCard
|
|
from mlflow.recipes.step import BaseStep, StepClass
|
|
from mlflow.recipes.utils.execution import get_step_output_path
|
|
from mlflow.recipes.utils.step import get_pandas_data_profiles, validate_classification_config
|
|
from mlflow.recipes.utils.tracking import TrackingConfig, get_recipe_tracking_config
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
_USER_DEFINED_TRANSFORM_STEP_MODULE = "steps.transform"
|
|
|
|
|
|
def _generate_feature_names(num_features):
|
|
max_length = len(str(num_features))
|
|
return ["f_" + str(i).zfill(max_length) for i in range(num_features)]
|
|
|
|
|
|
def _get_output_feature_names(transformer, num_features, input_features):
|
|
import sklearn
|
|
|
|
# `get_feature_names_out` was introduced in scikit-learn 1.0.0.
|
|
if Version(sklearn.__version__) < Version("1.0.0"):
|
|
return _generate_feature_names(num_features)
|
|
|
|
try:
|
|
# `get_feature_names_out` fails if `transformer` contains a transformer that doesn't
|
|
# implement `get_feature_names_out`. For example, `FunctionTransformer` only implements
|
|
# `get_feature_names_out` when it's instantiated with `feature_names_out`.
|
|
# In scikit-learn >= 1.1.0, all transformers implement `get_feature_names_out`.
|
|
# In scikit-learn == 1.0.*, some transformers implement `get_feature_names_out`.
|
|
return transformer.get_feature_names_out(input_features)
|
|
except Exception as e:
|
|
_logger.warning(
|
|
f"Failed to get output feature names with `get_feature_names_out`: {e}. "
|
|
"Falling back to using auto-generated feature names."
|
|
)
|
|
return _generate_feature_names(num_features)
|
|
|
|
|
|
def _validate_user_code_output(transformer_fn):
|
|
transformer = transformer_fn()
|
|
if transformer is not None and not (hasattr(transformer, "fit") and callable(transformer.fit)):
|
|
raise MlflowException(
|
|
message="The transformer provided doesn't have a fit method."
|
|
) from None
|
|
|
|
if transformer is not None and not (
|
|
hasattr(transformer, "transform") and callable(transformer.transform)
|
|
):
|
|
raise MlflowException(
|
|
message="The transformer provided doesn't have a transform method."
|
|
) from None
|
|
|
|
return transformer
|
|
|
|
|
|
class TransformStep(BaseStep):
|
|
def __init__(self, step_config, recipe_root):
|
|
super().__init__(step_config, recipe_root)
|
|
self.tracking_config = TrackingConfig.from_dict(self.step_config)
|
|
|
|
def _validate_and_apply_step_config(self):
|
|
self.target_col = self.step_config.get("target_col")
|
|
self.positive_class = self.step_config.get("positive_class")
|
|
if self.target_col is None:
|
|
raise MlflowException(
|
|
"Missing target_col config in recipe config.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
if "using" in self.step_config:
|
|
if self.step_config["using"] not in ["custom"]:
|
|
raise MlflowException(
|
|
f"Invalid transform step configuration value {self.step_config['using']} for "
|
|
f"key 'using'. Supported values are: ['custom']",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
else:
|
|
self.step_config["using"] = "custom"
|
|
self.run_end_time = None
|
|
self.execution_duration = None
|
|
self.skip_data_profiling = self.step_config.get("skip_data_profiling", False)
|
|
|
|
def _run(self, output_directory):
|
|
import pandas as pd
|
|
|
|
run_start_time = time.time()
|
|
|
|
train_data_path = get_step_output_path(
|
|
recipe_root_path=self.recipe_root,
|
|
step_name="split",
|
|
relative_path="train.parquet",
|
|
)
|
|
train_df = pd.read_parquet(train_data_path)
|
|
validate_classification_config(self.task, self.positive_class, train_df, self.target_col)
|
|
|
|
validation_data_path = get_step_output_path(
|
|
recipe_root_path=self.recipe_root,
|
|
step_name="split",
|
|
relative_path="validation.parquet",
|
|
)
|
|
validation_df = pd.read_parquet(validation_data_path)
|
|
|
|
sys.path.append(self.recipe_root)
|
|
|
|
def get_identity_transformer():
|
|
from sklearn.pipeline import Pipeline
|
|
from sklearn.preprocessing import FunctionTransformer
|
|
|
|
return Pipeline(steps=[("identity", FunctionTransformer())])
|
|
|
|
if "transformer_method" not in self.step_config and self.step_config["using"] == "custom":
|
|
raise MlflowException(
|
|
"Missing 'transformer_method' configuration in the transform step, "
|
|
"which is using 'custom'.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
method_config = self.step_config.get("transformer_method")
|
|
transformer = None
|
|
if method_config and self.step_config["using"] == "custom":
|
|
transformer_fn = getattr(
|
|
importlib.import_module(_USER_DEFINED_TRANSFORM_STEP_MODULE), method_config
|
|
)
|
|
transformer = _validate_user_code_output(transformer_fn)
|
|
transformer = transformer if transformer else get_identity_transformer()
|
|
transformer.fit(train_df.drop(columns=[self.target_col]), train_df[self.target_col])
|
|
|
|
def transform_dataset(dataset):
|
|
features = dataset.drop(columns=[self.target_col])
|
|
transformed_features = transformer.transform(features)
|
|
if not isinstance(transformed_features, pd.DataFrame):
|
|
num_features = transformed_features.shape[1]
|
|
columns = _get_output_feature_names(transformer, num_features, features.columns)
|
|
transformed_features = pd.DataFrame(transformed_features, columns=columns)
|
|
transformed_features[self.target_col] = dataset[self.target_col].values
|
|
return transformed_features
|
|
|
|
train_transformed = transform_dataset(train_df)
|
|
validation_transformed = transform_dataset(validation_df)
|
|
|
|
with open(os.path.join(output_directory, "transformer.pkl"), "wb") as f:
|
|
cloudpickle.dump(transformer, f)
|
|
|
|
train_transformed.to_parquet(
|
|
os.path.join(output_directory, "transformed_training_data.parquet")
|
|
)
|
|
validation_transformed.to_parquet(
|
|
os.path.join(output_directory, "transformed_validation_data.parquet")
|
|
)
|
|
|
|
self.run_end_time = time.time()
|
|
self.execution_duration = self.run_end_time - run_start_time
|
|
|
|
return self._build_profiles_and_card(train_df, train_transformed, transformer)
|
|
|
|
def _build_profiles_and_card(self, train_df, train_transformed, transformer) -> BaseCard:
|
|
# Build card
|
|
card = BaseCard(self.recipe_name, self.name)
|
|
|
|
if not self.skip_data_profiling:
|
|
# Tab 1: build profiles for train_transformed
|
|
train_transformed_profile = get_pandas_data_profiles(
|
|
[["Profile of Train Transformed Dataset", train_transformed]]
|
|
)
|
|
card.add_tab("Data Profile (Train Transformed)", "{{PROFILE}}").add_pandas_profile(
|
|
"PROFILE", train_transformed_profile
|
|
)
|
|
|
|
# Tab 3: transformer diagram
|
|
from sklearn import set_config
|
|
from sklearn.utils import estimator_html_repr
|
|
|
|
set_config(display="diagram")
|
|
transformer_repr = estimator_html_repr(transformer)
|
|
card.add_tab("Transformer", "{{TRANSFORMER}}").add_html("TRANSFORMER", transformer_repr)
|
|
|
|
# Tab 4: transformer input schema
|
|
card.add_tab("Input Schema", "{{INPUT_SCHEMA}}").add_html(
|
|
"INPUT_SCHEMA",
|
|
BaseCard.render_table({"Name": n, "Type": t} for n, t in train_df.dtypes.items()),
|
|
)
|
|
|
|
# Tab 5: transformer output schema
|
|
try:
|
|
card.add_tab("Output Schema", "{{OUTPUT_SCHEMA}}").add_html(
|
|
"OUTPUT_SCHEMA",
|
|
BaseCard.render_table(
|
|
{"Name": n, "Type": t} for n, t in train_transformed.dtypes.items()
|
|
),
|
|
)
|
|
except Exception as e:
|
|
card.add_tab("Output Schema", "{{OUTPUT_SCHEMA}}").add_html(
|
|
"OUTPUT_SCHEMA", f"Failed to extract transformer schema. Error: {e}"
|
|
)
|
|
|
|
# Tab 6: transformer output data preview
|
|
card.add_tab("Data Preview", "{{DATA_PREVIEW}}").add_html(
|
|
"DATA_PREVIEW", BaseCard.render_table(train_transformed.head())
|
|
)
|
|
|
|
# Tab 7: run summary
|
|
(
|
|
card.add_tab(
|
|
"Run Summary",
|
|
"""
|
|
{{ EXE_DURATION }}
|
|
{{ LAST_UPDATE_TIME }}
|
|
""",
|
|
)
|
|
)
|
|
|
|
return card
|
|
|
|
@classmethod
|
|
def from_recipe_config(cls, recipe_config, recipe_root):
|
|
step_config = {}
|
|
if recipe_config.get("steps", {}).get("transform", {}) is not None:
|
|
step_config.update(recipe_config.get("steps", {}).get("transform", {}))
|
|
step_config["target_col"] = recipe_config.get("target_col")
|
|
step_config["recipe"] = recipe_config.get("recipe", "regression/v1")
|
|
if "positive_class" in recipe_config:
|
|
step_config["positive_class"] = recipe_config.get("positive_class")
|
|
step_config.update(
|
|
get_recipe_tracking_config(
|
|
recipe_root_path=recipe_root,
|
|
recipe_config=recipe_config,
|
|
).to_dict()
|
|
)
|
|
return cls(step_config, recipe_root)
|
|
|
|
@property
|
|
def name(self):
|
|
return "transform"
|
|
|
|
def get_artifacts(self):
|
|
return [
|
|
DataframeArtifact(
|
|
"transformed_training_data",
|
|
self.recipe_root,
|
|
self.name,
|
|
"transformed_training_data.parquet",
|
|
),
|
|
DataframeArtifact(
|
|
"transformed_validation_data",
|
|
self.recipe_root,
|
|
self.name,
|
|
"transformed_validation_data.parquet",
|
|
),
|
|
TransformerArtifact(
|
|
"transformer", self.recipe_root, self.name, self.tracking_config.tracking_uri
|
|
),
|
|
]
|
|
|
|
def step_class(self):
|
|
return StepClass.TRAINING
|