from transformers import DPRReader, DPRReaderTokenizer
from typing import List, Dict, Tuple
from dotenv import load_dotenv

from src.readers.base_reader import Reader


load_dotenv()


class DprReader(Reader):
    def __init__(self) -> None:
        self._tokenizer = DPRReaderTokenizer.from_pretrained(
            "facebook/dpr-reader-single-nq-base")
        self._model = DPRReader.from_pretrained(
            "facebook/dpr-reader-single-nq-base")

    def read(self,
             query: str,
             context: Dict[str, List[str]],
             num_answers=5) -> List[Tuple]:
        encoded_inputs = self._tokenizer(
            questions=query,
            titles=context['titles'],
            texts=context['texts'],
            return_tensors='pt',
            truncation=True,
            padding=True
        )
        outputs = self._model(**encoded_inputs)

        predicted_spans = self._tokenizer.decode_best_spans(
            encoded_inputs,
            outputs,
            num_spans=num_answers,
            num_spans_per_passage=2,
            max_answer_length=512)

        return predicted_spans