File size: 680 Bytes
224cd33
d182e77
 
f68bac7
d182e77
 
4372d3f
d182e77
 
 
4372d3f
d182e77
 
4372d3f
d182e77
 
 
 
4372d3f
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import torch

tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
model = AutoModelForQuestionAnswering.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")

def answer_question(context, question):
    inputs = tokenizer.encode_plus(question, context, return_tensors="pt")
    answer_start_scores, answer_end_scores = model(**inputs).values()

    start = torch.argmax(answer_start_scores)
    end = torch.argmax(answer_end_scores) + 1

    answer = tokenizer.convert_tokens_to_string(
        tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][start:end])
    )
    return answer