BERTでTOEICの問題を解いてみる③

import torch
from transformers import BertTokenizer, BertForPreTraining

tokenizer = BertTokenizer.from_pretrained('bert-large-cased-whole-word-masking-finetuned-squad')
model = BertForQuestionAnswering.from_pretrained('bert-large-cased-whole-word-masking-finetuned-squad')

text = """
Used Car For Sale. Six-year-old Carlisle Custom.
Only one owner.
Low mileage.
Car used to commute short distances to town.
Brakes and tires replaced six months ago.
Struts replaced two weeks ago.
Air conditioning works well, but heater takes a while to warm up.
Brand new spare tire included.
Priced to sell.
Owner going overseas at the end of this month and must sell the car.
Call Firoozeh Ghorbani at (848) 555-0132.
"""
question1 = "What is suggested about the car?"
question2 = "According to the advertisement, why is Ms. Ghorbani selling her car?"
candidate1 = [
    "It was recently repaired.",
    "It has had more than one owner.",
    "It is very fuel efficient.",
    "It has been on sale for six months.",
]
candidate2 = [
    "She cannot repair the car’s temperature control.",
    "She finds it difficult to maintain.",
    "She would like to have a newer model.",
    "She is leaving for another country.",
]

tokens = ["[CLS]"] + tokenizer.tokenize(question1) + ["[SEP]"] + tokenizer.tokenize(text) + ["[SEP]"]
input_ids = tokenizer.convert_tokens_to_ids(tokens)

token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))]

with torch.no_grad():
    start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))

all_tokens = tokenizer.convert_ids_to_tokens(input_ids)
answer_token = all_tokens[torch.argmax(start_scores):torch.argmax(end_scores)+1]
print("answer token:", " ".join(answer_token))

ans = ""
for a in answer_token:
    if "#" in a:
        ans += a.replace("#", "")
    else:
        ans += " " + a
print("extract answer:", ans)
# answer token: Low mile ##age . Car used to com ##mu ##te short distances to town
# extract answer: Low mile age. Car used to commute short distances to town 

# 文章ベクトルを取得する関数
def get_bert_vec(c):
    tokens = tokenizer.tokenize(c)
    ids = tokenizer.convert_tokens_to_ids(tokens)
    ids = torch.tensor(ids).reshape(1,-1)
    with torch.no_grad():
        outputs, _ = model.bert(ids)
    vec = outputs[0].numpy().mean(0)
    return vec
    
answer_vec = get_bert_vec(ans)
c_vecs = [get_bert_vec(c) for c in candidate]

candidate1[cosine_similarity([answer_vec], c_vecs).argmax()]
# It was recently repaired.

def part7_slover(text, question, candidate):
    tokens = ["[CLS]"] + tokenizer.tokenize(question) + ["[SEP]"] + tokenizer.tokenize(text.replace("\n", " ")) + ["[SEP]"]
    input_ids = tokenizer.convert_tokens_to_ids(tokens)
    
    token_type_ids = [0 if i <= input_ids.index(102) else 1 for i in range(len(input_ids))]
    with torch.no_grad():
        start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))
        
    all_tokens = tokenizer.convert_ids_to_tokens(input_ids)
    answer_token = all_tokens[torch.argmax(start_scores):torch.argmax(end_scores)+1]
    print("answer token:", " ".join(answer_token))

    ans = ""
    for a in answer_token:
        if "#" in a:
            ans += a.replace("#", "")
        else:
            ans += " "+a
    print("extract answer:", ans)
    
    answer_vec = get_bert_vec(ans)
    c_vecs = [get_bert_vec(c) for c in candidate]

    return candidate[cosine_similarity([answer_vec], c_vecs).argmax()]
 
part7_slover(text, question2, candidate2)
#answer token: Own ##er going overseas
#extract answer:  Owner going overseas
#'She finds it difficult to maintain.'