workbench 0.8.213__py3-none-any.whl → 0.8.217__py3-none-any.whl

workbench/model_scripts/custom_models/uq_models/feature_space_proximity.py

@@ -0,0 +1,194 @@
+import pandas as pd
+import numpy as np
+from sklearn.preprocessing import StandardScaler
+from sklearn.neighbors import NearestNeighbors
+from typing import List, Optional
+import logging
+
+# Workbench Imports
+from workbench.algorithms.dataframe.proximity import Proximity
+from workbench.algorithms.dataframe.projection_2d import Projection2D
+
+# Set up logging
+log = logging.getLogger("workbench")
+
+
+class FeatureSpaceProximity(Proximity):
+    """Proximity computations for numeric feature spaces using Euclidean distance."""
+
+    def __init__(
+        self,
+        df: pd.DataFrame,
+        id_column: str,
+        features: List[str],
+        target: Optional[str] = None,
+        include_all_columns: bool = False,
+    ):
+        """
+        Initialize the FeatureSpaceProximity class.
+
+        Args:
+            df: DataFrame containing data for neighbor computations.
+            id_column: Name of the column used as the identifier.
+            features: List of feature column names to be used for neighbor computations.
+            target: Name of the target column. Defaults to None.
+            include_all_columns: Include all DataFrame columns in neighbor results. Defaults to False.
+        """
+        # Validate and filter features before calling parent init
+        self._raw_features = features
+        super().__init__(
+            df, id_column=id_column, features=features, target=target, include_all_columns=include_all_columns
+        )
+
+    def _prepare_data(self) -> None:
+        """Filter out non-numeric features and drop NaN rows."""
+        # Validate features
+        self.features = self._validate_features(self.df, self._raw_features)
+
+        # Drop NaN rows for the features we're using
+        self.df = self.df.dropna(subset=self.features).copy()
+
+    def _validate_features(self, df: pd.DataFrame, features: List[str]) -> List[str]:
+        """Remove non-numeric features and log warnings."""
+        non_numeric = [f for f in features if f not in df.select_dtypes(include=["number"]).columns]
+        if non_numeric:
+            log.warning(f"Non-numeric features {non_numeric} aren't currently supported, excluding them")
+        return [f for f in features if f not in non_numeric]
+
+    def _build_model(self) -> None:
+        """Standardize features and fit Nearest Neighbors model."""
+        self.scaler = StandardScaler()
+        X = self.scaler.fit_transform(self.df[self.features])
+        self.nn = NearestNeighbors().fit(X)
+
+    def _transform_features(self, df: pd.DataFrame) -> np.ndarray:
+        """Transform features using the fitted scaler."""
+        return self.scaler.transform(df[self.features])
+
+    def _project_2d(self) -> None:
+        """Project the numeric features to 2D for visualization."""
+        if len(self.features) >= 2:
+            self.df = Projection2D().fit_transform(self.df, features=self.features)
+
+
+# Testing the FeatureSpaceProximity class
+if __name__ == "__main__":
+
+    pd.set_option("display.max_columns", None)
+    pd.set_option("display.width", 1000)
+
+    # Create a sample DataFrame
+    data = {
+        "ID": [1, 2, 3, 4, 5],
+        "Feature1": [0.1, 0.2, 0.3, 0.4, 0.5],
+        "Feature2": [0.5, 0.4, 0.3, 0.2, 0.1],
+        "Feature3": [2.5, 2.4, 2.3, 2.3, np.nan],
+    }
+    df = pd.DataFrame(data)
+
+    # Test the FeatureSpaceProximity class
+    features = ["Feature1", "Feature2", "Feature3"]
+    prox = FeatureSpaceProximity(df, id_column="ID", features=features)
+    print(prox.neighbors(1, n_neighbors=2))
+
+    # Test the neighbors method with radius
+    print(prox.neighbors(1, radius=2.0))
+
+    # Test with Features list
+    prox = FeatureSpaceProximity(df, id_column="ID", features=["Feature1"])
+    print(prox.neighbors(1))
+
+    # Create a sample DataFrame
+    data = {
+        "id": ["a", "b", "c", "d", "e"],  # Testing string IDs
+        "Feature1": [0.1, 0.2, 0.3, 0.4, 0.5],
+        "Feature2": [0.5, 0.4, 0.3, 0.2, 0.1],
+        "target": [1, 0, 1, 0, 5],
+    }
+    df = pd.DataFrame(data)
+
+    # Test with String Ids
+    prox = FeatureSpaceProximity(
+        df,
+        id_column="id",
+        features=["Feature1", "Feature2"],
+        target="target",
+        include_all_columns=True,
+    )
+    print(prox.neighbors(["a", "b"]))
+
+    # Test duplicate IDs
+    data = {
+        "id": ["a", "b", "c", "d", "d"],  # Duplicate ID (d)
+        "Feature1": [0.1, 0.2, 0.3, 0.4, 0.5],
+        "Feature2": [0.5, 0.4, 0.3, 0.2, 0.1],
+        "target": [1, 0, 1, 0, 5],
+    }
+    df = pd.DataFrame(data)
+    prox = FeatureSpaceProximity(df, id_column="id", features=["Feature1", "Feature2"], target="target")
+    print(df.equals(prox.df))
+
+    # Test on real data from Workbench
+    from workbench.api import FeatureSet, Model
+
+    fs = FeatureSet("aqsol_features")
+    model = Model("aqsol-regression")
+    features = model.features()
+    df = fs.pull_dataframe()
+    prox = FeatureSpaceProximity(df, id_column=fs.id_column, features=model.features(), target=model.target())
+    print("\n" + "=" * 80)
+    print("Testing Neighbors...")
+    print("=" * 80)
+    test_id = df[fs.id_column].tolist()[0]
+    print(f"\nNeighbors for ID {test_id}:")
+    print(prox.neighbors(test_id))
+
+    print("\n" + "=" * 80)
+    print("Testing isolated_compounds...")
+    print("=" * 80)
+
+    # Test isolated data in the top 1%
+    isolated_1pct = prox.isolated(top_percent=1.0)
+    print(f"\nTop 1% most isolated compounds (n={len(isolated_1pct)}):")
+    print(isolated_1pct)
+
+    # Test isolated data in the top 5%
+    isolated_5pct = prox.isolated(top_percent=5.0)
+    print(f"\nTop 5% most isolated compounds (n={len(isolated_5pct)}):")
+    print(isolated_5pct)
+
+    print("\n" + "=" * 80)
+    print("Testing target_gradients...")
+    print("=" * 80)
+
+    # Test with different parameters
+    gradients_1pct = prox.target_gradients(top_percent=1.0, min_delta=1.0)
+    print(f"\nTop 1% target gradients (min_delta=5.0) (n={len(gradients_1pct)}):")
+    print(gradients_1pct)
+
+    gradients_5pct = prox.target_gradients(top_percent=5.0, min_delta=5.0)
+    print(f"\nTop 5% target gradients (min_delta=5.0) (n={len(gradients_5pct)}):")
+    print(gradients_5pct)
+
+    # Test proximity_stats
+    print("\n" + "=" * 80)
+    print("Testing proximity_stats...")
+    print("=" * 80)
+    stats = prox.proximity_stats()
+    print(stats)
+
+    # Plot the distance distribution using pandas
+    print("\n" + "=" * 80)
+    print("Plotting distance distribution...")
+    print("=" * 80)
+    prox.df["nn_distance"].hist(bins=50, figsize=(10, 6), edgecolor="black")
+
+    # Visualize the 2D projection
+    print("\n" + "=" * 80)
+    print("Visualizing 2D Projection...")
+    print("=" * 80)
+    from workbench.web_interface.components.plugin_unit_test import PluginUnitTest
+    from workbench.web_interface.components.plugins.scatter_plot import ScatterPlot
+
+    unit_test = PluginUnitTest(ScatterPlot, input_data=prox.df[:1000], x="x", y="y", color=model.target())
+    unit_test.run()

