Update app.py

app.py

@@ -19,6 +19,10 @@ def gradio_predict(input_text):
         max_length=512         # Ensure the sequence doesn't exceed the model's max length
     )
 
+    # Check that we have all necessary inputs
+    if "input_ids" not in tokenized_input or "attention_mask" not in tokenized_input:
+        return "Error: Missing required tokenizer outputs."
+
     # Convert tokenized inputs to numpy arrays and ensure correct shape
     input_ids = np.array(tokenized_input["input_ids"], dtype=np.int64)  # Shape should be [1, 512]
     attention_mask = np.array(tokenized_input["attention_mask"], dtype=np.int64)  # Shape should be [1, 512]
@@ -30,6 +34,9 @@ def gradio_predict(input_text):
     # Prepare decoder input ids if required by your model
     decoder_input_ids = input_ids  # Adjust as needed based on model requirements
 
+    # Print the shapes to check if they are correct
+    print(f"input_ids shape: {input_ids.shape}, attention_mask shape: {attention_mask.shape}")
+
     # Perform inference with ONNX model
     try:
         outputs = session.run(