from transformers import AutoTokenizer

class RerankerTokenizer():
    def __init__(self, total_maxlen, base):
        self.total_maxlen = total_maxlen
        self.tok = AutoTokenizer.from_pretrained(base)

    def tensorize(self, questions, passages):
        assert type(questions) in [list, tuple], type(questions)
        assert type(passages) in [list, tuple], type(passages)

        encoding = self.tok(questions, passages, padding='longest', truncation='longest_first',
                            return_tensors='pt', max_length=self.total_maxlen, add_special_tokens=True)

        return encoding