prefab 1.2.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

prefab/predict.py

@@ -1,37 +1,134 @@
-"""Prediction functions for ndarrays of device geometries."""
+"""
+Serverless prediction interface for nanofabrication modeling.
+
+This module provides functions for predicting nanofabrication outcomes using machine
+learning models hosted on a serverless platform. It supports multiple input formats
+(ndarrays, polygons, GDSII files) and model types (prediction, correction,
+segmentation). Gradient computation is available for inverse design applications
+using automatic differentiation.
+"""
 
 import base64
 import io
 import json
 import os
+from typing import Any
 
 import gdstk
 import numpy as np
+import numpy.typing as npt
 import requests
 import toml
 from autograd import primitive
 from autograd.extend import defvjp
 from PIL import Image
-from tqdm import tqdm
 
 from .geometry import binarize_hard
 from .models import Model
 
 BASE_ENDPOINT_URL = "https://prefab-photonics--predict"
-ENDPOINT_VERSION = "2"
+ENDPOINT_VERSION = "3"
+
+
+def predict_array(
+    device_array: npt.NDArray[Any],
+    model: Model,
+    model_type: str,
+    binarize: bool,
+) -> npt.NDArray[Any]:
+    """
+    Predict the nanofabrication outcome of a device array using a specified model.
+
+    This function sends the device array to a serverless prediction service, which uses
+    a specified machine learning model to predict the outcome of the nanofabrication
+    process.
+
+    Parameters
+    ----------
+    device_array : np.ndarray
+        A 2D array representing the planar geometry of the device. This array undergoes
+        various transformations to predict the nanofabrication process.
+    model : Model
+        The model to use for prediction, representing a specific fabrication process and
+        dataset. This model encapsulates details about the fabrication foundry and
+        process, as defined in `models.py`. Each model is associated with a version and
+        dataset that detail its creation and the data it was trained on, ensuring the
+        prediction is tailored to specific fabrication parameters.
+    model_type : str
+        The type of model to use (e.g., 'p' for prediction or 'c' for correction).
+    binarize : bool
+        If True, the predicted device geometry will be binarized using a threshold
+        method. This is useful for converting probabilistic predictions into binary
+        geometries.
+
+    Returns
+    -------
+    np.ndarray
+        The predicted output array. For single-level predictions, returns shape
+        (h, w, 1). For multi-level predictions, returns shape (h, w, n) where n is the
+        number of levels.
+
+    Raises
+    ------
+    RuntimeError
+        If the request to the prediction service fails.
+    ValueError
+        If the server returns an error or invalid response.
+    """
+    endpoint_url = f"{BASE_ENDPOINT_URL}-v{ENDPOINT_VERSION}.modal.run"
+    predict_data = {
+        "device_array": _encode_array(np.squeeze(device_array)),
+        "model": model.to_json(),
+        "model_type": model_type,
+    }
+    headers = _prepare_headers()
+
+    try:
+        response = requests.post(
+            url=endpoint_url, data=json.dumps(predict_data), headers=headers
+        )
+        response.raise_for_status()
+
+        if not response.content:
+            raise ValueError("Empty response received from server")
+
+        response_data = response.json()
+
+        if "error" in response_data:
+            raise ValueError(f"Prediction error: {response_data['error']}")
+
+        results = response_data["results"]
+        result_arrays = [
+            _decode_array(results[key])
+            for key in sorted(results.keys())
+            if key.startswith("result")
+        ]
+
+        prediction_array = np.stack(result_arrays, axis=-1)
+
+        if binarize:
+            prediction_array = binarize_hard(prediction_array)
+
+        return prediction_array
+
+    except requests.exceptions.RequestException as e:
+        raise RuntimeError(f"Request failed: {e}") from e
+    except json.JSONDecodeError as e:
+        raise ValueError(f"JSON decode error: {e}") from e
 
 
 def _predict_poly(
-    polygon_points: list,
+    polygon_points: list[Any],
     model: Model,
     model_type: str,
     eta: float = 0.5,
-) -> list:
+) -> list[Any]:
     """
-    Predict the nanofabrication outcome for a list of polygons.
+    Predict the nanofabrication outcome for a geometry (list of polygons).
 
-    This function sends polygon data to the server, which uses a specified machine
-    learning model to predict the outcome of the nanofabrication process.
+    This function sends the device array to a serverless prediction service, which uses
+    a specified machine learning model to predict the outcome of the nanofabrication
+    process.
 
     Parameters
     ----------
@@ -39,13 +136,12 @@ def _predict_poly(
         List of polygon points, where each polygon is a list of [x, y] coordinates.
     model : Model
         The model to use for prediction, representing a specific fabrication process and
-        dataset. This model encapsulates details about the fabrication foundry, process,
-        material, technology, thickness, and sidewall presence, as defined in
-        `models.py`. Each model is associated with a version and dataset that detail its
-        creation and the data it was trained on, ensuring the prediction is tailored to
-        specific fabrication parameters.
+        dataset. This model encapsulates details about the fabrication foundry and
+        process, as defined in `models.py`. Each model is associated with a version and
+        dataset that detail its creation and the data it was trained on, ensuring the
+        prediction is tailored to specific fabrication parameters.
     model_type : str
-        The type of model to use ('p' for prediction, 'c' for correction).
+        The type of model to use (e.g., 'p' for prediction or 'c' for correction).
     eta : float
         The threshold value for binarization. Defaults to 0.5. Because intermediate
         values cannot be preserved in the polygon data, the predicted polygons are
@@ -54,17 +150,15 @@ def _predict_poly(
     Returns
     -------
     list
-        List of predicted polygon points with channel information. Each polygon is a
-        dict with 'points' (list of coordinates) and 'channel' (int) keys.
+        List of predicted polygon points with level information. Each polygon is a dict
+        with 'points' (list of coordinates) and 'level' (int) keys.
 
     Raises
     ------
-    ValueError
-        If the server returns an error or empty response.
-    requests.exceptions.RequestException
+    RuntimeError
         If the request to the prediction service fails.
-    json.JSONDecodeError
-        If the response cannot be parsed as JSON.
+    ValueError
+        If the server returns an error or invalid response.
     """
     predict_data = {
         "polygons": polygon_points,
@@ -99,11 +193,89 @@ def _predict_poly(
             return []
 
     except requests.exceptions.RequestException as e:
-        print(f"Request failed: {str(e)}")
-        raise
+        raise RuntimeError(f"Request failed: {e}") from e
     except json.JSONDecodeError as e:
-        print(f"JSON decode error: {str(e)}")
-        raise
+        raise ValueError(f"JSON decode error: {e}") from e
+
+
+def predict_gds(
+    gds_path: str,
+    cell_name: str,
+    model: Model,
+    model_type: str,
+    gds_layer: tuple[int, int] = (1, 0),
+    eta: float = 0.5,
+    output_path: str | None = None,
+) -> None:
+    """
+    Predict the nanofabrication outcome for a GDS file and cell.
+
+    This function loads a GDS file, extracts the specified cell, and predicts the
+    nanofabrication outcome using the specified model. The predicted cell is
+    automatically added to the original GDS library and the file is written to the
+    specified output path (or overwrites the original if no output path is provided).
+
+    Parameters
+    ----------
+    gds_path : str
+        The file path to the GDS file.
+    cell_name : str
+        The name of the cell within the GDS file to predict.
+    model : Model
+        The model to use for prediction, representing a specific fabrication process and
+        dataset. This model encapsulates details about the fabrication foundry and
+        process, as defined in `models.py`. Each model is associated with a version and
+        dataset that detail its creation and the data it was trained on, ensuring the
+        prediction is tailored to specific fabrication parameters.
+    model_type : str
+        The type of model to use (e.g., 'p' for prediction or 'c' for correction).
+    gds_layer : tuple[int, int]
+        The layer and datatype to use within the GDS file. Defaults to (1, 0).
+    eta : float
+        The threshold value for binarization. Defaults to 0.5. Because intermediate
+        values cannot be preserved in the polygon data, the predicted polygons are
+        binarized using a threshold value of eta.
+    output_path : str, optional
+        The file path where the updated GDS file will be written. If None, the
+        original file will be overwritten. Defaults to None.
+
+    Raises
+    ------
+    RuntimeError
+        If the request to the prediction service fails.
+    ValueError
+        If the GDS file cannot be read, the specified cell is not found, or the server
+        returns an error or invalid response.
+    """
+    gdstk_library = gdstk.read_gds(gds_path)
+    cells = [
+        cell
+        for cell in gdstk_library.cells
+        if isinstance(cell, gdstk.Cell) and cell.name == cell_name
+    ]
+    if not cells:
+        raise ValueError(f"Cell '{cell_name}' not found in GDS file")
+    gdstk_cell = cells[0]
+
+    predicted_cell = predict_gdstk(
+        gdstk_cell=gdstk_cell,
+        model=model,
+        model_type=model_type,
+        gds_layer=gds_layer,
+        eta=eta,
+    )
+
+    base_name = predicted_cell.name
+    counter = 1
+    while predicted_cell.name in [cell.name for cell in gdstk_library.cells]:
+        predicted_cell.name = f"{base_name}_{counter}"
+        counter += 1
+
+    gdstk_library.add(predicted_cell)
+
+    write_path = output_path if output_path is not None else gds_path
+    print(f"Writing to {write_path}")
+    gdstk_library.write_gds(write_path, max_points=8190)
 
 
 def predict_gdstk(
@@ -116,8 +288,8 @@ def predict_gdstk(
     """
     Predict the nanofabrication outcome of a gdstk cell using a specified model.
 
-    This function extracts polygons from a gdstk cell, sends them to the prediction
-    server, and returns a new cell containing the predicted polygons.
+    This function extracts polygons from a gdstk cell, sends them to a serverless
+    prediction service, and returns a new cell containing the predicted polygons.
 
     Parameters
     ----------
@@ -125,13 +297,12 @@ def predict_gdstk(
         The gdstk.Cell object containing polygons to predict.
     model : Model
         The model to use for prediction, representing a specific fabrication process and
-        dataset. This model encapsulates details about the fabrication foundry, process,
-        material, technology, thickness, and sidewall presence, as defined in
-        `models.py`. Each model is associated with a version and dataset that detail its
-        creation and the data it was trained on, ensuring the prediction is tailored to
-        specific fabrication parameters.
+        dataset. This model encapsulates details about the fabrication foundry and
+        process, as defined in `models.py`. Each model is associated with a version and
+        dataset that detail its creation and the data it was trained on, ensuring the
+        prediction is tailored to specific fabrication parameters.
     model_type : str
-        The type of model to use ('p' for prediction, 'c' for correction).
+        The type of model to use (e.g., 'p' for prediction or 'c' for correction).
     gds_layer : tuple[int, int]
         The layer and datatype to use within the GDSTK cell. Defaults to (1, 0).
     eta : float
@@ -149,14 +320,17 @@ def predict_gdstk(
 
     Raises
     ------
+    RuntimeError
+        If the request to the prediction service fails.
     ValueError
-        If no polygons are found in the specified layer.
+        If no polygons are found in the specified layer, or the server returns an error
+        or invalid response.
     """
     polygons = gdstk_cell.get_polygons(layer=gds_layer[0], datatype=gds_layer[1])
     if not polygons:
         raise ValueError("No polygons found in the specified layer")
 
-    polygon_points = [polygon.points.tolist() for polygon in polygons]
+    polygon_points = [polygon.points.tolist() for polygon in polygons]  # pyright: ignore[reportAttributeAccessIssue]
 
     predicted_polygon_data = _predict_poly(
         polygon_points=polygon_points,
@@ -165,7 +339,8 @@ def predict_gdstk(
         eta=eta,
     )
 
-    result_cell = gdstk.Cell(f"{gdstk_cell.name}_predicted")
+    suffix = "corrected" if model_type == "c" else "predicted"
+    result_cell = gdstk.Cell(f"{gdstk_cell.name}_{suffix}")
 
     polygons_by_channel = {}
     for polygon_data in predicted_polygon_data:
@@ -183,129 +358,82 @@ def predict_gdstk(
 
         for points in points_list:
             points_array = np.array(points)
-            polygon = gdstk.Polygon(points_array, layer=layer, datatype=datatype)
+            polygon = gdstk.Polygon(points_array, layer=layer, datatype=datatype)  # pyright: ignore[reportArgumentType]
             result_cell.add(polygon)
 
     return result_cell
 
 
-def predict_array(
-    device_array: np.ndarray,
-    model: Model,
-    model_type: str,
-    binarize: bool,
-    gpu: bool = False,
-) -> np.ndarray:
+def _predict_array_with_grad(
+    device_array: npt.NDArray[Any], model: Model
+) -> tuple[npt.NDArray[Any], npt.NDArray[Any]]:
     """
-    Predict the nanofabrication outcome of a device array using a specified model.
+    Predict the nanofabrication outcome of a device array and compute its gradient.
 
-    This function sends the device array to a serverless prediction service, which uses
-    a specified machine learning model to predict the outcome of the nanofabrication
-    process. The prediction can be performed on a GPU if specified.
+    This function predicts the outcome of the nanofabrication process for a given
+    device array using a specified model. It also computes the gradient of the
+    prediction with respect to the input device array.
 
     Parameters
     ----------
     device_array : np.ndarray
-        A 2D array representing the planar geometry of the device. This array undergoes
-        various transformations to predict the nanofabrication process.
+        A 2D array representing the planar geometry of the device.
     model : Model
-        The model to use for prediction, representing a specific fabrication process and
-        dataset. This model encapsulates details about the fabrication foundry, process,
-        material, technology, thickness, and sidewall presence, as defined in
-        `models.py`. Each model is associated with a version and dataset that detail its
-        creation and the data it was trained on, ensuring the prediction is tailored to
-        specific fabrication parameters.
-    model_type : str
-        The type of model to use (e.g., 'p' for prediction, 'c' for correction, or 's'
-        for SEMulate).
-    binarize : bool
-        If True, the predicted device geometry will be binarized using a threshold
-        method. This is useful for converting probabilistic predictions into binary
-        geometries.
-    gpu : bool
-        If True, the prediction will be performed on a GPU. Defaults to False. Note: The
-        GPU option has more startup overhead and will take longer for small devices, but
-        will be faster for larger devices.
+        The model to use for prediction, representing a specific fabrication process.
 
     Returns
     -------
-    np.ndarray
-        The predicted output array.
+    tuple[np.ndarray, np.ndarray]
+        The predicted output array and gradient array.
 
     Raises
     ------
     RuntimeError
         If the request to the prediction service fails.
+    ValueError
+        If the server returns an error or invalid response.
     """
     headers = _prepare_headers()
-    predict_data = _prepare_predict_data(device_array, model, model_type, binarize)
-    endpoint_url = (
-        f"{BASE_ENDPOINT_URL}-gpu-v{ENDPOINT_VERSION}.modal.run"
-        if gpu
-        else f"{BASE_ENDPOINT_URL}-v{ENDPOINT_VERSION}.modal.run"
-    )
+    predict_data = {
+        "device_array": _encode_array(np.squeeze(device_array)),
+        "model": model.to_json(),
+        "model_type": "p",
+        "binary": False,
+    }
+    endpoint_url = f"{BASE_ENDPOINT_URL}-with-grad-v{ENDPOINT_VERSION}.modal.run"
 
     try:
-        with requests.post(
-            endpoint_url,
-            data=json.dumps(predict_data),
-            headers=headers,
-            stream=True,
-        ) as response:
-            response.raise_for_status()
-            result = _process_response(response, model_type, binarize)
-            if result is None:
-                raise RuntimeError("No prediction result received.")
-            return result
-    except requests.RequestException as e:
-        raise RuntimeError(f"Request failed: {e}") from e
-
-
-def _predict_array_with_grad(
-    device_array: np.ndarray, model: Model
-) -> tuple[np.ndarray, np.ndarray]:
-    """
-    Predict the nanofabrication outcome of a device array and compute its gradient.
+        response = requests.post(
+            endpoint_url, data=json.dumps(predict_data), headers=headers
+        )
+        response.raise_for_status()
 
-    This function predicts the outcome of the nanofabrication process for a given
-    device array using a specified model. It also computes the gradient of the
-    prediction with respect to the input device array.
+        if not response.content:
+            raise ValueError("Empty response received from server")
 
-    Notes
-    -----
-    This function is currently not used in the main `predict_array` function as
-    the main `predict_array` function (e.g., GPU support and progress bar) for now.
+        response_data = response.json()
 
-    Parameters
-    ----------
-    device_array : np.ndarray
-        A 2D array representing the planar geometry of the device.
-    model : Model
-        The model to use for prediction, representing a specific fabrication process.
+        if "error" in response_data:
+            raise ValueError(f"Prediction error: {response_data['error']}")
 
-    Returns
-    -------
-    tuple[np.ndarray, np.ndarray]
-        The predicted output array and gradient array.
-    """
-    headers = _prepare_headers()
-    predict_data = _prepare_predict_data(device_array, model, "p", False)
-    endpoint_url = f"{BASE_ENDPOINT_URL}-with-grad-v{ENDPOINT_VERSION}.modal.run"
+        prediction_array = _decode_array(response_data["prediction_array"])
+        gradient_array = _decode_array(response_data["gradient_array"])
+        gradient_min = response_data["gradient_min"]
+        gradient_max = response_data["gradient_max"]
+        gradient_range = gradient_max - gradient_min
+        gradient_array = gradient_array * gradient_range + gradient_min
+        return (prediction_array, gradient_array)
 
-    response = requests.post(
-        endpoint_url, data=json.dumps(predict_data), headers=headers
-    )
-    prediction_array = _decode_array(response.json()["prediction_array"])
-    gradient_array = _decode_array(response.json()["gradient_array"])
-    gradient_min = response.json()["gradient_min"]
-    gradient_max = response.json()["gradient_max"]
-    gradient_range = gradient_max - gradient_min
-    gradient_array = gradient_array * gradient_range + gradient_min
-    return (prediction_array, gradient_array)
+    except requests.exceptions.RequestException as e:
+        raise RuntimeError(f"Request failed: {e}") from e
+    except json.JSONDecodeError as e:
+        raise ValueError(f"JSON decode error: {e}") from e
 
 
 @primitive
-def predict_array_with_grad(device_array: np.ndarray, model: Model) -> np.ndarray:
+def predict_array_with_grad(
+    device_array: npt.NDArray[Any], model: Model
+) -> npt.NDArray[Any]:
     """
     Predict the nanofabrication outcome of a device array and compute its gradient.
 
@@ -325,15 +453,24 @@ def predict_array_with_grad(device_array: np.ndarray, model: Model) -> np.ndarra
     -------
     np.ndarray
         The predicted output array.
+
+    Raises
+    ------
+    RuntimeError
+        If the request to the prediction service fails.
+    ValueError
+        If the server returns an error or invalid response.
     """
     prediction_array, gradient_array = _predict_array_with_grad(
         device_array=device_array, model=model
     )
-    predict_array_with_grad.gradient_array = gradient_array  # type: ignore
+    predict_array_with_grad.gradient_array = gradient_array  # pyright: ignore[reportFunctionMemberAccess]
     return prediction_array
 
 
-def predict_array_with_grad_vjp(ans: np.ndarray, device_array: np.ndarray, *args):
+def predict_array_with_grad_vjp(
+    ans: npt.NDArray[Any], device_array: npt.NDArray[Any], *args: Any
+) -> Any:
     """
     Define the vector-Jacobian product (VJP) for the prediction function.
 
@@ -351,10 +488,10 @@ def predict_array_with_grad_vjp(ans: np.ndarray, device_array: np.ndarray, *args
     function
         A function that computes the VJP given an upstream gradient `g`.
     """
-    grad_x = predict_array_with_grad.gradient_array  # type: ignore
+    grad_x = predict_array_with_grad.gradient_array  # pyright: ignore[reportFunctionMemberAccess]
 
-    def vjp(g: np.ndarray) -> np.ndarray:
-        return g * grad_x
+    def vjp(g: npt.NDArray[Any]) -> npt.NDArray[Any]:
+        return g * grad_x  # type: ignore[no-any-return]
 
     return vjp
 
@@ -362,7 +499,7 @@ def predict_array_with_grad_vjp(ans: np.ndarray, device_array: np.ndarray, *args
 defvjp(predict_array_with_grad, predict_array_with_grad_vjp)
 
 
-def _encode_array(array):
+def _encode_array(array: npt.NDArray[Any]) -> str:
     """Encode an ndarray as a base64 encoded image for transmission."""
     image = Image.fromarray(np.uint8(array * 255))
     buffered = io.BytesIO()
@@ -371,15 +508,15 @@ def _encode_array(array):
     return encoded_png
 
 
-def _decode_array(encoded_png):
+def _decode_array(encoded_png: str) -> npt.NDArray[Any]:
     """Decode a base64 encoded image and return an ndarray."""
     binary_data = base64.b64decode(encoded_png)
     image = Image.open(io.BytesIO(binary_data))
-    return np.array(image) / 255
+    return np.array(image) / 255  # type: ignore[no-any-return]
 
 
-def _read_tokens():
-    """Read access and refresh tokens from the configuration file."""
+def _prepare_headers() -> dict[str, str]:
+    """Prepare HTTP headers for a server request."""
     token_file_path = os.path.expanduser("~/.prefab.toml")
     try:
         with open(token_file_path) as file:
@@ -388,119 +525,14 @@ def _read_tokens():
             refresh_token = tokens.get("refresh_token")
             if not access_token or not refresh_token:
                 raise ValueError("Tokens not found in the configuration file.")
-            return access_token, refresh_token
+            return {
+                "Authorization": f"Bearer {access_token}",
+                "X-Refresh-Token": refresh_token,
+            }
     except FileNotFoundError:
         raise FileNotFoundError(
             "Could not validate user.\n"
-            "Please update prefab using: pip install --upgrade prefab.\n"
-            "Signup/login and generate a new token.\n"
-            "See https://docs.prefabphotonics.com/."
+            + "Please update prefab using: pip install --upgrade prefab.\n"
+            + "Signup/login and generate a new token.\n"
+            + "See https://docs.prefabphotonics.com/."
         ) from None
-
-
-def _prepare_headers():
-    """Prepare HTTP headers for the request."""
-    access_token, refresh_token = _read_tokens()
-    return {
-        "Authorization": f"Bearer {access_token}",
-        "X-Refresh-Token": refresh_token,
-    }
-
-
-def _prepare_predict_data(device_array, model, model_type, binarize):
-    """Prepare the data payload for the prediction request."""
-    return {
-        "device_array": _encode_array(np.squeeze(device_array)),
-        "model": model.to_json(),
-        "model_type": model_type,
-        "binary": binarize,
-    }
-
-
-def _process_response(response, model_type, binarize):
-    """Process the streaming response from the prediction request."""
-    event_type = None
-    model_descriptions = {
-        "p": "Prediction",
-        "c": "Correction",
-        "s": "SEMulate",
-    }
-    progress_bar = tqdm(
-        total=100,
-        desc=model_descriptions.get(model_type, "Processing"),
-        unit="%",
-        colour="green",
-        bar_format="{l_bar}{bar:30}{r_bar}{bar:-10b}",
-    )
-
-    for line in response.iter_lines():
-        if line:
-            decoded_line = line.decode("utf-8").strip()
-            if decoded_line.startswith("event:"):
-                event_type = decoded_line.split(":", 1)[1].strip()
-            elif decoded_line.startswith("data:"):
-                data_content = _parse_data_line(decoded_line)
-                result = _handle_event(event_type, data_content, progress_bar, binarize)
-                if result is not None:
-                    progress_bar.close()
-                    return result
-    progress_bar.close()
-
-
-def _parse_data_line(decoded_line):
-    """Parse a data line from the response stream."""
-    data_line = decoded_line.split(":", 1)[1].strip()
-    try:
-        return json.loads(data_line)
-    except json.JSONDecodeError:
-        raise ValueError(f"Failed to decode JSON: {data_line}") from None
-
-
-def _handle_event(event_type, data_content, progress_bar, binarize):
-    """Handle different types of events received from the server."""
-    if event_type == "progress":
-        _update_progress(progress_bar, data_content)
-    elif event_type == "result":
-        return _process_result(data_content, binarize)
-    elif event_type == "end":
-        print("Stream ended.")
-    elif event_type == "auth":
-        _update_tokens(data_content.get("auth", {}))
-    elif event_type == "error":
-        raise ValueError(f"{data_content['error']}")
-
-
-def _update_progress(progress_bar, data_content):
-    """Update the progress bar based on the progress event."""
-    progress = round(100 * data_content.get("progress", 0))
-    progress_bar.update(progress - progress_bar.n)
-
-
-def _process_result(data_content, binarize):
-    """Process the result event and return the prediction."""
-    results = [
-        _decode_array(data_content[key])
-        for key in sorted(data_content.keys())
-        if key.startswith("result")
-    ]
-    if results:
-        prediction = np.stack(results, axis=-1)
-        if binarize:
-            prediction = binarize_hard(prediction)
-        return prediction
-
-
-def _update_tokens(auth_data):
-    """Update tokens if new tokens are provided in the auth event."""
-    new_access_token = auth_data.get("new_access_token")
-    new_refresh_token = auth_data.get("new_refresh_token")
-    if new_access_token and new_refresh_token:
-        prefab_file_path = os.path.expanduser("~/.prefab.toml")
-        with open(prefab_file_path, "w", encoding="utf-8") as toml_file:
-            toml.dump(
-                {
-                    "access_token": new_access_token,
-                    "refresh_token": new_refresh_token,
-                },
-                toml_file,
-            )