.

app.py

@@ -23,48 +23,31 @@ def predict_product(num1, num2):
     model.to('cuda' if torch.cuda.is_available() else 'cpu')
 
     generated_ids = inputs['input_ids']
-    prediction = ""
-    correct_product = ""
-    valid_input = True
+    outputs = model.generate(generated_ids, max_new_tokens=40, do_sample=False)
+    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    prediction = postprocess(full_output[len(input_text):])
 
     try:
         num1_int = int(num1)
         num2_int = int(num2)
         correct_product = str(num1_int * num2_int)
     except ValueError:
-        valid_input = False
+        return [], "Invalid input. Could not evaluate correctness."
 
-    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)
-        
-        # 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
+    # Create the diff for HighlightedText
+    diff = []
+    max_len = max(len(prediction), len(correct_product))
+    for i in range(max_len):
+        if i < len(prediction) and i < len(correct_product) and prediction[i] == correct_product[i]:
+            diff.append((prediction[i], None))  # No highlight for correct digits
+        elif i < len(prediction) and (i >= len(correct_product) or prediction[i] != correct_product[i]):
+            diff.append((prediction[i], "+"))  # Highlight incorrect digits in red
+        if i < len(correct_product) and (i >= len(prediction) or prediction[i] != correct_product[i]):
+            diff.append((correct_product[i], "-"))  # Highlight missing/incorrect digits in green
 
-        yield diff, ""
+    result_message = "Correct!" if prediction == correct_product else f"Incorrect! The correct product is {correct_product}."
 
-    if valid_input:
-        is_correct = prediction == correct_product
-        result_message = "Correct!" if is_correct else f"Incorrect! The correct product is {correct_product}."
-    else:
-        result_message = "Invalid input. Could not evaluate correctness."
-
-    # 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
+    return diff, result_message
 
 demo = gr.Interface(
     fn=predict_product,
@@ -73,7 +56,7 @@ demo = gr.Interface(
         gr.Textbox(label='Second Number (up to 12 digits)', value='67890'),
     ],
     outputs=[
-        gr.HighlightedText(label='Predicted Product with Matching Digits Highlighted', combine_adjacent=True, show_legend=True, color_map={"+": "red"}),
+        gr.HighlightedText(label='Predicted Product with Matching and Unmatching Digits Highlighted', combine_adjacent=True, show_legend=True, color_map={"-": "green", "+": "red"}),
         gr.HTML(label='Result Message')
     ],
     title='GPT2 Direct Multiplication Calculator (Without Using Chain-of-Thought)',