add param count

app.py

@@ -1,7 +1,17 @@
+from typing import Tuple
+
 import gradio as gr
 
 
-def deepmind_flops(n_layer, d_model, d_ff, d_attn, n_ctx, n_vocab, n_heads):
+def deepmind_flops(
+    n_layer: int,
+    d_model: int,
+    d_ff: int,
+    d_attn: int,
+    n_ctx: int,
+    n_vocab: int,
+    n_heads: int,
+) -> int:
     embeddings = 2 * n_ctx * n_vocab * d_model
     attn_qkv = 2 * n_ctx * 3 * d_model * (d_attn * n_heads)
     attn_logits = 2 * n_ctx * n_ctx * (d_attn * n_heads)
@@ -11,6 +21,12 @@ def deepmind_flops(n_layer, d_model, d_ff, d_attn, n_ctx, n_vocab, n_heads):
     ff = 2 * n_ctx * (d_model * d_ff + d_model * d_ff)
     logits = 2 * n_ctx * d_model * n_vocab
 
+    params = (
+        embeddings / n_ctx / 2,
+        (n_layer * (attn_qkv + attn_project + ff)) / n_ctx / 2,
+        logits / n_ctx / 2,
+    )
+
     return (
         embeddings,
         attn_qkv * n_layer,
@@ -20,25 +36,35 @@ def deepmind_flops(n_layer, d_model, d_ff, d_attn, n_ctx, n_vocab, n_heads):
         attn_project * n_layer,
         ff * n_layer,
         logits,
-    )
+    ), params
 
 
-def calculator(n_layer, d_model, n_heads, n_vocab, n_ctx, ff_ratio, incl_embed):
+def calculator(
+    n_layer: int,
+    d_model: int,
+    n_heads: int,
+    n_vocab: int,
+    n_ctx: int,
+    ff_ratio: int,
+    incl_embed: bool,
+) -> Tuple[int, int, int]:
     d_attn = d_model // n_heads
     if d_model % n_heads != 0:
         raise gr.Error("d_model must be divisible by n_heads")
     d_ff = d_model * ff_ratio
 
-    flops_terms = deepmind_flops(
+    flops_terms, params = deepmind_flops(
         n_layer, d_model, d_ff, d_attn, n_ctx, n_vocab, n_heads
     )
 
     if incl_embed:
         flops_per_sequence = sum(flops_terms)
+        params = sum(params)
     else:
-        flops_per_sequence = sum(flops_terms[1:-1])
+        flops_per_sequence = sum(flops_terms[1:3])
+        params = sum(params[1:3])
 
-    return flops_per_sequence, flops_per_sequence / n_ctx
+    return params, flops_per_sequence, flops_per_sequence / n_ctx
 
 
 with gr.Blocks() as iface:
@@ -54,20 +80,19 @@ with gr.Blocks() as iface:
             n_vocab = gr.Number(label="Vocabulary size (n_vocab)")
             n_ctx = gr.Number(label="Sequence length")
             ff_ratio = gr.Number(value=4, label="Feedforward ratio")
-            incl_embed = gr.Checkbox(
-                value=True, label="Include embedding and logits FLOPs"
-            )
+            incl_embed = gr.Checkbox(value=True, label="Include embeddings")
 
             btn = gr.Button(value="Submit", variant="primary")
 
         with gr.Column():
+            params = gr.Number(label="Model parameters")
             flops_per_sequence = gr.Number(label="FLOPs per sequence")
             flops_per_token = gr.Number(label="FLOPs per token")
 
     btn.click(
         calculator,
         inputs=[n_layer, d_model, n_heads, n_vocab, n_ctx, ff_ratio, incl_embed],
-        outputs=[flops_per_sequence, flops_per_token],
+        outputs=[params, flops_per_sequence, flops_per_token],
     )
 
     gr.Markdown("### GPT-3 model family examples")
@@ -87,7 +112,7 @@ with gr.Blocks() as iface:
             [96, 12288, 96, 50257, 4096, 4, True],
         ],
         [n_layer, d_model, n_heads, n_vocab, n_ctx, ff_ratio, incl_embed],
-        [flops_per_sequence, flops_per_token],
+        [params, flops_per_sequence, flops_per_token],
         calculator,
         cache_examples=False,
     )