Update app.py

app.py

@@ -18,6 +18,95 @@ if not os.path.exists(MODEL_PATH):
     print("Download complete!")
 
 # --- DEFINE YOUR MODEL ARCHITECTURE (TETAP SAMA) ---
+class LookThemLayer(nn.Module):
+    def __init__(self, num_tokens, in_features, hidden_dim):
+        super().__init__()
+        self.num_tokens = num_tokens
+        self.mod1_w1 = nn.Parameter(torch.randn(num_tokens, in_features, hidden_dim))
+        self.mod1_b1 = nn.Parameter(torch.zeros(num_tokens, hidden_dim))
+        self.mod1_w2 = nn.Parameter(torch.randn(num_tokens, hidden_dim, 1))
+        self.mod1_b2 = nn.Parameter(torch.zeros(num_tokens, 1))
+        self.mod2_w1 = nn.Parameter(torch.randn(num_tokens, in_features, hidden_dim))
+        self.mod2_b1 = nn.Parameter(torch.zeros(num_tokens, hidden_dim))
+        self.mod2_w2 = nn.Parameter(torch.randn(num_tokens, hidden_dim, 1))
+        self.mod2_b2 = nn.Parameter(torch.zeros(num_tokens, 1))
+        self.trans_w = nn.Parameter(torch.randn(num_tokens, 1, 1))
+        self.trans_b = nn.Parameter(torch.zeros(num_tokens, 1))
+        self._init_weights()
+
+    def _init_weights(self):
+        for w in [self.mod1_w1, self.mod2_w1, self.mod1_w2, self.mod2_w2]:
+            nn.init.xavier_uniform_(w)
+
+    def forward(self, x):
+        N = self.num_tokens
+        h1 = torch.einsum("bti,tij->btj", x, self.mod1_w1) + self.mod1_b1
+        out_m1 = torch.einsum("btj,tjk->btk", F.gelu(h1), self.mod1_w2) + self.mod1_b2
+        h2 = torch.einsum("bti,tij->btj", x, self.mod2_w1) + self.mod2_b1
+        out_m2 = torch.einsum("btj,tjk->btk", F.gelu(h2), self.mod2_w2) + self.mod2_b2
+
+        out_m2_safe = torch.sign(out_m2) * torch.clamp(torch.abs(out_m2), min=1e-6)
+        compare = torch.tanh(out_m1.unsqueeze(2) / out_m2_safe.unsqueeze(1))
+        compare2 = torch.tanh(out_m1.unsqueeze(1) / out_m2_safe.unsqueeze(2))
+
+        trans_compare = torch.einsum("bije,jef->bijf", compare, self.trans_w) + self.trans_b.view(1, 1, N, 1)
+        trans_compare2 = torch.einsum("bije,jef->bijf", compare2, self.trans_w) + self.trans_b.view(1, 1, N, 1)
+
+        interaksi = (trans_compare * x.unsqueeze(2) + trans_compare2 * x.unsqueeze(1)) / 2
+        mask = (1.0 - torch.eye(N, device=x.device)).view(1, N, N, 1)
+        return (interaksi * mask).sum(dim=2) / (N - 1.0)
+
+class LiteResidualBlock(nn.Module):
+    def __init__(self, dim, dropout=0.05):
+        super().__init__()
+        self.block = nn.Sequential(nn.Linear(dim, dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim, dim))
+        self.norm = nn.LayerNorm(dim)
+    def forward(self, x):
+        return self.norm(x + self.block(x))
+
+class LookThemV8MNIST(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.stream_a = nn.Sequential(
+                nn.Conv2d(1, 4, 3, 2, 1),
+                nn.BatchNorm2d(4), nn.GELU(),
+                nn.Conv2d(4, 8, 3, 2, 1),
+                nn.BatchNorm2d(8), nn.GELU(),
+                nn.AdaptiveMaxPool2d((8, 8)))
+        self.stream_b = nn.Sequential(
+                nn.Conv2d(1, 4, 3, 1, 1),
+                nn.BatchNorm2d(4), nn.GELU(),
+                nn.Conv2d(4, 8, 3, 1, 1),
+                nn.BatchNorm2d(8), nn.GELU(),
+                nn.AdaptiveMaxPool2d((8, 8)))
+
+        self.lookthemA = LookThemLayer(64, 8, 32)
+        self.lookthemB = LookThemLayer(64, 8, 32)
+        self.lookthem_comb = LookThemLayer(64, 16, 32)
+        self.comb_norm = nn.LayerNorm(16)
+
+        self.FFN1 = nn.Conv1d(16, 8, 1)
+        self.lookthem2 = LookThemLayer(64, 8, 32)
+        self.FFN2 = nn.Conv1d(8, 8, 1)
+
+        self.compressor = nn.Conv1d(8, 4, 1)
+        self.input_proj = nn.Linear(64 * 4, 128)
+        self.res_blocks = nn.Sequential(LiteResidualBlock(128), LiteResidualBlock(128))
+        self.head = nn.Sequential(nn.Linear(128, 128), nn.GELU(), nn.Linear(128, 10))
+
+    def forward(self, x):
+        b = x.size(0)
+        fa = self.lookthemA(self.stream_a(x).view(b, 8, 64).transpose(1, 2))
+        fb = self.lookthemB(self.stream_b(x).view(b, 8, 64).transpose(1, 2))
+        x = self.comb_norm(self.lookthem_comb(torch.cat([fa, fb], dim=2)))
+        x = x.transpose(1, 2)
+        x = self.FFN1(x).transpose(1, 2)
+        res = x
+        x = self.lookthem2(x).transpose(1, 2)
+        x = self.FFN2(x) + res.transpose(1, 2)
+        x = self.compressor(x).flatten(1)
+        x = self.res_blocks(self.input_proj(x))
+        return self.head(x)
 # ... (Salin definisi kelas LookThemLayer, LiteResidualBlock, dan LookThemV8MNIST Anda di sini) ...
 
 # --- LOAD WEIGHTS ON CPU/GPU ---