# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import bisect
import functools
import inspect
import io
import itertools
import json
import os
import re
import unicodedata
from collections import OrderedDict
from dataclasses import asdict, dataclass
from functools import lru_cache
from typing import Any, Dict, List, Literal, Optional, Tuple, Union

import numpy as np

from .....utils import logging
from .....utils.deps import class_requires_deps, is_dep_available
from .tokenizer_utils_base import (
    CHAT_TEMPLATE_CONFIG_NAME,
    AddedToken,
    BatchEncoding,
    EncodedInput,
    EncodedInputPair,
    PaddingStrategy,
    PreTokenizedInput,
    PreTokenizedInputPair,
    PretrainedTokenizerBase,
    TensorType,
    TextInput,
    TextInputPair,
    TruncationStrategy,
)
from .utils import convert_to_dict_message, fn_args_to_dict
from .vocab import Vocab

if is_dep_available("Jinja2"):
    from jinja2 import Template
    from jinja2.exceptions import TemplateError, TemplateSyntaxError
    from jinja2.sandbox import ImmutableSandboxedEnvironment

__all__ = [
    "ChatTemplate",
    "Trie",
    "ChatTemplateMixin",
    "PretrainedTokenizer",
    "InitTrackerMeta",
]


@class_requires_deps("Jinja2")
@dataclass
class ChatTemplate:
    conversation: Union[List[str], None] = None
    system: Union[str, None] = None
    query: str = None

    @staticmethod
    @lru_cache()
    def _compile_jinja_template(chat_template) -> "Template":
        def raise_exception(message):
            raise TemplateError(message)

        jinja_env = ImmutableSandboxedEnvironment(
            trim_blocks=True, lstrip_blocks=True, keep_trailing_newline=True
        )
        jinja_env.globals["raise_exception"] = raise_exception
        return jinja_env.from_string(chat_template)

    def render_conversation(
        self,
        conversation_data: Union[List[str], Dict[str, str]],
        index: int = 0,
        context_data: Dict[str, Any] = {},
    ) -> List[str]:
        """
        Args:
            conversation_data (list[str]): the conversation data which must be two parts
            index (int): the index of current conversation

        Returns:
            list[str]: the rendered conversation data
        """
        if self.conversation is None:
            raise ValueError(
                "The template for multi-turns is invalid, please check `conversation` filed in your chat-template."
            )

        if isinstance(conversation_data, (list, tuple)):
            assert (
                len(conversation_data) == 2
            ), "Each round/turn of conversation must be two participants, eg: [user-query, bot-query]"

            conversation_data = {
                "user": conversation_data[0],
                "bot": conversation_data[1],
                "index": index,
            }
        conversation_data.update(context_data)

        one_turn_conversation = []
        for conversation in self.conversation:
            template = self._compile_jinja_template(conversation)
            result = template.render(conversation_data)
            one_turn_conversation.append(result)
        return one_turn_conversation

    def render_query(
        self, query: str, index: int = 0, context_data: Dict[str, Union[int, str]] = {}
    ):
        if self.query is None:
            return query

        template = self._compile_jinja_template(self.query)
        return template.render(query=query, index=index, **context_data)

    def _init_context_data(
        self, context_data: Dict[str, Union[int, str]] = {}
    ) -> Dict[str, Union[int, str]]:
        """init the context data for chat-template"""
        context_data["is_training"] = context_data.get("is_training", False)
        return context_data

    def render_system(self, context_data: Dict[str, Union[int, str]] = {}) -> str:
        if self.system is None:
            return ""

        template = self._compile_jinja_template(self.system)
        return template.render(**context_data)

    def __call__(
        self,
        conversations: Union[List[List[str]], str],
        context_data: Dict[str, Union[int, str]] = {},
    ) -> str:
        """render the conversations by chat-template

        Args:
            conversations (list[list[str]]): the conversations of use and bot

        Returns:
            str: the result of conversation
        """
        if isinstance(conversations, str):
            conversations = [[conversations]]

        # [1 ... n-1] conversation
        final_query = self.render_system(context_data=context_data)
        context_data["length"] = len(conversations)
        for index, conversation in enumerate(conversations[:-1]):
            context_data["is_first"] = index == 0
            context_data["is_last"] = False
            final_query += "".join(
                self.render_conversation(
                    conversation, index=index, context_data=context_data
                )
            )

        if not isinstance(conversations[-1], list) and not len(conversations[-1]) != 1:
            raise ValueError(
                "The length of last conversation must be one, eg: [[user-query, bot-answer], [user-query, bot-answer], ..., [user-query]]"
            )
        if len(conversations[-1]) > 1:
            logging.warning(
                f"The last conversation is not a single-round, chat-template will skip the conversation: {conversations[-1][1:]}"
            )

        final_query += self.render_query(
            conversations[-1][0],
            index=len(conversations) - 1,
            context_data=context_data,
        )
        return final_query

    @classmethod
    def from_dict(cls, config: Dict):
        return cls(**config)

    @classmethod
    def from_file(cls, file: str):
        with open(file, "r", encoding="utf-8") as f:
            config = json.load(f)
        return cls.from_dict(config)


