.

app.py

@@ -23,31 +23,57 @@ def predict_product(num1, num2):
     model.to('cuda' if torch.cuda.is_available() else 'cpu')
 
     generated_ids = inputs['input_ids']
-    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):])
+    prediction = ""
+    correct_product = ""
+    valid_input = True
 
     try:
         num1_int = int(num1)
         num2_int = int(num2)
         correct_product = str(num1_int * num2_int)
     except ValueError:
-        return [], "Invalid input. Could not evaluate correctness."
+        valid_input = False
 
-    # Create the diff for HighlightedText
-    diff = []
-    max_len = max(len(prediction), len(correct_product))
-    for i in range(max_len):
+    eos_token_id = tokenizer.eos_token_id
+    for _ in range(100):  # Set a maximum limit to prevent infinite loops
+        outputs = model.generate(generated_ids, max_new_tokens=1, do_sample=False)
+        generated_ids = torch.cat((generated_ids, outputs[:, -1:]), dim=-1)
+        
+        if outputs[0, -1].item() == eos_token_id:
+            break
+        
+        output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+        prediction = postprocess(output_text[len(input_text):])
+        
+        # Create the diff for HighlightedText
+        diff = []
+        for i in range(max(len(prediction), len(correct_product))):
+            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, ""
+
+    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(max(len(prediction), len(correct_product))):
         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
+            final_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
+            final_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
-
-    result_message = "Correct!" if prediction == correct_product else f"Incorrect! The correct product is {correct_product}."
+            final_diff.append((correct_product[i], "-"))  # Highlight missing/incorrect digits in green
 
-    return diff, result_message
+    yield final_diff, result_message
 
 demo = gr.Interface(
     fn=predict_product,