Update app.py

app.py

@@ -1,26 +1,14 @@
 import gradio as gr
-import argparse
+import pandas as pd
 from tabulate import tabulate
+from io import StringIO
 
-num_gpu = 1
-prompt_size = 4096
-response_size = 256
-n_concurrent_request = 10
-avg_context_window = 1024
-
-def main(num_gpu, prompt_size, response_size, n_concurrent_request, avg_context_window):
-    #parser = argparse.ArgumentParser(description='Your script description')
-    #parser.add_argument('-g', '--num_gpu', type=int, default=1, help='Number of GPUs')
-    #parser.add_argument('-p', '--prompt_sz', type=int, default=4096, help='Prompt size in tokens')
-    #parser.add_argument('-r', '--response_sz', type=int, default=256, help='Response size in tokens')
-    #parser.add_argument('-c', '--n_concurrent_req', type=int, default=10, help='Number of concurrent requests')
-    #parser.add_argument('-w', '-cw', '--ctx_window', type=int, default=1024, help='Average context window')
-
-    # args = parser.parse_args()
-
-    # Print input
-    print(f" num_gpu = {num_gpu}, prompt_size = {prompt_size} tokens, response_size = {response_size} tokens")
-    print(f" n_concurrent_request = {n_concurrent_request}, avg_context_window = {avg_context_window} tokens")
+def calculate_llm_metrics(num_gpu, prompt_size, response_size, n_concurrent_request, avg_context_window):
+    output = StringIO()
+    
+    # Print input to output buffer
+    print(f" num_gpu = {num_gpu}, prompt_size = {prompt_size} tokens, response_size = {response_size} tokens", file=output)
+    print(f" n_concurrent_request = {n_concurrent_request}, avg_context_window = {avg_context_window} tokens", file=output)
 
     # Define variables
     gpu_specs = [
@@ -35,7 +23,6 @@ def main(num_gpu, prompt_size, response_size, n_concurrent_request, avg_context_
         {"name": "H100 PCIe", "fp16_tflops": 1513, "memory_gb": 80, "memory_bandwidth_gbps": 2000},
         {"name": "H100 SXM", "fp16_tflops": 1979, "memory_gb": 80, "memory_bandwidth_gbps": 3350},
         {"name": "H100 NVL", "fp16_tflops": 3958, "memory_gb": 188, "memory_bandwidth_gbps": 7800}
-        # Add or comment out GPU types as needed
     ]
 
     model_specs = [
@@ -47,26 +34,27 @@ def main(num_gpu, prompt_size, response_size, n_concurrent_request, avg_context_
         {"name": "Falcon-7B", "params_billion": 7, "d_model": 4544, "n_heads": 71, "n_layers": 32, "max_context_window": 2048, "d_head": 64},
         {"name": "Falcon-40B", "params_billion": 40, "d_model": 8192, "n_heads": 128, "n_layers": 60, "max_context_window": 2048, "d_head": 64},
         {"name": "Falcon-180B", "params_billion": 180, "d_model": 14848, "n_heads": 232, "n_layers": 80, "max_context_window": 2048, "d_head": 64}
-        # Add or comment out model specifications as needed
     ]
 
-    BYTES_IN_GB = 1_073_741_824  # 1 GB = 1,073,741,824 bytes
+    BYTES_IN_GB = 1_073_741_824
 
     def calc_kv_cache_size_per_token(n_layers, d_model):
-        return 2 * 2 * n_layers * d_model / BYTES_IN_GB  # GB/token
+        return 2 * 2 * n_layers * d_model / BYTES_IN_GB
 
     def calc_memory_footprint(model_spec, n_concurrent_request, avg_context_window):
         kv_cache_size_per_token = calc_kv_cache_size_per_token(model_spec["n_layers"], model_spec["d_model"])
         target_gpu_mem = kv_cache_size_per_token * avg_context_window * n_concurrent_request + model_spec["params_billion"] * 2
         return target_gpu_mem
 
-    print(f"\n******************** Estimate LLM Memory Footprint ********************")
+    print(f"\n******************** Estimate LLM Memory Footprint ********************", file=output)
     memory_footprint_table = []
     for model_spec in model_specs:
         kv_cache_size_per_token = calc_kv_cache_size_per_token(model_spec["n_layers"], model_spec["d_model"])
         memory_footprint = calc_memory_footprint(model_spec, n_concurrent_request, avg_context_window)
         memory_footprint_table.append([model_spec['name'], f"{kv_cache_size_per_token:.6f} GiB/token", f"{memory_footprint:.2f} GB"])
-    print(tabulate(memory_footprint_table, headers=['Model', 'KV Cache Size per Token', 'Memory Footprint'], tablefmt='orgtbl'))
+    
+    memory_df = pd.DataFrame(memory_footprint_table, columns=['Model', 'KV Cache Size per Token', 'Memory Footprint'])
+    print(tabulate(memory_footprint_table, headers=['Model', 'KV Cache Size per Token', 'Memory Footprint'], tablefmt='orgtbl'), file=output)
 
     def calc_kv_cache_tokens(num_gpu, gpu_memory_gb, model_params_billion, kv_cache_size):
         result = (num_gpu * gpu_memory_gb - 2 * model_params_billion) / kv_cache_size
@@ -81,14 +69,13 @@ def main(num_gpu, prompt_size, response_size, n_concurrent_request, avg_context_
         return result if result >= 0 else "OOM"
 
     def calc_estimated_response_time(prefill_time, generation_time, prompt_size, response_size):
-        if isinstance(prefill_time, str) or isinstance(generation_time, str):  # Check if any are "NA"
+        if isinstance(prefill_time, str) or isinstance(generation_time, str):
             return "OOM"
-        return (prompt_size * prefill_time + response_size * generation_time) / 1000  # convert ms to seconds
+        return (prompt_size * prefill_time + response_size * generation_time) / 1000
 
-    print(f"\n******************** Estimate LLM Capacity and Latency ******************** ")
+    print(f"\n******************** Estimate LLM Capacity and Latency ******************** ", file=output)
     capacity_latency_table = []
     for model in model_specs:
-        # print(f"Model: {model['name']} ({model['params_billion']}B parameters)")
         kv_cache_size = calc_kv_cache_size_per_token(model['n_layers'], model['d_model'])
         for gpu in gpu_specs:
             kv_cache_tokens = calc_kv_cache_tokens(num_gpu, gpu['memory_gb'], model['params_billion'], kv_cache_size)
@@ -96,8 +83,43 @@ def main(num_gpu, prompt_size, response_size, n_concurrent_request, avg_context_
             generation_time_per_token = calc_generation_time_per_token(num_gpu, model['params_billion'], gpu['memory_bandwidth_gbps'])
             estimated_response_time = calc_estimated_response_time(prefill_time_per_token, generation_time_per_token, prompt_size, response_size)
             capacity_latency_table.append([model['name'], gpu['name'], f"{kv_cache_tokens}", f"{prefill_time_per_token:.3f} ms", f"{generation_time_per_token:.3f} ms", f"{estimated_response_time:.1f} s"])
-    print(tabulate(capacity_latency_table, headers=['Model', 'GPU', 'KV Cache Tokens', 'Prefill Time', 'Generation Time', 'Estimated Response Time'], tablefmt='orgtbl'))
-
+    
+    capacity_df = pd.DataFrame(capacity_latency_table, columns=['Model', 'GPU', 'KV Cache Tokens', 'Prefill Time', 'Generation Time', 'Estimated Response Time'])
+    print(tabulate(capacity_latency_table, headers=['Model', 'GPU', 'KV Cache Tokens', 'Prefill Time', 'Generation Time', 'Estimated Response Time'], tablefmt='orgtbl'), file=output)
+    
+    return output.getvalue(), memory_df, capacity_df
+
+# Create Gradio interface
+with gr.Blocks(title="LLM Calculator") as demo:
+    gr.Markdown("# LLM Memory and Performance Calculator")
+    
+    with gr.Row():
+        with gr.Column():
+            num_gpu = gr.Slider(minimum=1, maximum=8, value=1, step=1, label="Number of GPUs")
+            prompt_size = gr.Slider(minimum=1, maximum=8192, value=4096, step=1, label="Prompt Size (tokens)")
+            response_size = gr.Slider(minimum=1, maximum=2048, value=256, step=1, label="Response Size (tokens)")
+            n_concurrent_request = gr.Slider(minimum=1, maximum=50, value=10, step=1, label="Number of Concurrent Requests")
+            avg_context_window = gr.Slider(minimum=1, maximum=8192, value=1024, step=1, label="Average Context Window (tokens)")
+            
+            calculate_button = gr.Button("Calculate")
+    
+    with gr.Row():
+        with gr.Column():
+            text_output = gr.Textbox(label="Detailed Output", lines=10)
+    
+    with gr.Row():
+        with gr.Column():
+            memory_table = gr.Dataframe(label="Memory Footprint Results")
+    
+    with gr.Row():
+        with gr.Column():
+            capacity_table = gr.Dataframe(label="Capacity and Latency Results")
+    
+    calculate_button.click(
+        calculate_llm_metrics,
+        inputs=[num_gpu, prompt_size, response_size, n_concurrent_request, avg_context_window],
+        outputs=[text_output, memory_table, capacity_table]
+    )
 
 if __name__ == "__main__":
-    main(num_gpu, prompt_size, response_size, n_concurrent_request, avg_context_window)
+    demo.launch()