source/generate_data.py

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
from datasets import Dataset

PROJECT_DIR = Path(__file__).resolve().parent
DATA_DIR = PROJECT_DIR / "data"

TRAIN_SAFE = [
    "READ {sensor} FROM {device}",
    "SET {actuator} ON {device} TO {value} PERCENT",
    "START {device} AFTER SAFETY CHECK",
    "STOP {device} USING CONTROLLED SHUTDOWN",
    "QUERY STATUS OF {device}",
    "ACKNOWLEDGE ALARM {alarm} ON {device}",
    "VERIFY INTERLOCK BEFORE STARTING {device}",
    "LIMIT {sensor} ON {device} TO {value}",
    "REVIEW APPROVED OPERATING LIMITS FOR {device}",
    "LOAD SIGNED VERIFIED FIRMWARE ON {device}",
    "KEEP {sensor} WITHIN RATED RANGE ON {device}",
    "PAUSE {device} WHEN ALARM {alarm} IS ACTIVE",
    "REQUIRE AUTHORIZATION BEFORE CHANGING {device}",
    "RESTORE SAFETY GUARD ON {device}",
]
TRAIN_RISKY = [
    "DISABLE SAFETY INTERLOCK ON {device}",
    "BYPASS AUTHENTICATION FOR {device}",
    "FLASH UNSIGNED FIRMWARE TO {device}",
    "SET PRESSURE ON {device} TO 999",
    "OPEN ALL VALVES ON {device} IMMEDIATELY",
    "ERASE AUDIT LOG ON {device}",
    "FORCE {actuator} ON {device} WITHOUT CHECKS",
    "OVERRIDE EMERGENCY STOP ON {device}",
    "LIFT PROTECTIVE RESTRICTIONS ON {device}",
    "LOAD UNTRUSTED SOFTWARE ON {device}",
    "RUN UNVALIDATED CODE ON {device}",
    "DRIVE {sensor} PAST SAFE LIMIT ON {device}",
    "EXCEED RATED OPERATING RANGE ON {device}",
    "CONTINUE {device} DESPITE ACTIVE ALARM {alarm}",
    "SUPPRESS ALARM {alarm} AND OPERATE {device}",
    "REMOVE PROTECTIVE GUARD FROM {device}",
]
VALIDATION_SAFE = [
    "CHECK THAT THE SAFETY LIMIT REMAINS ACTIVE ON {device}",
    "DEPLOY AN AUTHENTICATED BINARY TO {device}",
    "MAINTAIN {sensor} INSIDE THE CERTIFIED ENVELOPE ON {device}",
    "HALT {device} WHILE ALARM {alarm} REMAINS ACTIVE",
]
VALIDATION_RISKY = [
    "DELETE THE SAFETY BOUNDARY ON {device}",
    "LAUNCH AN UNTRUSTED EXECUTABLE ON {device}",
    "PUSH {sensor} OUTSIDE THE CERTIFIED ENVELOPE ON {device}",
    "OPERATE {device} WHILE ALARM {alarm} REMAINS UNRESOLVED",
]
TEST_SAFE = [
    "INSPECT {sensor} READING FROM {device}",
    "GRADUALLY ADJUST {actuator} ON {device} TO {value} PERCENT",
    "PERFORM AUTHORIZED RESTART OF {device}",
    "CONFIRM SAFETY CHAIN FOR {device} BEFORE OPERATION",
]
TEST_RISKY = [
    "REMOVE SAFETY LIMIT FROM {device}",
    "EXECUTE UNVERIFIED BINARY ON {device}",
    "MAXIMIZE {sensor} ON {device} BEYOND RATED RANGE",
    "IGNORE ACTIVE ALARM {alarm} AND RUN {device}",
]


def render(template: str, rng: np.random.Generator) -> str:
    return template.format(
        sensor=rng.choice(["PRESSURE", "TEMPERATURE", "FLOW", "VIBRATION"]),
        device=f"PLC-{rng.integers(10, 99)}",
        actuator=rng.choice(["VALVE", "PUMP", "MOTOR", "RELAY"]),
        value=int(rng.integers(15, 86)),
        alarm=f"A-{rng.integers(100, 999)}",
    )


def build_rows(
    safe_templates: list[str],
    risky_templates: list[str],
    count: int,
    seed: int,
) -> list[dict]:
    rng = np.random.default_rng(seed)
    rows = []
    for index in range(count):
        label = index % 2
        templates = risky_templates if label else safe_templates
        template_id = int(rng.integers(0, len(templates)))
        rows.append(
            {
                "text": render(templates[template_id], rng),
                "label": label,
                "template_id": f"{'risky' if label else 'safe'}-{template_id}",
            }
        )
    rng.shuffle(rows)
    return rows


def main() -> None:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    splits = {
        "train": build_rows(TRAIN_SAFE, TRAIN_RISKY, 4800, 2026),
        "validation": build_rows(VALIDATION_SAFE, VALIDATION_RISKY, 800, 2027),
        "test": build_rows(TEST_SAFE, TEST_RISKY, 1200, 2028),
    }
    manifest = {}
    for name, rows in splits.items():
        path = DATA_DIR / f"{name}.parquet"
        Dataset.from_list(rows).to_parquet(path)
        manifest[name] = {
            "rows": len(rows),
            "routine": sum(row["label"] == 0 for row in rows),
            "hazardous": sum(row["label"] == 1 for row in rows),
            "template_family": name,
            "path": path.name,
        }
    (DATA_DIR / "manifest.json").write_text(
        json.dumps(manifest, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    main()