Update app.py

app.py

@@ -44,20 +44,36 @@ def answer_question(text, question, max_length=512):
     Returns:
         The answer extracted from the text using the model.
     """
-    qa_model_name = "bert-base-uncased"  # Replace with your chosen model
+    
+    qa_model_name = "bert-base-uncased"  # Replace with your model
 
     qa_model = TFBertForQuestionAnswering.from_pretrained(qa_model_name)
     tokenizer = AutoTokenizer.from_pretrained(qa_model_name)
 
-    # Truncate text if necessary
+    # Truncate text if necessary:
     if len(text) > max_length:
         text = text[:max_length]
 
-    inputs = tokenizer(question, text, return_tensors="tf")  # Tokenize inputs for TensorFlow
+    # Add special tokens and tokenize:
+    inputs = tokenizer(
+        question, text, return_tensors="tf", padding="max_length", truncation=True
+    )
 
     start_logits, end_logits = qa_model(inputs)
 
-    answer_start = tf.math.argmax(start_logits, axis=1)  # Get predicted start index
+    # Check and ensure data type of start_logits:
+    if start_logits.dtype not in (tf.float32, tf.int32):
+        start_logits = tf.cast(start_logits, tf.float32)  # Example casting to float32
+
+    # Verify axis type:
+    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
+
+    # Ensure compatibility for argmax (e.g., non-empty tensor):
+    if start_logits.shape[0] == 0:
+        raise ValueError("start_logits tensor is empty")
+
+    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]]