from transformers import BertTokenizer, BertForQuestionAnswering import torch # 加载BERT模型和分词器 model_name = 'bert-base-uncased' tokenizer = BertTokenizer.from_pretrained(model_name) model = BertForQuestionAnswering.from_pretrained(model_name) # 输入文本和问题 context = "The Apollo program, also known as Project Apollo, was the third United States human spaceflight program carried out by the National Aeronautics and Space Administration (NASA), which succeeded in landing the first humans on the Moon from 1969 to 1972. Apollo was first conceived during the Eisenhower administration in early 1960 as a follow-up to Project Mercury. It was dedicated to President John F. Kennedy's national goal of landing Americans on the Moon before the end of the 1960s." question = "What was the goal of the Apollo program?" # 对输入进行编码 encoding = tokenizer.encode_plus(question, context, max_length=512, padding='max_length', truncation=True, return_tensors='pt') # 获取输入ids和注意力掩码 input_ids = encoding['input_ids'] attention_mask = encoding['attention_mask'] # 使用BERT模型进行问答 outputs = model(input_ids=input_ids, attention_mask=attention_mask) 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(tokenizer.convert_ids_to_tokens(input_ids[0][start_index:end_index+1])) print(answer)
时间: 2024-04-13 19:25:26 浏览: 208
基于torch transformers 直接加载bert预训练模型计算句子相似度
5星 · 资源好评率100%
运行这段代码时,你加载了预训练的 BERT 模型和分词器,并对文本和问题进行了编码。然后,你使用 BERT 模型进行问答,通过输入输入ids和注意力掩码,模型输出了起始和结束位置的得分。接着,通过找到得分最高的位置,你解码出了答案,并将其打印出来。
请注意,这段代码还需要安装 Transformers 库和 PyTorch 库,并确保已经下载了预训练的 BERT 模型。
阅读全文