def adapt_stale_fwd_patch(self, name, value):
    """
    Since there are some monkey patches for forward of PretrainedModel, such as
    model compression, we make these patches compatible with the latest forward
    method.
    """

    if name == "forward":
        # NOTE(guosheng): In dygraph to static, `layer.forward` would be patched
        # by an instance of `StaticFunction`. And use string compare to avoid to
        # import fluid.
        if type(value).__name__.endswith(
            "StaticFunction"
        ) or self.forward.__class__.__name__.endswith("StaticFunction"):
            return value
        (
            patch_spec_args,
            patch_spec_varargs,
            patch_spec_varkw,
            patch_spec_defaults,
            _,
            _,
            _,
        ) = inspect.getfullargspec(value)
        (spec_args, spec_varargs, spec_varkw, spec_defaults, _, _, _) = (
            inspect.getfullargspec(self.forward)
        )
        new_args = [
            arg
            for arg in ("output_hidden_states", "output_attentions", "return_dict")
            if arg not in patch_spec_args and arg in spec_args
        ]

        if new_args:
            import paddle

            if self.__module__.startswith("paddlenlp"):
                logging.warning(
                    f"The `forward` method of {self.__class__ if isinstance(self, paddle.nn.Layer) else self} is patched and the patch "
                    "might be based on an old oversion which missing some "
                    f"arguments compared with the latest, such as {new_args}. "
                    "We automatically add compatibility on the patch for "
                    "these arguments, and maybe the patch should be updated."
                )
            else:
                logging.warning(
                    f"The `forward` method of {self.__class__ if isinstance(self, paddle.nn.Layer) else self} "
                    "is patched and the patch might be conflict with patches made "
                    f"by paddlenlp which seems have more arguments such as {new_args}. "
                    "We automatically add compatibility on the patch for "
                    "these arguments, and maybe the patch should be updated."
                )
            if isinstance(self, paddle.nn.Layer) and inspect.isfunction(value):

                @functools.wraps(value)
                def wrap_fwd(*args, **kwargs):
                    for arg in new_args:
                        kwargs.pop(arg, None)
                    return value(self, *args, **kwargs)

            else:

                @functools.wraps(value)
                def wrap_fwd(*args, **kwargs):
                    for arg in new_args:
                        kwargs.pop(arg, None)
                    return value(*args, **kwargs)

            return wrap_fwd
    return value


