class Tokenizer:
    """tokenizing and encoding/decoding text using SentencePiece."""
    def __init__(self, model_path: str):
        # 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()
        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

# Encode with BOS token for prompt start
tokens = tokenizer.encode("Hello world", bos=True, eos=False)
# Output: [1, 15043, 3186]  # 1 is BOS token

# Encode with both BOS and EOS for complete sequences
tokens = tokenizer.encode("Hello world", bos=True, eos=True)
# Output: [1, 15043, 3186, 2]  # 2 is EOS token

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)

B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"

SPECIAL_TAGS = [B_INST, E_INST, "<<SYS>>", "<</SYS>>"]

[INST] User message here [/INST]

[INST] <<SYS>>
You are a helpful assistant.
<</SYS>>

User message here [/INST]

[INST] <<SYS>>
You are a helpful assistant.
<</SYS>>

First user message [/INST] First assistant response [INST] Second user message [/INST] Second assistant response

if dialog[0]["role"] == "system":
    dialog = [
        {
            "role": dialog[1]["role"],
            "content": B_SYS
            + dialog[0]["content"]
            + E_SYS
            + dialog[1]["content"],
        }
    ] + dialog[2:]

dialog_tokens: List[int] = sum(
    [
        self.tokenizer.encode(
            f"{B_INST} {(prompt['content']).strip()} {E_INST} {(answer['content']).strip()} ",
            bos=True,
            eos=True,
        )
        for prompt, answer in zip(
            dialog[::2],
            dialog[1::2],
        )
    ],
    [],
)

dialog_tokens += self.tokenizer.encode(
    f"{B_INST} {(dialog[-1]['content']).strip()} {E_INST}",
    bos=True,
    eos=False,
)

dialog = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "What is the capital of France?"},
]

[INST] <<SYS>>
You are a helpful assistant.
<</SYS>>

What is the capital of France? [/INST]

prompt_tokens = tokenizer.encode(
    "[INST] <<SYS>>\nYou are a helpful assistant.\n<</SYS>>\n\nWhat is the capital of France? [/INST]",
    bos=True,
    eos=False
)

prompt_tokens = [self.tokenizer.encode(x, bos=True, eos=False) for x in prompts]

assert all([msg["role"] == "user" for msg in dialog[::2]]) and all(
    [msg["role"] == "assistant" for msg in dialog[1::2]]
), (
    "model only supports 'system', 'user' and 'assistant' roles, "
    "starting with 'system', then 'user' and alternating (u/a/u/a/u...)"
)

unsafe_requests = [
    any([tag in msg["content"] for tag in SPECIAL_TAGS for msg in dialog])
]

UNSAFE_ERROR = "Error: special tags are not allowed as part of the prompt."