.

app.py

@@ -2,6 +2,7 @@ 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)
@@ -18,41 +19,43 @@ def postprocess(raw_output):
 
 @spaces.GPU
 def predict_product(num1, num2):
-    # Reverse input digits and add spaces
     input_text = f'{preprocess(num1)} * {preprocess(num2)} ='
-
     inputs = tokenizer(input_text, return_tensors='pt').to('cuda' if torch.cuda.is_available() else 'cpu')
     model.to('cuda' if torch.cuda.is_available() else 'cpu')
 
-    # Generate output
-    outputs = model.generate(**inputs, max_new_tokens=40)
-
-    output = outputs[0][inputs['input_ids'].shape[-1]:]
-    raw_output = tokenizer.decode(output, skip_special_tokens=True)
-    prediction = postprocess(raw_output)
+    generated_ids = inputs['input_ids']
+    prediction = ""
+    correct_product = ""
+    valid_input = True
 
-    # Evaluate the correctness of the result
     try:
         num1_int = int(num1)
         num2_int = int(num2)
-        valid_input = True
+        correct_product = str(num1_int * num2_int)
     except ValueError:
         valid_input = False
+
+    for _ in range(40):  # Adjust the range to control the maximum number of generated tokens
+        outputs = model.generate(generated_ids, max_new_tokens=1, do_sample=False)
+        generated_ids = torch.cat((generated_ids, outputs[:, -1:]), dim=-1)
+        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>"
+
+        yield result_html, ""
+
     if valid_input:
-        correct_product = num1_int * num2_int
-        try:
-            prediction_int = int(prediction)
-            is_correct = (prediction_int == correct_product)
-        except ValueError:
-            is_correct = False
-        result_color = "green" if is_correct else "red"
+        is_correct = prediction == correct_product
         result_message = "Correct!" if is_correct else f"Incorrect! The correct product is {correct_product}."
     else:
-        result_color = "black"
         result_message = "Invalid input. Could not evaluate correctness."
-    result_html = f"<div style='color: {result_color};'>{result_message}</div>"
 
-    return prediction, result_html
+    yield result_html, result_message
 
 demo = gr.Interface(
     fn=predict_product,
@@ -61,7 +64,7 @@ demo = gr.Interface(
         gr.Textbox(label='Second Number (up to 12 digits)', value='67890'),
     ],
     outputs=[
-        gr.Textbox(label='Predicted Product'),
+        gr.HTML(label='Predicted Product with Matching Digits Highlighted'),
         gr.HTML(label='Result Message')
     ],
     title='GPT2 Direct Multiplication Calculator (Without Using Chain-of-Thought)',
@@ -73,7 +76,7 @@ demo = gr.Interface(
     """,
     clear_btn=None,
     submit_btn="Multiply!",
-    live=False
+    live=True
 )
 
 demo.launch()