Upload folder using huggingface_hub

.gitattributes

@@ -1,2 +1,4 @@
 benchmark_results.png filter=lfs diff=lfs merge=lfs -text
 benchmark_results.jpg filter=lfs diff=lfs merge=lfs -text
+model.safetensors filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

.ipynb_checkpoints/chat_template-checkpoint.jinja

@@ -0,0 +1,33 @@
+{# Lumma schema-conditioned extraction chat template.
+#
+# INFERENCE: pass input_text + template (type hints) or input_text + json_schema (null schema).
+#   tokenizer.apply_chat_template(
+#       [],
+#       input_text="John Smith joined Acme.",
+#       template={"name": "string", "company": "string"},
+#       add_generation_prompt=True,
+#   )
+# NandiTokenizer converts template -> null schema via json.dumps before rendering.
+#
+# TRAINING: pass full messages (system / user / assistant) as usual.
+# #}
+{%- set extraction_system = "You are Lumma, an AI model developed by FrontiersMind. You will be provided with an input text and a JSON schema. Your task is to extract information from the input text strictly according to the provided schema. Populate only the fields defined in the schema. If a field cannot be determined from the input text, set its value to null. Do not infer, assume, rewrite, rephrase, summarize, or generate any information that is not explicitly present in the input text. Return only a valid JSON object that exactly follows the provided schema." -%}
+{%- if input_text is defined and json_schema is defined -%}
+<|im_start|><|system|>{{ '\n' }}{{ extraction_system }}<|endoftext|>{{ '\n' }}
+<|im_start|><|user|>{{ '\n' }}Input-Text: {{ input_text }}
+
+JSON-Schema: {%- if json_schema is string -%}{{ json_schema }}{%- else -%}{{ json_schema | tojson }}{%- endif -%}<|endoftext|>{{ '\n' }}
+{%- else -%}
+{%- for message in messages -%}
+{%- if message["role"] == "system" -%}
+<|im_start|><|system|>{{ '\n' }}{{ message["content"] }}<|endoftext|>{{ '\n' }}
+{%- elif message["role"] == "user" -%}
+<|im_start|><|user|>{{ '\n' }}{{ message["content"] }}<|endoftext|>{{ '\n' }}
+{%- elif message["role"] == "assistant" -%}
+<|im_start|><|assistant|>{{ '\n' }}{%- generation -%}{{ message["content"] }}<|endoftext|>{{ '\n' }}{%- endgeneration -%}
+{%- endif -%}
+{%- endfor -%}
+{%- endif -%}
+{%- if add_generation_prompt -%}
+<|im_start|><|assistant|>{{ '\n' }}
+{%- endif -%}

.ipynb_checkpoints/config-checkpoint.json

@@ -0,0 +1,45 @@
+{
+  "architectures": [
+    "NandiForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_nandi.NandiConfig",
+    "AutoModel": "modeling_nandi.NandiModel",
+    "AutoModelForCausalLM": "modeling_nandi.NandiForCausalLM"
+  },
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "embedding_rank": 512,
+  "eos_token_id": 0,
+  "factorized_embedding": true,
+  "head_dim": 90,
+  "hidden_act": "silu",
+  "hidden_size": 1440,
+  "initializer_range": 0.02,
+  "intermediate_size": 3280,
+  "kv_cache_mode": "shared",
+  "layer_sharing": false,
+  "layer_sharing_repeats": 1,
+  "max_position_embeddings": 12288,
+  "mlp_bias": false,
+  "model_type": "nandi",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 30,
+  "num_key_value_heads": 8,
+  "pad_token_id": 3,
+  "pretraining_tp": 1,
+  "q_norm": true,
+  "qk_norm": false,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000.0,
+    "rope_type": "default"
+  },
+  "shared_kv": true,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.14.1",
+  "use_cache": false,
+  "vocab_size": 131072
+}

.ipynb_checkpoints/generation_config-checkpoint.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": [
+    0
+  ],
+  "pad_token_id": 3,
+  "transformers_version": "5.14.1",
+  "use_cache": true
+}

