Update app.py

app.py

@@ -5,6 +5,14 @@ from tabulate import tabulate
 def greet(name):
     return f"Ciao, {name}!"
 
+# Move estimate_capacity_latency outside of main()
+def estimate_capacity_latency(model, gpu):
+    kv_cache_tokens = calc_kv_cache_tokens(num_gpu, gpu['memory_gb'], model['params_billion'], kv_cache_size_per_token)
+    prefill_time_per_token = calc_prefill_time_per_token(num_gpu, model['params_billion'], gpu['fp16_tflops'])
+    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)
+    return f"{prefill_time_per_token:.3f} ms", f"{generation_time_per_token:.3f} ms", f"{estimated_response_time:.1f} s"
+
 def main():
     parser = argparse.ArgumentParser(description='Your script description')
     parser.add_argument('-g', '--num_gpu', type=int, default=1, help='Number of GPUs')
@@ -87,12 +95,6 @@ def main():
         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'))
 
-    def estimate_capacity_latency(model, gpu):
-        kv_cache_tokens = calc_kv_cache_tokens(num_gpu, gpu['memory_gb'], model['params_billion'], kv_cache_size_per_token)
-        prefill_time_per_token = calc_prefill_time_per_token(num_gpu, model['params_billion'], gpu['fp16_tflops'])
-        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)
-        return f"{prefill_time_per_token:.3f} ms", f"{generation_time_per_token:.3f} ms", f"{estimated_response_time:.1f} s"
 
     capacity_latency_table = []
     for model in model_specs:
@@ -106,6 +108,7 @@ def main():
 if __name__ == '__main__':
     main()
 
+# Modify create_gradio_interface to use the global estimate_capacity_latency
 def create_gradio_interface():
     demo = gr.Interface(
         fn=estimate_capacity_latency,
@@ -123,8 +126,8 @@ def create_gradio_interface():
     
     return demo
 
-# Creare l'interfaccia Gradio
+# Create the Gradio interface
 gr_interface = create_gradio_interface()
 
-# Avviare l'interfaccia
-gr_interface.launch()
+# Start the interface
+gr_interface.launch()