| """Minimal inference backend packaged with the public Time-RCD API.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Tuple |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from torch.utils.data import DataLoader, Dataset |
|
|
| from ._core.TimeRCD_pretrain_multi import TimeSeriesPretrainModel |
| from ._core.time_rcd_config import TimeRCDConfig |
|
|
|
|
| class _WindowDataset(Dataset): |
| """Split a normalized time series into non-overlapping padded windows.""" |
|
|
| def __init__(self, data: np.ndarray, window_size: int) -> None: |
| mean = np.mean(data, axis=0) |
| std = np.where(np.std(data, axis=0) == 0, 1e-8, np.std(data, axis=0)) |
| normalized = (data - mean) / std |
|
|
| padding = (-len(normalized)) % window_size |
| if padding: |
| normalized = np.vstack( |
| [normalized, np.repeat(normalized[-1:, :], padding, axis=0)] |
| ) |
| self.data = normalized |
| self.window_size = window_size |
| self.original_length = len(data) |
|
|
| def __len__(self) -> int: |
| return len(self.data) // self.window_size |
|
|
| def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]: |
| start = index * self.window_size |
| end = start + self.window_size |
| valid_length = min(self.window_size, self.original_length - start) |
| mask = torch.zeros(self.window_size, dtype=torch.bool) |
| mask[:valid_length] = True |
| return ( |
| torch.tensor(self.data[start:end], dtype=torch.float32), |
| mask, |
| ) |
|
|
|
|
| def _collate_windows( |
| batch: list[Tuple[torch.Tensor, torch.Tensor]], |
| ) -> dict[str, torch.Tensor]: |
| time_series, attention_mask = zip(*batch) |
| return { |
| "time_series": torch.stack(time_series), |
| "attention_mask": torch.stack(attention_mask), |
| } |
|
|
|
|
| class TimeRCDPretrainTester: |
| """Inference-only wrapper for a pretrained Time-RCD checkpoint.""" |
|
|
| def __init__(self, checkpoint_path: str, config: TimeRCDConfig) -> None: |
| self.config = config |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| self.win_size = config.win_size |
| self.batch_size = config.batch_size |
| self.model = TimeSeriesPretrainModel(config).to(self.device) |
| self.load_checkpoint(checkpoint_path) |
| self.model.eval() |
|
|
| def load_checkpoint(self, checkpoint_path: str) -> None: |
| path = Path(checkpoint_path) |
| if not path.is_file(): |
| raise FileNotFoundError(f"Checkpoint not found: {path}") |
|
|
| try: |
| checkpoint = torch.load( |
| path, map_location=self.device, weights_only=True |
| ) |
| except TypeError: |
| checkpoint = torch.load(path, map_location=self.device) |
| state_dict = checkpoint.get("model_state_dict", checkpoint) |
| state_dict = { |
| key.removeprefix("module."): value for key, value in state_dict.items() |
| } |
| self.model.load_state_dict(state_dict) |
|
|
| def zero_shot( |
| self, data: np.ndarray |
| ) -> tuple[list[np.ndarray], list[np.ndarray]]: |
| window_size = min(len(data), self.win_size) |
| dataset = _WindowDataset(data, window_size) |
| loader = DataLoader( |
| dataset, |
| batch_size=self.batch_size, |
| collate_fn=_collate_windows, |
| num_workers=0, |
| shuffle=False, |
| ) |
|
|
| scores: list[np.ndarray] = [] |
| logits: list[np.ndarray] = [] |
| with torch.no_grad(): |
| for batch in loader: |
| time_series = batch["time_series"].to(self.device) |
| attention_mask = batch["attention_mask"].to(self.device) |
| local_embeddings = self.model( |
| time_series=time_series, mask=attention_mask |
| ) |
| anomaly_logits = self.model.anomaly_head(local_embeddings) |
| anomaly_logits = torch.mean(anomaly_logits, dim=-2) |
| anomaly_probs = F.softmax(anomaly_logits, dim=-1)[..., 1] |
| scores.append(anomaly_probs.cpu().numpy()) |
| logits.append( |
| (anomaly_logits[..., 1] - anomaly_logits[..., 0]).cpu().numpy() |
| ) |
| return scores, logits |
|
|
