dragon-ml-toolbox 2.3.0__py3-none-any.whl → 3.0.0__py3-none-any.whl

ml_tools/ML_trainer.py

@@ -0,0 +1,344 @@
+from typing import List, Literal, Union, Optional
+from pathlib import Path
+from torch.utils.data import DataLoader, Dataset
+import torch
+from torch import nn
+import numpy as np
+
+from .ML_callbacks import Callback, History, TqdmProgressBar
+from .ML_evaluation import classification_metrics, regression_metrics, plot_losses, shap_summary_plot
+from .utilities import _script_info, LogKeys
+from .logger import _LOGGER
+
+
+__all__ = [
+    "MyTrainer"
+]
+
+
+class MyTrainer:
+    def __init__(self, model: nn.Module, train_dataset: Dataset, test_dataset: Dataset, 
+                 kind: Literal["regression", "classification"],
+                 criterion: nn.Module, optimizer: torch.optim.Optimizer, 
+                 device: Union[Literal['cuda', 'mps', 'cpu'],str], dataloader_workers: int = 2, callbacks: Optional[List[Callback]] = None):
+        """
+        Automates the training process of a PyTorch Model.
+        
+        Built-in Callbacks: `History`, `TqdmProgressBar`
+
+        Args:
+            model (nn.Module): The PyTorch model to train.
+            train_dataset (Dataset): The training dataset.
+            test_dataset (Dataset): The testing/validation dataset.
+            kind (str): The type of task, 'regression' or 'classification'.
+            criterion (nn.Module): The loss function.
+            optimizer (torch.optim.Optimizer): The optimizer.
+            device (str): The device to run training on ('cpu', 'cuda', 'mps').
+            dataloader_workers (int): Subprocesses for data loading. Defaults to 2.
+            callbacks (List[Callback] | None): A list of callbacks to use during training.
+            
+        Note:
+            For **regression** tasks, suggested criterions include `nn.MSELoss` or `nn.L1Loss`.
+            
+            For **classification** tasks, `nn.CrossEntropyLoss` (multi-class) or `nn.BCEWithLogitsLoss` (binary) are common choices.
+        """
+        if kind not in ["regression", "classification"]:
+            raise TypeError("Kind must be 'regression' or 'classification'.")
+
+        self.model = model
+        self.train_dataset = train_dataset
+        self.test_dataset = test_dataset
+        self.kind = kind
+        self.criterion = criterion
+        self.optimizer = optimizer
+        self.device = self._validate_device(device)
+        self.dataloader_workers = dataloader_workers
+        
+        # Callback handler - History and TqdmProgressBar are added by default
+        default_callbacks = [History(), TqdmProgressBar()]
+        user_callbacks = callbacks if callbacks is not None else []
+        self.callbacks = default_callbacks + user_callbacks
+        self._set_trainer_on_callbacks()
+
+        # Internal state
+        self.train_loader = None
+        self.test_loader = None
+        self.history = {}
+        self.epoch = 0
+        self.epochs = 0 # Total epochs for the fit run
+        self.stop_training = False
+
+    def _validate_device(self, device: str) -> torch.device:
+        """Validates the selected device and returns a torch.device object."""
+        device_lower = device.lower()
+        if "cuda" in device_lower and not torch.cuda.is_available():
+            _LOGGER.warning("CUDA not available, switching to CPU.")
+            device = "cpu"
+        elif device_lower == "mps" and not torch.backends.mps.is_available():
+            _LOGGER.warning("Apple Metal Performance Shaders (MPS) not available, switching to CPU.")
+            device = "cpu"
+        return torch.device(device)
+
+    def _set_trainer_on_callbacks(self):
+        """Gives each callback a reference to this trainer instance."""
+        for callback in self.callbacks:
+            callback.set_trainer(self)
+
+    def _create_dataloaders(self, batch_size: int, shuffle: bool):
+        """Initializes the DataLoaders."""
+        # Ensure stability on MPS devices by setting num_workers to 0
+        loader_workers = 0 if self.device.type == 'mps' else self.dataloader_workers
+        
+        self.train_loader = DataLoader(
+            dataset=self.train_dataset, 
+            batch_size=batch_size, 
+            shuffle=shuffle, 
+            num_workers=loader_workers, 
+            pin_memory=(self.device.type == "cuda")
+        )
+        self.test_loader = DataLoader(
+            dataset=self.test_dataset, 
+            batch_size=batch_size, 
+            shuffle=False, 
+            num_workers=loader_workers, 
+            pin_memory=(self.device.type == "cuda")
+        )
+
+    def fit(self, epochs: int = 10, batch_size: int = 32, shuffle: bool = True):
+        """
+        Starts the training-validation process of the model.
+
+        Args:
+            epochs (int): The total number of epochs to train for.
+            batch_size (int): The number of samples per batch.
+            shuffle (bool): Whether to shuffle the training data at each epoch.
+        """
+        self.epochs = epochs
+        self._create_dataloaders(batch_size, shuffle)
+        self.model.to(self.device)
+        
+        # Reset stop_training flag on the trainer
+        self.stop_training = False
+
+        self.callbacks_hook('on_train_begin')
+        
+        for epoch in range(1, self.epochs + 1):
+            self.epoch = epoch
+            epoch_logs = {}
+            self.callbacks_hook('on_epoch_begin', epoch, logs=epoch_logs)
+
+            train_logs = self._train_step()
+            epoch_logs.update(train_logs)
+
+            val_logs = self._validation_step()
+            epoch_logs.update(val_logs)
+            
+            self.callbacks_hook('on_epoch_end', epoch, logs=epoch_logs)
+            
+            # Check the early stopping flag
+            if self.stop_training:
+                break
+
+        self.callbacks_hook('on_train_end')
+        return self.history
+    
+    def _train_step(self):
+        self.model.train()
+        running_loss = 0.0
+        # Enumerate to get batch index
+        for batch_idx, (features, target) in enumerate(self.train_loader): # type: ignore
+            # Create a log dictionary for the batch
+            batch_logs = {
+                LogKeys.BATCH_INDEX: batch_idx, 
+                LogKeys.BATCH_SIZE: features.size(0)
+            }
+            self.callbacks_hook('on_batch_begin', batch_idx, logs=batch_logs)
+
+            features, target = features.to(self.device), target.to(self.device)
+            self.optimizer.zero_grad()
+            output = self.model(features)
+            if isinstance(self.criterion, (nn.MSELoss, nn.L1Loss)):
+                output = output.view_as(target)
+            loss = self.criterion(output, target)
+            loss.backward()
+            self.optimizer.step()
+
+            # Calculate batch loss and update running loss for the epoch
+            batch_loss = loss.item()
+            running_loss += batch_loss * features.size(0)
+
+            # Add the batch loss to the logs and call the end-of-batch hook
+            batch_logs[LogKeys.BATCH_LOSS] = batch_loss
+            self.callbacks_hook('on_batch_end', batch_idx, logs=batch_logs)
+
+        # Return the average loss for the entire epoch
+        return {LogKeys.TRAIN_LOSS: running_loss / len(self.train_loader.dataset)} # type: ignore
+
+    def _validation_step(self):
+        self.model.eval()
+        running_loss = 0.0
+        with torch.no_grad():
+            for features, target in self.test_loader: # type: ignore
+                features, target = features.to(self.device), target.to(self.device)
+                output = self.model(features)
+                if isinstance(self.criterion, (nn.MSELoss, nn.L1Loss)):
+                    output = output.view_as(target)
+                loss = self.criterion(output, target)
+                running_loss += loss.item() * features.size(0)
+        logs = {LogKeys.VAL_LOSS: running_loss / len(self.test_loader.dataset)} # type: ignore
+        return logs
+    
+    def predict(self, dataloader: DataLoader):
+        """
+        Yields model predictions batch by batch, avoids loading all predictions into memory at once.
+
+        Args:
+            dataloader (DataLoader): The dataloader to predict on.
+
+        Yields:
+            tuple: A tuple containing (y_pred_batch, y_prob_batch, y_true_batch).
+                   y_prob_batch is None for regression tasks.
+        """
+        self.model.eval()
+        self.model.to(self.device)
+        with torch.no_grad():
+            for features, target in dataloader:
+                features = features.to(self.device)
+                output = self.model(features).cpu()
+                y_true_batch = target.numpy()
+
+                if self.kind == "classification":
+                    probs = nn.functional.softmax(output, dim=1)
+                    preds = torch.argmax(probs, dim=1)
+                    y_pred_batch = preds.numpy()
+                    y_prob_batch = probs.numpy()
+                else:
+                    y_pred_batch = output.numpy()
+                    y_prob_batch = None
+                
+                yield y_pred_batch, y_prob_batch, y_true_batch
+    
+    def evaluate(self, data: Optional[Union[DataLoader, Dataset]] = None, save_dir: Optional[Union[str,Path]] = None):
+        """
+        Evaluates the model on the given data.
+
+        Args:
+            data (DataLoader | Dataset | None ): The data to evaluate on.
+                Can be a DataLoader or a Dataset. If None, defaults to the trainer's internal test_dataset.
+            save_dir (str | Path | None): Directory to save all reports and plots. If None, metrics are shown but not saved.
+        """
+        eval_loader = None
+        if isinstance(data, DataLoader):
+            eval_loader = data
+        else:
+            # Determine which dataset to use (the one passed in, or the default test_dataset)
+            dataset_to_use = data if data is not None else self.test_dataset
+            if not isinstance(dataset_to_use, Dataset):
+                raise ValueError("Cannot evaluate. No valid DataLoader or Dataset was provided, "
+                                 "and no test_dataset is available in the trainer.")
+
+            # Create a new DataLoader from the dataset
+            eval_loader = DataLoader(
+                dataset=dataset_to_use,
+                batch_size=32,  # A sensible default for evaluation
+                shuffle=False,
+                num_workers=0 if self.device.type == 'mps' else self.dataloader_workers,
+                pin_memory=(self.device.type == "cuda")
+            )
+            
+        print("\n--- Model Evaluation ---")
+
+        # Collect results from the predict generator
+        all_preds, all_probs, all_true = [], [], []
+        for y_pred_b, y_prob_b, y_true_b in self.predict(eval_loader):
+            all_preds.append(y_pred_b)
+            if y_prob_b is not None:
+                all_probs.append(y_prob_b)
+            all_true.append(y_true_b)
+
+        y_pred = np.concatenate(all_preds)
+        y_true = np.concatenate(all_true)
+        y_prob = np.concatenate(all_probs) if self.kind == "classification" else None
+
+        if self.kind == "classification":
+            classification_metrics(y_true, y_pred, y_prob, save_dir=save_dir)
+        else:
+            regression_metrics(y_true.flatten(), y_pred.flatten(), save_dir=save_dir)
+
+        print("\n--- Training History ---")
+        plot_losses(self.history, save_dir=save_dir)
+    
+    def explain(self, explain_dataset: Optional[Dataset] = None, n_samples: int = 100, 
+                feature_names: Optional[List[str]] = None, save_dir: Optional[str] = None):
+        """
+        Explains model predictions using SHAP and saves all artifacts.
+
+        The background data is automatically sampled from the trainer's training dataset.
+
+        Args:
+            explain_dataset (Dataset, optional): A specific dataset to explain. 
+                                                 If None, the trainer's test dataset is used.
+            n_samples (int): The number of samples to use for both background and explanation.
+            feature_names (List[str], optional): Names for the features.
+            save_dir (str, optional): Directory to save all SHAP artifacts.
+        """
+        # Internal helper to create a dataloader and get a random sample
+        def _get_random_sample(dataset: Dataset, num_samples: int):
+            if dataset is None:
+                return None
+            
+            # For MPS devices, num_workers must be 0 to ensure stability
+            loader_workers = 0 if self.device.type == 'mps' else self.dataloader_workers
+            
+            loader = DataLoader(
+                dataset, 
+                batch_size=64,
+                shuffle=False,
+                num_workers=loader_workers
+            )
+            
+            all_features = [features for features, _ in loader]
+            if not all_features:
+                return None
+            
+            full_data = torch.cat(all_features, dim=0)
+            
+            if num_samples >= full_data.size(0):
+                return full_data
+            
+            rand_indices = torch.randperm(full_data.size(0))[:num_samples]
+            return full_data[rand_indices]
+
+        print(f"\n--- Preparing SHAP Data (sampling up to {n_samples} instances) ---")
+
+        # 1. Get background data from the trainer's train_dataset
+        background_data = _get_random_sample(self.train_dataset, n_samples)
+        if background_data is None:
+            print("Warning: Trainer's train_dataset is empty or invalid. Skipping SHAP analysis.")
+            return
+
+        # 2. Determine target dataset and get explanation instances
+        target_dataset = explain_dataset if explain_dataset is not None else self.test_dataset
+        instances_to_explain = _get_random_sample(target_dataset, n_samples)
+        if instances_to_explain is None:
+            print("Warning: Explanation dataset is empty or invalid. Skipping SHAP analysis.")
+            return
+
+        # 3. Call the plotting function
+        shap_summary_plot(
+            model=self.model,
+            background_data=background_data,
+            instances_to_explain=instances_to_explain,
+            feature_names=feature_names,
+            save_dir=save_dir
+        )
+    
+
+    def callbacks_hook(self, method_name: str, *args, **kwargs):
+        """Calls the specified method on all callbacks."""
+        for callback in self.callbacks:
+            method = getattr(callback, method_name)
+            method(*args, **kwargs)
+
+def info():
+    _script_info(__all__)

