Update index.html

index.html

@@ -12,12 +12,15 @@
             display: flex;
             flex-direction: row;
             height: 100vh;
+            background: linear-gradient(120deg, #f8b7d4, #d1c1f2);
+            color: #4a154b;
         }
 
         .console {
             width: 30%;
             padding: 20px;
-            background-color: #f4f4f4;
+            background-color: #fef4fc;
+            border-right: 2px solid #d1c1f2;
             overflow-y: auto;
         }
 
@@ -30,6 +33,7 @@
         table {
             width: 100%;
             border-collapse: collapse;
+            margin-bottom: 20px;
         }
 
         th, td {
@@ -39,17 +43,65 @@
         }
 
         th {
-            background-color: #f2f2f2;
+            background-color: #d1c1f2;
+            color: #4a154b;
         }
 
         .highlight {
-            background-color: #d4edda;
-            color: #155724;
+            background-color: #f8d7f1;
+            color: #4a154b;
             padding: 10px;
             margin: 20px 0;
-            border: 1px solid #c3e6cb;
+            border: 1px solid #d1c1f2;
             border-radius: 4px;
         }
+
+        h3 {
+            color: #4a154b;
+        }
+
+        label, input, select {
+            display: block;
+            margin-bottom: 10px;
+            width: 100%;
+        }
+
+        input, select {
+            padding: 8px;
+            border: 1px solid #d1c1f2;
+            border-radius: 4px;
+            background: #fef4fc;
+            color: #4a154b;
+        }
+
+        button {
+            background-color: #d1c1f2;
+            color: white;
+            padding: 10px 20px;
+            border: none;
+            border-radius: 4px;
+            cursor: pointer;
+        }
+
+        button:hover {
+            background-color: #bfa8e3;
+        }
+
+        .section-title {
+            background-color: #e9d8f5;
+            color: #4a154b;
+            text-align: center;
+            padding: 5px 0;
+            margin-bottom: 5px;
+        }
+
+        .sub-table {
+            margin-top: 10px;
+        }
+
+        .no-bias {
+            display: none;
+        }
     </style>
     <script>
         function calculateParameters() {
@@ -63,10 +115,12 @@
 
             // Attention calculations
             const attentionQ = hiddenSize * hiddenSize;
+            const attentionK = hiddenSize * numKeyValueHeads;
             const attentionV = hiddenSize * numKeyValueHeads;
             const attentionBq = includeBias ? hiddenSize : 0;
+            const attentionBk = includeBias ? numKeyValueHeads : 0;
             const attentionBv = includeBias ? numKeyValueHeads : 0;
-            const attentionTotal = 2 * (attentionQ + attentionBq) + 2 * (attentionV + attentionBv);
+            const attentionTotal = 2 * (attentionQ + attentionBq) + 2 * (attentionK + attentionBk) + 2 * (attentionV + attentionBv);
 
             // Feed Forward calculations
             const switchW = hiddenSize * intermediateSize;
@@ -90,78 +144,24 @@
             const fullSize = fullLayersParams + embeddingTotal;
 
             // Display results
-            const outputTable = document.getElementById('output_table');
-            outputTable.innerHTML = `
-                <tr>
-                    <th>Section</th>
-                    <th>Parameter</th>
-                    <th>Value</th>
-                </tr>
-                <tr>
-                    <td>Attention</td>
-                    <td>Q</td>
-                    <td>${attentionQ.toLocaleString()}</td>
-                </tr>
-                <tr>
-                    <td>Attention</td>
-                    <td>Bq</td>
-                    <td>${attentionBq.toLocaleString()}</td>
-                </tr>
-                <tr>
-                    <td>Attention</td>
-                    <td>V</td>
-                    <td>${attentionV.toLocaleString()}</td>
-                </tr>
-                <tr>
-                    <td>Attention</td>
-                    <td>Bv</td>
-                    <td>${attentionBv.toLocaleString()}</td>
-                </tr>
-                <tr>
-                    <td>Attention</td>
-                    <td>Total</td>
-                    <td>${attentionTotal.toLocaleString()}</td>
-                </tr>
-                <tr>
-                    <td>Feed Forward</td>
-                    <td>Switch W</td>
-                    <td>${switchW.toLocaleString()}</td>
-                </tr>
-                <tr>
-                    <td>Feed Forward</td>
-                    <td>Switch b</td>
-                    <td>${switchB.toLocaleString()}</td>
-                </tr>
-                <tr>
-                    <td>Feed Forward</td>
-                    <td>LU W</td>
-                    <td>${luW.toLocaleString()}</td>
-                </tr>
-                <tr>
-                    <td>Feed Forward</td>
-                    <td>LU b</td>
-                    <td>${luB.toLocaleString()}</td>
-                </tr>
-                <tr>
-                    <td>Feed Forward</td>
-                    <td>Proj W</td>
-                    <td>${projW.toLocaleString()}</td>
-                </tr>
-                <tr>
-                    <td>Feed Forward</td>
-                    <td>Proj b</td>
-                    <td>${projB.toLocaleString()}</td>
-                </tr>
-                <tr>
-                    <td>Feed Forward</td>
-                    <td>Total</td>
-                    <td>${feedForwardTotal.toLocaleString()}</td>
-                </tr>
-                <tr>
-                    <td>Embedding</td>
-                    <td>Embedding</td>
-                    <td>${embeddingTotal.toLocaleString()}</td>
-                </tr>
+            const attentionTable = document.getElementById('attention_table');
+            attentionTable.innerHTML = `
+                <tr><th>Type</th><th>Q</th><th>K</th><th>V</th><th>Total</th></tr>
+                <tr><td>Weights</td><td>${attentionQ.toLocaleString()}</td><td>${attentionK.toLocaleString()}</td><td>${attentionV.toLocaleString()}</td><td>${attentionTotal.toLocaleString()}</td></tr>
+                ${includeBias ? `<tr class="bias-row"><td>Bias</td><td>${attentionBq.toLocaleString()}</td><td>${attentionBk.toLocaleString()}</td><td>${attentionBv.toLocaleString()}</td><td></td></tr>` : ''}
+            `;
+
+            const feedForwardTable = document.getElementById('feed_forward_table');
+            feedForwardTable.innerHTML = `
+                <tr><th>Type</th><th>Switch W</th><th>LU W</th><th>Proj W</th><th>Total</th></tr>
+                <tr><td>Weights</td><td>${switchW.toLocaleString()}</td><td>${luW.toLocaleString()}</td><td>${projW.toLocaleString()}</td><td>${feedForwardTotal.toLocaleString()}</td></tr>
+                ${includeBias ? `<tr class="bias-row"><td>Bias</td><td>${switchB.toLocaleString()}</td><td>${luB.toLocaleString()}</td><td>${projB.toLocaleString()}</td><td></td></tr>` : ''}
+            `;
+
+            const embeddingTable = document.getElementById('embedding_table');
+            embeddingTable.innerHTML = `
+                <tr><th>Type</th><th>Total</th></tr>
+                <tr><td>Embedding</td><td>${embeddingTotal.toLocaleString()}</td></tr>
             `;
 
             document.getElementById('one_layer_params').innerText = oneLayerParams.toLocaleString();
@@ -174,29 +174,29 @@
 
 <div class="console">
     <h3>Input Parameters</h3>
-    <label for="hidden_size">Hidden size:</label><br>
-    <input type="number" id="hidden_size" value="2048"><br><br>
+    <label for="hidden_size">Hidden size:</label>
+    <input type="number" id="hidden_size" value="2048">
 
-    <label for="intermediate_size">Intermediate size:</label><br>
-    <input type="number" id="intermediate_size" value="16384"><br><br>
+    <label for="intermediate_size">Intermediate size:</label>
+    <input type="number" id="intermediate_size" value="16384">
 
-    <label for="vocab_size">Vocab size:</label><br>
-    <input type="number" id="vocab_size" value="64000"><br><br>
+    <label for="vocab_size">Vocab size:</label>
+    <input type="number" id="vocab_size" value="64000">
 
-    <label for="num_key_value_heads">Number of key-value heads:</label><br>
-    <input type="number" id="num_key_value_heads" value="80"><br><br>
+    <label for="num_key_value_heads">Number of key-value heads:</label>
+    <input type="number" id="num_key_value_heads" value="80">
 
-    <label for="num_attention_heads">Number of attention heads:</label><br>
-    <input type="number" id="num_attention_heads" value="8"><br><br>
+    <label for="num_attention_heads">Number of attention heads:</label>
+    <input type="number" id="num_attention_heads" value="8">
 
-    <label for="num_hidden_layers">Number of hidden layers:</label><br>
-    <input type="number" id="num_hidden_layers" value="64"><br><br>
+    <label for="num_hidden_layers">Number of hidden layers:</label>
+    <input type="number" id="num_hidden_layers" value="64">
 
-    <label for="include_bias">Include bias?</label><br>
+    <label for="include_bias">Include bias?</label>
     <select id="include_bias">
         <option value="yes">Yes</option>
         <option value="no">No</option>
-    </select><br><br>
+    </select>
 
     <button onclick="calculateParameters()">Calculate</button>
 </div>
@@ -204,13 +204,14 @@
 <div class="output">
     <h3>Model Parameter Results</h3>
 
-    <table id="output_table">
-        <tr>
-            <th>Section</th>
-            <th>Parameter</th>
-            <th>Value</th>
-        </tr>
-    </table>
+    <div class="section-title">Attention Parameters</div>
+    <table id="attention_table" class="sub-table"></table>
+
+    <div class="section-title">Feed Forward Parameters</div>
+    <table id="feed_forward_table" class="sub-table"></table>
+
+    <div class="section-title">Embedding Parameters</div>
+    <table id="embedding_table" class="sub-table"></table>
 
     <div class="highlight">
         <strong>1 Layer Parameters:</strong> <span id="one_layer_params">0</span><br>