tokenization_nandi.py

# Copyright 2026 The HuggingFace Inc. team. 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.
"""Tokenization classes for the Nandi family."""

from __future__ import annotations

import json
from typing import Any

from tokenizers import Regex, Tokenizer, decoders, normalizers, pre_tokenizers
from tokenizers.models import BPE

from transformers.tokenization_utils_tokenizers import TokenizersBackend
from transformers.utils import logging


logger = logging.get_logger(__name__)

PRETOKENIZE_REGEX = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?(?:\p{L}\p{M}*)+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""

ALLOWED_TEMPLATE_LEAF_TYPES = frozenset({"string", "number", "integer", "boolean", "null"})
_IM_START = "<|im_start|>"


def normalize_extraction_template(template: Any) -> dict[str, Any]:
    """Convert a type-hint template to null-placeholder schema (matches SFT/DPO training)."""
    if isinstance(template, str):
        try:
            template = json.loads(template)
        except json.JSONDecodeError as exc:
            raise ValueError(f"template must be valid JSON: {exc}") from exc

    if not isinstance(template, dict):
        raise ValueError("template root must be a JSON object")

    return _nullify_template_node(template)


def _nullify_template_node(node: Any) -> Any:
    if node is None:
        return None

    if isinstance(node, dict):
        return {key: _nullify_template_node(value) for key, value in node.items()}

    if isinstance(node, list):
        if len(node) == 0:
            return []
        if len(node) == 1:
            item = node[0]
            if isinstance(item, str) and item in ALLOWED_TEMPLATE_LEAF_TYPES:
                return []
            return [_nullify_template_node(item)]
        raise ValueError(
            'array template must be [] or a one-element type list, e.g. ["string"]'
        )

    if isinstance(node, str):
        if node in ALLOWED_TEMPLATE_LEAF_TYPES:
            return None
        raise ValueError(
            f"invalid template leaf {node!r}; use a type name like 'string' or null"
        )

    if isinstance(node, bool):
        raise ValueError("template leaf must be a type name, not a boolean literal")

    if isinstance(node, (int, float)):
        raise ValueError("template leaf must be a type name, not a numeric literal")

    raise ValueError(f"unsupported template value: {node!r}")


def _maybe_add_im_start_prefix(text: str) -> str:
    stripped = text.lstrip()
    if stripped.startswith(_IM_START):
        return text
    return f"{_IM_START} {text}"


class NandiTokenizer(TokenizersBackend):
    model_input_names = ["input_ids", "attention_mask"]
    model = BPE

    def __init__(
        self,
        vocab: str | dict[str, int] | None = None,
        merges: str | list[str] | None = None,
        vocab_file=None,
        merges_file=None,
        unk_token: str = "<|endoftext|>",
        bos_token: str = "<|im_start|>",
        eos_token: str = "<|endoftext|>",
        pad_token: str = "<|pad|>",
        add_prefix_space: bool | None = None,
        **kwargs,
    ):
        self._vocab = (
            vocab
            if vocab is not None
            else {
                "<|endoftext|>": 0,
            }
        )
        self._merges = merges or []

        self._tokenizer = Tokenizer(
            BPE(
                vocab=self._vocab,
                merges=self._merges,
                dropout=None,
                unk_token=None,
                continuing_subword_prefix="",
                end_of_word_suffix="",
                fuse_unk=False,
                byte_fallback=False,
            )
        )
        self._tokenizer.decoder = decoders.ByteLevel()
        self._tokenizer.normalizer = normalizers.NFC()
        self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
            [
                pre_tokenizers.Split(
                    Regex(PRETOKENIZE_REGEX),
                    behavior="isolated",
                    invert=False,
                ),
                pre_tokenizers.ByteLevel(
                    add_prefix_space=False,
                    trim_offsets=True,
                    use_regex=False,
                ),
            ]
        )

        super().__init__(
            vocab_file=vocab_file,
            merges_file=merges_file,
            unk_token=unk_token,
            bos_token=bos_token,
            eos_token=eos_token,
            pad_token=pad_token,
            add_prefix_space=add_prefix_space,
            **kwargs,
        )

    def apply_chat_template(self, conversation=None, *args, **kwargs):
        """Support extraction inference via input_text + template (type hints)."""
        template = kwargs.pop("template", None)
        input_text = kwargs.get("input_text")
        json_schema = kwargs.get("json_schema")

        extraction_mode = (
            template is not None or input_text is not None or json_schema is not None
        )

        if template is not None:
            if json_schema is not None:
                raise ValueError("Pass either template or json_schema, not both.")
            if input_text is None:
                raise ValueError("input_text is required when template is provided.")
            null_schema = normalize_extraction_template(template)
            kwargs["json_schema"] = json.dumps(null_schema, ensure_ascii=False)
        elif json_schema is not None and not isinstance(json_schema, str):
            if isinstance(json_schema, (dict, list)):
                kwargs["json_schema"] = json.dumps(json_schema, ensure_ascii=False)

        if extraction_mode:
            if not conversation:
                conversation = None
        elif conversation is None:
            raise ValueError(
                "conversation is required unless using extraction kwargs "
                "(input_text + template, or input_text + json_schema)."
            )

        return super().apply_chat_template(conversation, *args, **kwargs)

    normalize_template = staticmethod(normalize_extraction_template)

    def __call__(self, text, *args, **kwargs):
        add_special_tokens = kwargs.get("add_special_tokens", False)

        if not add_special_tokens:
            if isinstance(text, list):
                text = [_maybe_add_im_start_prefix(t) if isinstance(t, str) else t for t in text]
            elif isinstance(text, str):
                text = _maybe_add_im_start_prefix(text)

        return super().__call__(text, *args, **kwargs)

    def encode(
        self,
        text,
        text_pair=None,
        add_special_tokens: bool = True,
        padding=False,
        truncation=None,
        max_length=None,
        stride: int = 0,
        padding_side=None,
        return_tensors=None,
        **kwargs,
    ):
        if isinstance(text, str):
            text = _maybe_add_im_start_prefix(text)
        return super().encode(
            text,
            text_pair=text_pair,
            add_special_tokens=add_special_tokens,
            padding=padding,
            truncation=truncation,
            max_length=max_length,
            stride=stride,
            padding_side=padding_side,
            return_tensors=return_tensors,
            **kwargs,
        )


__all__ = ["NandiTokenizer", "normalize_extraction_template"]