xx

app.py

@@ -81,10 +81,12 @@ def model_inference(input_dict, history):
     streamer = TextIteratorStreamer(
         processor, skip_prompt=True, skip_special_tokens=True
     )
-    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=2048)
+    # start timer just before generation begins
+    start_time = time.time()
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
-    buffer = ""
+    buffer = "Baseline Model Response: "
     yield progress_bar_html("Processing...")
     for new_text in streamer:
         escaped_new_text = html.escape(new_text)
@@ -93,6 +95,13 @@ def model_inference(input_dict, history):
         time.sleep(0.001)
         yield buffer
 
+    # Ensure generation thread has finished and measure elapsed time
+    thread.join()
+    elapsed = time.time() - start_time
+    elapsed_text = f"\nBaseline generation time: {elapsed:.2f} s"
+    buffer += html.escape(elapsed_text)
+    yield buffer
+
 
 examples = [
     [