Upload folder using huggingface_hub

README.md

@@ -0,0 +1,9 @@
+# Atomight-V1-GRAFT-0.8M
+
+Custom from-scratch hybrid discriminative+generative model for abstract reasoning (RAVEN / PGM / I-RAVEN-X style matrix puzzles).
+
+Architecture: GRAFT (Grid Relational Attention Fusion Transformer), ~0.79M params.
+
+Narration is templated (captions the classifier's decision, not an independent derivation from ground-truth attribute labels).
+
+Load with the GRAFT class definition (see modeling.py) + model.pt state_dict.

config.json

@@ -0,0 +1 @@
+{"D": 128, "D2": 96, "vocab_size": 49, "architecture": "GRAFT"}

model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3c126b5fdfecb5dbf8d8ed177005fa6d843ac03c759697281a84d1c63fb1c13
+size 3173304

modeling.py

@@ -0,0 +1,73 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+D = 128
+D2 = 96
+
+class DepthwiseSeparableConv(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.dw = nn.Conv2d(in_ch, in_ch, 3, padding=1, groups=in_ch)
+        self.pw = nn.Conv2d(in_ch, out_ch, 1)
+    def forward(self, x):
+        return self.pw(self.dw(x))
+
+class AtomEncoder(nn.Module):
+    def __init__(self, embed_dim=D):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(1, 16, 3, stride=2, padding=1), nn.GELU(),
+            nn.Conv2d(16, 32, 3, stride=2, padding=1), nn.GELU(),
+            DepthwiseSeparableConv(32, 64), nn.GELU(),
+            nn.AdaptiveAvgPool2d(1),
+        )
+        self.proj = nn.Linear(64, embed_dim)
+    def forward(self, x):
+        return self.proj(self.net(x).flatten(1))
+
+class GridRelationalAttention(nn.Module):
+    def __init__(self, dim=D, heads=4, n_relations=4):
+        super().__init__()
+        self.missing_token = nn.Parameter(torch.randn(dim) * 0.02)
+        self.row_attn = nn.MultiheadAttention(dim, heads, batch_first=True)
+        self.col_attn = nn.MultiheadAttention(dim, heads, batch_first=True)
+        self.relation_queries = nn.Parameter(torch.randn(n_relations, dim) * 0.02)
+        self.cross_attn = nn.MultiheadAttention(dim, heads, batch_first=True)
+        self.rule_proj = nn.Linear(n_relations * dim, dim)
+        self.n_relations, self.dim = n_relations, dim
+
+    def forward(self, context_panels):
+        B = context_panels.shape[0]
+        missing = self.missing_token.expand(B, 1, self.dim)
+        grid = torch.cat([context_panels, missing], dim=1)
+        rows_idx = [[0,1,2],[3,4,5],[6,7,8]]
+        cols_idx = [[0,3,6],[1,4,7],[2,5,8]]
+        row_vecs = torch.cat([
+            self.row_attn(grid[:, idx, :], grid[:, idx, :], grid[:, idx, :])[0].mean(1, keepdim=True)
+            for idx in rows_idx
+        ], dim=1)
+        col_vecs = torch.cat([
+            self.col_attn(grid[:, idx, :], grid[:, idx, :], grid[:, idx, :])[0].mean(1, keepdim=True)
+            for idx in cols_idx
+        ], dim=1)
+        combined_ctx = torch.cat([row_vecs, col_vecs], dim=1)
+        rq = self.relation_queries.unsqueeze(0).expand(B, -1, -1)
+        rel_out, _ = self.cross_attn(rq, combined_ctx, combined_ctx)
+        return self.rule_proj(rel_out.flatten(1))
+
+class CandidateScorer(nn.Module):
+    def __init__(self, dim=D):
+        super().__init__()
+        self.score_proj = nn.Linear(dim, dim)
+    def forward(self, rule_embedding, candidate_embeds):
+        proj_rule = self.score_proj(rule_embedding)
+        return torch.einsum("bd,bkd->bk", proj_rule, candidate_embeds)
+
+class GRAFT(nn.Module):
+    def __init__(self, vocab_size, narrator_cls):
+        super().__init__()
+        self.encoder = AtomEncoder()
+        self.gra = GridRelationalAttention()
+        self.scorer = CandidateScorer()
+        self.narrator = narrator_cls(vocab_size=vocab_size)

vocab.json

@@ -0,0 +1 @@
+["<pad>", "<bos>", "<eos>", " ", "'", "(", ")", ",", "-", ".", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", ":", "?", "_", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z"]