Update app.py

app.py

@@ -16,19 +16,22 @@ MAX_RESULT_TOKENS = 10
 
 @spaces.GPU
 def predict_answer(question):
-    input_text = ' '.join(question.split()).strip() + ' ' + tokenizer.eos_token
-    print (input_text)
-    inputs = tokenizer(input_text, return_tensors='pt').to('cuda' if torch.cuda.is_available() else 'cpu')
-    implicit_cot_model.to('cuda' if torch.cuda.is_available() else 'cpu')
-
-    input_ids = inputs['input_ids']
-    print (input_ids)
-    outputs = implicit_cot_model.generate(input_ids=input_ids,
-            max_new_tokens=MAX_RESULT_TOKENS,
-            do_sample=False)
-    print (outputs)
-
-    prediction = tokenizer.decode(outputs[0, input_ids.shape[-1]:], skip_special_tokens=True)
+    try:
+        input_text = ' '.join(question.split()).strip() + ' ' + tokenizer.eos_token
+        print (input_text)
+        inputs = tokenizer(input_text, return_tensors='pt').to('cuda' if torch.cuda.is_available() else 'cpu')
+        implicit_cot_model.to('cuda' if torch.cuda.is_available() else 'cpu')
+    
+        input_ids = inputs['input_ids']
+        #print (input_ids)
+        outputs = implicit_cot_model.generate(input_ids=input_ids,
+                max_new_tokens=MAX_RESULT_TOKENS,
+                do_sample=False)
+        #print (outputs)
+    
+        prediction = tokenizer.decode(outputs[0, input_ids.shape[-1]:], skip_special_tokens=True)
+    except Exception as e:
+        prediction = f'{e}'
 
     return prediction