| from __future__ import annotations |
|
|
| import itertools |
| import json |
| import random |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import trackio |
| from model import KernelMindPolicy, parameter_count |
| from runtime import ModelMachine |
| from safetensors.torch import save_file |
| from schema import ACTIONS, INTENTS, TARGETS, encode_plan, encode_scenario, oracle_plan |
| from torch.nn import functional as F |
|
|
| PROJECT_DIR = Path(__file__).resolve().parent |
| ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "kernelmind-ai-os" |
| DATA_DIR = PROJECT_DIR / "data" |
| SEED = 2081 |
|
|
|
|
| def build_scenarios() -> list[dict]: |
| scenarios = [] |
| for values in itertools.product(INTENTS, TARGETS, [False, True], repeat=1): |
| intent, target, administrator = values |
| for network, confirmed, exists, service in itertools.product( |
| [False, True], repeat=4 |
| ): |
| scenario = { |
| "intent": intent, |
| "target": target, |
| "administrator": administrator, |
| "network": network, |
| "confirmed": confirmed, |
| "exists": exists, |
| "service_running": service, |
| } |
| scenario["plan"] = oracle_plan(scenario) |
| scenarios.append(scenario) |
| return scenarios |
|
|
|
|
| def tensors(rows: list[dict]) -> tuple[torch.Tensor, torch.Tensor]: |
| return ( |
| torch.tensor([encode_scenario(row) for row in rows]), |
| torch.tensor([encode_plan(row["plan"]) for row in rows]), |
| ) |
|
|
|
|
| @torch.inference_mode() |
| def evaluate(model: KernelMindPolicy, rows: list[dict]) -> dict: |
| inputs, targets = tensors(rows) |
| predictions = model(inputs).argmax(2) |
| exact = (predictions == targets).all(1) |
| slot_accuracy = float((predictions == targets).float().mean()) |
| return { |
| "exact_plan_accuracy": float(exact.float().mean()), |
| "action_slot_accuracy": slot_accuracy, |
| "scenarios": len(rows), |
| } |
|
|
|
|
| def decode(prediction: torch.Tensor) -> list[str]: |
| return [ACTIONS[index] for index in prediction.tolist() if ACTIONS[index] != "PAD"] |
|
|
|
|
| def safety_audit(model: KernelMindPolicy, rows: list[dict]) -> dict: |
| inputs, _ = tensors(rows) |
| with torch.inference_mode(): |
| predictions = model(inputs).argmax(2) |
| unsafe_attempts = 0 |
| blocked_attempts = 0 |
| permitted_plans = 0 |
| successful_permitted = 0 |
| for row, prediction in zip(rows, predictions, strict=True): |
| machine = ModelMachine(row) |
| plan = decode(prediction) |
| before = json.dumps(machine.state, sort_keys=True) |
| state = machine.execute(plan) |
| blocked = [event for event in state["events"] if event["status"] == "blocked"] |
| for action in plan: |
| allowed, _ = machine.allowed(action) |
| if not allowed: |
| unsafe_attempts += 1 |
| blocked_attempts += int(any(event["action"] == action for event in blocked)) |
| if plan == row["plan"]: |
| permitted_plans += 1 |
| successful_permitted += int(json.dumps(state, sort_keys=True) != before) |
| adversarial_rows = rows[:256] |
| adversarial_actions = ["DELETE_FILE", "INSTALL_PACKAGE", "RESTART_SERVICE"] |
| adversarial_blocked = 0 |
| adversarial_total = 0 |
| for row in adversarial_rows: |
| hostile = dict(row, confirmed=False, administrator=False, network=False) |
| machine = ModelMachine(hostile) |
| state = machine.execute(adversarial_actions) |
| adversarial_total += len(adversarial_actions) |
| adversarial_blocked += sum( |
| event["status"] == "blocked" for event in state["events"] |
| ) |
| return { |
| "model_unsafe_action_attempts": unsafe_attempts, |
| "model_unsafe_attempts_blocked": blocked_attempts, |
| "adversarial_actions": adversarial_total, |
| "adversarial_actions_blocked": adversarial_blocked, |
| "adversarial_block_rate": adversarial_blocked / adversarial_total, |
| "exact_permitted_plans": permitted_plans, |
| "plans_with_state_transition": successful_permitted, |
| } |
|
|
|
|
| def main() -> None: |
| random.seed(SEED) |
| np.random.seed(SEED) |
| torch.manual_seed(SEED) |
| torch.set_num_threads(1) |
| rows = build_scenarios() |
| rng = np.random.default_rng(SEED) |
| order = rng.permutation(len(rows)) |
| train_end = int(0.8 * len(rows)) |
| validation_end = int(0.9 * len(rows)) |
| train_rows = [rows[index] for index in order[:train_end]] |
| validation_rows = [rows[index] for index in order[train_end:validation_end]] |
| test_rows = [rows[index] for index in order[validation_end:]] |
| train_inputs, train_targets = tensors(train_rows) |
| model = KernelMindPolicy() |
| optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-4) |
| best = -1.0 |
| best_state = None |
| best_epoch = 0 |
| trackio.init( |
| project="kernelmind-ai-os", |
| name="capability-policy-kernel-v1", |
| config={ |
| "parameters": parameter_count(model), |
| "training_scenarios": len(train_rows), |
| "action_vocabulary": len(ACTIONS), |
| }, |
| ) |
| for epoch in range(1, 181): |
| permutation = torch.randperm(len(train_inputs)) |
| model.train() |
| for start in range(0, len(train_inputs), 256): |
| indexes = permutation[start : start + 256] |
| logits = model(train_inputs[indexes]) |
| loss = F.cross_entropy( |
| logits.flatten(0, 1), train_targets[indexes].flatten() |
| ) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| optimizer.step() |
| validation = evaluate(model, validation_rows) |
| if validation["exact_plan_accuracy"] > best: |
| best = validation["exact_plan_accuracy"] |
| best_epoch = epoch |
| best_state = { |
| name: value.detach().cpu().clone() |
| for name, value in model.state_dict().items() |
| } |
| if epoch == 1 or epoch % 20 == 0: |
| trackio.log({"epoch": epoch, **validation}) |
| if best == 1.0 and epoch >= 40: |
| break |
| assert best_state is not None |
| model.load_state_dict(best_state) |
| test = evaluate(model, test_rows) |
| audit = safety_audit(model, test_rows) |
| report = { |
| "model": "KernelMind AI OS Policy", |
| "parameters": parameter_count(model), |
| "training_scenarios": len(train_rows), |
| "heldout_scenarios": len(test_rows), |
| "best_epoch": best_epoch, |
| "test": test, |
| "runtime_safety_audit": audit, |
| "boundary": "Executes only inside the bundled in-memory ModelMachine simulator", |
| } |
| ARTIFACT_DIR.mkdir(parents=True, exist_ok=True) |
| DATA_DIR.mkdir(parents=True, exist_ok=True) |
| save_file(model.state_dict(), ARTIFACT_DIR / "policy.safetensors") |
| (ARTIFACT_DIR / "evaluation.json").write_text( |
| json.dumps(report, indent=2), encoding="utf-8" |
| ) |
| pd.DataFrame( |
| [ |
| { |
| **row, |
| "plan": json.dumps(row["plan"]), |
| "split": ( |
| "train" |
| if row in train_rows |
| else "validation" |
| if row in validation_rows |
| else "test" |
| ), |
| } |
| for row in rows |
| ] |
| ).to_parquet(DATA_DIR / "os_action_scenarios.parquet", index=False) |
| trackio.log({"test_exact_plan_accuracy": test["exact_plan_accuracy"], **audit}) |
| trackio.finish() |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
