使用Lightning-Transformers问答示例代码触发AssertionError求助

运行Lightning-Transformers官方文档提供的SQuAD问答任务示例代码时，触发如下断言错误：

AssertionError: Got 2 predictions and 10784 features.

解决方案 #

方法一：禁用验证步骤

错误触发于训练后验证阶段的预测后处理逻辑，直接在训练器中禁用验证即可绕过该问题，适合仅验证训练流程的场景：

trainer = pl.Trainer(
    accelerator="auto", 
    devices="auto", 
    max_epochs=1,
    limit_val_batches=0  # 关闭验证步骤
)

方法二：移除冗余的数据集配置参数

示例中设置的dataset_config_name="plain_text"是不必要的，SQuAD数据集的默认配置已适配Lightning-Transformers的处理逻辑，移除该参数后可正常完成训练和验证：

dm = SquadDataModule(
    batch_size=1,
    max_length=384,
    version_2_with_negative=False,
    null_score_diff_threshold=0.0,
    doc_stride=128,
    n_best_size=20,
    max_answer_length=30,
    tokenizer=tokenizer,
)

错误原因 #

该错误源于验证阶段的预测后处理步骤：模型输出的预测样本数量与数据集生成的特征数量不匹配。plain_text配置会改变SQuAD数据集的加载和预处理方式，导致特征生成逻辑与模型预测输出的处理流程不兼容；而禁用验证则直接跳过了触发错误的后处理环节。

内容的提问来源于stack exchange，提问作者thistleknot