export_models.py

from __future__ import annotations

import json
import os
import random
import shutil
import sys
from importlib.metadata import version
from pathlib import Path

os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0")

import numpy as np
import torch
from torch import nn

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT / "src"))

from forge_lab.models.numpy_runtime import (  # noqa: E402
    DEFAULT_CENTER,
    DEFAULT_SCALE,
    WINDOW_SIZE,
)
from forge_lab.models.pytorch_model import TinyDriftNet  # noqa: E402

ARTIFACTS = ROOT
MODEL_VERSION = "forge-tiny-drift-v0.1"


def synthetic_windows(count: int = 2048, seed: int = 17) -> tuple[np.ndarray, np.ndarray]:
    rng = np.random.default_rng(seed)
    windows = np.empty((count, WINDOW_SIZE, 2), dtype=np.float32)
    labels = np.empty(count, dtype=np.float32)
    x = np.arange(WINDOW_SIZE, dtype=np.float32)
    for index in range(count):
        slope = rng.uniform(-0.01, 0.13)
        force = (
            rng.uniform(96.0, 104.0)
            + slope * x
            + rng.normal(0.0, rng.uniform(0.03, 0.18), WINDOW_SIZE)
        )
        deviation = (
            0.07 + np.maximum(0.0, force - 101.0) * 0.011 + rng.normal(0.0, 0.002, WINDOW_SIZE)
        )
        windows[index, :, 0] = force
        windows[index, :, 1] = deviation
        shift = force[-10:].mean() - force[:10].mean()
        labels[index] = float(slope >= 0.07 and shift >= 2.0)
    return windows, labels


def train_model() -> tuple[TinyDriftNet, dict]:
    random.seed(17)
    np.random.seed(17)
    torch.manual_seed(17)
    x, y = synthetic_windows()
    model = TinyDriftNet()
    optimizer = torch.optim.AdamW(model.classifier.parameters(), lr=0.04, weight_decay=0.001)
    loss_fn = nn.BCELoss()
    x_tensor = torch.from_numpy(x)
    y_tensor = torch.from_numpy(y)
    model.train()
    for _ in range(240):
        optimizer.zero_grad(set_to_none=True)
        loss = loss_fn(model(x_tensor), y_tensor)
        loss.backward()
        optimizer.step()
    model.eval()
    with torch.no_grad():
        probabilities = model(x_tensor)
        predictions = probabilities >= 0.5
        labels = y_tensor.bool()
        accuracy = float((predictions == labels).float().mean())
        false_positives = ((predictions == 1) & (labels == 0)).sum()
        false_negatives = ((predictions == 0) & (labels == 1)).sum()
        false_positive_rate = float(false_positives / (labels == 0).sum())
        false_negative_rate = float(false_negatives / (labels == 1).sum())
    metrics = {
        "synthetic_examples": len(x),
        "accuracy": round(accuracy, 6),
        "false_positive_rate": round(false_positive_rate, 6),
        "false_negative_rate": round(false_negative_rate, 6),
        "training_seed": 17,
    }
    return model, metrics


def export_pytorch_and_onnx(model: TinyDriftNet) -> None:
    torch.save(
        {
            "model_version": MODEL_VERSION,
            "state_dict": model.state_dict(),
        },
        ARTIFACTS / "tiny_drift_pytorch.pt",
    )
    example = torch.zeros(2, WINDOW_SIZE, 2, dtype=torch.float32)
    batch = torch.export.Dim("batch", min=1, max=512)
    onnx_program = torch.onnx.export(
        model,
        (example,),
        input_names=["telemetry"],
        output_names=["probability"],
        dynamic_shapes=({0: batch},),
        dynamo=True,
        verify=True,
        external_data=False,
    )
    onnx_program.save(ARTIFACTS / "tiny_drift.onnx")


def export_tensorflow_and_litert(weight: np.ndarray, bias: float) -> None:
    import tensorflow as tf

    from forge_lab.models.tensorflow_model import TensorFlowDriftModel

    model = TensorFlowDriftModel(weight, bias)
    saved_model = ARTIFACTS / "tensorflow_saved_model"
    if saved_model.exists():
        shutil.rmtree(saved_model)
    concrete = model.__call__.get_concrete_function()
    tf.saved_model.save(model, saved_model, signatures={"serving_default": concrete})
    converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete], model)
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    (ARTIFACTS / "tiny_drift.tflite").write_bytes(converter.convert())


def write_metadata(model: TinyDriftNet, metrics: dict) -> None:
    weight = model.classifier.weight.detach().cpu().numpy().reshape(-1)
    bias = float(model.classifier.bias.detach().cpu().item())
    weights = {
        "model_version": MODEL_VERSION,
        "feature_center": DEFAULT_CENTER.tolist(),
        "feature_scale": DEFAULT_SCALE.tolist(),
        "linear_weight": weight.tolist(),
        "linear_bias": bias,
    }
    (ARTIFACTS / "weights.json").write_text(json.dumps(weights, indent=2) + "\n")
    manifest = {
        "model_version": MODEL_VERSION,
        "input": {"name": "telemetry", "shape": ["batch", WINDOW_SIZE, 2], "dtype": "float32"},
        "channels": ["peak_force_kn", "part_deviation_mm"],
        "outputs": ["probability"],
        "runtimes": ["numpy", "pytorch", "onnxruntime", "tensorflow", "litert"],
        "framework_versions": {
            "python": sys.version.split()[0],
            "numpy": np.__version__,
            "pytorch": torch.__version__,
            "pytorch_cuda_build": torch.version.cuda,
            "onnx": version("onnx"),
            "onnxruntime": version("onnxruntime"),
            "tensorflow": version("tensorflow-cpu"),
            "litert": version("ai-edge-litert"),
        },
        "metrics": metrics,
        "data": "deterministic synthetic telemetry; no plant data",
        "safety": "L0 decision support only; output cannot actuate equipment",
        "cuda_validation": "not executed; no GPU or nvcc in the build workspace",
    }
    (ARTIFACTS / "model_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")


def main() -> None:
    ARTIFACTS.mkdir(parents=True, exist_ok=True)
    model, metrics = train_model()
    write_metadata(model, metrics)
    weight = model.classifier.weight.detach().cpu().numpy().reshape(-1)
    bias = float(model.classifier.bias.detach().cpu().item())
    export_pytorch_and_onnx(model)
    export_tensorflow_and_litert(weight, bias)
    print(json.dumps(metrics, indent=2))


if __name__ == "__main__":
    main()