safetycage 0.0.1__tar.gz

safetycage-0.0.1/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2026 [OUR AUTHORS ...]
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

safetycage-0.0.1/PKG-INFO

@@ -0,0 +1,88 @@
+Metadata-Version: 2.4
+Name: safetycage
+Version: 0.0.1
+Summary: The official safety cage package for misclassification detection in predictive models.
+Author: Pål Vegard Bun Johnsen, Joel Bjervig, Julia Qiu
+Author-email: Pål Vegard Bun Johnsen <pal.johnsen@sintef.no>, Joel Bjervig <Joel.Bjervig@sintef.no>, Julia Qiu <Julia.Qiu@sintef.no>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Dist: joblib==1.3.2
+Requires-Dist: matplotlib==3.8.2
+Requires-Dist: numpy>=1.25.2
+Requires-Dist: omegaconf==2.3.0
+Requires-Dist: pyrootutils==1.0.4
+Requires-Dist: scikit-learn==1.1.3
+Requires-Dist: scipy==1.9.3
+Requires-Dist: statsmodels==0.13.5
+Requires-Dist: tqdm>=4.67.3
+Requires-Python: ==3.11.7
+Project-URL: Homepage, https://github.com/safety-cage/safetycage
+Project-URL: Issues, https://github.com/safety-cage/safetycage/issues
+Description-Content-Type: text/markdown
+
+# safetycage(-pypi) (SafetyCage Test Package)
+
+This is the git repo that hosts the development and testing of the safetycage python package (src & tests respectfully).
+
+## Background?
+
+## Description
+Let people know what your project can do specifically. Provide context and add a link to any reference visitors might be unfamiliar with. A list of Features or a Background subsection can also be added here. If there are alternatives to your project, this is a good place to list differentiating factors.
+
+## Visuals
+- could put a gif on installing and using the package
+
+## Installation
+TBA
+
+## Requirements
+- Python 3.11.7
+
+## Usage
+Use examples liberally, and show the expected output if you can. It's helpful to have inline the smallest example of usage that you can demonstrate, while providing links to more sophisticated examples if they are too long to reasonably include in the README.
+
+## Support
+Tell people where they can go to for help. It can be any combination of an issue tracker, a chat room, an email address, etc.
+
+## Integrate with your tools
+
+* [Set up project integrations](https://gitlab.sintef.no/aai-internal/safetycage-pypi/-/settings/integrations)
+
+## Collaborate with your team
+
+* [Invite team members and collaborators](https://docs.gitlab.com/ee/user/project/members/)
+* [Create a new merge request](https://docs.gitlab.com/ee/user/project/merge_requests/creating_merge_requests.html)
+* [Automatically close issues from merge requests](https://docs.gitlab.com/ee/user/project/issues/managing_issues.html#closing-issues-automatically)
+* [Enable merge request approvals](https://docs.gitlab.com/ee/user/project/merge_requests/approvals/)
+* [Set auto-merge](https://docs.gitlab.com/user/project/merge_requests/auto_merge/)
+
+## Test and Deploy
+
+Use the built-in continuous integration in GitLab.
+
+* [Get started with GitLab CI/CD](https://docs.gitlab.com/ee/ci/quick_start/)
+* [Analyze your code for known vulnerabilities with Static Application Security Testing (SAST)](https://docs.gitlab.com/ee/user/application_security/sast/)
+* [Deploy to Kubernetes, Amazon EC2, or Amazon ECS using Auto Deploy](https://docs.gitlab.com/ee/topics/autodevops/requirements.html)
+* [Use pull-based deployments for improved Kubernetes management](https://docs.gitlab.com/ee/user/clusters/agent/)
+* [Set up protected environments](https://docs.gitlab.com/ee/ci/environments/protected_environments.html)
+
+## Roadmap
+If you have ideas for releases in the future, it is a good idea to list them in the README.
+
+## Contributing
+State if you are open to contributions and what your requirements are for accepting them.
+
+For people who want to make changes to your project, it's helpful to have some documentation on how to get started. Perhaps there is a script that they should run or some environment variables that they need to set. Make these steps explicit. These instructions could also be useful to your future self.
+
+You can also document commands to lint the code or run tests. These steps help to ensure high code quality and reduce the likelihood that the changes inadvertently break something. Having instructions for running tests is especially helpful if it requires external setup, such as starting a Selenium server for testing in a browser.
+
+## Authors and acknowledgment
+Show your appreciation to those who have contributed to the project.
+
+## License
+MIT License
+
+## Project status
+If you have run out of energy or time for your project, put a note at the top of the README saying that development has slowed down or stopped completely. Someone may choose to fork your project or volunteer to step in as a maintainer or owner, allowing your project to keep going. You can also make an explicit request for maintainers.

safetycage-0.0.1/README.md

@@ -0,0 +1,64 @@
+# safetycage(-pypi) (SafetyCage Test Package)
+
+This is the git repo that hosts the development and testing of the safetycage python package (src & tests respectfully).
+
+## Background?
+
+## Description
+Let people know what your project can do specifically. Provide context and add a link to any reference visitors might be unfamiliar with. A list of Features or a Background subsection can also be added here. If there are alternatives to your project, this is a good place to list differentiating factors.
+
+## Visuals
+- could put a gif on installing and using the package
+
+## Installation
+TBA
+
+## Requirements
+- Python 3.11.7
+
+## Usage
+Use examples liberally, and show the expected output if you can. It's helpful to have inline the smallest example of usage that you can demonstrate, while providing links to more sophisticated examples if they are too long to reasonably include in the README.
+
+## Support
+Tell people where they can go to for help. It can be any combination of an issue tracker, a chat room, an email address, etc.
+
+## Integrate with your tools
+
+* [Set up project integrations](https://gitlab.sintef.no/aai-internal/safetycage-pypi/-/settings/integrations)
+
+## Collaborate with your team
+
+* [Invite team members and collaborators](https://docs.gitlab.com/ee/user/project/members/)
+* [Create a new merge request](https://docs.gitlab.com/ee/user/project/merge_requests/creating_merge_requests.html)
+* [Automatically close issues from merge requests](https://docs.gitlab.com/ee/user/project/issues/managing_issues.html#closing-issues-automatically)
+* [Enable merge request approvals](https://docs.gitlab.com/ee/user/project/merge_requests/approvals/)
+* [Set auto-merge](https://docs.gitlab.com/user/project/merge_requests/auto_merge/)
+
+## Test and Deploy
+
+Use the built-in continuous integration in GitLab.
+
+* [Get started with GitLab CI/CD](https://docs.gitlab.com/ee/ci/quick_start/)
+* [Analyze your code for known vulnerabilities with Static Application Security Testing (SAST)](https://docs.gitlab.com/ee/user/application_security/sast/)
+* [Deploy to Kubernetes, Amazon EC2, or Amazon ECS using Auto Deploy](https://docs.gitlab.com/ee/topics/autodevops/requirements.html)
+* [Use pull-based deployments for improved Kubernetes management](https://docs.gitlab.com/ee/user/clusters/agent/)
+* [Set up protected environments](https://docs.gitlab.com/ee/ci/environments/protected_environments.html)
+
+## Roadmap
+If you have ideas for releases in the future, it is a good idea to list them in the README.
+
+## Contributing
+State if you are open to contributions and what your requirements are for accepting them.
+
+For people who want to make changes to your project, it's helpful to have some documentation on how to get started. Perhaps there is a script that they should run or some environment variables that they need to set. Make these steps explicit. These instructions could also be useful to your future self.
+
+You can also document commands to lint the code or run tests. These steps help to ensure high code quality and reduce the likelihood that the changes inadvertently break something. Having instructions for running tests is especially helpful if it requires external setup, such as starting a Selenium server for testing in a browser.
+
+## Authors and acknowledgment
+Show your appreciation to those who have contributed to the project.
+
+## License
+MIT License
+
+## Project status
+If you have run out of energy or time for your project, put a note at the top of the README saying that development has slowed down or stopped completely. Someone may choose to fork your project or volunteer to step in as a maintainer or owner, allowing your project to keep going. You can also make an explicit request for maintainers.

safetycage-0.0.1/pyproject.toml

@@ -0,0 +1,47 @@
+[build-system]
+requires = ["uv_build >= 0.10.0, <0.11.0"]
+build-backend = "uv_build"
+
+[project]
+name = "safetycage"
+version = "0.0.1"
+authors = [
+    { name="Pål Vegard Bun Johnsen", email="pal.johnsen@sintef.no" },
+    { name="Joel Bjervig", email="Joel.Bjervig@sintef.no" },
+    { name="Julia Qiu", email="Julia.Qiu@sintef.no" },
+]
+description = "The official safety cage package for misclassification detection in predictive models."
+readme = "README.md"
+requires-python = "==3.11.7"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: OS Independent",
+]
+license = "MIT"
+license-files = ["LICEN[CS]E*"]
+dependencies = [
+    "joblib==1.3.2",
+    "matplotlib==3.8.2",
+    "numpy>=1.25.2",
+    "omegaconf==2.3.0",
+    "pyrootutils==1.0.4",
+    "scikit-learn==1.1.3",
+    "scipy==1.9.3",
+    "statsmodels==0.13.5",
+    "tqdm>=4.67.3",
+]
+
+[project.urls]
+Homepage = "https://github.com/safety-cage/safetycage"
+Issues = "https://github.com/safety-cage/safetycage/issues"
+
+[tool.uv]
+package = true
+
+[tool.uv.sources]
+safetycage-testing = { workspace = true }
+
+[dependency-groups]
+dev = [
+    "safetycage-testing",
+]

safetycage-0.0.1/src/safetycage/ABC/datamodule.py

@@ -0,0 +1,101 @@
+from abc import ABC, abstractmethod
+from typing import List, Any
+from pathlib import Path
+
+class DataModule(ABC):
+    """
+    Abstract base class for handling batched data regardless of source.
+
+    This class defines a common interface for loading, transforming, and splitting
+    data for the use by safety cage methods and model modules. It also provides 
+    access to dataset properties such as the class labels and dataset size.
+
+    Subclasses must implement methods for data loading, preprocessing, and splitting.
+    """
+    
+    def __init__(
+        self,
+        data_dir:str=None,
+        from_cache:bool = False,
+        batch_size: int = 32,
+        device:str="cpu"
+        ) -> None:
+        """
+        Initialize the data module.
+
+        Sets configuration parameters for loading and processing data, such as
+        caching behavior, batch size, and device placement.
+        
+        The data directory is created if it does not already exist.
+
+        Args:
+            data_dir (str, optional): Path to the data directory. (default: None).
+            from_cache (bool, optional): Whether to load data from cached files. (default: False).
+            batch_size (int, optional): Number of samples per batch. (default: 32)
+            device (str, optional): Device used for data processing (e.g., "cpu", "cuda") (default: "cpu"). 
+        """
+
+        # Data parameters
+        self.from_cache = from_cache
+        self.batch_size = batch_size
+        
+        # Device configuration
+        self.device = device
+        
+        self.data_dir = Path(data_dir)
+        self.data_dir.mkdir(parents=True, exist_ok=True)
+
+
+    @property
+    @abstractmethod
+    def num_classes(self) -> int:
+        """Returns the number of classes in the dataset"""
+        raise NotImplementedError("Subclasses should implement this method to return the number of classes.")
+
+
+    @property
+    @abstractmethod
+    def classes(self) -> List[Any]:
+        """Returns the class names in the dataset"""
+        raise NotImplementedError("Subclasses should implement this method to return the class names.")
+
+
+    @property
+    @abstractmethod
+    def dataset_name(self) -> str:
+        """Get the name of the dataset."""
+        raise NotImplementedError("Subclasses should implement this method to return the dataset name.")
+
+
+    @abstractmethod
+    def setup(self) -> None:
+        """Setup the data module, should be implemented by subclasses"""
+        raise NotImplementedError("Subclasses should implement this method to setup the data module.")
+
+
+    @abstractmethod
+    def _load_data(self, filepath:str) -> None:
+        """Load data from the source or cache"""
+        raise NotImplementedError("Subclasses should implement this method to load data.")
+
+
+    @abstractmethod
+    def _transform(self,x,y):
+        """Transform the data, should be implemented by subclasses"""
+        raise NotImplementedError("Subclasses should implement this method to transform the data.")
+
+
+    @abstractmethod
+    def _split(self,x,y,split):
+        """Split the data into training and validation sets."""
+        raise NotImplementedError("Subclasses should implement this method to split the data.")
+    
+    @abstractmethod
+    def to_joblib(self, path: str) -> None:
+        """Save the data module to a joblib file."""
+        raise NotImplementedError("Subclasses should implement this method to save the data module to a joblib file.")
+    
+    @abstractmethod
+    def from_joblib(self, path: str) -> None:
+        """Load the data module from a joblib file."""
+        raise NotImplementedError("Subclasses should implement this method to load the data module from a joblib file.")

safetycage-0.0.1/src/safetycage/ABC/modelmodule.py

@@ -0,0 +1,83 @@
+from abc import ABC, abstractmethod
+from typing import Tuple, Union, List, Any, Dict
+import numpy as np
+import omegaconf
+
+class ModelModule(ABC):
+    """
+    Abstract base class for handling models used by safety cage methods.
+
+    This class defines a common interface for generating predictions, probabilities,
+    and intermediate activations from a model. It enables safety cage methods to
+    interact with different model types in a consistent way.
+
+    Subclasses must implement methods for prediction, probability estimation, and 
+    activation extraction.
+    """
+
+    def __init__(
+        self,
+        selected_layers:Union[str, List[str]],
+        use_onehot_encoder:bool,
+        model:Any,
+        **kwargs
+        ):
+        """
+        Initialize the model module.
+
+        Sets configuration parameters for interacting with the model, such as
+        which layers to extract activations from and how outputs are formatted.
+
+        Args:
+            model (Any, optional): Underlying model object. (default: None).
+            selected_layers (list[str], optional): List of layer names used for feature extraction. (default: None).
+            use_onehot_encoder (bool, optional): Whether model outputs are one-hot encoded. (default: False).
+            device (str, optional): Device used for model computation (e.g., "cpu", "cuda"). (default: "cpu").
+        """
+        super(ModelModule, self).__init__()
+        
+        self.model = model
+        self.use_onehot_encoder = use_onehot_encoder
+
+        # Handle different types of selected_layers input
+        if isinstance(selected_layers, str):
+            self.selected_layers = [selected_layers]
+        elif isinstance(selected_layers, (list, omegaconf.listconfig.ListConfig)):
+            if all(isinstance(layer, str) for layer in selected_layers):
+                self.selected_layers = list(selected_layers)  # Convert to regular list if it's ListConfig
+            else:
+                raise ValueError("All elements in selected_layers must be strings")
+        else:
+            raise ValueError(f"selected_layers must be a string or list of strings, got {type(selected_layers)}")
+
+
+    @abstractmethod
+    def _get_predictions(self, x: np.ndarray) -> np.ndarray:
+        """Get model predictions for input x."""
+        raise NotImplementedError("Implement based on your model architecture")
+    
+    
+    @abstractmethod
+    def _get_activations(self, x: np.ndarray) -> List[np.ndarray]:
+        """Calculate activations for each layer given input x."""
+        # This method should be implemented based on your specific model architecture
+        raise NotImplementedError("Implement based on your model architecture")
+
+
+    @abstractmethod
+    def _get_pre_activations(self, x: np.ndarray) -> List[np.ndarray]:
+        """Calculate pre-activation values for each layer given input x."""
+        raise NotImplementedError("Implement based on your model architecture")
+
+
+    @abstractmethod
+    def _calc_model_shape(self) -> Dict[str,int]:
+        """
+        Get the shape of each layer in the model.
+        Returns: List of integers representing the number of neurons in each layer
+        """
+        raise NotImplementedError("Implement based on your model architecture")
+    
+if __name__ == '__main__':
+    #Quick test to confirm it builds
+    model_module = ModelModule()

safetycage-0.0.1/src/safetycage/ABC/safetycage.py

@@ -0,0 +1,271 @@
+from typing import Any
+import pyrootutils
+from abc import ABC, abstractmethod
+import numpy as np
+import os
+import json
+import joblib
+pyrootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
+
+class SafetyCage(ABC):
+    """
+    Abstract base class for safety cage methods.
+
+    A safety cage detects misclassification in classification tasks. Safety cage methods
+    evaluate predictive models by computing statistics that indicate whether a sample is 
+    likely to be misclassified and comparing to some optimal threshold statistic, called "alpha". 
+    Concrete base classes define how these statistics are computed, how predictions are 
+    evaluated, and how to find the optimal alpha value.
+
+    Subclasses must implement training, prediction, and the method to compute the statistic.
+    Implementations on how to flag misclassifications, find the best threshold, and save/load the 
+    safetycage are provided for simplicity, but should be overridden if they do not meet the requirements
+    of the specific safety cage method.
+
+    Attributes:
+        model_module: Reference to model module object for making predictions.
+        data_module: Reference to data module object for handling data.
+        num_classes (int): Number of classes. Retrieved from the data module.
+        selected_classes (list): List of classes. Retrieved from the data module.
+        alpha (float): Threshold statistic value used for flagging misclassifications.
+        layer_params (dict, optional): Dictionary to store parameters for specific layers, if needed by 
+            safety cage.
+        leq (bool, optional): If True, samples with statistic less than or equal to alpha are flagged as 
+            misclassified. Only required if the default flag method is used.
+    """
+
+    def __init__(
+        self,
+        model_module,
+        data_module: Any,
+        **kwargs
+        ) -> None:
+        """
+        Initialize the safety cage.
+
+        Stores references to the model and data modules and initializes shared attributes 
+        used across all safety cage methods.
+
+        Args:
+            model_module: Model module used for predictions and activations.
+            data_module (Any): Data module providing datasets and class information.
+        """
+
+        self.model_module = model_module
+        self.data_module = data_module
+        
+        self.num_classes = data_module.num_classes
+        self.selected_classes = data_module.classes
+        
+        self.alpha = None
+        
+    #Train the parameters of the specified SafetyCage
+    @abstractmethod
+    def train_cage(self) -> None:
+        """
+        Train the safety cage.
+
+        Learns parameters from training data that are later used to evaluate
+        whether predictions are reliable.
+
+        Returns:
+            None
+        """
+        pass
+    
+    #Apply the SafetyCage on unseen test samples
+    @abstractmethod
+    def predict(self, x, y) -> None:
+        """
+        Evaluate input samples using the trained safety cage.
+
+        Computes statistics that indicate how likely each prediction is a misclassification.
+
+        Args:
+            x: Input data samples
+            y: True labels
+
+        Returns:
+            numpy.ndarray: Computed statistics for each sample
+        """
+        pass
+
+    #Compute the statistics to evaluate whether each test sample is wrongly predicted
+    @abstractmethod
+    def _compute_statistics(self, x, y):
+        """
+        Compute per-sample statistics used to evaluate prediction reliability.
+
+        Args:
+            x: Input data samples
+            y: True labels
+
+        Returns:
+            numpy.ndarray: Per-sample statistics
+        """
+        pass
+
+    #Flag predictions as being correct (0) or wrong (1)
+    def flag(self, statistics: float | np.ndarray, alpha: float | None = None) -> float | np.ndarray:
+        """
+        Flag samples with probability less than or equal (safetycage.leq = True) to alpha as incorrect
+        or probability more than or equal (safetycage.leq = False) to alpha as incorrect.
+        
+        This method identifies samples where the maximum/minimum probability is below/above a
+        specified threshold (alpha), marking them as potentially incorrect classifications.
+
+        *Requires safetycage.leq to be defined, not None.*
+
+        Args:
+            statistics (numpy.ndarray): Array of probability values to evaluate
+            alpha (float): Threshold value for flagging samples (0 to 1)
+        Returns:
+            numpy.ndarray: Boolean array where True indicates probabilities below/above the alpha threshold
+                depending on safetycage.leq
+        """
+                
+        # Check priority of alpha parameter
+        if alpha is None:
+            # If not provided as input, try to use self.alpha
+            if hasattr(self, 'alpha') and self.alpha is not None:
+                alpha = self.alpha
+            else:
+                # If neither source is available, raise an error
+                raise ValueError("Missing alpha parameter: must be provided as input or set as class attribute")
+            
+        if self.leq:
+            flags = statistics <= alpha
+        elif not self.leq and self.leq is not None:
+            flags = statistics >= alpha
+        else:
+            raise ValueError("Define safetycage.leq")
+
+        return flags
+
+    def find_best_threshold_flag(self, y_true, y_probs, metric_fn, greater_is_better=True) -> float | np.ndarray:
+        """
+        Find the optimal threshold for flagging samples by calling self.flag().
+
+        Evaluates thresholds t (from 1000 samples between min to max) and selects the one that maximizes 
+        the given metric.
+
+        Args:
+            y_true (numpy.ndarray): Ground-truth misclassification labels
+            y_probs (numpy.ndarray): Computed statistics or probabilities
+            metric_fn (callable): Function to evaluate performance
+            greater_is_better (bool, optional): Whether greater metric values are better (default: True)
+
+        Returns:
+            dict: Dictionary containing the optimal threshold and the corresponding best metric value
+        """
+        if greater_is_better:
+            compare = lambda x, y: x > y
+        else:
+            compare = lambda x, y: x < y
+
+        best_metric = -np.inf 
+        best_alpha = -np.inf
+        for t in np.linspace(min(y_probs), max(y_probs), num=1000):
+
+            flag = self.flag(y_probs, t)
+
+            # Flag is true when misclassification occurs
+            tps = np.sum(flag & y_true)
+            fps = np.sum(flag & (1 - y_true))
+            
+            total_pos = y_true.sum()
+            total_neg = y_true.size - total_pos
+            
+            fns = total_pos - tps
+            tns = total_neg - fps
+            
+            metric = metric_fn(TP=tps, TN=tns, FP=fps, FN=fns)
+            
+            if compare(metric, best_metric):
+                best_metric = metric
+                best_alpha = t
+                
+        return {
+            "alpha_opt": best_alpha,
+            "metric_max": best_metric,
+        }
+    
+    def save_cage(self, path):
+        """
+        Save the safety cage parameters, alpha and/or layer_params, to a specified folder path
+        in a file called "parameters.json" and/or "parameters.joblib".
+
+        Args:
+            path (str): Folder path where the safety cage parameters should be saved
+        """
+
+        os.makedirs(path, exist_ok=True)
+
+        if getattr(self, "alpha", None) is not None:
+            if getattr(self, "layer_params", None) is not None:
+                
+                parameters_path = os.path.join(path, 'parameters.joblib')
+                parameters = {
+                    'alpha': self.alpha,
+                    'layer_params': self.layer_params
+                }
+                joblib.dump(parameters, parameters_path)
+            
+            else:
+
+                parameters_path = os.path.join(path, 'parameters.json')
+
+                with open(parameters_path, 'w') as f:
+                    json.dump({'alpha': str(self.alpha)}, f)
+
+        else:
+            raise ValueError("alpha is not set.")
+
+    def load_cage(self, path):
+        """
+        Load the safety cage parameters, alpha and/or layer_params, from a specified folder path.
+        
+        Looks for a parameters.json or parameters.joblib file respectively.
+        Load cage assumes only one of these two files are present in the folder, and gives priority 
+        to parameters.joblib if both are present.
+
+        Args:
+            path (str): Path to the parameters file or directory where the safety cage parameters should be loaded.
+        """
+        joblib_path = os.path.join(path, "parameters.joblib")
+        json_path = os.path.join(path, "parameters.json")
+
+        if os.path.isfile(path):
+            if path.endswith(".joblib"):
+                
+                parameters = joblib.load(path)
+                self.alpha = parameters["alpha"]
+                self.layer_params = parameters["layer_params"]
+
+            elif path.endswith(".json"):
+
+                with open(path, "r") as f:
+                    data = json.load(f)
+                self.alpha = float(data["alpha"])
+
+            else:
+                raise ValueError(f"Unsupported file type: {path}")
+            
+        elif os.path.exists(joblib_path):
+
+            parameters = joblib.load(joblib_path)
+            
+            # Restore parameters
+            self.alpha = parameters['alpha']
+            self.layer_params = parameters['layer_params']
+
+        elif os.path.exists(json_path):
+
+            with open(json_path, 'r') as f:
+                data = json.load(f)
+                self.alpha = float(data['alpha'])
+        else:
+            raise FileNotFoundError(f"Safety cage parameters file not found at {path}")
+
+if __name__ == "__main__":
+    SafetyCage(None, None, None)