Handle input errors

app.py

@@ -1,6 +1,7 @@
 from transformers import AutoModelForCausalLM
 import torch
 import gradio as gr
+import re
 
 model = AutoModelForCausalLM.from_pretrained("Manuel2011/addition_model")
 
@@ -22,7 +23,14 @@ class NumberTokenizer:
 tokenizer = NumberTokenizer(13)
 
 def generate_solution(input, solution_length=6, model=model):
+  try:
+    parsed_input = re.search(r'(\d)\s*\+\s*(\d)', input)
+    first_number = int(parsed_input.group(1))
+    second_number = int(parsed_input.group(2))
+  except:
+    return 'Invalid input'
   model.eval()
+  input = f'{first_number} + {second_number} ='
   input = torch.tensor(tokenizer(input))
   input = input
   solution = []