.

app.py

@@ -18,6 +18,13 @@ def postprocess(raw_output):
 
 @spaces.GPU
 def predict_product(num1, num2):
+    try:
+        num1_int = int(num1)
+        num2_int = int(num2)
+        valid_input = True
+    except ValueError:
+        valid_input = False
+
     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')
@@ -26,13 +33,6 @@ def predict_product(num1, num2):
     raw_output = tokenizer.decode(output, skip_special_tokens=True)
     prediction = postprocess(raw_output)
 
-    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)
@@ -42,33 +42,29 @@ def predict_product(num1, num2):
         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>")
+    return input_text, raw_output, prediction, result_message
 
 demo = gr.Interface(
     fn=predict_product,
-    inputs=[gr.Textbox(label='First Number (up to 9 digits)', value='12345'), gr.Textbox(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.HTML(label='Result Message')
     ],
-    title='GPT-2 Multiplication Predictor',
-    description='Enter two numbers up to 9 digits each and get the predicted product.',
+    title='GPT-2 Multiplication Calculator',
+    description='This demo uses GPT-2 to directly predict the product of two numbers without using any intermediate steps.',
     article="""
     ### Additional Resources
     - [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)
     """,
-    css="""
-    .output-html {
-        font-size: 1.5em;
-    }
-    """
+    live=False
 )
 
 demo.launch()