Update index.html

index.html

@@ -12,20 +12,18 @@
             display: flex;
             flex-direction: row;
             height: 100vh;
-            background: linear-gradient(to right, #ff9a9e, #fad0c4);
-            color: #4b0082;
         }
 
         .console {
-            width: 30%;
+            width: 20%;
             padding: 20px;
-            background-color: #fdfdfd;
-            border-right: 2px solid #ff9a9e;
+            background-color: #f6b5b5;
             overflow-y: auto;
+            border-right: 2px solid #ee4a4f;
         }
 
         .output {
-            width: 70%;
+            width: 80%;
             padding: 20px;
             overflow-y: auto;
         }
@@ -34,66 +32,112 @@
             width: 100%;
             border-collapse: collapse;
             margin-bottom: 20px;
-            background-color: #ffffff;
-            border-radius: 8px;
-            overflow: hidden;
-            box-shadow: 0px 4px 6px rgba(0, 0, 0, 0.1);
+            background-color: #fff7f9;
+            border: 1px solid #f6b5b5;
         }
 
         th, td {
-            border: 1px solid #ddd;
+            border: 1px solid #f6b5b5;
             padding: 8px;
             text-align: center;
         }
 
         th {
-            background-color: #ff9a9e;
-            color: #fff;
-            text-transform: uppercase;
+            background-color: #f6d4d4;
         }
 
         .highlight {
-            background-color: #fdf2ff;
-            color: #800080;
+            background-color: #ffe8e8;
+            color: #ad0d0d;
             padding: 10px;
             margin: 20px 0;
-            border: 1px solid #ff9a9e;
-            border-radius: 8px;
+            border: 1px solid #f6b5b5;
+            border-radius: 4px;
+            text-align: center;
+        }
+
+        h3 {
+            color: #ad0d0d;
         }
 