.ipynb_checkpoints/tokenization_nandi-checkpoint.py

@@ -0,0 +1,219 @@
+# 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, *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")
+
+        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)
+
+        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"]

.ipynb_checkpoints/tokenizer_config-checkpoint.json

@@ -0,0 +1,9 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|im_start|>",
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 12288,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|endoftext|>"
+}

chat_template.jinja

@@ -0,0 +1,33 @@
+{# Lumma schema-conditioned extraction chat template.
+#
+# INFERENCE: pass input_text + template (type hints) or input_text + json_schema (null schema).
+#   tokenizer.apply_chat_template(
+#       [],
+#       input_text="John Smith joined Acme.",
+#       template={"name": "string", "company": "string"},
+#       add_generation_prompt=True,
+#   )
+# NandiTokenizer converts template -> null schema via json.dumps before rendering.
+#
+# TRAINING: pass full messages (system / user / assistant) as usual.
+# #}
+{%- set extraction_system = "You are Lumma, an AI model developed by FrontiersMind. You will be provided with an input text and a JSON schema. Your task is to extract information from the input text strictly according to the provided schema. Populate only the fields defined in the schema. If a field cannot be determined from the input text, set its value to null. Do not infer, assume, rewrite, rephrase, summarize, or generate any information that is not explicitly present in the input text. Return only a valid JSON object that exactly follows the provided schema." -%}
+{%- if input_text is defined and json_schema is defined -%}
+<|im_start|><|system|>{{ '\n' }}{{ extraction_system }}<|endoftext|>{{ '\n' }}
+<|im_start|><|user|>{{ '\n' }}Input-Text: {{ input_text }}
+
+JSON-Schema: {%- if json_schema is string -%}{{ json_schema }}{%- else -%}{{ json_schema | tojson }}{%- endif -%}<|endoftext|>{{ '\n' }}
+{%- else -%}
+{%- for message in messages -%}
+{%- if message["role"] == "system" -%}
+<|im_start|><|system|>{{ '\n' }}{{ message["content"] }}<|endoftext|>{{ '\n' }}
+{%- elif message["role"] == "user" -%}
+<|im_start|><|user|>{{ '\n' }}{{ message["content"] }}<|endoftext|>{{ '\n' }}
+{%- elif message["role"] == "assistant" -%}
+<|im_start|><|assistant|>{{ '\n' }}{%- generation -%}{{ message["content"] }}<|endoftext|>{{ '\n' }}{%- endgeneration -%}
+{%- endif -%}
+{%- endfor -%}
+{%- endif -%}
+{%- if add_generation_prompt -%}
+<|im_start|><|assistant|>{{ '\n' }}
+{%- endif -%}

config.json

@@ -0,0 +1,45 @@
+{
+  "architectures": [
+    "NandiForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_nandi.NandiConfig",
+    "AutoModel": "modeling_nandi.NandiModel",
+    "AutoModelForCausalLM": "modeling_nandi.NandiForCausalLM"
+  },
+  "bos_token_id": 1,
+  "dtype": "bfloat16",
+  "embedding_rank": 512,
+  "eos_token_id": 0,
+  "factorized_embedding": true,
+  "head_dim": 90,
+  "hidden_act": "silu",
+  "hidden_size": 1440,
+  "initializer_range": 0.02,
+  "intermediate_size": 3280,
+  "kv_cache_mode": "shared",
+  "layer_sharing": false,
+  "layer_sharing_repeats": 1,
+  "max_position_embeddings": 12288,
+  "mlp_bias": false,
+  "model_type": "nandi",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 30,
+  "num_key_value_heads": 8,
+  "pad_token_id": 3,
+  "pretraining_tp": 1,
+  "q_norm": true,
+  "qk_norm": false,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000.0,
+    "rope_type": "default"
+  },
+  "shared_kv": true,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.14.1",
+  "use_cache": false,
+  "vocab_size": 131072
+}

configuration_nandi.py

