Update app_dialogue.py

app_dialogue.py

@@ -14,6 +14,20 @@ from gradio_client.client import DEFAULT_TEMP_DIR
 from text_generation import Client
 from transformers import AutoProcessor
 
+def timer_decorator(func):
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        elapsed_time = end_time - start_time
+        print(f"{func.__name__} took {elapsed_time:.5f} seconds to run.")
+        return result
+    return wrapper
+
+gr.Blocks.preprocess_data = timer_decorator(gr.Blocks.preprocess_data)
+gr.Blocks.call_function = timer_decorator(gr.Blocks.call_function)
+gr.Blocks.postprocess_data = timer_decorator(gr.Blocks.postprocess_data)
+
 
 MODELS = [  # TODO uncomment
     "HuggingFaceM4/idefics-9b-instruct",
@@ -519,9 +533,7 @@ with gr.Blocks(title="IDEFICS Playground", theme=gr.themes.Base()) as demo:
         stream = client.generate_stream(prompt=query, **generation_args)
 
         acc_text = ""
-        import time
         for idx, response in enumerate(stream):
-            start = time.time()
             text_token = response.token.text
 
             if response.details:
@@ -539,7 +551,6 @@ with gr.Blocks(title="IDEFICS Playground", theme=gr.themes.Base()) as demo:
             last_turn = chat_history.pop(-1)
             last_turn[-1] += acc_text
             chat_history.append(last_turn)
-            print(">", time.time() - start)
             yield "", None, chat_history
             acc_text = ""