Upload 4 files

modeling_jetoncount.py

@@ -79,20 +79,41 @@ def _engineer_features_tensor(base: torch.Tensor) -> torch.Tensor:
 
 
 class JetonCountMLP(nn.Module):
+    """
+    IMPORTANT: num_layers means total Linear layers, exactly like the training script.
+
+    For num_layers == 1:
+        Linear(in_features -> 1)
+
+    For num_layers > 1:
+        Linear(in_features -> hidden_dim)
+        [num_layers - 2] x Linear(hidden_dim -> hidden_dim)
+        Linear(hidden_dim -> 1)
+    """
+
     def __init__(self, input_dim: int, hidden_dim: int, num_layers: int, dropout: float, activation: str):
         super().__init__()
         if num_layers < 1:
             raise ValueError("num_layers must be >= 1")
 
         layers = []
-        in_dim = input_dim
-        for _ in range(num_layers):
-            layers.append(nn.Linear(in_dim, hidden_dim))
+
+        if num_layers == 1:
+            layers.append(nn.Linear(input_dim, 1))
+        else:
+            layers.append(nn.Linear(input_dim, hidden_dim))
             layers.append(_get_activation(activation))
             if dropout > 0:
                 layers.append(nn.Dropout(dropout))
-            in_dim = hidden_dim
-        layers.append(nn.Linear(in_dim, 1))
+
+            for _ in range(num_layers - 2):
+                layers.append(nn.Linear(hidden_dim, hidden_dim))
+                layers.append(_get_activation(activation))
+                if dropout > 0:
+                    layers.append(nn.Dropout(dropout))
+
+            layers.append(nn.Linear(hidden_dim, 1))
+
         self.net = nn.Sequential(*layers)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor: