Update app.py

app.py

@@ -59,23 +59,20 @@ def answer_question(text, question, max_length=512):
         question, text, return_tensors="tf", padding="max_length", truncation=True
     )
 
-    start_logits, end_logits = qa_model(inputs)
+    outputs = qa_model(inputs)
 
-    # Ensure start_logits is a tensor
-    start_logits = tf.convert_to_tensor(start_logits)
+    start_logits = outputs.start_logits
+    end_logits = outputs.end_logits
 
-    # 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
+    # Ensure start_logits and end_logits are tensors
+    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
 
-    # 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)
 
@@ -84,7 +81,6 @@ def answer_question(text, question, max_length=512):
     return answer if answer else "No answer found."
 
 
-
 ## Streamlit app
 
 st.set_page_config(page_title="Enhanced PDF Summarizer")
@@ -119,4 +115,3 @@ if uploaded_file is not None:
                     st.write(answer)
         else:
             st.error("No text found in the PDF.")
-