| """GAT symbol-spotting model (T020, research R4). |
| |
| A two-layer Graph Attention Network over the attributed drawing graph: |
| node = entity, output = semantic type logits. Weights are pinned artifacts — |
| loaded in eval mode and hash-checked against models/meta.json so a silent |
| weight swap cannot change outputs unnoticed (Constitution IV, R11). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| from pathlib import Path |
|
|
| import torch |
| from torch_geometric.nn import GATConv |
|
|
| from cadmorph.classify.features import CLASSES, FEATURE_DIM |
| from cadmorph.determinism import check_weights, file_sha256 |
|
|
| CLASSIFIER_WEIGHTS = "gat_classifier.pt" |
| META_FILE = "meta.json" |
|
|
|
|
| def models_dir() -> Path: |
| override = os.environ.get("CADMORPH_MODELS_DIR") |
| if override: |
| return Path(override) |
| return Path(__file__).resolve().parents[3] / "models" |
|
|
|
|
| class GATClassifier(torch.nn.Module): |
| """GAT with a self-feature skip path. |
| |
| Pure stacked attention over kNN-proximity edges over-smooths: an |
| entity's own vector-exact features (size, stroke, segment structure) |
| are what identify its symbol class, while attention supplies |
| neighborhood context. The head sees both, concatenated. |
| """ |
|
|
| def __init__(self, in_dim: int = FEATURE_DIM, hidden: int = 32, heads: int = 4) -> None: |
| super().__init__() |
| self.lin_self = torch.nn.Linear(in_dim, hidden * heads) |
| self.conv1 = GATConv(in_dim, hidden, heads=heads) |
| self.conv2 = GATConv(hidden * heads, hidden, heads=1) |
| self.head = torch.nn.Linear(hidden * heads + hidden, len(CLASSES)) |
|
|
| def forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor: |
| h_self = torch.nn.functional.relu(self.lin_self(x)) |
| h_graph = torch.nn.functional.elu(self.conv1(x, edge_index)) |
| h_graph = torch.nn.functional.elu(self.conv2(h_graph, edge_index)) |
| return self.head(torch.cat([h_self, h_graph], dim=1)) |
|
|
|
|
| def save_pinned(model: torch.nn.Module, weights_name: str, directory: Path | None = None) -> str: |
| """Save weights and record their hash in meta.json; returns the hash.""" |
| directory = directory or models_dir() |
| directory.mkdir(parents=True, exist_ok=True) |
| path = directory / weights_name |
| torch.save(model.state_dict(), path) |
| digest = file_sha256(path) |
| meta_path = directory / META_FILE |
| meta = json.loads(meta_path.read_text(encoding="utf-8")) if meta_path.exists() else {} |
| meta[weights_name] = digest |
| meta_path.write_text(json.dumps(meta, indent=2, sort_keys=True), encoding="utf-8") |
| return digest |
|
|
|
|
| def pinned_hash(weights_name: str, directory: Path | None = None) -> str: |
| directory = directory or models_dir() |
| meta = json.loads((directory / META_FILE).read_text(encoding="utf-8")) |
| return meta[weights_name] |
|
|
|
|
| def load_pinned( |
| model: torch.nn.Module, weights_name: str, directory: Path | None = None |
| ) -> torch.nn.Module: |
| """Load hash-checked weights into ``model`` and freeze it in eval mode.""" |
| directory = directory or models_dir() |
| path = directory / weights_name |
| check_weights(path, pinned_hash(weights_name, directory)) |
| model.load_state_dict(torch.load(path, map_location="cpu", weights_only=True)) |
| model.eval() |
| for param in model.parameters(): |
| param.requires_grad_(False) |
| return model |
|
|
|
|
| def load_classifier(directory: Path | None = None) -> GATClassifier: |
| return load_pinned(GATClassifier(), CLASSIFIER_WEIGHTS, directory) |
|
|