# NOTE:
# Modification:
#   class InitTrackerMeta(type(paddle.nn.Layer)) -> class InitTrackerMeta(type)
# Context:
#   1. In paddle 3.0rc, type(paddle.nn.Layer) == type
#   2. Solve the conflict between ultra-infer and paddle
class InitTrackerMeta(type):
    """
    This metaclass wraps the `__init__` method of a class to add `init_config`
    attribute for instances of that class, and `init_config` use a dict to track
    the initial configuration. If the class has `_pre_init` or `_post_init`
    method, it would be hooked before or after `__init__` and called as
    `_pre_init(self, init_fn, init_args)` or `_post_init(self, init_fn, init_args)`.
    Since InitTrackerMeta would be used as metaclass for pretrained model classes,
    which always are Layer and `type(Layer)` is not `type`, thus use `type(Layer)`
    rather than `type` as base class for it to avoid inheritance metaclass
    conflicts.
    """

    def __init__(cls, name, bases, attrs):
        init_func = cls.__init__
        # If attrs has `__init__`, wrap it using accessible `_pre_init, _post_init`.
        # Otherwise, no need to wrap again since the super cls has been wrapped.
        # TODO: remove reduplicated tracker if using super cls `__init__`
        pre_init_func = getattr(cls, "_pre_init", None) if "__init__" in attrs else None
        post_init_func = (
            getattr(cls, "_post_init", None) if "__init__" in attrs else None
        )
        cls.__init__ = InitTrackerMeta.init_and_track_conf(
            init_func, pre_init_func, post_init_func
        )
        super(InitTrackerMeta, cls).__init__(name, bases, attrs)

    @staticmethod
    def init_and_track_conf(init_func, pre_init_func=None, post_init_func=None):
        """
        wraps `init_func` which is `__init__` method of a class to add `init_config`
        attribute for instances of that class.
        Args:
            init_func (callable): It should be the `__init__` method of a class.
                warning: `self` always is the class type of down-stream model, eg: BertForTokenClassification
            pre_init_func (callable, optional): If provided, it would be hooked after
                `init_func` and called as `pre_init_func(self, init_func, *init_args, **init_args)`.
                Default None.
            post_init_func (callable, optional): If provided, it would be hooked after
                `init_func` and called as `post_init_func(self, init_func, *init_args, **init_args)`.
                Default None.

        Returns:
            function: the wrapped function
        """

        @functools.wraps(init_func)
        def __impl__(self, *args, **kwargs):
            # registered helper by `pre_init_func`
            if pre_init_func:
                pre_init_func(self, init_func, *args, **kwargs)
            # keep full configuration
            init_func(self, *args, **kwargs)
            # registered helper by `post_init_func`
            if post_init_func:
                post_init_func(self, init_func, *args, **kwargs)
            self.init_config = kwargs
            if args:
                kwargs["init_args"] = args
            kwargs["init_class"] = self.__class__.__name__

        return __impl__

    def __setattr__(self, name, value):
        value = adapt_stale_fwd_patch(self, name, value)
        return super(InitTrackerMeta, self).__setattr__(name, value)


