import numpy
from transformers import DPRContextEncoder, DPRContextEncoderTokenizer

from towhee import register
from towhee.operator import NNOperator

import warnings
import logging


warnings.filterwarnings('ignore')
logging.getLogger("transformers").setLevel(logging.ERROR)
log = logging.getLogger()


@register(output_schema=['vec'])
class Dpr(NNOperator):
    """
    This class uses Dense Passage Retrieval to generate embedding.
    Dense Passage Retrieval (DPR) is a set of tools and models for state-of-the-art open-domain Q&A research.
    It was introduced in Dense Passage Retrieval for Open-Domain Question Answering by Vladimir Karpukhin,
    Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, Wen-tau Yih.
    Ref: https://huggingface.co/docs/transformers/v4.16.2/en/model_doc/dpr

    Args:
        model_name (`str`):
            Which model to use for the embeddings.
    """
    def __init__(self, model_name: str = "facebook/dpr-ctx_encoder-single-nq-base") -> None:
        self.model_name = model_name
        try:
            self.tokenizer = DPRContextEncoderTokenizer.from_pretrained(model_name)
        except Exception as e:
            log.error(f'Fail to load tokenizer by name: {model_name}')
            raise e
        try:
            self.model = DPRContextEncoder.from_pretrained(model_name)
        except Exception as e:
            log.error(f'Fail to load model by name: {model_name}')
            raise e

    def __call__(self, txt: str) -> numpy.ndarray:
        try:
            input_ids = self.tokenizer(txt, return_tensors="pt")["input_ids"]
        except Exception as e:
            log.error(f'Invalid input for the tokenizer: {self.model_name}')
            raise e
        try:
            embeddings = self.model(input_ids).pooler_output
        except Exception as e:
            log.error(f'Invalid input for the model: {self.model_name}')
            raise e
        vec = embeddings.squeeze(0).detach().numpy()
        return vec


def get_model_list():
    full_list = [
        "facebook/dpr-ctx_encoder-single-nq-base",
        "facebook/dpr-ctx_encoder-multiset-base",
    ]
    full_list.sort()
    return full_list