workbench/model_scripts/custom_models/uq_models/meta_uq.template

@@ -18,7 +18,7 @@ import pandas as pd
 from typing import List, Tuple
 
 # Local Imports
-from proximity import Proximity
+from feature_space_proximity import FeatureSpaceProximity
 
 
 # Template Placeholders
@@ -28,7 +28,7 @@ TEMPLATE_PARAMS = {
     "features": "{{feature_list}}",
     "compressed_features": "{{compressed_features}}",
     "train_all_data": "{{train_all_data}}",
-    "track_columns": "{{track_columns}}",
+    "include_all_columns": "{{include_all_columns}}",
 }
 
 
@@ -166,7 +166,7 @@ if __name__ == "__main__":
     orig_features = features.copy()
     compressed_features = TEMPLATE_PARAMS["compressed_features"]
     train_all_data = TEMPLATE_PARAMS["train_all_data"]
-    track_columns = TEMPLATE_PARAMS["track_columns"]  # Can be None
+    include_all_columns = TEMPLATE_PARAMS["include_all_columns"]  # Defaults to False
     validation_split = 0.2
 
     # Script arguments for input/output directories
@@ -260,8 +260,8 @@ if __name__ == "__main__":
     with open(os.path.join(args.model_dir, "feature_columns.json"), "w") as fp:
         json.dump(orig_features, fp)  # We save the original features, not the decompressed ones
 
