Update app.py

app.py

@@ -68,15 +68,12 @@ def answer_question(text, question, max_length=512):
     start_logits = tf.convert_to_tensor(start_logits)
     end_logits = tf.convert_to_tensor(end_logits)
 
-    # Verify axis type:
-    axis = 1  # Assuming axis 1 for argmax
-    if not isinstance(axis, tf.Tensor) or axis.dtype not in (tf.int32, tf.int64):
-        axis = tf.constant(1, dtype=tf.int32)  # Replace with correct axis if needed
+    # Find the indices of the start and end positions
+    answer_start = tf.argmax(start_logits, axis=1).numpy()[0]
+    answer_end = (tf.argmax(end_logits, axis=1) + 1).numpy()[0]  # Increment by 1 for exclusive end index
 
-    answer_start = tf.math.argmax(start_logits, axis=axis)
-    answer_end = tf.math.argmax(end_logits, axis=1) + 1  # Get predicted end index (exclusive)
-
-    answer = tf.gather(text, answer_start, axis=1).numpy()[0][answer_start[0]:answer_end[0]]
+    # Extract the answer text from the original text
+    answer = text[answer_start:answer_end].strip()
 
     return answer if answer else "No answer found."