Time generation. Align responses.

app.py

@@ -22,13 +22,17 @@ print(f'transformers=={transformers.__version__}')
 print(f'accelerate=={accelerate.__version__}')
 print()
 
+# Initialize logging.
+logging.basicConfig(format='%(levelname)s:%(name)s: %(message)s', level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # gr.DataFrame is currently bugged for updating values,
 # so we must use raw HTML.
 # https://github.com/gradio-app/gradio/issues/8160
 def make_html_table(headers, data):
     rows = ['<tr>' + ''.join(f'<th style="width: 50%">{h}</th>' for h in headers) + '</tr>\n']
     for row in data:
-        rows.append('<tr>' + ''.join(f'<td style="width: 50%; font-family: monospace; white-space: pre-wrap;">{v}</td>' for v in row) + '</tr>\n')
+        rows.append('<tr>' + ''.join(f'<td style="width: 50%; font-family: monospace; white-space: pre-wrap; text-align: left; vertical-align: top;">{v}</td>' for v in row) + '</tr>\n')
     return '<table style="width: 100%; table-layout: fixed">\n' + ''.join(rows) + '</table>\n'
 
 def highlight_prefix(tokens, prefix_len):
@@ -56,7 +60,7 @@ def format_response_pair(tokens_a, tokens_b):
 HEADERS = ['Response (Left)', 'Response (Right)']
 repo_id = "Qwen/Qwen2-0.5B-Instruct"
 
-DRY_RUN = False
+DRY_RUN = True
 
 if DRY_RUN:
     from load import load_tokenizer
@@ -64,13 +68,16 @@ if DRY_RUN:
     tokenizer = load_tokenizer(repo_id)
 
     def fn(max_tokens, num_responses, prompt_x, prompt_y):
+        logger.info('Starting generation...')
+        generation_start = time.perf_counter()
+
         rows = [['']*2 for i in range(num_responses)]
         
         yield make_html_table(HEADERS, rows)
 
         for j in range(num_responses):
             response_raw_a = f'Sure!\n\n1 2 3 4 & 5.'
-            response_raw_b = f'Sure!\n\n1 2 3 4 5 & 6.'
+            response_raw_b = f'Sure!\n\n1 2 3 4 5 &\n\n\n\n6.'
 
             response_tok_a = tokenizer.encode(response_raw_a, add_special_tokens=False, return_tensors='np')[0]
             response_tok_b = tokenizer.encode(response_raw_b, add_special_tokens=False, return_tensors='np')[0]
@@ -78,7 +85,7 @@ if DRY_RUN:
             steps = 1 + max(len(response_tok_a), len(response_tok_b))
 
             for i in range(steps):
-                time.sleep(0.1)
+                time.sleep(0.01)
                 prefix_tok_a = response_tok_a[:i]
                 prefix_tok_b = response_tok_b[:i]
 
@@ -88,12 +95,13 @@ if DRY_RUN:
                 rows[j][1] = content_b
 
                 yield make_html_table(HEADERS, rows)
+
+        generation_end = time.perf_counter()
+        logger.info(f'Generation took {(generation_end - generation_start):.3f} s')
 else:
     from load import load_model
     import algorithms
-
-    logging.basicConfig(format='%(levelname)s:%(name)s: %(message)s')
-    algorithms.logger.setLevel(logging.INFO)
+    #algorithms.logger.setLevel(logging.DEBUG)
 
     model, tokenizer = load_model(repo_id)
 
@@ -112,6 +120,9 @@ else:
 
     @spaces.GPU
     def fn(max_tokens, num_responses, prompt_x, prompt_y):
+        logger.info('Starting generation...')
+        generation_start = time.perf_counter()
+
         # Is this necessary with ZeroGPU?
         torch.use_deterministic_algorithms(True)
 
@@ -151,6 +162,9 @@ else:
                 
                 yield make_html_table(HEADERS, rows)
 
+        generation_end = time.perf_counter()
+        logger.info(f'Generation took {(generation_end - generation_start):.3f} s')
+
 demo = gr.Interface(
     fn=fn,
     inputs=[