# Copyright (c) Meta Platforms, Inc. and affiliates.
# Source from: https://github.com/facebookresearch/llama/blob/main/llama/tokenizer.py

import os
from logging import getLogger
from typing import List

from sentencepiece import SentencePieceProcessor


logger = getLogger()


class TextTokenizer:
    """Tokenizing and encoding/decoding text using SentencePiece."""

    def __init__(self, model_path=None):
        """
        Initializes the Tokenizer with a SentencePiece model.

        Args:
            model_path (str): The path to the SentencePiece model file.
        """
        if model_path is None:
            model_path = os.path.join(
                os.path.dirname(os.path.abspath(__file__)), "text_tokenizer.model"
            )
        # reload tokenizer
        assert os.path.isfile(model_path), model_path
        self.sp_model = SentencePieceProcessor(model_file=model_path)
        logger.info(f"Reloaded SentencePiece model from {model_path}")
        # BOS / EOS token IDs
        self.n_words: int = self.sp_model.vocab_size()
        self.bos_id: int = self.sp_model.bos_id()
        self.eos_id: int = self.sp_model.eos_id()
        self.pad_id: int = self.sp_model.pad_id()
        self.pad_id += self.n_words if self.pad_id < 0 else 0
        logger.info(f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}")
        assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()

    def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
        """
        Encodes a string into a list of token IDs.

        Args:
            s (str): The input string to be encoded.
            bos (bool): Whether to prepend the beginning-of-sequence token.
            eos (bool): Whether to append the end-of-sequence token.

        Returns:
            List[int]: A list of token IDs.
        """
        assert type(s) is str
        t = self.sp_model.encode(s)
        if bos:
            t = [self.bos_id] + t
        if eos:
            t = t + [self.eos_id]
        return t

    def decode(self, t: List[int]) -> str:
        """
        Decodes a list of token IDs into a string.

        Args:
            t (List[int]): The list of token IDs to be decoded.

        Returns:
            str: The decoded string.
        """
        return self.sp_model.decode(t)

    def tokenize(self, texts, context_length=None):
        """Encode a list of string.

        Parameters
        ----------
        texts : Union[str, List[str]]
            The input text(s).
        context_length : int, optional
            The max token length.

        Returns
        -------
        List[List[int]]
            The encoded token indices.

        """
        if isinstance(texts, str):
            texts = [texts]
        tokens = [self.encode(text, bos=True, eos=True) for text in texts]
        if context_length is None:
            return tokens
        truncated_tokens = []
        for k, t in enumerate(tokens):
            if len(t) > context_length:
                t = t[:context_length]
                t[-1] = self.eos_id
            truncated_tokens.append(t)
        return truncated_tokens

    def detokenize(self, tokens):
        """Decode a list of string.

        Parameters
        ----------
        tokens : Union[List[List[int]], numpy.ndarray]
            The input tokens.

        Returns
        -------
        List[str]
            The decoded text strings.

        """
        if hasattr(tokens, "tolist"):
            tokens = tokens.tolist()
        texts = []
        for i in range(len(tokens)):
            t = tokens[i][1:]
            try:
                eot_idx = t.index(self.eos_id)
                t = t[:eot_idx]
            except ValueError:
                pass
            texts.append(self.decode(t))
        return texts