workbench 0.8.205__py3-none-any.whl → 0.8.212__py3-none-any.whl

workbench/model_scripts/pytorch_model/pytorch_utils.py

@@ -0,0 +1,395 @@
+"""PyTorch utilities for tabular data modeling.
+
+Provides a lightweight TabularMLP model with categorical embeddings and
+training utilities for use in Workbench model scripts.
+"""
+
+import json
+import os
+from typing import Optional
+
+import joblib
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, TensorDataset
+
+
+class FeatureScaler:
+    """Standard scaler for continuous features (zero mean, unit variance)."""
+
+    def __init__(self):
+        self.means: Optional[np.ndarray] = None
+        self.stds: Optional[np.ndarray] = None
+        self.feature_names: Optional[list[str]] = None
+
+    def fit(self, df: pd.DataFrame, continuous_cols: list[str]) -> "FeatureScaler":
+        """Fit the scaler on training data."""
+        self.feature_names = continuous_cols
+        data = df[continuous_cols].values.astype(np.float32)
+        self.means = np.nanmean(data, axis=0)
+        self.stds = np.nanstd(data, axis=0)
+        # Avoid division by zero for constant features
+        self.stds[self.stds == 0] = 1.0
+        return self
+
+    def transform(self, df: pd.DataFrame) -> np.ndarray:
+        """Transform data using fitted parameters."""
+        data = df[self.feature_names].values.astype(np.float32)
+        # Fill NaN with mean before scaling
+        for i, mean in enumerate(self.means):
+            data[np.isnan(data[:, i]), i] = mean
+        return (data - self.means) / self.stds
+
+    def fit_transform(self, df: pd.DataFrame, continuous_cols: list[str]) -> np.ndarray:
+        """Fit and transform in one step."""
+        self.fit(df, continuous_cols)
+        return self.transform(df)
+
+    def save(self, path: str) -> None:
+        """Save scaler parameters."""
+        joblib.dump(
+            {
+                "means": self.means.tolist(),
+                "stds": self.stds.tolist(),
+                "feature_names": self.feature_names,
+            },
+            path,
+        )
+
+    @classmethod
+    def load(cls, path: str) -> "FeatureScaler":
+        """Load scaler from saved parameters."""
+        data = joblib.load(path)
+        scaler = cls()
+        scaler.means = np.array(data["means"], dtype=np.float32)
+        scaler.stds = np.array(data["stds"], dtype=np.float32)
+        scaler.feature_names = data["feature_names"]
+        return scaler
+
+
+class TabularMLP(nn.Module):
+    """Feedforward neural network for tabular data with optional categorical embeddings.
+
+    Args:
+        n_continuous: Number of continuous input features
+        categorical_cardinalities: List of cardinalities for each categorical feature
+        embedding_dims: List of embedding dimensions for each categorical feature
+        hidden_layers: List of hidden layer sizes (e.g., [256, 128, 64])
+        n_outputs: Number of output units
+        task: "regression" or "classification"
+        dropout: Dropout rate
+        use_batch_norm: Whether to use batch normalization
+    """
+
+    def __init__(
+        self,
+        n_continuous: int,
+        categorical_cardinalities: list[int],
+        embedding_dims: list[int],
+        hidden_layers: list[int],
+        n_outputs: int,
+        task: str = "regression",
+        dropout: float = 0.1,
+        use_batch_norm: bool = True,
+    ):
+        super().__init__()
+        self.task = task
+        self.n_continuous = n_continuous
+        self.categorical_cardinalities = categorical_cardinalities
+
+        # Embedding layers for categorical features
+        self.embeddings = nn.ModuleList(
+            [nn.Embedding(n_cats, emb_dim) for n_cats, emb_dim in zip(categorical_cardinalities, embedding_dims)]
+        )
+
+        # Calculate input dimension
+        total_emb_dim = sum(embedding_dims)
+        input_dim = n_continuous + total_emb_dim
+
+        # Build MLP layers
+        layers = []
+        for hidden_dim in hidden_layers:
+            layers.append(nn.Linear(input_dim, hidden_dim))
+            if use_batch_norm:
+                layers.append(nn.BatchNorm1d(hidden_dim))
+            layers.append(nn.LeakyReLU())
+            layers.append(nn.Dropout(dropout))
+            input_dim = hidden_dim
+
+        self.mlp = nn.Sequential(*layers)
+        self.head = nn.Linear(input_dim, n_outputs)
+
+    def forward(self, x_cont: torch.Tensor, x_cat: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Forward pass.
+
+        Args:
+            x_cont: Continuous features tensor of shape (batch, n_continuous)
+            x_cat: Categorical features tensor of shape (batch, n_categoricals), optional
+
+        Returns:
+            Output tensor of shape (batch, n_outputs)
+        """
+        # Embed categorical features and concatenate with continuous
+        if x_cat is not None and len(self.embeddings) > 0:
+            embs = [emb(x_cat[:, i]) for i, emb in enumerate(self.embeddings)]
+            x = torch.cat([x_cont] + embs, dim=1)
+        else:
+            x = x_cont
+
+        x = self.mlp(x)
+        out = self.head(x)
+
+        if self.task == "classification":
+            out = torch.softmax(out, dim=1)
+
+        return out
+
+
+def compute_embedding_dims(cardinalities: list[int], max_dim: int = 50) -> list[int]:
+    """Compute embedding dimensions using the rule of thumb: min(50, (n+1)//2)."""
+    return [min(max_dim, (n + 1) // 2) for n in cardinalities]
+
+
+def prepare_data(
+    df: pd.DataFrame,
+    continuous_cols: list[str],
+    categorical_cols: list[str],
+    target_col: Optional[str] = None,
+    category_mappings: Optional[dict] = None,
+    scaler: Optional[FeatureScaler] = None,
+) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], dict, Optional[FeatureScaler]]:
+    """Prepare dataframe for model input.
+
+    Args:
+        df: Input dataframe
+        continuous_cols: List of continuous feature column names
+        categorical_cols: List of categorical feature column names
+        target_col: Target column name (optional, for training)
+        category_mappings: Existing category mappings (for inference)
+        scaler: Existing FeatureScaler (for inference), or None to fit a new one
+
+    Returns:
+        Tuple of (x_cont, x_cat, y, category_mappings, scaler)
+    """
+    # Continuous features with standardization
+    if scaler is None:
+        scaler = FeatureScaler()
+        cont_data = scaler.fit_transform(df, continuous_cols)
+    else:
+        cont_data = scaler.transform(df)
+    x_cont = torch.tensor(cont_data, dtype=torch.float32)
+
+    # Categorical features
+    x_cat = None
+    if categorical_cols:
+        if category_mappings is None:
+            category_mappings = {}
+            for col in categorical_cols:
+                unique_vals = df[col].unique().tolist()
+                category_mappings[col] = {v: i for i, v in enumerate(unique_vals)}
+
+        cat_indices = []
+        for col in categorical_cols:
+            mapping = category_mappings[col]
+            # Map values to indices, use 0 for unknown categories
+            indices = df[col].map(lambda x: mapping.get(x, 0)).values
+            cat_indices.append(indices)
+
+        x_cat = torch.tensor(np.column_stack(cat_indices), dtype=torch.long)
+
+    # Target
+    y = None
+    if target_col is not None:
+        y = torch.tensor(df[target_col].values, dtype=torch.float32)
+        if len(y.shape) == 1:
+            y = y.unsqueeze(1)
+
+    return x_cont, x_cat, y, category_mappings, scaler
+
+
+def create_model(
+    n_continuous: int,
+    categorical_cardinalities: list[int],
+    hidden_layers: list[int],
+    n_outputs: int,
+    task: str = "regression",
+    dropout: float = 0.1,
+    use_batch_norm: bool = True,
+) -> TabularMLP:
+    """Create a TabularMLP model with appropriate embedding dimensions."""
+    embedding_dims = compute_embedding_dims(categorical_cardinalities)
+    return TabularMLP(
+        n_continuous=n_continuous,
+        categorical_cardinalities=categorical_cardinalities,
+        embedding_dims=embedding_dims,
+        hidden_layers=hidden_layers,
+        n_outputs=n_outputs,
+        task=task,
+        dropout=dropout,
+        use_batch_norm=use_batch_norm,
+    )
+
+
+def train_model(
+    model: TabularMLP,
+    train_x_cont: torch.Tensor,
+    train_x_cat: Optional[torch.Tensor],
+    train_y: torch.Tensor,
+    val_x_cont: torch.Tensor,
+    val_x_cat: Optional[torch.Tensor],
+    val_y: torch.Tensor,
+    task: str = "regression",
+    max_epochs: int = 200,
+    patience: int = 20,
+    batch_size: int = 128,
+    learning_rate: float = 1e-3,
+    device: str = "cpu",
+) -> tuple[TabularMLP, dict]:
+    """Train the model with early stopping.
+
+    Returns:
+        Tuple of (trained model, training history dict)
+    """
+    model = model.to(device)
+
+    # Create dataloaders
+    if train_x_cat is not None:
+        train_dataset = TensorDataset(train_x_cont, train_x_cat, train_y)
+        val_dataset = TensorDataset(val_x_cont, val_x_cat, val_y)
+    else:
+        # Use dummy categorical tensor
+        dummy_cat = torch.zeros(train_x_cont.shape[0], 0, dtype=torch.long)
+        dummy_val_cat = torch.zeros(val_x_cont.shape[0], 0, dtype=torch.long)
+        train_dataset = TensorDataset(train_x_cont, dummy_cat, train_y)
+        val_dataset = TensorDataset(val_x_cont, dummy_val_cat, val_y)
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
+
+    # Loss and optimizer
+    if task == "classification":
+        criterion = nn.CrossEntropyLoss()
+    else:
+        criterion = nn.MSELoss()
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+
+    # Training loop with early stopping
+    best_val_loss = float("inf")
+    best_state = None
+    epochs_without_improvement = 0
+    history = {"train_loss": [], "val_loss": []}
+
+    for epoch in range(max_epochs):
+        # Training
+        model.train()
+        train_losses = []
+        for batch in train_loader:
+            x_cont, x_cat, y = [b.to(device) for b in batch]
+            x_cat = x_cat if x_cat.shape[1] > 0 else None
+
+            optimizer.zero_grad()
+            out = model(x_cont, x_cat)
+
+            if task == "classification":
+                loss = criterion(out, y.squeeze().long())
+            else:
+                loss = criterion(out, y)
+
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            train_losses.append(loss.item())
+
+        # Validation
+        model.eval()
+        val_losses = []
+        with torch.no_grad():
+            for batch in val_loader:
+                x_cont, x_cat, y = [b.to(device) for b in batch]
+                x_cat = x_cat if x_cat.shape[1] > 0 else None
+                out = model(x_cont, x_cat)
+
+                if task == "classification":
+                    loss = criterion(out, y.squeeze().long())
+                else:
+                    loss = criterion(out, y)
+                val_losses.append(loss.item())
+
+        train_loss = np.mean(train_losses)
+        val_loss = np.mean(val_losses)
+        history["train_loss"].append(train_loss)
+        history["val_loss"].append(val_loss)
+
+        # Early stopping check
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
+            epochs_without_improvement = 0
+        else:
+            epochs_without_improvement += 1
+
+        if (epoch + 1) % 10 == 0:
+            print(f"Epoch {epoch + 1}: train_loss={train_loss:.4f}, val_loss={val_loss:.4f}")
+
+        if epochs_without_improvement >= patience:
+            print(f"Early stopping at epoch {epoch + 1}")
+            break
+
+    # Load best weights
+    if best_state is not None:
+        model.load_state_dict(best_state)
+
+    model = model.to("cpu")
+    return model, history
+
+
+def predict(
+    model: TabularMLP,
+    x_cont: torch.Tensor,
+    x_cat: Optional[torch.Tensor] = None,
+    device: str = "cpu",
+) -> np.ndarray:
+    """Run inference with the model."""
+    model = model.to(device)
+    model.eval()
+
+    with torch.no_grad():
+        x_cont = x_cont.to(device)
+        if x_cat is not None:
+            x_cat = x_cat.to(device)
+        out = model(x_cont, x_cat)
+
+    return out.cpu().numpy()
+
+
+def save_model(model: TabularMLP, path: str, model_config: dict) -> None:
+    """Save model weights and configuration."""
+    os.makedirs(path, exist_ok=True)
+    torch.save(model.state_dict(), os.path.join(path, "model.pt"))
+    with open(os.path.join(path, "config.json"), "w") as f:
+        json.dump(model_config, f, indent=2)
+
+
+def load_model(path: str, device: str = "cpu") -> TabularMLP:
+    """Load model from saved weights and configuration."""
+    with open(os.path.join(path, "config.json")) as f:
+        config = json.load(f)
+
+    model = create_model(
+        n_continuous=config["n_continuous"],
+        categorical_cardinalities=config["categorical_cardinalities"],
+        hidden_layers=config["hidden_layers"],
+        n_outputs=config["n_outputs"],
+        task=config["task"],
+        dropout=config.get("dropout", 0.1),
+        use_batch_norm=config.get("use_batch_norm", True),
+    )
+
+    state_dict = torch.load(os.path.join(path, "model.pt"), map_location=device, weights_only=True)
+    model.load_state_dict(state_dict)
+    model.eval()
+
+    return model

workbench/model_scripts/pytorch_model/requirements.txt

@@ -1,2 +1,2 @@
-# Note: The training and inference images already have torch and pytorch-tabular installed.
+# Note: The training and inference images already have torch + supporting packages installed.
 #       So we only need to install packages that are not already included in the images.

workbench/model_scripts/pytorch_model/uq_harness.py

@@ -0,0 +1,278 @@
+"""UQ Harness: Uncertainty Quantification using MAPIE Conformalized Quantile Regression.
+
+This module provides a reusable UQ harness that can wrap any point predictor model
+(XGBoost, PyTorch, ChemProp, etc.) to provide calibrated prediction intervals.
+
+Usage:
+    # Training
+    uq_models, uq_metadata = train_uq_models(X_train, y_train, X_val, y_val)
+    save_uq_models(uq_models, uq_metadata, model_dir)
+
+    # Inference
+    uq_models, uq_metadata = load_uq_models(model_dir)
+    df = predict_intervals(df, X, uq_models, uq_metadata)
+    df = compute_confidence(df, uq_metadata["median_interval_width"])
+"""
+
+import json
+import os
+import numpy as np
+import pandas as pd
+import joblib
+from lightgbm import LGBMRegressor
+from mapie.regression import ConformalizedQuantileRegressor
+
+
+# Default confidence levels for prediction intervals
+DEFAULT_CONFIDENCE_LEVELS = [0.50, 0.68, 0.80, 0.90, 0.95]
+
+
+def train_uq_models(
+    X_train: pd.DataFrame | np.ndarray,
+    y_train: pd.Series | np.ndarray,
+    X_val: pd.DataFrame | np.ndarray,
+    y_val: pd.Series | np.ndarray,
+    confidence_levels: list[float] | None = None,
+) -> tuple[dict, dict]:
+    """Train MAPIE UQ models for multiple confidence levels.
+
+    Args:
+        X_train: Training features
+        y_train: Training targets
+        X_val: Validation features for conformalization
+        y_val: Validation targets for conformalization
+        confidence_levels: List of confidence levels (default: [0.50, 0.68, 0.80, 0.90, 0.95])
+
+    Returns:
+        Tuple of (uq_models dict, uq_metadata dict)
+    """
+    if confidence_levels is None:
+        confidence_levels = DEFAULT_CONFIDENCE_LEVELS
+
+    mapie_models = {}
+
+    for confidence_level in confidence_levels:
+        alpha = 1 - confidence_level
+        lower_q = alpha / 2
+        upper_q = 1 - alpha / 2
+
+        print(f"\nTraining quantile models for {confidence_level * 100:.0f}% confidence interval...")
+        print(f"  Quantiles: {lower_q:.3f}, {upper_q:.3f}, 0.500")
+
+        # Train three LightGBM quantile models for this confidence level
+        quantile_estimators = []
+        for q in [lower_q, upper_q, 0.5]:
+            print(f"    Training model for quantile {q:.3f}...")
+            est = LGBMRegressor(
+                objective="quantile",
+                alpha=q,
+                n_estimators=1000,
+                max_depth=6,
+                learning_rate=0.01,
+                num_leaves=31,
+                min_child_samples=20,
+                subsample=0.8,
+                colsample_bytree=0.8,
+                random_state=42,
+                verbose=-1,
+                force_col_wise=True,
+            )
+            est.fit(X_train, y_train)
+            quantile_estimators.append(est)
+
+        # Create MAPIE CQR model for this confidence level
+        print(f"  Setting up MAPIE CQR for {confidence_level * 100:.0f}% confidence...")
+        mapie_model = ConformalizedQuantileRegressor(
+            quantile_estimators, confidence_level=confidence_level, prefit=True
+        )
+
+        # Conformalize the model with validation data
+        print("  Conformalizing with validation data...")
+        mapie_model.conformalize(X_val, y_val)
+
+        # Store the model
+        model_name = f"mapie_{confidence_level:.2f}"
+        mapie_models[model_name] = mapie_model
+
+        # Validate coverage for this confidence level
+        y_pred, y_pis = mapie_model.predict_interval(X_val)
+        coverage = np.mean((y_val >= y_pis[:, 0, 0]) & (y_val <= y_pis[:, 1, 0]))
+        print(f"  Coverage: Target={confidence_level * 100:.0f}%, Empirical={coverage * 100:.1f}%")
+
+    # Compute median interval width for confidence calculation (using 80% CI = q_10 to q_90)
+    print("\nComputing normalization statistics for confidence scores...")
+    model_80 = mapie_models["mapie_0.80"]
+    _, y_pis_80 = model_80.predict_interval(X_val)
+    interval_width = np.abs(y_pis_80[:, 1, 0] - y_pis_80[:, 0, 0])
+    median_interval_width = float(np.median(interval_width))
+    print(f"  Median interval width (q_10-q_90): {median_interval_width:.6f}")
+
+    # Analyze interval widths across confidence levels
+    print("\nInterval Width Analysis:")
+    for conf_level in confidence_levels:
+        model = mapie_models[f"mapie_{conf_level:.2f}"]
+        _, y_pis = model.predict_interval(X_val)
+        widths = y_pis[:, 1, 0] - y_pis[:, 0, 0]
+        print(f"  {conf_level * 100:.0f}% CI: Mean width={np.mean(widths):.3f}, Std={np.std(widths):.3f}")
+
+    uq_metadata = {
+        "confidence_levels": confidence_levels,
+        "median_interval_width": median_interval_width,
+    }
+
+    return mapie_models, uq_metadata
+
+
+def save_uq_models(uq_models: dict, uq_metadata: dict, model_dir: str) -> None:
+    """Save UQ models and metadata to disk.
+
+    Args:
+        uq_models: Dictionary of MAPIE models keyed by name (e.g., "mapie_0.80")
+        uq_metadata: Dictionary with confidence_levels and median_interval_width
+        model_dir: Directory to save models
+    """
+    # Save each MAPIE model
+    for model_name, model in uq_models.items():
+        joblib.dump(model, os.path.join(model_dir, f"{model_name}.joblib"))
+
+    # Save median interval width
+    with open(os.path.join(model_dir, "median_interval_width.json"), "w") as fp:
+        json.dump(uq_metadata["median_interval_width"], fp)
+
+    # Save UQ metadata
+    with open(os.path.join(model_dir, "uq_metadata.json"), "w") as fp:
+        json.dump(uq_metadata, fp, indent=2)
+
+    print(f"Saved {len(uq_models)} UQ models to {model_dir}")
+
+
+def load_uq_models(model_dir: str) -> tuple[dict, dict]:
+    """Load UQ models and metadata from disk.
+
+    Args:
+        model_dir: Directory containing saved models
+
+    Returns:
+        Tuple of (uq_models dict, uq_metadata dict)
+    """
+    # Load UQ metadata
+    uq_metadata_path = os.path.join(model_dir, "uq_metadata.json")
+    if os.path.exists(uq_metadata_path):
+        with open(uq_metadata_path) as fp:
+            uq_metadata = json.load(fp)
+    else:
+        # Fallback for older models that only have median_interval_width.json
+        uq_metadata = {"confidence_levels": DEFAULT_CONFIDENCE_LEVELS}
+        median_width_path = os.path.join(model_dir, "median_interval_width.json")
+        if os.path.exists(median_width_path):
+            with open(median_width_path) as fp:
+                uq_metadata["median_interval_width"] = json.load(fp)
+
+    # Load all MAPIE models
+    uq_models = {}
+    for conf_level in uq_metadata["confidence_levels"]:
+        model_name = f"mapie_{conf_level:.2f}"
+        model_path = os.path.join(model_dir, f"{model_name}.joblib")
+        if os.path.exists(model_path):
+            uq_models[model_name] = joblib.load(model_path)
+
+    return uq_models, uq_metadata
+
+
+def predict_intervals(
+    df: pd.DataFrame,
+    X: pd.DataFrame | np.ndarray,
+    uq_models: dict,
+    uq_metadata: dict,
+) -> pd.DataFrame:
+    """Add prediction intervals to a DataFrame.
+
+    Args:
+        df: DataFrame to add interval columns to
+        X: Features for prediction (must match training features)
+        uq_models: Dictionary of MAPIE models
+        uq_metadata: Dictionary with confidence_levels
+
+    Returns:
+        DataFrame with added quantile columns (q_025, q_05, ..., q_975)
+    """
+    confidence_levels = uq_metadata["confidence_levels"]
+
+    for conf_level in confidence_levels:
+        model_name = f"mapie_{conf_level:.2f}"
+        model = uq_models[model_name]
+
+        # Get conformalized predictions
+        y_pred, y_pis = model.predict_interval(X)
+
+        # Map confidence levels to quantile column names
+        if conf_level == 0.50:  # 50% CI
+            df["q_25"] = y_pis[:, 0, 0]
+            df["q_75"] = y_pis[:, 1, 0]
+            df["q_50"] = y_pred  # Median prediction
+        elif conf_level == 0.68:  # 68% CI (~1 std)
+            df["q_16"] = y_pis[:, 0, 0]
+            df["q_84"] = y_pis[:, 1, 0]
+        elif conf_level == 0.80:  # 80% CI
+            df["q_10"] = y_pis[:, 0, 0]
+            df["q_90"] = y_pis[:, 1, 0]
+        elif conf_level == 0.90:  # 90% CI
+            df["q_05"] = y_pis[:, 0, 0]
+            df["q_95"] = y_pis[:, 1, 0]
+        elif conf_level == 0.95:  # 95% CI
+            df["q_025"] = y_pis[:, 0, 0]
+            df["q_975"] = y_pis[:, 1, 0]
+
+    # Calculate pseudo-standard deviation from the 68% interval width
+    if "q_84" in df.columns and "q_16" in df.columns:
+        df["prediction_std"] = (df["q_84"] - df["q_16"]).abs() / 2.0
+
+    # Reorder quantile columns for easier reading
+    quantile_cols = ["q_025", "q_05", "q_10", "q_16", "q_25", "q_50", "q_75", "q_84", "q_90", "q_95", "q_975"]
+    existing_q_cols = [c for c in quantile_cols if c in df.columns]
+    other_cols = [c for c in df.columns if c not in quantile_cols]
+    df = df[other_cols + existing_q_cols]
+
+    return df
+
+
+def compute_confidence(
+    df: pd.DataFrame,
+    median_interval_width: float,
+    lower_q: str = "q_10",
+    upper_q: str = "q_90",
+    alpha: float = 1.0,
+    beta: float = 1.0,
+) -> pd.DataFrame:
+    """Compute confidence scores (0.0 to 1.0) based on prediction interval width.
+
+    Uses exponential decay based on:
+    1. Interval width relative to median (alpha weight)
+    2. Distance from median prediction (beta weight)
+
+    Args:
+        df: DataFrame with 'prediction', 'q_50', and quantile columns
+        median_interval_width: Pre-computed median interval width from training data
+        lower_q: Lower quantile column name (default: 'q_10')
+        upper_q: Upper quantile column name (default: 'q_90')
+        alpha: Weight for interval width term (default: 1.0)
+        beta: Weight for distance from median term (default: 1.0)
+
+    Returns:
+        DataFrame with added 'confidence' column
+    """
+    # Interval width
+    interval_width = (df[upper_q] - df[lower_q]).abs()
+
+    # Distance from median, normalized by interval width
+    distance_from_median = (df["prediction"] - df["q_50"]).abs()
+    normalized_distance = distance_from_median / (interval_width + 1e-6)
+
+    # Cap the distance penalty at 1.0
+    normalized_distance = np.minimum(normalized_distance, 1.0)
+
+    # Confidence using exponential decay
+    interval_term = interval_width / median_interval_width
+    df["confidence"] = np.exp(-(alpha * interval_term + beta * normalized_distance))
+
+    return df

workbench/model_scripts/script_generation.py

@@ -110,18 +110,15 @@ def generate_model_script(template_params: dict) -> str:
     if template_params.get("model_class"):
         template_name = "scikit_learn.template"
         model_script_dir = "scikit_learn"
-    elif template_params["model_framework"] == ModelFramework.PYTORCH_TABULAR:
+    elif template_params["model_framework"] == ModelFramework.PYTORCH:
         template_name = "pytorch.template"
         model_script_dir = "pytorch_model"
     elif template_params["model_framework"] == ModelFramework.CHEMPROP:
         template_name = "chemprop.template"
         model_script_dir = "chemprop"
-    elif template_params["model_type"] in [ModelType.REGRESSOR, ModelType.CLASSIFIER]:
+    elif template_params["model_type"] in [ModelType.REGRESSOR, ModelType.UQ_REGRESSOR, ModelType.CLASSIFIER]:
         template_name = "xgb_model.template"
         model_script_dir = "xgb_model"
-    elif template_params["model_type"] == ModelType.UQ_REGRESSOR:
-        template_name = "mapie.template"
-        model_script_dir = "uq_models"
     elif template_params["model_type"] == ModelType.ENSEMBLE_REGRESSOR:
         template_name = "ensemble_xgb.template"
         model_script_dir = "ensemble_xgb"