content: # Apache Software License 2.0
  #
  # Copyright (c) ZenML GmbH 2025. All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  # http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  #
  
  from typing import List, Optional, Tuple
  
  import pandas as pd
  from sklearn.pipeline import Pipeline
  from sklearn.preprocessing import MinMaxScaler
  from typing_extensions import Annotated
  from utils.preprocess import ColumnsDropper, DataFrameCaster, NADropper
  
  from zenml import log_metadata, step
  
  
  @step
  def data_preprocessor(
      random_state: int,
      dataset_trn: pd.DataFrame,
      dataset_tst: pd.DataFrame,
      drop_na: Optional[bool] = None,
      normalize: Optional[bool] = None,
      drop_columns: Optional[List[str]] = None,
      target: Optional[str] = "target",
  ) -> Tuple[
      Annotated[pd.DataFrame, "dataset_trn"],
      Annotated[pd.DataFrame, "dataset_tst"],
      Annotated[Pipeline, "preprocess_pipeline"],
  ]:
      """Data preprocessor step.
  
      This is an example of a data processor step that prepares the data so that
      it is suitable for model training. It takes in a dataset as an input step
      artifact and performs any necessary preprocessing steps like cleaning,
      feature engineering, feature selection, etc. It then returns the processed
      dataset as a step output artifact.
  
      This step is parameterized, which allows you to configure the step
      independently of the step code, before running it in a pipeline.
      In this example, the step can be configured to drop NA values, drop some
      columns and normalize numerical columns. See the documentation for more
      information:
  
          https://docs.zenml.io/how-to/build-pipelines/use-pipeline-step-parameters
  
      Args:
          random_state: Random state for sampling.
          dataset_trn: The train dataset.
          dataset_tst: The test dataset.
          drop_na: If `True` all NA rows will be dropped.
          normalize: If `True` all numeric fields will be normalized.
          drop_columns: List of column names to drop.
          target: Name of target column in dataset.
  
      Returns:
          The processed datasets (dataset_trn, dataset_tst) and fitted `Pipeline` object.
      """
      # We use the sklearn pipeline to chain together multiple preprocessing steps
      preprocess_pipeline = Pipeline([("passthrough", "passthrough")])
      if drop_na:
          preprocess_pipeline.steps.append(("drop_na", NADropper()))
      if drop_columns:
          # Drop columns
          preprocess_pipeline.steps.append(
              ("drop_columns", ColumnsDropper(drop_columns))
          )
      if normalize:
          # Normalize the data
          preprocess_pipeline.steps.append(("normalize", MinMaxScaler()))
      preprocess_pipeline.steps.append(
          ("cast", DataFrameCaster(dataset_trn.columns))
      )
      dataset_trn = preprocess_pipeline.fit_transform(dataset_trn)
      dataset_tst = preprocess_pipeline.transform(dataset_tst)
  
      # Log metadata so we can load it in the inference pipeline
      log_metadata(
          metadata={"random_state": random_state, "target": target},
          artifact_name="preprocess_pipeline",
          infer_artifact=True,
      )
      return dataset_trn, dataset_tst, preprocess_pipeline
  