.

app.py

@@ -18,13 +18,6 @@ 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')
@@ -33,6 +26,13 @@ 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,7 +42,9 @@ 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_html = f"<div style='color: {result_color};'>{result_message}</div>"
+
+    return input_text, raw_output, prediction, result_html
 
 demo = gr.Interface(
     fn=predict_product,