class Trie:
    """
    Trie in Python. Creates a Trie out of a list of words. The trie is used to split on `added_tokens` in one pass
    Loose reference https://en.wikipedia.org/wiki/Trie
    """

    def __init__(self):
        self.data = {}

    def add(self, word: str):
        """
        Passes over every char (utf-8 char) on word and recursively adds it to the internal `data` trie representation.
        The special key `""` is used to represent termination.

        This function is idempotent, adding twice the same word will leave the trie unchanged

        Example:

        ```python
        >>> trie = Trie()
        >>> trie.add("Hello 友達")
        >>> trie.data
        {"H": {"e": {"l": {"l": {"o": {" ": {"友": {"達": {"": 1}}}}}}}}}

        >>> trie.add("Hello")
        >>> trie.data
        {"H": {"e": {"l": {"l": {"o": {"": 1, " ": {"友": {"達": {"": 1}}}}}}}}}
        ```
        """
        if not word:
            # Prevent empty string
            return
        ref = self.data
        for char in word:
            ref[char] = char in ref and ref[char] or {}
            ref = ref[char]
        ref[""] = 1

    def split(self, text: str) -> List[str]:
        """
        Will look for the words added to the trie within `text`. Output is the original string splitted along the
        boundaries of the words found.

        This trie will match the longest possible word first !

        Example:

        ```python
        >>> trie = Trie()
        >>> trie.split("[CLS] This is a extra_id_100")
        ["[CLS] This is a extra_id_100"]

        >>> trie.add("[CLS]")
        >>> trie.add("extra_id_1")
        >>> trie.add("extra_id_100")
        >>> trie.split("[CLS] This is a extra_id_100")
        ["[CLS]", " This is a ", "extra_id_100"]
        ```
        """
        # indexes are counted left of the chars index.
        # "hello", index 0, is left of h, index 1 is between h and e.
        # index 5 is right of the "o".

        # States are going to capture every possible start (indexes as above)
        # as keys, and have as values, a pointer to the position in the trie
        # where we're at. This is a partial match for now.
        # This enables to keep track of multiple matches while we're iterating
        # the string
        # If the trie contains, "blowing", and "lower" and we encounter the
        # string "blower", we need to split into ["b", "lower"].
        # This is where we need to keep track of multiple possible starts.
        states = OrderedDict()

        # This will contain every indices where we need
        # to cut.
        # We force to cut at offset 0 and len(text) (added later)
        offsets = [0]

        # This is used by the lookahead which needs to skip over
        # some text where the full match exceeded the place in the initial
        # for loop
        skip = 0
        # Main loop, Giving this algorithm O(n) complexity
        for current, current_char in enumerate(text):
            if skip and current < skip:
                # Prevents the lookahead for matching twice
                # like extra_id_100 and id_100
                continue

            # This will track every state
            # that stop matching, we need to stop tracking them.
            # If we look at "lowball", we're going to match "l" (add it to states), "o", "w", then
            # fail on "b", we need to remove 0 from the valid states.
            to_remove = set()
            # Whenever we found a match, we need to drop everything
            # this is a greedy algorithm, it will match on the first found token
            reset = False

            # In this case, we already have partial matches (But unfinished)
            for start, trie_pointer in states.items():
                if "" in trie_pointer:
                    # This is a final match, we need to reset and
                    # store the results in `offsets`.

                    # Lookahead to match longest first
                    # Important in case of extra_id_1 vs extra_id_100
                    # Here we are also actively looking for other earlier partial
                    # matches
                    # "[CLS]", "L", we need to match CLS even if L is special
                    for lookstart, looktrie_pointer in states.items():
                        if lookstart > start:
                            # This partial match is later, we can stop looking
                            break
                        elif lookstart < start:
                            # This partial match is earlier, the trie pointer
                            # was already updated, so index is + 1
                            lookahead_index = current + 1
                            end = current + 1
                        else:
                            # Here lookstart == start and
                            #      looktrie_pointer == trie_pointer
                            # It wasn't updated yet so indices are current ones
                            lookahead_index = current
                            end = current
                        next_char = (
                            text[lookahead_index]
                            if lookahead_index < len(text)
                            else None
                        )
                        if "" in looktrie_pointer:
                            start = lookstart
                            end = lookahead_index
                            skip = lookahead_index

                        while next_char in looktrie_pointer:
                            looktrie_pointer = looktrie_pointer[next_char]
                            lookahead_index += 1
                            if "" in looktrie_pointer:
                                start = lookstart
                                end = lookahead_index
                                skip = lookahead_index

                            if lookahead_index == len(text):
                                # End of string
                                break
                            next_char = text[lookahead_index]
                        # End lookahead

                        # Storing and resetting
                    offsets.append(start)
                    offsets.append(end)
                    reset = True
                    break
                elif current_char in trie_pointer:
                    # The current character being looked at has a match within the trie
                    # update the pointer (it will be stored back into states later).
                    trie_pointer = trie_pointer[current_char]

                    # Storing back the new pointer into the states.
                    # Partial matches got longer by one.
                    states[start] = trie_pointer
                else:
                    # The new character has not match in the trie, we need
                    # to stop keeping track of this partial match.
                    # We can't do it directly within the loop because of how
                    # python iteration works
                    to_remove.add(start)

            # Either clearing the full start (we found a real match)
            # Or clearing only the partial matches that didn't work.
            if reset:
                states = {}
            else:
                for start in to_remove:
                    del states[start]

            # If this character is a starting character within the trie
            # start keeping track of this partial match.
            if current >= skip and current_char in self.data:
                states[current] = self.data[current_char]

        # We have a cut at the end with states.
        for start, trie_pointer in states.items():
            if "" in trie_pointer:
                # This is a final match, we need to reset and
                # store the results in `offsets`.
                end = len(text)
                offsets.append(start)
                offsets.append(end)
                # Longest cut is always the one with lower start so the first
                # item so we need to break.
                break

        return self.cut_text(text, offsets)

    def cut_text(self, text, offsets):
        # We have all the offsets now, we just need to do the actual splitting.
        # We need to eventually add the first part of the string and the eventual
        # last part.
        offsets.append(len(text))
        tokens = []
        start = 0
        for end in offsets:
            if start > end:
                logging.error(
                    "There was a bug in Trie algorithm in tokenization. Attempting to recover. Please report it anyway."
                )
                continue
            elif start == end:
                # This might happen if there's a match at index 0
                # we're also preventing zero-width cuts in case of two
                # consecutive matches
                continue
            tokens.append(text[start:end])
            start = end

        return tokens


def _insert_one_token_to_ordered_list(token_list: List[str], new_token: str):
    """
    Inserts one token to an ordered list if it does not already exist. Note: token_list must be sorted.
    """
    insertion_idx = bisect.bisect_left(token_list, new_token)
    # Checks if new_token is already in the ordered token_list
    if insertion_idx < len(token_list) and token_list[insertion_idx] == new_token:
        # new_token is in token_list, don't add
        return
    else:
        token_list.insert(insertion_idx, new_token)