-    # Now the Proximity model
-    model = Proximity(df_train, id_column, features, target, track_columns=track_columns)
+    # Now the FeatureSpaceProximity model
+    model = FeatureSpaceProximity(df_train, id_column=id_column, features=features, target=target, include_all_columns=include_all_columns)
 
     # Now serialize the model
     model.serialize(args.model_dir)
@@ -282,7 +282,7 @@ def model_fn(model_dir) -> dict:
     ngb_model = joblib.load(os.path.join(model_dir, "ngb_model.joblib"))
 
     # Deserialize the proximity model
-    prox_model = Proximity.deserialize(model_dir)
+    prox_model = FeatureSpaceProximity.deserialize(model_dir)
 
     return {"xgboost": xgb_model, "ngboost": ngb_model, "proximity": prox_model}
 

workbench/model_scripts/meta_model/generated_model_script.py

@@ -0,0 +1,209 @@
+# Meta Model Template for Workbench
+#
+# NOTE: This is called a "meta model" but it's really a "meta endpoint" - it aggregates
+# predictions from multiple child endpoints. We call it a "model" because Workbench
+# creates Model artifacts that get deployed as Endpoints, so this follows that pattern.
+#
+# Assumptions/Shortcuts:
+# - All child endpoints are regression models
+# - All child endpoints output 'prediction' and 'confidence' columns
+# - Aggregation uses model weights (provided at meta model creation time)
+#
+# This template:
+# - Has no real training phase (just saves metadata including model weights)
+# - At inference time, calls child endpoints and aggregates their predictions
+
+import argparse
+import json
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from io import StringIO
+
+import pandas as pd
+
+from workbench_bridges.endpoints.fast_inference import fast_inference
+
+# Template parameters (filled in by Workbench)
+TEMPLATE_PARAMS = {
+    "child_endpoints": ['logd-reg-pytorch', 'logd-reg-chemprop'],
+    "target_column": "logd",
+    "model_weights": {'logd-reg-pytorch': 0.4228205813233993, 'logd-reg-chemprop': 0.5771794186766008},
+    "model_metrics_s3_path": "s3://sandbox-sageworks-artifacts/models/logd-meta/training",
+    "aws_region": "us-west-2",
+}
+
+
+def invoke_endpoints_parallel(endpoint_names: list[str], df: pd.DataFrame) -> dict[str, pd.DataFrame]:
+    """Call multiple child endpoints in parallel and collect their results.
+
+    Args:
+        endpoint_names: List of endpoint names to call
+        df: Input DataFrame to send to each endpoint
+
+    Returns:
+        Dict mapping endpoint_name -> result DataFrame (or None if failed)
+    """
+    results = {}
+
+    def call_endpoint(name: str) -> tuple[str, pd.DataFrame | None]:
+        try:
+            return name, fast_inference(name, df)
+        except Exception as e:
+            print(f"Error calling endpoint {name}: {e}")
+            return name, None
+
+    with ThreadPoolExecutor(max_workers=len(endpoint_names)) as executor:
+        futures = {executor.submit(call_endpoint, name): name for name in endpoint_names}
+        for future in as_completed(futures):
+            name, result = future.result()
+            results[name] = result
+
+    return results
+
+
+def aggregate_predictions(results: dict[str, pd.DataFrame], model_weights: dict[str, float]) -> pd.DataFrame:
+    """Aggregate predictions from multiple endpoints using model weights.
+
+    Args:
+        results: Dict mapping endpoint_name -> predictions DataFrame
+                 Each DataFrame must have 'prediction' and 'confidence' columns
+        model_weights: Dict mapping endpoint_name -> weight
+
+    Returns:
+        DataFrame with aggregated prediction, prediction_std, and confidence
+    """
+    # Filter out failed endpoints
+    valid_results = {k: v for k, v in results.items() if v is not None}
+    if not valid_results:
+        raise ValueError("All child endpoints failed")
+
+    # Use first result as base (for id columns, etc.)
+    first_df = list(valid_results.values())[0]
+    output_df = first_df.drop(columns=["prediction", "confidence", "prediction_std"], errors="ignore").copy()
+
+    # Build DataFrames of predictions and confidences from all endpoints
+    pred_df = pd.DataFrame({name: df["prediction"] for name, df in valid_results.items()})
+    conf_df = pd.DataFrame({name: df["confidence"] for name, df in valid_results.items()})
+
+    # Apply model weights (renormalize for valid endpoints only)
+    valid_weights = {k: model_weights.get(k, 1.0) for k in valid_results}
+    weight_sum = sum(valid_weights.values())
+    normalized_weights = {k: v / weight_sum for k, v in valid_weights.items()}
+
+    # Weighted average
+    output_df["prediction"] = sum(pred_df[name] * w for name, w in normalized_weights.items())
+
+    # Ensemble std across child endpoints
+    output_df["prediction_std"] = pred_df.std(axis=1)
+
+    # Aggregated confidence: weighted mean of child confidences
+    output_df["confidence"] = sum(conf_df[name] * w for name, w in normalized_weights.items())
+
+    return output_df
+
+
+# =============================================================================
+# Model Loading (for SageMaker inference)
+# =============================================================================
+def model_fn(model_dir: str) -> dict:
+    """Load meta model configuration."""
+    with open(os.path.join(model_dir, "meta_config.json")) as f:
+        config = json.load(f)
+
+    # Set AWS_REGION for fast_inference (baked in at training time)
+    if config.get("aws_region"):
+        os.environ["AWS_REGION"] = config["aws_region"]
+
+    print(f"Meta model loaded: {len(config['child_endpoints'])} child endpoints")
+    print(f"Model weights: {config.get('model_weights')}")
+    print(f"AWS region: {config.get('aws_region')}")
+    return config
+
+
+def input_fn(input_data, content_type):
+    """Parse input data and return a DataFrame."""
+    if not input_data:
+        raise ValueError("Empty input data is not supported!")
+
+    # Decode bytes to string if necessary
+    if isinstance(input_data, bytes):
+        input_data = input_data.decode("utf-8")
+
+    if "text/csv" in content_type:
+        return pd.read_csv(StringIO(input_data))
+    elif "application/json" in content_type:
+        return pd.DataFrame(json.loads(input_data))
+    else:
+        raise ValueError(f"{content_type} not supported!")
+
+
+def output_fn(output_df, accept_type):
+    """Supports both CSV and JSON output formats."""
+    if "text/csv" in accept_type:
+        return output_df.to_csv(index=False), "text/csv"
+    elif "application/json" in accept_type:
+        return output_df.to_json(orient="records"), "application/json"
+    else:
+        raise RuntimeError(f"{accept_type} accept type is not supported by this script.")
+
+
+# =============================================================================
+# Inference (for SageMaker inference)
+# =============================================================================
+def predict_fn(df: pd.DataFrame, config: dict) -> pd.DataFrame:
+    """Run inference by calling child endpoints and aggregating results."""
+    child_endpoints = config["child_endpoints"]
+    model_weights = config.get("model_weights", {})
+
+    print(f"Calling {len(child_endpoints)} child endpoints: {child_endpoints}")
+
+    # Call all child endpoints
+    results = invoke_endpoints_parallel(child_endpoints, df)
+
+    # Report status
+    for name, result in results.items():
+        status = f"{len(result)} rows" if result is not None else "FAILED"
+        print(f"  {name}: {status}")
+
+    # Aggregate predictions using model weights
+    output_df = aggregate_predictions(results, model_weights)
+
+    print(f"Aggregated {len(output_df)} predictions from {len(results)} endpoints")
+    return output_df
+
+
+# =============================================================================
+# Training (just saves configuration - no actual training)
+# =============================================================================
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
+    parser.add_argument("--output-data-dir", type=str, default=os.environ.get("SM_OUTPUT_DATA_DIR", "/opt/ml/output/data"))
+    parser.add_argument("--train", type=str, default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
+    args = parser.parse_args()
+
+    child_endpoints = TEMPLATE_PARAMS["child_endpoints"]
+    target_column = TEMPLATE_PARAMS["target_column"]
+    model_weights = TEMPLATE_PARAMS["model_weights"]
+    aws_region = TEMPLATE_PARAMS["aws_region"]
+
+    print("=" * 60)
+    print("Meta Model Configuration")
+    print("=" * 60)
+    print(f"Child endpoints: {child_endpoints}")
+    print(f"Target column: {target_column}")
+    print(f"Model weights: {model_weights}")
+    print(f"AWS region: {aws_region}")
+
+    # Save configuration for inference
+    config = {
+        "child_endpoints": child_endpoints,
+        "target_column": target_column,
+        "model_weights": model_weights,
+        "aws_region": aws_region,
+    }
+
+    with open(os.path.join(args.model_dir, "meta_config.json"), "w") as f:
+        json.dump(config, f, indent=2)
+
+    print(f"\nMeta model configuration saved to {args.model_dir}")

workbench/model_scripts/meta_model/meta_model.template

@@ -0,0 +1,209 @@
+# Meta Model Template for Workbench
+#
+# NOTE: This is called a "meta model" but it's really a "meta endpoint" - it aggregates
+# predictions from multiple child endpoints. We call it a "model" because Workbench
+# creates Model artifacts that get deployed as Endpoints, so this follows that pattern.
+#
+# Assumptions/Shortcuts:
+# - All child endpoints are regression models
+# - All child endpoints output 'prediction' and 'confidence' columns
+# - Aggregation uses model weights (provided at meta model creation time)
+#
+# This template:
+# - Has no real training phase (just saves metadata including model weights)
+# - At inference time, calls child endpoints and aggregates their predictions
+
+import argparse
+import json
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from io import StringIO
+
+import pandas as pd
+
+from workbench_bridges.endpoints.fast_inference import fast_inference
+
+# Template parameters (filled in by Workbench)
+TEMPLATE_PARAMS = {
+    "child_endpoints": "{{child_endpoints}}",
+    "target_column": "{{target_column}}",
+    "model_weights": "{{model_weights}}",
+    "model_metrics_s3_path": "{{model_metrics_s3_path}}",
+    "aws_region": "{{aws_region}}",
+}
+
+
+def invoke_endpoints_parallel(endpoint_names: list[str], df: pd.DataFrame) -> dict[str, pd.DataFrame]:
+    """Call multiple child endpoints in parallel and collect their results.
+
+    Args:
+        endpoint_names: List of endpoint names to call
+        df: Input DataFrame to send to each endpoint
+
+    Returns:
+        Dict mapping endpoint_name -> result DataFrame (or None if failed)
+    """
+    results = {}
+
+    def call_endpoint(name: str) -> tuple[str, pd.DataFrame | None]:
+        try:
+            return name, fast_inference(name, df)
+        except Exception as e:
+            print(f"Error calling endpoint {name}: {e}")
+            return name, None
+
+    with ThreadPoolExecutor(max_workers=len(endpoint_names)) as executor:
+        futures = {executor.submit(call_endpoint, name): name for name in endpoint_names}
+        for future in as_completed(futures):
+            name, result = future.result()
+            results[name] = result
+
+    return results
+
+
+def aggregate_predictions(results: dict[str, pd.DataFrame], model_weights: dict[str, float]) -> pd.DataFrame:
+    """Aggregate predictions from multiple endpoints using model weights.
+
+    Args:
+        results: Dict mapping endpoint_name -> predictions DataFrame
+                 Each DataFrame must have 'prediction' and 'confidence' columns
+        model_weights: Dict mapping endpoint_name -> weight
+
+    Returns:
+        DataFrame with aggregated prediction, prediction_std, and confidence
+    """
+    # Filter out failed endpoints
+    valid_results = {k: v for k, v in results.items() if v is not None}
+    if not valid_results:
+        raise ValueError("All child endpoints failed")
+
+    # Use first result as base (for id columns, etc.)
+    first_df = list(valid_results.values())[0]
+    output_df = first_df.drop(columns=["prediction", "confidence", "prediction_std"], errors="ignore").copy()
+
+    # Build DataFrames of predictions and confidences from all endpoints
+    pred_df = pd.DataFrame({name: df["prediction"] for name, df in valid_results.items()})
+    conf_df = pd.DataFrame({name: df["confidence"] for name, df in valid_results.items()})
+
+    # Apply model weights (renormalize for valid endpoints only)
+    valid_weights = {k: model_weights.get(k, 1.0) for k in valid_results}
+    weight_sum = sum(valid_weights.values())
+    normalized_weights = {k: v / weight_sum for k, v in valid_weights.items()}
+
+    # Weighted average
+    output_df["prediction"] = sum(pred_df[name] * w for name, w in normalized_weights.items())
+
+    # Ensemble std across child endpoints
+    output_df["prediction_std"] = pred_df.std(axis=1)
+
+    # Aggregated confidence: weighted mean of child confidences
+    output_df["confidence"] = sum(conf_df[name] * w for name, w in normalized_weights.items())
+
+    return output_df
+
+
+# =============================================================================
+# Model Loading (for SageMaker inference)
+# =============================================================================
+def model_fn(model_dir: str) -> dict:
+    """Load meta model configuration."""
+    with open(os.path.join(model_dir, "meta_config.json")) as f:
+        config = json.load(f)
+
+    # Set AWS_REGION for fast_inference (baked in at training time)
+    if config.get("aws_region"):
+        os.environ["AWS_REGION"] = config["aws_region"]
+
+    print(f"Meta model loaded: {len(config['child_endpoints'])} child endpoints")
+    print(f"Model weights: {config.get('model_weights')}")
+    print(f"AWS region: {config.get('aws_region')}")
+    return config
+
+
+def input_fn(input_data, content_type):
+    """Parse input data and return a DataFrame."""
+    if not input_data:
+        raise ValueError("Empty input data is not supported!")
+
+    # Decode bytes to string if necessary
+    if isinstance(input_data, bytes):
+        input_data = input_data.decode("utf-8")
+
+    if "text/csv" in content_type:
+        return pd.read_csv(StringIO(input_data))
+    elif "application/json" in content_type:
+        return pd.DataFrame(json.loads(input_data))
+    else:
+        raise ValueError(f"{content_type} not supported!")
+
+
+def output_fn(output_df, accept_type):
+    """Supports both CSV and JSON output formats."""
+    if "text/csv" in accept_type:
+        return output_df.to_csv(index=False), "text/csv"
+    elif "application/json" in accept_type:
+        return output_df.to_json(orient="records"), "application/json"
+    else:
+        raise RuntimeError(f"{accept_type} accept type is not supported by this script.")
+
+
+# =============================================================================
+# Inference (for SageMaker inference)
+# =============================================================================
+def predict_fn(df: pd.DataFrame, config: dict) -> pd.DataFrame:
+    """Run inference by calling child endpoints and aggregating results."""
+    child_endpoints = config["child_endpoints"]
+    model_weights = config.get("model_weights", {})
+
+    print(f"Calling {len(child_endpoints)} child endpoints: {child_endpoints}")
+
+    # Call all child endpoints
+    results = invoke_endpoints_parallel(child_endpoints, df)
+
+    # Report status
+    for name, result in results.items():
+        status = f"{len(result)} rows" if result is not None else "FAILED"
+        print(f"  {name}: {status}")
+
+    # Aggregate predictions using model weights
+    output_df = aggregate_predictions(results, model_weights)
+
+    print(f"Aggregated {len(output_df)} predictions from {len(results)} endpoints")
+    return output_df
+
+
+# =============================================================================
+# Training (just saves configuration - no actual training)
+# =============================================================================
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
+    parser.add_argument("--output-data-dir", type=str, default=os.environ.get("SM_OUTPUT_DATA_DIR", "/opt/ml/output/data"))
+    parser.add_argument("--train", type=str, default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
+    args = parser.parse_args()
+
+    child_endpoints = TEMPLATE_PARAMS["child_endpoints"]
+    target_column = TEMPLATE_PARAMS["target_column"]
+    model_weights = TEMPLATE_PARAMS["model_weights"]
+    aws_region = TEMPLATE_PARAMS["aws_region"]
+
+    print("=" * 60)
+    print("Meta Model Configuration")
+    print("=" * 60)
+    print(f"Child endpoints: {child_endpoints}")
+    print(f"Target column: {target_column}")
+    print(f"Model weights: {model_weights}")
+    print(f"AWS region: {aws_region}")
+
+    # Save configuration for inference
+    config = {
+        "child_endpoints": child_endpoints,
+        "target_column": target_column,
+        "model_weights": model_weights,
+        "aws_region": aws_region,
+    }
+
+    with open(os.path.join(args.model_dir, "meta_config.json"), "w") as f:
+        json.dump(config, f, indent=2)
+
+    print(f"\nMeta model configuration saved to {args.model_dir}")