.

app.py

@@ -6,6 +6,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 model_name = 'yuntian-deng/gpt2-implicit-cot-multiplication'
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
+MAX_PRODUCT_DIGITS = 100
 
 def preprocess(num):
     num = str(num).strip().replace(' ', '')
@@ -21,8 +22,10 @@ def predict_product(num1, num2):
     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')
+    eos_token_id = tokenizer.eos_token_id
 
-    generated_ids = inputs['input_ids']
+    input_ids = inputs['input_ids']
+    input_len = input_ids.shape[-1]
     prediction = ""
     correct_product = ""
     valid_input = True
@@ -34,55 +37,74 @@ def predict_product(num1, num2):
     except ValueError:
         valid_input = False
 
-    eos_token_id = tokenizer.eos_token_id
+    generated_ids = inputs['input_ids']
     past_key_values = None
-    for _ in range(100):  # Set a maximum limit to prevent infinite loops
+    for _ in range(MAX_PRODUCT_DIGITS):  # Set a maximum limit to prevent infinite loops
         outputs = model(
             input_ids=generated_ids,
             past_key_values=past_key_values,
             use_cache=True
         )
         logits = outputs.logits
-        past_key_values = outputs.past_key_values
         
         next_token_id = torch.argmax(logits[:, -1, :], dim=-1)
-        generated_ids = torch.cat((generated_ids, next_token_id.unsqueeze(-1)), dim=-1)
+        generated_ids = torch.cat((generated_ids, next_token_id.view(1,-1)), dim=-1)
         
         if next_token_id.item() == eos_token_id:
             break
+        past_key_values = outputs.past_key_values
         
-        output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-        prediction = postprocess(output_text[len(input_text):])
+        output_text = tokenizer.decode(generated_ids[0, input_len:], skip_special_tokens=True)
+        #prediction = postprocess(output_text)
+        predicted_digits_reversed = output_text.strip().split(' ')
+        correct_digits_reversed = ' '.join(correct_product)[::-1]
         
         # 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
+        correct_digits = []
+        is_correct_sofar = True
+        for i in range(len(predicted_digits_reversed)):
+            predicted_digit = predicted_digits_reversed[i]
+            correct_digit = correct_digits_reversed[i]
+            correct_digits.append((correct_digit, None))
+            if i >= len(correct_digits_reversed):
+                if predicted_digit == '0' and is_correct_sofar:
+                    is_correct_digit = True
+                else:
+                    is_correct_digit = True
+            else:
+                if predicted_digit == correct_digit:
+                    is_correct_digit = True
+                else:
+                    is_correct_digit = False
+            if not is_correct_digit:
+                is_correct_sofar = False
+            if is_correct_digit:
+                diff.append((correct_product[i], "-"))
+            else:
+                diff.append((predicted_digit, "+"))
+        diff = diff[::-1]
+        correct_digits = correct_digits[::-1]
 
-        yield diff, ""
+        yield correct_digits, 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."
+    #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]:
-            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]):
-            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]):
-            final_diff.append((correct_product[i], "-"))  # Highlight missing/incorrect digits in green
+    ## 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]:
+    #        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]):
+    #        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]):
+    #        final_diff.append((correct_product[i], "-"))  # Highlight missing/incorrect digits in green
 
-    yield final_diff, result_message
+    #yield final_diff, result_message
 
 demo = gr.Interface(
     fn=predict_product,
@@ -91,7 +113,8 @@ demo = gr.Interface(
         gr.Textbox(label='Second Number (up to 12 digits)', value='67890'),
     ],
     outputs=[
-        gr.HighlightedText(label='Predicted Product with Matching and Unmatching Digits Highlighted', combine_adjacent=True, show_legend=True, color_map={"-": "green", "+": "red"}),
+        gr.Textbox(label='Ground Truth Product'),
+        gr.HighlightedText(label='Predicted Product', combine_adjacent=False, show_legend=False, color_map={"-": "green", "+": "red"}),
         gr.HTML(label='Result Message')
     ],
     title='GPT2 Direct Multiplication Calculator (Without Using Chain-of-Thought)',