def _is_control(char):
    """Checks whether `chars` is a control character."""
    # These are technically control characters but we count them as whitespace
    # characters.
    if char == "\t" or char == "\n" or char == "\r":
        return False
    cat = unicodedata.category(char)
    if cat.startswith("C"):
        return True
    return False


def _is_nonnormalized_char(char):
    """Check whether `chars` is a non-normalized character."""
    cp = ord(char)
    if (
        (0xFF00 <= cp <= 0xFFEF)
        or (0xFE50 <= cp <= 0xFE6B)  # Halfwidth and Fullwidth Forms
        or (0x3358 <= cp <= 0x33FF)  # Small Form Variants
        or (0x249C <= cp <= 0x24E9)  # CJK Compatibility
        or (0x3200 <= cp <= 0x32FF)  # Enclosed Alphanumerics: Ⓛ ⒰
    ):  # Enclosed CJK Letters and Months
        return True

    return False


def _is_nonnormalized_numeric(char):
    """Check whether `chars` is a non-normalized numeric character."""
    cp = ord(char)
    if (
        (0x2460 <= cp <= 0x249B)
        or (0x24EA <= cp <= 0x24FF)  #
        or (0x2776 <= cp <= 0x2793)  #
        or (0x2160 <= cp <= 0x217F)  # Enclosed Alphanumerics
    ):  # Number Forms
        return True

    return False


def normalize_chars(text):
    """
    Normalize the text for multiligual and chinese models. Unicode range:
    https://www.ling.upenn.edu/courses/Spring_2003/ling538/UnicodeRanges.html
    """
    output = []
    for char in text:
        if _is_nonnormalized_char(char):
            for c in unicodedata.normalize("NFKC", char):
                output.append(c)
        elif _is_nonnormalized_numeric(char):
            output.append(" ")
            for c in str(int(unicodedata.numeric(char))):
                output.append(c)
            output.append(" ")
        elif ord(char) == 0xF979:  # https://www.zhihu.com/question/20697984
            output.append("凉")
        else:
            output.append(char)
    return "".join(output)


