.

app.py

@@ -56,90 +56,93 @@ def predict_product(num1, num2):
     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_annotations_per_model = {}
-    for step in range(max(MAX_PRODUCT_DIGITS_PER_MODEL.values())):  # Set a maximum limit to prevent infinite loops
-        # Ground Truth
-        if not valid_input:
-            ground_truth_annotations = [('Invalid Input!', None)]
-        else:
-            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]
-            if finished_per_model[model_name]:
-                continue
-            if step >= MAX_PRODUCT_DIGITS_PER_MODEL[model_name]:
-                continue
-            generation_kwargs = {
-                'input_ids': generated_ids_per_model[model_name],
-                'max_new_tokens': 1,
-                'do_sample': False,
-                'past_key_values': past_key_values_per_model[model_name],
-                'return_dict_in_generate': True,
-                'use_cache': True
-            }
-            if step == 0:
-                del generation_kwargs['past_key_values']
-            outputs = model.generate(**generation_kwargs)
-            generated_ids = outputs.sequences
-            next_token_id = generated_ids[0, -1]
-            #print (next_token_id)
-            
-            if next_token_id.item() == tokenizer.eos_token_id:
-                finished_per_model[model_name] = True
-                if valid_input:
-                    if len([item for item in predicted_annotations_per_model[model_name] if item[1] is not None]) < len(ground_truth_digits_reversed):
-                        predicted_annotations_per_model[model_name].insert(0, ('⠀', 'wrong'))
-                continue
+    try:
+        for step in range(max(MAX_PRODUCT_DIGITS_PER_MODEL.values())):  # Set a maximum limit to prevent infinite loops
+            # Ground Truth
+            if not valid_input:
+                ground_truth_annotations = [('Invalid Input!', None)]
+            else:
+                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]
+                if finished_per_model[model_name]:
+                    continue
+                if step >= MAX_PRODUCT_DIGITS_PER_MODEL[model_name]:
+                    continue
+                generation_kwargs = {
+                    'input_ids': generated_ids_per_model[model_name],
+                    'max_new_tokens': 1,
+                    'do_sample': False,
+                    'past_key_values': past_key_values_per_model[model_name],
+                    'return_dict_in_generate': True,
+                    'use_cache': True
+                }
+                if step == 0:
+                    del generation_kwargs['past_key_values']
+                outputs = model.generate(**generation_kwargs)
+                generated_ids = outputs.sequences
+                next_token_id = generated_ids[0, -1]
+                #print (next_token_id)
+                
+                if next_token_id.item() == tokenizer.eos_token_id:
+                    finished_per_model[model_name] = True
+                    if valid_input:
+                        if len([item for item in predicted_annotations_per_model[model_name] if item[1] is not None]) < len(ground_truth_digits_reversed):
+                            predicted_annotations_per_model[model_name].insert(0, ('⠀', 'wrong'))
+                    continue
+                    
+                generated_ids_per_model[model_name] = generated_ids
+                past_key_values_per_model[model_name] = outputs.past_key_values
                 
-            generated_ids_per_model[model_name] = generated_ids
-            past_key_values_per_model[model_name] = outputs.past_key_values
+                output_text = tokenizer.decode(generated_ids[0, input_len:], skip_special_tokens=True)
+                predicted_digits_reversed = output_text.strip().split(' ')
             
-            output_text = tokenizer.decode(generated_ids[0, input_len:], skip_special_tokens=True)
-            predicted_digits_reversed = output_text.strip().split(' ')
-        
-            predicted_annotations = []
-            is_correct_sofar = True
-            if model_name == 'explicit':
-                if '=' not in predicted_digits_reversed:
-                    predicted_annotations = [(predicted_digit, None) for predicted_digit in predicted_digits_reversed]
-                    predicted_digits_reversed = []
-                else:
-                    equal_sign_position = predicted_digits_reversed.index('=')
-                    predicted_annotations = [(predicted_digit, None) for predicted_digit in predicted_digits_reversed[:equal_sign_position+1]]
-                    predicted_digits_reversed = predicted_digits_reversed[equal_sign_position+1:]
-
-            for i in range(len(predicted_digits_reversed)):
-                predicted_digit = predicted_digits_reversed[i]
-                if not valid_input:
-                    is_correct_digit = None
-                elif i >= len(ground_truth_digits_reversed):
-                    if predicted_digit == '0' and is_correct_sofar:
-                        is_correct_digit = True
+                predicted_annotations = []
+                is_correct_sofar = True
+                if model_name == 'explicit':
+                    if '=' not in predicted_digits_reversed:
+                        predicted_annotations = [(predicted_digit, None) for predicted_digit in predicted_digits_reversed]
+                        predicted_digits_reversed = []
                     else:
-                        is_correct_digit = False
-                else:
-                    ground_truth_digit = ground_truth_digits_reversed[i]
-                    if predicted_digit == ground_truth_digit:
-                        is_correct_digit = True
+                        equal_sign_position = predicted_digits_reversed.index('=')
+                        predicted_annotations = [(predicted_digit, None) for predicted_digit in predicted_digits_reversed[:equal_sign_position+1]]
+                        predicted_digits_reversed = predicted_digits_reversed[equal_sign_position+1:]
+
+                for i in range(len(predicted_digits_reversed)):
+                    predicted_digit = predicted_digits_reversed[i]
+                    if not valid_input:
+                        is_correct_digit = None
+                    elif i >= len(ground_truth_digits_reversed):
+                        if predicted_digit == '0' and is_correct_sofar:
+                            is_correct_digit = True
+                        else:
+                            is_correct_digit = False
                     else:
-                        is_correct_digit = False
-                if not is_correct_digit:
-                    is_correct_sofar = False
-                if is_correct_digit is None:
-                    predicted_annotations.append((predicted_digit, None))
-                elif is_correct_digit:
-                    predicted_annotations.append((predicted_digit, "correct"))
-                else:
-                    predicted_annotations.append((predicted_digit, "wrong"))
-            predicted_annotations = predicted_annotations[::-1]
-            predicted_annotations_per_model[model_name] = predicted_annotations
-
-        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_annotations, predicted_annotations_implicit_cot, predicted_annotations_nocot, predicted_annotations_explicit_cot
+                        ground_truth_digit = ground_truth_digits_reversed[i]
+                        if predicted_digit == ground_truth_digit:
+                            is_correct_digit = True
+                        else:
+                            is_correct_digit = False
+                    if not is_correct_digit:
+                        is_correct_sofar = False
+                    if is_correct_digit is None:
+                        predicted_annotations.append((predicted_digit, None))
+                    elif is_correct_digit:
+                        predicted_annotations.append((predicted_digit, "correct"))
+                    else:
+                        predicted_annotations.append((predicted_digit, "wrong"))
+                predicted_annotations = predicted_annotations[::-1]
+                predicted_annotations_per_model[model_name] = predicted_annotations
+
+            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_annotations, predicted_annotations_implicit_cot, predicted_annotations_nocot, predicted_annotations_explicit_cot
+    except Exception as e:
+        pass
 
 color_map = {"correct": "green", "wrong": "red"}