Update app.py

app.py

@@ -49,7 +49,9 @@ def run_generation(user_text, use_assistant, temperature, max_new_tokens):
     model_output = ""
     for new_text in streamer:
         model_output += new_text
-        yield [model_output, round(time.time() - start, 3)]
+        time_so_far = round(time.time() - start, 3)
+        tokens_so_far = tokenizer(model_output, return_tensors="pt").input_ids.shape[1]
+        yield [model_output, tokens_so_far/time_so_far]
 
 
 def reset_textbox():
@@ -82,11 +84,11 @@ with gr.Blocks() as demo:
             temperature = gr.Slider(
                 minimum=0.0, maximum=2.0, value=0.6, step=0.05, interactive=True, label="Temperature (0.0 = Greedy)",
             )
-            gr.Markdown("### Generation time (seconds)")
-            generation_time = gr.Textbox(lines=1, interactive=False, show_label=False)
+            gr.Markdown("### Tokens per secon")
+            tokens_per_second = gr.Textbox(lines=1, interactive=False, show_label=False)
 
     generate_inputs = [user_text, use_assistant, temperature, max_new_tokens]
-    generate_outputs = [model_output, generation_time]
+    generate_outputs = [model_output, tokens_per_second]
     user_text.submit(run_generation, generate_inputs, generate_outputs)
     button_submit.click(run_generation, generate_inputs, generate_outputs)