class ChatTemplateMixin:
    chat_template: Optional[ChatTemplate] = None

    def apply_chat_template(
        self,
        conversation: Union[List[List[str]], Dict[str, str], str],
        tokenize: bool = True,
        context_data: Dict[str, Any] = {},
        **tokenizer_kwargs,
    ):
        """apply chat_template rules to conversation which should not be batched data

        Args:
            conversation (List[List[str]] , str): the conversation messages between user and bot
            context_data (Dict[str, Any]): the context data for chat_template.json
            tokenize (bool, optional): whether do tokenization. Defaults to True.

        Returns:
            str | dict[str, Union[numpy.ndarray, paddle.Tensor]]: return the result of applied data
        """
        if not self.chat_template:
            raise ValueError(
                "chat_template is not set, please set chat_template first."
            )
        elif isinstance(self.chat_template, Template):
            add_generation_prompt = tokenizer_kwargs.pop("add_generation_prompt", True)
            query = self._apply_chat_template(
                conversation, add_generation_prompt=add_generation_prompt
            )
        elif isinstance(self.chat_template, ChatTemplate):
            query = self._apply_chat_template_paddle(conversation, context_data)

        if not tokenize:
            return query

        # chat_template should not add special tokens
        tokenizer_kwargs["add_special_tokens"] = False
        return self(query, **tokenizer_kwargs)

    def _apply_chat_template_paddle(
        self,
        conversation: Union[List[List[str]], str],
        context_data: Dict[str, Any] = {},
    ):
        context_data = self.chat_template._init_context_data(context_data)

        if isinstance(conversation, str):
            conversation = [[conversation]]
        elif isinstance(conversation, list) and isinstance(conversation[0], str):
            raise ValueError(
                "apply_chat_template do not support applying batch conversations, "
                "so you should apply the conversation one by one."
            )

        query = self.chat_template(conversation, context_data=context_data)
        return query

    def _apply_chat_template(
        self,
        conversation: Union[List[List[str]], Dict[str, str], str],
        add_generation_prompt=True,
    ):
        if isinstance(conversation, str):
            conversations = [{"role": "user", "content": conversation}]
        elif isinstance(conversation, list):
            assert len(conversation) > 0, "empty conversation is not allowed"
            if isinstance(conversation[0], list):
                conversations = convert_to_dict_message(conversation)
            elif isinstance(conversation[0], dict):
                conversations = conversation
            else:
                raise ValueError(
                    "apply_chat_template do not support applying batch conversations, "
                    "so you should apply the conversation one by one."
                )
        try:
            query = self.chat_template.render(
                messages=conversations,
                **self.special_tokens_map,
                add_generation_prompt=add_generation_prompt,
            )
        except TypeError:
            for i in range(len(conversations)):
                content = conversations[i]["content"]
                if isinstance(content, list):
                    new_content = ""
                    for part in content:
                        if part.get("type") == "text":
                            new_content = part["text"]
                            break
                conversations[i]["content"] = new_content
            query = self.chat_template.render(
                messages=conversations,
                **self.special_tokens_map,
                add_generation_prompt=add_generation_prompt,
            )
        return query

    def encode_chat_inputs(
        self,
        conversations: List[List[str]],
        context_data: Dict[str, Any] = {},
        **kwargs,
    ):
        """Encodes conversation to pairs of token ids.
        Turn 0: bos + system + sep + user     bot + eos
        Turn t: sep + bot + query             bot + eos

        Args:
            conversation (List[List[str]]): the conversation of data
            context_data (Dict[str, Any]): the context data of conversation

        Returns:
            List[list[int], list[int]]: the pair of input_ids and target_ids
        """
        if not self.chat_template:
            raise ValueError(
                "chat_template is not set, please set chat_template first."
            )
        elif isinstance(self.chat_template, Template):
            add_generation_prompt = kwargs.pop("add_generation_prompt", True)
            query = self._encode_chat_inputs(
                conversations, context_data, add_generation_prompt=add_generation_prompt
            )
        elif isinstance(self.chat_template, ChatTemplate):
            query = self._encode_chat_inputs_paddle(conversations, context_data)
        return query

    def _encode_chat_inputs_paddle(
        self, conversations: List[List[str]], context_data: Dict[str, Any] = {}
    ):
        context_data = self.chat_template._init_context_data(context_data)
        # encode system
        result = {}
        if self.chat_template.system:
            system = self.chat_template.render_system(context_data)
            result["system"] = self.encode(system, add_special_tokens=False)[
                "input_ids"
            ]

        # encode conversation
        conversation_ids = []
        for index, conversation in enumerate(conversations):
            # give more control to chat_template
            context_data["is_first"] = index == 0
            context_data["is_last"] = index == len(conversations) - 1

            user_input, bot_output = self.chat_template.render_conversation(
                conversation, index=index, context_data=context_data
            )
            user_ids = self.encode(user_input, add_special_tokens=False)["input_ids"]
            bot_ids = self.encode(bot_output, add_special_tokens=False)["input_ids"]
            conversation_ids.append([user_ids, bot_ids])

        result["conversations"] = conversation_ids
        return result

    def _encode_chat_inputs(
        self,
        conversations: List[List[str]],
        context_data: Dict[str, Any] = {},
        system: str = None,
        add_generation_prompt=True,
    ):
        result = {}

        # Some template do not support system msg, so we need to check it first.
        if system:
            try:
                self.chat_template.render(
                    messages={"role": "system", "content": system}
                )
            except Exception as e:
                raise ValueError("System is not supported in this tokenizer.", e)

        # convert list msg to role dict msg
        conversation_dict = []
        origin_msg = []
        for round in conversations:
            round_role = [
                {"role": "user", "content": round[0]},
                {"role": "assistant", "content": round[1]},
            ]
            origin_msg.extend(round_role)
            conversation_dict.append(round_role)
        ans = []

        # get answer in single round, then compile the chat entirely and split by single round ans
        # attention: answer should include end token!
        for conv in conversation_dict:
            roundi = [system] + conv if system else conv
            roundi_str = self.chat_template.render(
                messages=roundi, add_generation_prompt=False, **self.special_tokens_map
            )
            roundi_no_ans = [system] + [conv[0]] if system else [conv[0]]
            roundi_no_ans_str = self.chat_template.render(
                messages=roundi_no_ans,
                add_generation_prompt=add_generation_prompt,
                **self.special_tokens_map,
            )
            ans_roundi = roundi_str[len(roundi_no_ans_str) :]
            ans.append(ans_roundi)

        non_learnable_parts = self._extract_non_learnable_parts(origin_msg, ans)
        assert len(non_learnable_parts) == len(
            ans