-        .section-title {
-            margin-top: 30px;
-            font-size: 1.2em;
+        .section {
+            margin-bottom: 20px;
+        }
+        label {
             font-weight: bold;
-            color: #800080;
-            text-align: center;
+            font-size: 0.9rem;
+        }
+
+        input[type="number"], select {
+            width: 100%;
+            padding: 8px;
+            border: 1px solid #ddd;
+            border-radius: 5px;
+            box-sizing: border-box;
+            font-size: 0.9rem;
+        }
+
+        input[type="number"]:focus, select:focus {
+            border-color: #ff6666;
+            outline: none;
+            box-shadow: 0px 0px 5px rgba(255, 102, 102, 0.5);
+        }
+
+        button {
+            background-color: #ff6666;
+            color: white;
+            border: none;
+            border-radius: 5px;
+            padding: 5px;
+            width: 100%;
+            font-size: 1rem;
+            cursor: pointer;
+            transition: background-color 0.3s ease;
+        }
+
+        button:hover {
+            background-color: #e60000;
         }
     </style>
     <script>
         function calculateParameters() {
             const hiddenSize = parseInt(document.getElementById('hidden_size').value);
+            const num_kv_heads = parseInt(document.getElementById('num_key_value_heads').value);
+            const num_heads = parseInt(document.getElementById('num_heads').value);
             const intermediateSize = parseInt(document.getElementById('intermediate_size').value);
             const vocabSize = parseInt(document.getElementById('vocab_size').value);
-            const numKeyValueHeads = parseInt(document.getElementById('num_key_value_heads').value);
             const numHiddenLayers = parseInt(document.getElementById('num_hidden_layers').value);
             const includeBias = document.getElementById('include_bias').value === 'yes';
 
+            const calcTotal = (input, output) => input * output;
+            const calcKVdim = (hiddenSize, num_heads,num_kv_heads) => hiddenSize * num_kv_heads / num_heads;
+
             // Attention calculations
-            const attentionQW = [hiddenSize, hiddenSize, hiddenSize * hiddenSize];
-            const attentionV = [hiddenSize, numKeyValueHeads, hiddenSize * numKeyValueHeads];
-            const attentionBq = includeBias ? [hiddenSize, 1, hiddenSize] : [0, 0, 0];
-            const attentionBv = includeBias ? [numKeyValueHeads, 1, numKeyValueHeads] : [0, 0, 0];
-            const attentionTotal = attentionQW[2] + attentionV[2] + attentionBq[2] + attentionBv[2];
+            const attention = [
+                { name: 'Query', input: hiddenSize, output: hiddenSize },
+                { name: 'Key', input: hiddenSize, output: calcKVdim(hiddenSize, num_heads,num_kv_heads) },
+                { name: 'Value', input: hiddenSize, output: calcKVdim(hiddenSize, num_heads,num_kv_heads) },
+                { name: 'Projection', input: hiddenSize, output: hiddenSize }
+            ].map(entry => ({
+                ...entry,
+                wTotal: calcTotal(entry.input, entry.output),
+                bTotal: includeBias ? entry.output : 0
+            }));
+
+            const attentionTotal = attention.reduce((sum, entry) => sum + entry.wTotal + entry.bTotal, 0);
 
             // Feed Forward calculations
-            const switchW = [hiddenSize, intermediateSize, hiddenSize * intermediateSize];
-            const switchB = includeBias ? [intermediateSize, 1, intermediateSize] : [0, 0, 0];
-            const projW = [intermediateSize, hiddenSize, intermediateSize * hiddenSize];
-            const projB = includeBias ? [hiddenSize, 1, hiddenSize] : [0, 0, 0];
-            const feedForwardTotal = switchW[2] + switchB[2] + projW[2] + projB[2];
+            const switchW = calcTotal(hiddenSize, intermediateSize);
+            const switchB = includeBias ? intermediateSize : 0;
+            const luW = calcTotal(hiddenSize, intermediateSize);
+            const luB = includeBias ? intermediateSize : 0;
+            const projW = calcTotal(intermediateSize, hiddenSize);
+            const projB = includeBias ? hiddenSize : 0;
+
+            const feedForwardTotal = switchW + switchB + luW + luB + projW + projB;
 
             // Embedding
-            const embedding = [vocabSize, hiddenSize, vocabSize * hiddenSize];
+            const embeddingTotal = calcTotal(vocabSize, hiddenSize);
 
             // 1 layer parameters
             const oneLayerParams = attentionTotal + feedForwardTotal;
@@ -102,107 +146,119 @@
             const fullLayersParams = oneLayerParams * numHiddenLayers;
 
             // Full size (includes embedding)
-            const fullSize = fullLayersParams + embedding[2];
-
-            // Generate tables
-            const attentionTable = `
-                <table>
-                    <tr>
-                        <th>Attention</th>
-                        <th>Input Size</th>
-                        <th>Output Size</th>
-                        <th>Total Size</th>
-                    </tr>
-                    <tr>
-                        <td>Q / Wq</td>
-                        <td>${attentionQW[0]}</td>
-                        <td>${attentionQW[1]}</td>
-                        <td>${attentionQW[2].toLocaleString()}</td>
-                    </tr>
-                    <tr>
-                        <td>V / Wv</td>
-                        <td>${attentionV[0]}</td>
-                        <td>${attentionV[1]}</td>
-                        <td>${attentionV[2].toLocaleString()}</td>
-                    </tr>
-                    ${includeBias ? `
-                    <tr>
-                        <td>Bq</td>
-                        <td>${attentionBq[0]}</td>
-                        <td>${attentionBq[1]}</td>
-                        <td>${attentionBq[2].toLocaleString()}</td>
-                    </tr>
-                    <tr>
-                        <td>Bv</td>
-                        <td>${attentionBv[0]}</td>
-                        <td>${attentionBv[1]}</td>
-                        <td>${attentionBv[2].toLocaleString()}</td>
-                    </tr>
-                    ` : ''}
-                </table>`;
-
-            const feedForwardTable = `
-                <table>
-                    <tr>
-                        <th>Feed Forward</th>
-                        <th>Input Size</th>
-                        <th>Output Size</th>
-                        <th>Total Size</th>
-                    </tr>
-                    <tr>
-                        <td>Switch / W</td>
-                        <td>${switchW[0]}</td>
-                        <td>${switchW[1]}</td>
-                        <td>${switchW[2].toLocaleString()}</td>
-                    </tr>
-                    ${includeBias ? `<tr>
-                        <td>Switch / b</td>
-                        <td>${switchB[0]}</td>
-                        <td>${switchB[1]}</td>
-                        <td>${switchB[2].toLocaleString()}</td>
-                    </tr>` : ''}
-                    <tr>
-                        <td>Proj / W</td>
-                        <td>${projW[0]}</td>
-                        <td>${projW[1]}</td>
-                        <td>${projW[2].toLocaleString()}</td>
-                    </tr>
-                    ${includeBias ? `<tr>
-                        <td>Proj / b</td>
-                        <td>${projB[0]}</td>
-                        <td>${projB[1]}</td>
-                        <td>${projB[2].toLocaleString()}</td>
-                    </tr>` : ''}
-                </table>`;
-
-            const embeddingTable = `
-                <table>
-                    <tr>
-                        <th>Embedding</th>
-                        <th>Input Size</th>
-                        <th>Output Size</th>
-                        <th>Total Size</th>
-                    </tr>
-                    <tr>
-                        <td>Embedding</td>
-                        <td>${embedding[0]}</td>
-                        <td>${embedding[1]}</td>
-                        <td>${embedding[2].toLocaleString()}</td>
-                    </tr>
-                </table>`;
-
-            // Update output
-            document.getElementById('output').innerHTML = `
+            const fullSize = fullLayersParams + embeddingTotal;
+
+            // Display results
+            const outputDiv = document.getElementById('output');
+            outputDiv.innerHTML = `
+                <h1>Model Parameter Calculator</h1>
+                <div class="section">
+                    <h3>Attention</h3>
+                    <table>
+                        <tr>
+                            <th>Parameter</th>
+                            <th>Input Size</th>
+                            <th>Output Size</th>
+                            <th>Total parameters (input*ouput)</th>
+                        </tr>
+                        ${attention.map(entry => `
+                            <tr>
+                                <td>${entry.name} - W</td>
+                                <td>${entry.input.toLocaleString()}</td>
+                                <td>${entry.output.toLocaleString()}</td>
+                                <td>${entry.wTotal.toLocaleString()}</td>
+                            </tr>
+                            ${includeBias ? `
+                            <tr>
+                                <td>${entry.name} - b</td>
+                                <td>-</td>
+                                <td>${entry.output.toLocaleString()}</td>
+                                <td>${entry.bTotal.toLocaleString()}</td>
+                            </tr>
+                            ` : ''}
+                        `).join('')}
+                        <tr>
+                            <th colspan="3">Total attention parameters</th>
+                            <td colspan="2">${attentionTotal.toLocaleString()}</td>
+                        </tr>
+                    </table>
+                </div>
+                <div class="section">
+                    <h3>Feed Forward</h3>
+                    <table>
+                        <tr>
+                            <th>Parameter</th>
+                            <th>Input Size</th>
+                            <th>Output Size</th>
+                            <th>Total (input*output)</th>
+                        </tr>
+                        <tr>
+                            <td>Swish - W</td>
+                            <td>${hiddenSize.toLocaleString()}</td>
+                            <td>${intermediateSize.toLocaleString()}</td>
+                            <td>${switchW.toLocaleString()}</td>
+                        </tr>
+                        ${includeBias ? `<tr>
+                            <td>Swish - b</td>
+                            <td>-</td>
+                            <td>${intermediateSize.toLocaleString()}</td>
+                            <td>${switchB.toLocaleString()}</td>
+                        </tr>` : ''}
+                        <tr>
+                            <td>GLU W</td>
+                            <td>${hiddenSize.toLocaleString()}</td>
+                            <td>${intermediateSize.toLocaleString()}</td>
+                            <td>${luW.toLocaleString()}</td>
+                        </tr>
+                        ${includeBias ? `<tr>
+                            <td>GLU - b</td>
+                            <td>-</td>
+                            <td>${intermediateSize.toLocaleString()}</td>
+                            <td>${luB.toLocaleString()}</td>
+                        </tr>` : ''}
+                        <tr>
+                            <td>Projection - W</td>
+                            <td>${intermediateSize.toLocaleString()}</td>
+                            <td>${hiddenSize.toLocaleString()}</td>
+                            <td>${projW.toLocaleString()}</td>
+                        </tr>
+                        ${includeBias ? `<tr>
+                            <td>Projection - b</td>
+                            <td>-</td>
+                            <td>${hiddenSize.toLocaleString()}</td>
+                            <td>${projB.toLocaleString()}</td>
+                        </tr>` : ''}
+                        <tr>
+                            <th colspan="3">Total Feed Forward parameters</th>
+                            <td>${feedForwardTotal.toLocaleString()}</td>
+                        </tr>
+                    </table>
+                </div>
                 <div class="highlight">
-                    <strong>1 Layer Parameters:</strong> ${oneLayerParams.toLocaleString()}<br>
-                    <strong>Full Layers Parameters:</strong> ${fullLayersParams.toLocaleString()}<br>
+                    <strong>1 Layer Parameters (sum of attention and FFN):</strong> ${oneLayerParams.toLocaleString()}<br>
+                    <strong>Full Layers Parameters (1 layer parameters * num layers):</strong> ${fullLayersParams.toLocaleString()}<br>
+                </div>
+                <div class="section">
+                    <h3>Embedding</h3>
+                    <table>
+                        <tr>
+                            <th>Parameter</th>
+                            <th>Vocab Size</th>
+                            <th>Hidden Size</th>
+                            <th>Total (vocab * hidden)</th>
+                        </tr>
+                        <tr>
+                            <td>Embedding</td>
+                            <td>${vocabSize.toLocaleString()}</td>
+                            <td>${hiddenSize.toLocaleString()}</td>
+                            <td>${embeddingTotal.toLocaleString()}</td>
+                        </tr>
+                    </table>
                 </div>
-                ${attentionTable}
-                ${feedForwardTable}
                 <div class="highlight">
-                    <strong>Complete Model Size:</strong> ${fullSize.toLocaleString()}
+                    <strong>Complete Model Parmeters (embedding size + full layers size):</strong> ${fullSize.toLocaleString()}
                 </div>
-                ${embeddingTable}`;
+            `;
         }
     </script>
 </head>
@@ -211,31 +267,35 @@
 <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>
+    <input type="number" id="hidden_size" value="896"><br><br>
 
     <label for="intermediate_size">Intermediate size:</label><br>
-    <input type="number" id="intermediate_size" value="16384"><br><br>
+    <input type="number" id="intermediate_size" value="4864"><br><br>
 
     <label for="vocab_size">Vocab size:</label><br>
-    <input type="number" id="vocab_size" value="64000"><br><br>
+    <input type="number" id="vocab_size" value="151646"><br><br>
 
     <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>
+    <input type="number" id="num_key_value_heads" value="2"><br><br>
+
+    <label for="num_heads">Number of attention (query) heads:</label><br>
+    <input type="number" id="num_heads" value="14"><br><br>
 
     <label for="num_hidden_layers">Number of hidden layers:</label><br>
-    <input type="number" id="num_hidden_layers" value="64"><br><br>
+    <input type="number" id="num_hidden_layers" value="24"><br><br>
 
     <label for="include_bias">Include bias?</label><br>
     <select id="include_bias">
-        <option value="yes">Yes</option>
         <option value="no">No</option>
+        <option value="yes">Yes</option>
     </select><br><br>
 
     <button onclick="calculateParameters()">Calculate</button>
 </div>
 
 <div class="output" id="output">
-    <h3>Model Parameter Results</h3>
+    <h1>Model Parameter Calculator</h1>
+    <h3>Enter model parameters in the console and press calculate</h3>
 </div>
 
 </body>