@@ -0,0 +1,123 @@
+from transformers.configuration_utils import PretrainedConfig
+
+
+class NandiConfig(PretrainedConfig):
+    r"""
+    Configuration class for the Nandi model.
+
+    Example:
+
+    ```python
+    >>> from transformers import AutoConfig, AutoModelForCausalLM
+
+    >>> configuration = AutoConfig.from_pretrained("Rta-AILabs/Nandi-500M-remote", trust_remote_code=True)
+
+    >>> model = AutoModelForCausalLM.from_pretrained("Rta-AILabs/Nandi-500M-remote", trust_remote_code=True)
+
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "nandi"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+
+    def __init__(
+        self,
+        vocab_size=131072,
+        hidden_size=1248,
+        intermediate_size=3556,
+        num_hidden_layers=28,
+        num_attention_heads=16,
+        num_key_value_heads=8,
+        head_dim=None,
+        hidden_act="silu",
+        max_position_embeddings=2048,
+        initializer_range=0.008,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=1,
+        eos_token_id=0,
+        pretraining_tp=1,
+        tie_word_embeddings=True,
+        rope_parameters=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        mlp_bias=False,
+        factorized_embedding=False,
+        embedding_rank=768,
+        layer_sharing=False,
+        layer_sharing_repeats=1,
+        qk_norm=True,
+        q_norm=False,
+        shared_kv=True,
+        kv_cache_mode="shared",
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads if num_key_value_heads is not None else num_attention_heads
+        self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.pretraining_tp = pretraining_tp
+        self.rope_parameters = rope_parameters if rope_parameters is not None else {"rope_theta": 1000000.0}
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.mlp_bias = mlp_bias
+        self.factorized_embedding = factorized_embedding
+        self.embedding_rank = embedding_rank
+        self.layer_sharing = layer_sharing
+
+        self.layer_sharing_repeats = max(1, int(layer_sharing_repeats or 1))
+        if qk_norm and q_norm:
+            raise ValueError("Use either qk_norm or q_norm, not both")
+        self.qk_norm = qk_norm
+        self.q_norm = q_norm
+
+        self.shared_kv = shared_kv
+
+        if kv_cache_mode not in ("shared", "vanilla"):
+            raise ValueError(
+                f"`kv_cache_mode` must be 'shared' or 'vanilla', got {kv_cache_mode!r}."
+            )
+        self.kv_cache_mode = kv_cache_mode
+
+        if self.factorized_embedding and self.embedding_rank <= 0:
+            raise ValueError(
+                f"`embedding_rank` must be positive when `factorized_embedding=True`, got {self.embedding_rank}."
+            )
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError(
+                f"`hidden_size` ({self.hidden_size}) must be divisible by "
+                f"`num_attention_heads` ({self.num_attention_heads})."
+            )
+        if self.layer_sharing_repeats < 1:
+            raise ValueError(f"`layer_sharing_repeats` must be >= 1, got {self.layer_sharing_repeats}.")
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["NandiConfig"]

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": [
+    0
+  ],
+  "pad_token_id": 3,
+  "transformers_version": "5.14.1",
+  "use_cache": true
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3424f17dd80518a25878cf1c8f9816a9276ce0422ce38232648d975dbb32a085
+size 1298775408

modeling_nandi.py

@@ -0,0 +1,541 @@
+from collections.abc import Callable
+
+import torch
+import torch.nn as nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache, DynamicLayer
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring
+from transformers.utils.deprecation import deprecate_kwarg
+from transformers.utils.generic import can_return_tuple, merge_with_config_defaults
+from transformers.utils.output_capturing import capture_outputs
+from .configuration_nandi import NandiConfig
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class NandiRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class NandiRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor
+
+    def __init__(self, config: NandiConfig, device=None):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_type = self.config.rope_parameters.get("rope_type", "default")
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != "default":
+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: NandiConfig | None = None,
+        device: torch.device | None = None,
+        seq_len: int | None = None,
+    ) -> tuple[torch.Tensor, float]:
+        del seq_len
+        base = config.rope_parameters["rope_theta"]
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+        attention_factor = 1.0
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, attention_factor
+
+    @torch.no_grad()
+    @dynamic_rope_update
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    del position_ids
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    del kwargs
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class NandiAttention(nn.Module):
+    def __init__(self, config: NandiConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = config.head_dim
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        self.shared_kv = getattr(config, "shared_kv", False)
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        if self.shared_kv:
+            self.v_proj = None
+        else:
+            self.v_proj = nn.Linear(
+                config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+            )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+
+        self.qk_norm = getattr(config, "qk_norm", False)
+        self.q_norm_enabled = getattr(config, "q_norm", False)
+        if self.qk_norm:
+            self.q_norm = NandiRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.k_norm = NandiRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        elif self.q_norm_enabled:
+            self.q_norm = NandiRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.k_norm = None
+        else:
+            self.q_norm = None
+            self.k_norm = None
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        k_raw = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        if self.shared_kv:
+            kv_cache_mode = getattr(self.config, "kv_cache_mode", "shared")
+
+            if self.qk_norm:
+                query_states = self.q_norm(query_states)
+            elif self.q_norm_enabled:
+                query_states = self.q_norm(query_states)
+
+            if kv_cache_mode == "shared":
+                if past_key_values is not None:
+                    empty_v = torch.empty(
+                        k_raw.shape[0],
+                        k_raw.shape[1],
+                        0,
+                        k_raw.shape[3],
+                        device=k_raw.device,
+                        dtype=k_raw.dtype,
+                    )
+                    k_raw_full, _ = past_key_values.update(k_raw, empty_v, self.layer_idx)
+                else:
+                    k_raw_full = k_raw
+
+                value_states = k_raw_full
+                key_states = self.k_norm(k_raw_full) if self.qk_norm else k_raw_full
+
+                cos, sin = position_embeddings
+                q_len = query_states.shape[-2]
+                cos_q = cos[..., -q_len:, :]
+                sin_q = sin[..., -q_len:, :]
+                query_states, _ = apply_rotary_pos_emb(query_states, query_states, cos_q, sin_q)
+                _, key_states = apply_rotary_pos_emb(key_states, key_states, cos, sin)
+
+            else:  
+                key_states = self.k_norm(k_raw) if self.qk_norm else k_raw
+                value_states = k_raw  
+
+                cos, sin = position_embeddings  
+                query_states, key_states = apply_rotary_pos_emb(
+                    query_states, key_states, cos, sin
+                )
+
+                if past_key_values is not None:
+                    key_states, value_states = past_key_values.update(
+                        key_states, value_states, self.layer_idx
+                    )
+
+        else:
+            value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+            key_states = k_raw
+
+            if self.qk_norm:
+                query_states = self.q_norm(query_states)
+                key_states = self.k_norm(key_states)
+            elif self.q_norm_enabled:
+                query_states = self.q_norm(query_states)
+
+            cos, sin = position_embeddings
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+            if past_key_values is not None:
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class NandiMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_bias)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class NandiDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: NandiConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.self_attn = NandiAttention(config=config, layer_idx=layer_idx)
+        self.mlp = NandiMLP(config)
+        self.input_layernorm = NandiRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = NandiRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class _VirtualLayerCache:
+    """Proxy that shifts cache layer indices by `offset` to give each repeat its own virtual slots."""
+
+    def __init__(self, cache: Cache, offset: int):
+        self._cache = cache
+        self._offset = offset
+
+    def __getattr__(self, name):
+        return getattr(self._cache, name)
+
+    def update(self, key_states, value_states, layer_idx, cache_kwargs=None):
+        virtual_idx = layer_idx + self._offset
+        # grow the backing cache if generate() pre-allocated fewer slots than needed
+        while len(self._cache.layers) <= virtual_idx:
+            self._cache.layers.append(DynamicLayer())
+        return self._cache.update(key_states, value_states, virtual_idx, cache_kwargs)
+
+    def get_seq_length(self, layer_idx: int = 0) -> int:
+        return self._cache.get_seq_length(layer_idx + self._offset)
+
+
+@auto_docstring
+class NandiPreTrainedModel(PreTrainedModel):
+    config: NandiConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["NandiDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": NandiDecoderLayer,
+        "attentions": NandiAttention,
+    }
+
+    def __init__(self, config: NandiConfig):
+        super().__init__(config)
+
+
+@auto_docstring
+class NandiModel(NandiPreTrainedModel):
+    def __init__(self, config: NandiConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        embedding_dim = config.embedding_rank if config.factorized_embedding else config.hidden_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, embedding_dim, self.padding_idx)
+        self.embedding_proj = (
+            nn.Linear(config.embedding_rank, config.hidden_size, bias=False) if config.factorized_embedding else None
+        )
+        self.layers = nn.ModuleList(
+            [NandiDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = NandiRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = NandiRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if self.embedding_proj is not None:
+            inputs_embeds = self.embedding_proj(inputs_embeds)
+        repeats = max(1, int(getattr(self.config, "layer_sharing_repeats", 1) or 1))
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if position_ids is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            position_ids = position_ids.unsqueeze(0)
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+        hidden_states = inputs_embeds
+        kv_cache_mode = getattr(self.config, "kv_cache_mode", "shared")
+        if (
+            getattr(self.config, "shared_kv", False)
+            and kv_cache_mode == "shared"
+            and past_key_values is not None
+        ):
+            past_len = past_key_values.get_seq_length(0)
+            cur_len = inputs_embeds.shape[1]
+            full_position_ids = torch.arange(
+                past_len + cur_len, device=inputs_embeds.device
+            ).unsqueeze(0)
+            position_embeddings = self.rotary_emb(hidden_states, position_ids=full_position_ids)
+        else:
+            position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            for repeat_idx in range(repeats):
+
+                repeat_cache = (
+                    _VirtualLayerCache(past_key_values, repeat_idx * self.config.num_hidden_layers)
+                    if (past_key_values is not None and repeat_idx > 0)
+                    else past_key_values
+                )
+                hidden_states = decoder_layer(
+                    hidden_states,
+                    attention_mask=causal_mask,
+                    position_embeddings=position_embeddings,
+                    position_ids=position_ids,
+                    past_key_values=repeat_cache,
+                    use_cache=use_cache,
+                    **kwargs,
+                )
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+@auto_docstring
+class NandiForCausalLM(NandiPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_gather_output"}
+    _pp_plan = {
+        "lm_head_proj": (["hidden_states"], ["hidden_states"]),
+        "lm_head": (["hidden_states"], ["logits"]),
+    }
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = NandiModel(config)
+        self.vocab_size = config.vocab_size
+
+        lm_head_in_features = config.embedding_rank if config.factorized_embedding else config.hidden_size
+        self.lm_head_proj = (
+            nn.Linear(config.hidden_size, config.embedding_rank, bias=False) if config.factorized_embedding else None
+        )
+        self.lm_head = nn.Linear(lm_head_in_features, config.vocab_size, bias=False)
+
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        if self.lm_head_proj is not None:
+            hidden_states = self.lm_head_proj(hidden_states)
+
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = ["NandiPreTrainedModel", "NandiModel", "NandiForCausalLM"]

tokenization_nandi.py

@@ -0,0 +1,232 @@
+# 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"]

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:354154322dbfd5971fff276473422331b7e2f23073a835ebdaa5f5a84f11a340
+size 12460613

tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|im_start|>",
+  "bos_token_id": 1,
+  "eos_token": "<|endoftext|>",
+  "eos_token_id": 0,
+  "pad_token": "<|pad|>",
+  "pad_token_id": 3,
+  "model_max_length": 12288,
+  "tokenizer_class": "NandiTokenizer",
+  "unk_token": "<|endoftext|>",
+  "auto_map": {
+    "AutoTokenizer": ["tokenization_nandi.NandiTokenizer", null]
+  }
+}