使用transformers的BertForQuestionAnswering做抽取式QA解包报错如何解决

问题原因 #

Hugging Face Transformers库v3.0及之后的版本，模型默认返回的是输出dataclass对象，而非直接返回元组。你直接用start, end = model(**inputs)解包时，并不会拿到两个张量，第一个变量实际拿到的是输出对象的字符串表示，所以调用softmax/argmax时会报错提示str类型没有对应方法。

解决方法 #

方法1：指定return_dict=False获取元组输出 #

调用模型时传入return_dict=False参数，模型就会返回顺序排列的元组，第一个元素为起始位置得分张量，第二个为结束位置得分张量，原有解包逻辑即可正常运行：

import os
import torch
import torch.nn.functional as F
from transformers import BertForQuestionAnswering, BertTokenizer

# 优先使用经过SQuAD数据集微调的问答模型，不要直接用bert-base-uncased底座
tokenizer = BertTokenizer.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
model = BertForQuestionAnswering.from_pretrained(
    'bert-large-uncased-whole-word-masking-finetuned-squad',
    cache_dir=os.getenv("cache_dir", "../../models")
)
question = "What is the capital of Italy?"
text = "The capital of Italy is Rome."
inputs = tokenizer.encode_plus(question, text, return_tensors='pt')
# 加return_dict=False参数
start, end = model(**inputs, return_dict=False)
start_max = torch.argmax(F.softmax(start, dim = -1))
end_max = torch.argmax(F.softmax(end, dim = -1)) + 1
answer = tokenizer.decode(inputs["input_ids"][0][start_max : end_max])
print(answer)

方法2：通过输出对象的属性访问（更推荐，可读性更高） #

不需要修改模型调用参数，直接接收完整输出对象，通过属性名获取对应的得分张量：

encoding = tokenizer.encode_plus(text=question,text_pair=text, add_special_tokens=True)
inputs = encoding['input_ids']
sentence_embedding = encoding['token_type_ids']
tokens = tokenizer.convert_ids_to_tokens(inputs)
# 接收完整输出对象
outputs = model(input_ids=torch.tensor([inputs]), token_type_ids=torch.tensor([sentence_embedding]))
# 通过属性名获取得分
start_scores = outputs.start_logits
end_scores = outputs.end_logits
start_index = torch.argmax(start_scores)
end_index = torch.argmax(end_scores)
answer = tokenizer.convert_tokens_to_string(tokens[start_index:end_index+1])

额外注意 #

你原来使用的bert-base-uncased是通用预训练底座，没有经过抽取式问答任务微调，直接用它跑问答任务输出结果会完全不符合预期，建议替换为官方在SQuAD数据集上微调好的模型，比如bert-large-uncased-whole-word-masking-finetuned-squad、distilbert-base-cased-distilled-squad这类专门的问答模型。

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