from tokenizers import Tokenizer
from utils.config import TOKENIZER_DIR
import numpy as np


class ModelTokenizer:
    def __init__(self):
        TOKENIZER_FILE = TOKENIZER_DIR / "tokenizer.json"
        self.tokenizer = Tokenizer.from_file(str(TOKENIZER_FILE))
        self.MAX_LEN = 77

    def tokenize_text(self, text):
        encoding = self.tokenizer.encode(text)
        input_ids = encoding.ids
        # implementing padding
        ids = encoding.ids[:self.MAX_LEN]
        mask = [1]*len(ids)
        pad = self.MAX_LEN-len(ids)
        ids += [0]*pad
        mask += [0]*pad

        input_ids = np.array([ids], dtype=np.int64)
        attention_mask = np.array([mask], dtype=np.int64)

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask
        }