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