# Copyright (c) Gensyn Ltd.
# Licensed under the Apache License, Version 2.0.
"""Chat-aware tokenizer for open-1b.

Thin override of ``apply_chat_template`` that heals a BPE pretokenizer quirk:
encoding a string ending in whitespace (nothing after it) merges the trailing
run into different token(s) than when real content follows. The training
template's generation prompt is
``<|start_header|>assistant<|end_header|>\\n\\n`` — during SFT that ``\\n\\n``
is always training *context* immediately before the assistant's answer, never
a completion target (``sft.data.chat_template.render_conversation`` masks the
whole header, including the ``\\n\\n``, out of the loss). A plain encode of
the generation prompt on its own — nothing follows the ``\\n\\n`` — merges it
into a single token the model never saw at that position: its next-token
distribution reverts to the pretraining prior and it emits EOS immediately,
which looks fine on a single turn (the model usually recovers) but collapses
multi-turn chats to empty replies once a malformed turn enters the history.

We heal it the standard way: encode the prompt with a throwaway sentinel
character appended, then drop the sentinel's own token(s) — this forces real
content to follow the ``\\n\\n`` during encoding, so it tokenises the training
way (two separate tokens), and the sentinel's own tokens are cleanly
separable because it starts a new BPE pretoken. Same trick as
transformer-sft's ``sft.data.chat_template.encode_healed``, just applied
automatically so plain ``AutoTokenizer.apply_chat_template(...)`` users get
correct multi-turn behavior without knowing about the quirk.
"""

from __future__ import annotations

from transformers import PreTrainedTokenizerFast
from transformers.tokenization_utils_base import BatchEncoding

_HEAL_SENTINEL = "a"


class Open1BTokenizerFast(PreTrainedTokenizerFast):
    def apply_chat_template(
        self,
        conversation,
        tools=None,
        documents=None,
        chat_template=None,
        add_generation_prompt=False,
        continue_final_message=False,
        tokenize=True,
        padding=False,
        truncation=False,
        max_length=None,
        return_tensors=None,
        return_dict=True,
        return_assistant_tokens_mask=False,
        tokenizer_kwargs=None,
        **kwargs,
    ):
        # The quirk only bites the live-generation case: a single conversation,
        # tokenized, with a bare (no-content-yet) generation prompt. Anything
        # else — rendering text only, continuing/completing a final message,
        # batched conversations, assistant-token masking for training — falls
        # through to the standard implementation unchanged.
        is_batch = conversation and isinstance(conversation[0], list)
        if not (
            tokenize and add_generation_prompt and not continue_final_message
            and not return_assistant_tokens_mask and not is_batch
        ):
            return super().apply_chat_template(
                conversation,
                tools=tools,
                documents=documents,
                chat_template=chat_template,
                add_generation_prompt=add_generation_prompt,
                continue_final_message=continue_final_message,
                tokenize=tokenize,
                padding=padding,
                truncation=truncation,
                max_length=max_length,
                return_tensors=return_tensors,
                return_dict=return_dict,
                return_assistant_tokens_mask=return_assistant_tokens_mask,
                tokenizer_kwargs=tokenizer_kwargs,
                **kwargs,
            )

        text = super().apply_chat_template(
            conversation,
            tools=tools,
            documents=documents,
            chat_template=chat_template,
            add_generation_prompt=True,
            tokenize=False,
            **kwargs,
        )
        ids = self._encode_healed(text)
        encoded = BatchEncoding({"input_ids": [ids], "attention_mask": [[1] * len(ids)]})
        if return_tensors is not None:
            encoded.convert_to_tensors(tensor_type=return_tensors)
        if not return_dict:
            input_ids = encoded["input_ids"]
            return input_ids[0] if return_tensors is None else input_ids
        return encoded

    def _encode_healed(self, text: str, sentinel: str = _HEAL_SENTINEL) -> list[int]:
        full = self.encode(text + sentinel, add_special_tokens=False)
        tail = self.encode(sentinel, add_special_tokens=False)
        n = len(full) - len(tail)
        if n >= 0 and full[n:] == tail:
            return full[:n]
        return self.encode(text, add_special_tokens=False)