| from __future__ import annotations |
|
|
| import itertools |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import trackio |
| from model import NeuralModelMachine, parameter_count |
| from safetensors.torch import save_file |
| from schema import ACTIONS, CAPABILITY_NAMES, STATE_NAMES, reference_transition |
| from torch.nn import functional as F |
|
|
| PROJECT_DIR = Path(__file__).resolve().parent |
| ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "kernelmind-model-machine" |
| DATA_DIR = PROJECT_DIR / "data" |
| SEED = 2087 |
|
|
|
|
| def build_rows() -> list[dict]: |
| rows = [] |
| for state in itertools.product([0, 1], repeat=8): |
| for capabilities in itertools.product([0, 1], repeat=4): |
| for action_index, action in enumerate(ACTIONS): |
| next_state, blocked = reference_transition( |
| list(state), list(capabilities), action |
| ) |
| rows.append( |
| { |
| "state": list(state), |
| "capabilities": list(capabilities), |
| "action": action, |
| "action_index": action_index, |
| "next_state": next_state, |
| "blocked": blocked, |
| } |
| ) |
| return rows |
|
|
|
|
| def tensors(rows: list[dict]) -> tuple[torch.Tensor, ...]: |
| return ( |
| torch.tensor([row["state"] for row in rows], dtype=torch.float32), |
| torch.tensor([row["capabilities"] for row in rows], dtype=torch.float32), |
| torch.tensor([row["action_index"] for row in rows]), |
| torch.tensor( |
| [row["next_state"] + [row["blocked"]] for row in rows], |
| dtype=torch.float32, |
| ), |
| ) |
|
|
|
|
| @torch.inference_mode() |
| def evaluate(model: NeuralModelMachine, rows: list[dict]) -> dict: |
| state, capabilities, action, targets = tensors(rows) |
| probabilities = torch.sigmoid(model(state, capabilities, action)) |
| predictions = (probabilities >= 0.5).float() |
| exact = (predictions == targets).all(1) |
| return { |
| "exact_transition_accuracy": float(exact.float().mean()), |
| "state_bit_accuracy": float( |
| (predictions[:, :8] == targets[:, :8]).float().mean() |
| ), |
| "blocked_accuracy": float( |
| (predictions[:, 8] == targets[:, 8]).float().mean() |
| ), |
| "transitions": len(rows), |
| } |
|
|
|
|
| @torch.inference_mode() |
| def rollout_audit(model: NeuralModelMachine, samples: int = 2_000) -> dict: |
| rng = np.random.default_rng(SEED + 1) |
| exact_steps = 0 |
| exact_rollouts = 0 |
| for _ in range(samples): |
| state = rng.integers(0, 2, size=8).tolist() |
| capabilities = rng.integers(0, 2, size=4).tolist() |
| learned = torch.tensor([state], dtype=torch.float32) |
| reference = state |
| rollout_exact = True |
| for action_index in rng.integers(0, len(ACTIONS), size=6): |
| action = ACTIONS[action_index] |
| reference, reference_blocked = reference_transition( |
| reference, capabilities, action |
| ) |
| predicted, blocked_probability = model.transition( |
| learned, |
| torch.tensor([capabilities], dtype=torch.float32), |
| torch.tensor([action_index]), |
| ) |
| learned = predicted |
| step_exact = learned[0].int().tolist() == reference |
| step_exact = step_exact and ( |
| int(blocked_probability[0] >= 0.5) == reference_blocked |
| ) |
| exact_steps += int(step_exact) |
| rollout_exact = rollout_exact and step_exact |
| exact_rollouts += int(rollout_exact) |
| return { |
| "rollouts": samples, |
| "steps_per_rollout": 6, |
| "exact_step_fraction": exact_steps / (samples * 6), |
| "exact_six_step_rollout_fraction": exact_rollouts / samples, |
| } |
|
|
|
|
| def main() -> None: |
| torch.manual_seed(SEED) |
| torch.set_num_threads(1) |
| rows = build_rows() |
| 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:]] |
| state, capabilities, action, targets = tensors(train_rows) |
| model = NeuralModelMachine() |
| optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-5) |
| best = -1.0 |
| best_state = None |
| best_epoch = 0 |
| trackio.init( |
| project="kernelmind-model-machine", |
| name="neural-os-transition-runtime-v1", |
| config={ |
| "parameters": parameter_count(model), |
| "training_transitions": len(train_rows), |
| "state_bits": len(STATE_NAMES), |
| "capability_bits": len(CAPABILITY_NAMES), |
| }, |
| ) |
| for epoch in range(1, 81): |
| permutation = torch.randperm(len(state)) |
| model.train() |
| for start in range(0, len(state), 512): |
| indexes = permutation[start : start + 512] |
| logits = model(state[indexes], capabilities[indexes], action[indexes]) |
| loss = F.binary_cross_entropy_with_logits(logits, targets[indexes]) |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| optimizer.step() |
| validation = evaluate(model, validation_rows) |
| if validation["exact_transition_accuracy"] > best: |
| best = validation["exact_transition_accuracy"] |
| best_epoch = epoch |
| best_state = { |
| name: value.detach().cpu().clone() |
| for name, value in model.state_dict().items() |
| } |
| if epoch == 1 or epoch % 10 == 0: |
| trackio.log({"epoch": epoch, **validation}) |
| if best == 1.0 and epoch >= 20: |
| break |
| assert best_state is not None |
| model.load_state_dict(best_state) |
| test = evaluate(model, test_rows) |
| rollouts = rollout_audit(model) |
| report = { |
| "model": "KernelMind Neural Model Machine", |
| "parameters": parameter_count(model), |
| "training_transitions": len(train_rows), |
| "heldout_transitions": len(test_rows), |
| "best_epoch": best_epoch, |
| "test": test, |
| "multi_step_audit": rollouts, |
| "boundary": "Models an eight-bit virtual state; it does not control the host OS", |
| } |
| ARTIFACT_DIR.mkdir(parents=True, exist_ok=True) |
| DATA_DIR.mkdir(parents=True, exist_ok=True) |
| save_file(model.state_dict(), ARTIFACT_DIR / "model_machine.safetensors") |
| (ARTIFACT_DIR / "evaluation.json").write_text( |
| json.dumps(report, indent=2), encoding="utf-8" |
| ) |
| pd.DataFrame( |
| { |
| "state": [row["state"] for row in rows], |
| "capabilities": [row["capabilities"] for row in rows], |
| "action": [row["action"] for row in rows], |
| "next_state": [row["next_state"] for row in rows], |
| "blocked": [row["blocked"] for row in rows], |
| } |
| ).to_parquet(DATA_DIR / "transitions.parquet", index=False) |
| trackio.log({**test, **rollouts}) |
| trackio.finish() |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
