File size: 1,066 Bytes
cd6775d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | from transformers import DPRReader, DPRReaderTokenizer
import pathlib, os
os.environ["CUDA_VISIBLE_DEVICES"] = '1'
device = "cuda"
tokenizer = DPRReaderTokenizer.from_pretrained("facebook/dpr-reader-multiset-base")
model = DPRReader.from_pretrained("facebook/dpr-reader-multiset-base")
model.eval()
model.to(device)
def get_answer(query,texts,title):
encoded_inputs = tokenizer(
questions=[query],
titles=[title],
texts=[texts],
return_tensors="pt",
max_length=512,
truncation=True,
)
outputs = model(**encoded_inputs.to(device))
start_logits = outputs.start_logits
end_logits = outputs.end_logits
relevance_logits = outputs.relevance_logits
answer_start_index = outputs.start_logits.argmax()
answer_end_index = outputs.end_logits.argmax()
predict_answer_tokens = encoded_inputs.input_ids[0, answer_start_index : answer_end_index + 1]
#print(tokenizer.decode(predict_answer_tokens))
answer = tokenizer.decode(predict_answer_tokens)
return answer,relevance_logits
|