ml_tools/ML_tutorial.py

@@ -0,0 +1,300 @@
+import json
+from typing import Literal, Optional, Union
+from pathlib import Path
+from .logger import _LOGGER
+from .utilities import make_fullpath, sanitize_filename
+
+
+__all__ = [
+    "generate_notebook"
+]
+
+def _get_notebook_content(kind: str):
+    """Helper function to generate the cell content for the notebook."""
+    
+    # --- Common Cells ---
+    imports_cell = {
+        "cell_type": "code",
+        "source": [
+            "import torch\n",
+            "from torch import nn\n",
+            "from torch.utils.data import TensorDataset, DataLoader\n",
+            "import numpy as np\n",
+            "from pathlib import Path\n",
+            "\n",
+            "# Import from dragon_ml_toolbox\n",
+            "from ml_tools.ML_trainer import MyTrainer\n",
+            "from ml_tools.ML_callbacks import EarlyStopping, ModelCheckpoint"
+            "from ml_tools.utilities import LogKeys"
+        ]
+    }
+    
+    device_cell = {
+        "cell_type": "code",
+        "source": [
+            "import torch\\n",
+            "if torch.cuda.is_available():\\n",
+            "    device = 'cuda'\\n",
+            "elif torch.backends.mps.is_available():\\n",
+            "    device = 'mps'\\n",
+            "else:\\n",
+            "    device = 'cpu'\\n",
+            "\\n",
+            "print(f'Using device: {device}')"
+        ]
+    }
+    
+    model_definition_cell = {
+        "cell_type": "markdown",
+        "source": [
+            "### 3. Define the Model, Criterion, and Optimizer\n",
+            "Next, we define a simple neural network for our task. We also need to choose a loss function (`criterion`) and an `optimizer`."
+        ]
+    }
+
+    callbacks_cell = {
+        "cell_type": "code",
+        "source": [
+            "# Define callbacks for training\n",
+            "model_filepath = 'best_model.pth'\n",
+            "monitor_metric = LogKeys.VAL_LOSS\n",
+            "\n",
+            "model_checkpoint = ModelCheckpoint(\n",
+            "    filepath=model_filepath, \n",
+            "    save_best_only=True, \n",
+            "    monitor=monitor_metric, \n",
+            "    mode='min'\n",
+            ")\n",
+            "\n",
+            "early_stopping = EarlyStopping(\n",
+            "    patience=10, \n",
+            "    monitor=monitor_metric, \n",
+            "    mode='min'\n",
+            ")"
+        ]
+    }
+
+    trainer_instantiation_cell = {
+        "cell_type": "code",
+        "source": [
+            "trainer = MyTrainer(\n",
+            "    model=model,\n",
+            "    train_dataset=train_dataset,\n",
+            "    test_dataset=test_dataset,\n",
+            f"    kind='{kind}',\n",
+            "    criterion=criterion,\n",
+            "    optimizer=optimizer,\n",
+            "    device=device,\\n",
+            "    callbacks=[model_checkpoint, early_stopping]\n",
+            ")"
+        ]
+    }
+
+    fit_cell = {
+        "cell_type": "code",
+        "source": [
+            "history = trainer.fit(epochs=100, batch_size=16)"
+        ]
+    }
+
+    evaluation_cell = {
+        "cell_type": "code",
+        "source": [
+            "save_dir = Path('tutorial_results')\n",
+            "\n",
+            "# The evaluate method will automatically use the test_loader.\n",
+            "# First, we load the best weights saved by ModelCheckpoint.\n",
+            "model_path = Path(model_filepath)\n",
+            "if model_path.exists():\n",
+            "    print(f'Loading best model from {model_path}')\n",
+            "    trainer.model.load_state_dict(torch.load(model_path))\n",
+            "\n",
+            "print('\\n--- Evaluating Model ---')\n",
+            "# All evaluation artifacts will be saved in the 'evaluation' subdirectory.\n",
+            "trainer.evaluate(save_dir=save_dir / 'evaluation')"
+        ]
+    }
+    
+    explanation_cell = {
+        "cell_type": "code",
+        "source": [
+            "print('\\n--- Explaining Model ---')\n",
+            "# We can also generate SHAP plots to explain the model's predictions.\n",
+            "# All SHAP artifacts will be saved in the 'explanation' subdirectory.\n",
+            "trainer.explain(\n",
+            "    background_loader=trainer.train_loader,\n",
+            "    explain_loader=trainer.test_loader,\n",
+            "    save_dir=save_dir / 'explanation'\n",
+            ")"
+        ]
+    }
+
+
+    # --- Task-Specific Cells ---
+    if kind == 'classification':
+        title = "Classification Tutorial"
+        data_prep_source = [
+            "### 2. Prepare the Data\n",
+            "For this example, we'll generate some simple, linearly separable mock data for a binary classification task. We'll then wrap it in PyTorch `TensorDataset` objects."
+        ]
+        data_creation_source = [
+            "from sklearn.datasets import make_classification\n",
+            "from sklearn.model_selection import train_test_split\n",
+            "\n",
+            "X, y = make_classification(n_samples=200, n_features=10, n_informative=5, n_redundant=0, random_state=42)\n",
+            "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)\n",
+            "\n",
+            "# Convert to PyTorch tensors\n",
+            "X_train = torch.FloatTensor(X_train)\n",
+            "y_train = torch.LongTensor(y_train)\n",
+            "X_test = torch.FloatTensor(X_test)\n",
+            "y_test = torch.LongTensor(y_test)\n",
+            "\n",
+            "# Create TensorDatasets\n",
+            "train_dataset = TensorDataset(X_train, y_train)\n",
+            "test_dataset = TensorDataset(X_test, y_test)"
+        ]
+        model_creation_source = [
+            "class SimpleClassifier(nn.Module):\n",
+            "    def __init__(self, input_features, num_classes):\n",
+            "        super().__init__()\n",
+            "        self.layer_1 = nn.Linear(input_features, 32)\n",
+            "        self.layer_2 = nn.Linear(32, num_classes)\n",
+            "        self.relu = nn.ReLU()\n",
+            "    \n",
+            "    def forward(self, x):\n",
+            "        return self.layer_2(self.relu(self.layer_1(x)))\n",
+            "\n",
+            "model = SimpleClassifier(input_features=10, num_classes=2)\n",
+            "criterion = nn.CrossEntropyLoss()\n",
+            "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)"
+        ]
+
+    elif kind == 'regression':
+        title = "Regression Tutorial"
+        data_prep_source = [
+            "### 2. Prepare the Data\n",
+            "For this example, we'll generate some simple mock data for a regression task. We'll then wrap it in PyTorch `TensorDataset` objects."
+        ]
+        data_creation_source = [
+            "from sklearn.datasets import make_regression\n",
+            "from sklearn.model_selection import train_test_split\n",
+            "\n",
+            "X, y = make_regression(n_samples=200, n_features=5, noise=15, random_state=42)\n",
+            "y = y.reshape(-1, 1) # Reshape for compatibility with MSELoss\n",
+            "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)\n",
+            "\n",
+            "# Convert to PyTorch tensors\n",
+            "X_train = torch.FloatTensor(X_train)\n",
+            "y_train = torch.FloatTensor(y_train)\n",
+            "X_test = torch.FloatTensor(X_test)\n",
+            "y_test = torch.FloatTensor(y_test)\n",
+            "\n",
+            "# Create TensorDatasets\n",
+            "train_dataset = TensorDataset(X_train, y_train)\n",
+            "test_dataset = TensorDataset(X_test, y_test)"
+        ]
+        model_creation_source = [
+            "class SimpleRegressor(nn.Module):\n",
+            "    def __init__(self, input_features, output_features):\n",
+            "        super().__init__()\n",
+            "        self.layer_1 = nn.Linear(input_features, 32)\n",
+            "        self.layer_2 = nn.Linear(32, output_features)\n",
+            "        self.relu = nn.ReLU()\n",
+            "    \n",
+            "    def forward(self, x):\n",
+            "        return self.layer_2(self.relu(self.layer_1(x)))\n",
+            "\n",
+            "model = SimpleRegressor(input_features=5, output_features=1)\n",
+            "criterion = nn.MSELoss()\n",
+            "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)"
+        ]
+    else:
+        raise ValueError("kind must be 'classification' or 'regression'")
+
+    # --- Assemble Notebook ---
+    cells = [
+        {"cell_type": "markdown", "source": [f"# Dragon ML Toolbox - {title}\n", "This notebook demonstrates how to use the `MyTrainer` class for a complete training and evaluation workflow."]},
+        {"cell_type": "markdown", "source": ["### 1. Imports\n", "First, let's import all the necessary components."]},
+        imports_cell,
+        {"cell_type": "markdown", "source": data_prep_source},
+        {"cell_type": "code", "source": data_creation_source},
+        model_definition_cell,
+        {"cell_type": "code", "source": model_creation_source},
+        {"cell_type": "markdown", "source": ["### 4. Configure Callbacks\n", "We'll set up `ModelCheckpoint` to save the best model and `EarlyStopping` to prevent overfitting."]},
+        callbacks_cell,
+        {"cell_type": "markdown", "source": ["### 5. Initialize the Trainer\\n", "First, we'll determine the best device to run on. Then, we can instantiate `MyTrainer` with all our components."]}, 
+        device_cell,
+        trainer_instantiation_cell,
+        {"cell_type": "markdown", "source": ["### 6. Train the Model\n", "Call the `.fit()` method to start training."]},
+        fit_cell,
+        {"cell_type": "markdown", "source": ["### 7. Evaluate the Model\n", "Finally, call the `.evaluate()` method to see the performance report and save all plots and metrics."]},
+        evaluation_cell,
+        {"cell_type": "markdown", "source": ["### 8. Explain the Model\n", "We can also use the `.explain()` method to generate and save SHAP plots for model interpretability."]},
+        explanation_cell,
+    ]
+    
+    # Add execution counts to code cells
+    for cell in cells:
+        if cell['cell_type'] == 'code':
+            cell['execution_count'] = None
+            cell['metadata'] = {}
+            cell['outputs'] = []
+            
+    return cells
+
+
+def generate_notebook(kind: Literal['classification', 'regression'] = 'classification', filepath: Optional[Union[str,Path]] = None):
+    """
+    Generates a tutorial Jupyter Notebook (.ipynb) file.
+
+    This function creates a complete, runnable notebook with mock data,
+    a simple model, and a full training/evaluation cycle using MyTrainer.
+
+    Args:
+        kind (str): The type of tutorial to generate, either 'classification' or 'regression'.
+        filepath (str | Path | None): The path to save the notebook file. 
+                                  If None, defaults to 'classification_tutorial.ipynb' or
+                                  'regression_tutorial.ipynb' in the current directory.
+    """
+    if kind not in ["classification", "regression"]:
+        raise ValueError("kind must be 'classification' or 'regression'")    
+    
+    if filepath is None:
+        sanitized_filepath = f"{kind}_tutorial.ipynb"
+    else:
+        sanitized_filepath = sanitize_filename(str(filepath))
+    
+    # check suffix
+    if not sanitized_filepath.endswith(".ipynb"):
+        sanitized_filepath = sanitized_filepath + ".ipynb"
+        
+    new_filepath = make_fullpath(sanitized_filepath, make=True)
+
+    _LOGGER.info(f"Generating {kind} tutorial notebook at: {filepath}")
+
+    cells = _get_notebook_content(kind)
+
+    notebook = {
+        "cells": cells,
+        "metadata": {
+            "kernelspec": {
+                "display_name": "Python 3",
+                "language": "python",
+                "name": "python3"
+            },
+            "language_info": {
+                "name": "python",
+                "version": "3.10.0" 
+            }
+        },
+        "nbformat": 4,
+        "nbformat_minor": 2
+    }
+
+    try:
+        with open(new_filepath, 'w') as f:
+            json.dump(notebook, f, indent=2)
+        _LOGGER.info("Notebook generated successfully.")
+    except Exception as e:
+        _LOGGER.error(f"Error generating notebook: {e}")