.

app.py

@@ -25,15 +25,36 @@ def predict_product(num1, num2):
     output = outputs[0][inputs['input_ids'].shape[-1]:]
     raw_output = tokenizer.decode(output, skip_special_tokens=True)
     prediction = postprocess(raw_output)
-    return input_text, raw_output, prediction
+
+    try:
+        num1_int = int(num1)
+        num2_int = int(num2)
+        valid_input = True
+    except ValueError:
+        valid_input = False
+
+    if valid_input:
+        correct_product = str(num1_int * num2_int)
+        is_correct = (prediction == correct_product)
+        result_color = "green" if is_correct else "red"
+        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."
+
+    return input_text, raw_output, prediction, result_message, result_color
+
+def output_component(value, color):
+    return gr.HTML.update(value=f"<div style='color: {color};'>{value}</div>")
 
 demo = gr.Interface(
     fn=predict_product,
-    inputs=[gr.Number(label='First Number (up to 9 digits)', value=12345), gr.Number(label='Second Number (up to 9 digits)', value=67890)],
+    inputs=[gr.Textbox(label='First Number (up to 9 digits)', value='12345'), gr.Textbox(label='Second Number (up to 9 digits)', value='67890')],
     outputs=[
         gr.Textbox(label='Raw Input to GPT-2'),
         gr.Textbox(label='Raw Output from GPT-2'),
-        gr.Textbox(label='Predicted Product')
+        gr.Textbox(label='Predicted Product'),
+        gr.HTML(label='Result Message')
     ],
     title='GPT-2 Multiplication Predictor',
     description='Enter two numbers up to 9 digits each and get the predicted product.',
@@ -42,6 +63,12 @@ demo = gr.Interface(
     - [Paper: From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step](https://arxiv.org/pdf/2405.14838)
     - [Code Repository](https://github.com/da03/Internalize_CoT_Step_by_Step)
     - [Tweet Announcement](https://twitter.com/yuntiandeng/status/1795854740879774036)
+    """,
+    live=True,
+    css="""
+    .output-html {
+        font-size: 1.5em;
+    }
     """
 )