Update app.py

app.py

@@ -50,15 +50,24 @@ def predict_answer(context, question):
 # Load short model
 short_answer_model = pipeline(model="Nighter/QA_wiki_data_short_answer", from_tf=True)
 
-def answer_questions(context, question):
-    long_answer = predict_answer(context, question)
-    short_answer_result = short_answer_model(question=question, context=remove_parentheses(context))
-    return short_answer_result['answer'], long_answer
+# def answer_questions(context, question):
+#     long_answer = predict_answer(context, question)
+#     short_answer_result = short_answer_model(question=question, context=remove_parentheses(context))
+#     return short_answer_result['answer'], long_answer
 
 def answer_questions(context, question):
     long_score, long_answer = predict_answer(context, question)
-    short_answer_result = short_answer_model(question=question, context=remove_parentheses(context))
-    # short_answer_result = short_answer_model(question=question, context=remove_parentheses(long_answer))
+
+    # Check if the original context is longer than 512 tokens
+    if len(tokenizer.texts_to_sequences([context])[0]) > 512:
+        # If yes, use the long answer as the context for the short answer model
+        short_context = long_answer
+    else:
+        # If no, use the original context
+        short_context = remove_parentheses(context)
+
+    short_answer_result = short_answer_model(question=question, context=short_context)
+
     return short_answer_result['answer'], short_answer_result['score'], long_answer, long_score