| """Loader + server helper for the CBC Manufacturing RUL LSTM reference model. | |
| The model is a torch LSTM, so (unlike the other CBC reference models) it cannot be loaded | |
| with joblib alone — you need this class definition plus the saved state_dict and the | |
| normalization meta. Usage: | |
| from model import load_model, predict_rul | |
| model, meta = load_model(".") # dir holding manufacturing_lstm.pt + _meta.joblib | |
| rul = predict_rul(model, meta, raw_window) # raw_window: list[ list[float] ], shape window x n_sensors | |
| """ | |
| from __future__ import annotations | |
| import json | |
| from pathlib import Path | |
| import joblib | |
| import numpy as np | |
| import torch | |
| from torch import nn | |
| class LSTMReg(nn.Module): | |
| def __init__(self, n_feat: int, hidden: int = 64) -> None: | |
| super().__init__() | |
| self.lstm = nn.LSTM(n_feat, hidden, batch_first=True) | |
| self.head = nn.Linear(hidden, 1) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| out, _ = self.lstm(x) | |
| return self.head(out[:, -1, :]).squeeze(-1) | |
| def load_model(model_dir: str = "."): | |
| """Return (model, meta). model_dir holds manufacturing_lstm.pt + manufacturing_meta.joblib.""" | |
| d = Path(model_dir) | |
| meta = joblib.load(d / "manufacturing_meta.joblib") | |
| model = LSTMReg(meta["n_features"], meta["hidden"]) | |
| model.load_state_dict(torch.load(d / "manufacturing_lstm.pt", map_location="cpu")) | |
| model.eval() | |
| return model, meta | |
| def predict_rul(model, meta, raw_window) -> float: | |
| """raw_window: window x n_sensors of RAW sensor values (in meta['sensors'] order).""" | |
| sensors = meta["sensors"] | |
| mn = np.array([meta["mins"][s] for s in sensors], dtype=np.float32) | |
| mx = np.array([meta["maxs"][s] for s in sensors], dtype=np.float32) | |
| arr = np.asarray(raw_window, dtype=np.float32) | |
| if arr.shape != (meta["window"], len(sensors)): | |
| raise ValueError(f"expected window {meta['window']} x {len(sensors)} sensors, got {arr.shape}") | |
| norm = (arr - mn) / (mx - mn) | |
| with torch.no_grad(): | |
| pred = model(torch.from_numpy(norm[None, :, :]).float()).numpy() | |
| return float(np.clip(pred, 0, None)[0]) | |
| if __name__ == "__main__": | |
| m, meta = load_model(Path(__file__).parent) | |
| sample = json.load(open(Path(__file__).parent / "sample_input.json")) | |
| print("predicted RUL:", round(predict_rul(m, meta, sample["window"]), 2), | |
| "| true RUL:", sample.get("true_rul")) | |