.

app.py

@@ -2,7 +2,6 @@ import spaces
 import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import time
 
 model_name = 'yuntian-deng/gpt2-implicit-cot-multiplication'
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -41,13 +40,15 @@ def predict_product(num1, num2):
         output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
         prediction = postprocess(output_text)
         
-        result_html = "<div style='margin-bottom: 10px;'>Correct Result: " + " ".join(correct_product) + "</div><div>"
-        for i, pred_digit in enumerate(prediction):
-            color = "green" if i < len(correct_product) and pred_digit == correct_product[i] else "red"
-            result_html += f"<span style='color: {color};'>{pred_digit}</span>"
-        result_html += "</div>"
+        # Manually create the diff for HighlightedText
+        diff = []
+        for i in range(len(prediction)):
+            if i < len(correct_product) and prediction[i] == correct_product[i]:
+                diff.append((prediction[i], None))  # No highlight for correct digits
+            else:
+                diff.append((prediction[i], "+"))  # Highlight incorrect digits in red
 
-        yield result_html, ""
+        yield diff, ""
 
     if valid_input:
         is_correct = prediction == correct_product
@@ -55,7 +56,15 @@ def predict_product(num1, num2):
     else:
         result_message = "Invalid input. Could not evaluate correctness."
 
-    yield result_html, result_message
+    # Final diff for the complete prediction
+    final_diff = []
+    for i in range(len(prediction)):
+        if i < len(correct_product) and prediction[i] == correct_product[i]:
+            final_diff.append((prediction[i], None))  # No highlight for correct digits
+        else:
+            final_diff.append((prediction[i], "+"))  # Highlight incorrect digits in red
+
+    yield final_diff, result_message
 
 demo = gr.Interface(
     fn=predict_product,
@@ -64,7 +73,7 @@ demo = gr.Interface(
         gr.Textbox(label='Second Number (up to 12 digits)', value='67890'),
     ],
     outputs=[
-        gr.HTML(label='Predicted Product with Matching Digits Highlighted'),
+        gr.HighlightedText(label='Predicted Product with Matching Digits Highlighted', combine_adjacent=True, show_legend=True, color_map={"+": "red"}),
         gr.HTML(label='Result Message')
     ],
     title='GPT2 Direct Multiplication Calculator (Without Using Chain-of-Thought)',
@@ -76,7 +85,7 @@ demo = gr.Interface(
     """,
     clear_btn=None,
     submit_btn="Multiply!",
-    live=True
+    live=False
 )
 
 demo.launch()