Update app.py

app.py

@@ -348,7 +348,21 @@ def two_stage_qa(question, candidate_paragraphs_str, max_seq_len_mc=512, max_seq
                     logger.error(f"    !!!!!! CRITICAL: Key '{key_to_check}' in feature {i} contains non-integer {elem} (type: {type(elem)}) at index {elem_idx}!")
                     return f"錯誤: 特徵準備失敗，{key_to_check} 內部有非整數", "N/A", "N/A"
             logger.info(f"    Feature {i}, key '{key_to_check}' passed inspection. Length: {len(val_list)}")
+            
+    model_input_columns = ["input_ids", "attention_mask"]
+    if "token_type_ids" in qa_features_dataset.features:
+        model_input_columns.append("token_type_ids")
+    
+    try:
+        # 從 qa_features_dataset 中提取需要的列來創建新的 Dataset
+        # qa_features_dataset 本身就是一個 Dataset 對象
+        features_for_dataloader = qa_features_dataset.select_columns(model_input_columns)
+    except Exception as e:
+        logger.error(f"從 qa_features_dataset 選擇列時出錯: {e}. Features: {qa_features_dataset.features}")
+        return f"錯誤: 準備模型輸入時出錯 (列選擇)。 Error: {e}", "N/A", "N/A"
 
+    logger.info(f"--- In two_stage_qa, about to create DataLoader for question_id: {log_question_id} ---")
+    
     qa_dataloader = DataLoader(
         qa_features_dataset, collate_fn=default_data_collator, batch_size=8 # batch_size可以小一些
     )