.

app.py

@@ -51,14 +51,11 @@ def predict_product(num1, num2):
     generated_ids_per_model = {model_name: inputs['input_ids'].data.clone() for model_name in models}
     finished_per_model = {model_name: False for model_name in models}
     past_key_values_per_model = {model_name: None for model_name in models}
-    predicted_results_per_model = {}
+    predicted_annotations_per_model = {}
     for step in range(max(MAX_PRODUCT_DIGITS_PER_MODEL.values())):  # Set a maximum limit to prevent infinite loops
         # Ground Truth
-        ground_truth_results = []
-        for i in range(step+1):
-            ground_truth_digit = ground_truth_digits_reversed[i]
-            ground_truth_results.append((ground_truth_digit, None))
-        ground_truth_results = ground_truth_results[::-1]
+        ground_truth_annotations = [(ground_truth_digit, None) for ground_truth_digit in ground_truth_digits_reversed[:step+1]]
+        ground_truth_annotations = ground_truth_annotations[::-1]
         # Predicted
         for model_name in models:
             model = models[model_name]
@@ -91,7 +88,7 @@ def predict_product(num1, num2):
             output_text = tokenizer.decode(generated_ids[0, input_len:], skip_special_tokens=True)
             predicted_digits_reversed = output_text.strip().split(' ')
         
-            predicted_results = []
+            predicted_annotations = []
             is_correct_sofar = True
             for i in range(len(predicted_digits_reversed)):
                 predicted_digit = predicted_digits_reversed[i]
@@ -109,17 +106,17 @@ def predict_product(num1, num2):
                 if not is_correct_digit:
                     is_correct_sofar = False
                 if is_correct_digit:
-                    predicted_results.append((predicted_digit, "correct"))
+                    predicted_annotations.append((predicted_digit, "correct"))
                 else:
-                    predicted_results.append((predicted_digit, "wrong"))
-            predicted_results = predicted_results[::-1]
-            predicted_results_per_model[model_name] = predicted_results
+                    predicted_annotations.append((predicted_digit, "wrong"))
+            predicted_annotations = predicted_annotations[::-1]
+            predicted_annotations_per_model[model_name] = predicted_annotations
 
-        predicted_results_implicit_cot = predicted_results_per_model['implicit']
-        predicted_results_nocot = predicted_results_per_model['no']
-        predicted_results_explicit_cot = predicted_results_per_model['explicit']
+        predicted_annotations_implicit_cot = predicted_annotations_per_model['implicit']
+        predicted_annotations_nocot = predicted_annotations_per_model['no']
+        predicted_annotations_explicit_cot = predicted_annotations_per_model['explicit']
 
-        yield ground_truth_results, predicted_results_implicit_cot, predicted_results_nocot, predicted_results_explicit_cot
+        yield ground_truth_annotations, predicted_annotations_implicit_cot, predicted_annotations_nocot, predicted_annotations_explicit_cot
 
 color_map = {"correct": "green", "wrong": "red"}