Merge HF-format model into repo (config.json keeps JAX identity under jax_config)

config.json

@@ -1,5 +1,41 @@
 {
-  "library_name": "jax",
-  "model_type": "custom",
-  "architectures": ["SimpleAttentionNetwork"]
+  "architectures": [
+    "NeedleForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_needle.NeedleConfig",
+    "AutoModel": "modeling_needle.NeedleModel",
+    "AutoModelForCausalLM": "modeling_needle.NeedleForCausalLM",
+    "AutoModelForSeq2SeqLM": "modeling_needle.NeedleForCausalLM"
+  },
+  "bos_token_id": 2,
+  "d_model": 512,
+  "decoder_start_token_id": 1,
+  "eos_token_id": 1,
+  "hidden_size": 512,
+  "is_encoder_decoder": true,
+  "model_type": "needle",
+  "num_attention_heads": 8,
+  "num_decoder_layers": 8,
+  "num_encoder_layers": 12,
+  "num_heads": 8,
+  "num_hidden_layers": 8,
+  "num_key_value_heads": 4,
+  "num_kv_heads": 4,
+  "pad_token_id": 0,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "5.5.4",
+  "unk_token_id": 3,
+  "vocab_size": 8192,
+  "jax_config": {
+    "library_name": "jax",
+    "model_type": "custom",
+    "architectures": [
+      "SimpleAttentionNetwork"
+    ],
+    "checkpoint": "needle.pkl"
+  }
 }

configuration_needle.py

@@ -0,0 +1,63 @@
+"""Hugging Face configuration for Needle."""
+
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class NeedleConfig(PretrainedConfig):
+    model_type = "needle"
+
+    def __init__(
+        self,
+        vocab_size: int = 8192,
+        hidden_size: int | None = None,
+        d_model: int = 512,
+        num_attention_heads: int | None = None,
+        num_heads: int = 8,
+        num_key_value_heads: int | None = None,
+        num_kv_heads: int = 4,
+        num_encoder_layers: int = 12,
+        num_decoder_layers: int = 8,
+        rope_theta: float = 10000.0,
+        rms_norm_eps: float = 1e-6,
+        pad_token_id: int = 0,
+        eos_token_id: int = 1,
+        bos_token_id: int = 2,
+        unk_token_id: int = 3,
+        decoder_start_token_id: int | None = None,
+        tie_word_embeddings: bool = True,
+        torch_dtype: str = "bfloat16",
+        **kwargs,
+    ) -> None:
+        kwargs.pop("is_encoder_decoder", None)
+        hidden_size = int(hidden_size if hidden_size is not None else d_model)
+        num_attention_heads = int(num_attention_heads if num_attention_heads is not None else num_heads)
+        num_key_value_heads = int(num_key_value_heads if num_key_value_heads is not None else num_kv_heads)
+        decoder_start_token_id = eos_token_id if decoder_start_token_id is None else decoder_start_token_id
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            bos_token_id=bos_token_id,
+            unk_token_id=unk_token_id,
+            decoder_start_token_id=decoder_start_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            is_encoder_decoder=True,
+            torch_dtype=torch_dtype,
+            **kwargs,
+        )
+
+        self.vocab_size = int(vocab_size)
+        self.hidden_size = hidden_size
+        self.d_model = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.num_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.num_kv_heads = num_key_value_heads
+        self.num_encoder_layers = int(num_encoder_layers)
+        self.num_decoder_layers = int(num_decoder_layers)
+        self.num_hidden_layers = int(num_decoder_layers)
+        self.rope_theta = float(rope_theta)
+        self.rms_norm_eps = float(rms_norm_eps)
+        self.attention_head_dim = hidden_size // max(1, num_attention_heads)

model.safetensors

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

modeling_needle.py

@@ -0,0 +1,429 @@
+"""Minimal PyTorch Needle model for Cactus conversion."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
+
+from .configuration_needle import NeedleConfig
+
+
+class NeedleRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(hidden_size))
+        self.eps = float(eps)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        variance = x.float().pow(2).mean(dim=-1, keepdim=True)
+        x = x.float() * torch.rsqrt(variance + self.eps)
+        return x.to(dtype=dtype) * (1.0 + self.weight.to(dtype=dtype))
+
+
+def _padding_mask(input_ids: torch.Tensor, pad_token_id: int) -> torch.Tensor:
+    return (input_ids != int(pad_token_id))[:, None, None, :]
+
+
+def _causal_mask(seq_len: int, device: torch.device) -> torch.Tensor:
+    return torch.ones((seq_len, seq_len), dtype=torch.bool, device=device).tril()[None, None, :, :]
+
+
+def _build_inv_freq(head_dim: int, theta: float) -> torch.Tensor:
+    return 1.0 / (float(theta) ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / float(head_dim)))
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    half = x.shape[-1] // 2
+    return torch.cat((-x[..., half:], x[..., :half]), dim=-1)
+
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    cos = cos.unsqueeze(2)
+    sin = sin.unsqueeze(2)
+    return (x * cos) + (_rotate_half(x) * sin)
+
+
+def _rotary_tables(
+    inv_freq: torch.Tensor,
+    batch_size: int,
+    seq_len: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    position_ids = torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0).expand(batch_size, -1)
+    inv_freq = inv_freq[None, :, None].float().expand(batch_size, -1, 1).to(device)
+    freqs = (inv_freq @ position_ids[:, None, :].float()).transpose(1, 2)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos().to(dtype=dtype), emb.sin().to(dtype=dtype)
+
+
+def _rotary_tables_for_position_ids(
+    inv_freq: torch.Tensor,
+    position_ids: torch.Tensor,
+    dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    inv_freq = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(position_ids.device)
+    freqs = (inv_freq @ position_ids[:, None, :].float()).transpose(1, 2)
+    emb = torch.cat((freqs, freqs), dim=-1)
+    return emb.cos().to(dtype=dtype), emb.sin().to(dtype=dtype)
+
+
+def _add_clipped(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    return torch.clamp(a + b, min=-65500.0, max=65500.0)
+
+
+class NeedleAttention(nn.Module):
+    def __init__(self, config: NeedleConfig) -> None:
+        super().__init__()
+        self.hidden_size = int(config.hidden_size)
+        self.num_heads = int(config.num_attention_heads)
+        self.num_key_value_heads = int(config.num_key_value_heads)
+        self.head_dim = self.hidden_size // self.num_heads
+        kv_size = self.num_key_value_heads * self.head_dim
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, kv_size, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, kv_size, bias=False)
+        self.out_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.q_norm = NeedleRMSNorm(self.head_dim, config.rms_norm_eps)
+        self.k_norm = NeedleRMSNorm(self.head_dim, config.rms_norm_eps)
+        self.scale = 1.0 / math.sqrt(float(self.head_dim))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_value_states: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        rope: tuple[torch.Tensor, torch.Tensor] | None,
+    ) -> torch.Tensor:
+        batch, q_len, _ = hidden_states.shape
+        kv_len = key_value_states.shape[1]
+        q = self.q_proj(hidden_states).view(batch, q_len, self.num_heads, self.head_dim)
+        k = self.k_proj(key_value_states).view(batch, kv_len, self.num_key_value_heads, self.head_dim)
+        v = self.v_proj(key_value_states).view(batch, kv_len, self.num_key_value_heads, self.head_dim)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        if rope is not None:
+            cos, sin = rope
+            q = _apply_rope(q, cos, sin)
+            k = _apply_rope(k, cos, sin)
+        out = F.scaled_dot_product_attention(
+            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
+            attn_mask=attention_mask,
+            dropout_p=0.0,
+            is_causal=False,
+            scale=self.scale,
+            enable_gqa=self.num_heads != self.num_key_value_heads,
+        )
+        return self.out_proj(out.transpose(1, 2).contiguous().view(batch, q_len, self.hidden_size))
+
+    def project_kv(self, key_value_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        batch, kv_len, _ = key_value_states.shape
+        k = self.k_proj(key_value_states).view(batch, kv_len, self.num_key_value_heads, self.head_dim)
+        v = self.v_proj(key_value_states).view(batch, kv_len, self.num_key_value_heads, self.head_dim)
+        k = self.k_norm(k)
+        return k.contiguous(), v.contiguous()
+
+    def forward_with_kv(
+        self,
+        hidden_states: torch.Tensor,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        rope: tuple[torch.Tensor, torch.Tensor] | None,
+    ) -> torch.Tensor:
+        batch, q_len, _ = hidden_states.shape
+        q = self.q_proj(hidden_states).view(batch, q_len, self.num_heads, self.head_dim)
+        q = self.q_norm(q)
+        if rope is not None:
+            cos, sin = rope
+            q = _apply_rope(q, cos, sin)
+        out = F.scaled_dot_product_attention(
+            q.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2),
+            attn_mask=attention_mask,
+            dropout_p=0.0,
+            is_causal=False,
+            scale=self.scale,
+            enable_gqa=self.num_heads != self.num_key_value_heads,
+        )
+        return self.out_proj(out.transpose(1, 2).contiguous().view(batch, q_len, self.hidden_size))
+
+
+class NeedleEncoderLayer(nn.Module):
+    def __init__(self, config: NeedleConfig) -> None:
+        super().__init__()
+        self.input_layernorm = NeedleRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.self_attn = NeedleAttention(config)
+        self.attn_gate = nn.Parameter(torch.zeros(1))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        rope: tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        normed = self.input_layernorm(hidden_states)
+        attn = self.self_attn(normed, normed, attention_mask, rope)
+        return _add_clipped(hidden_states, torch.sigmoid(self.attn_gate).to(dtype=attn.dtype) * attn)
+
+
+class NeedleDecoderLayer(nn.Module):
+    def __init__(self, config: NeedleConfig) -> None:
+        super().__init__()
+        self.input_layernorm = NeedleRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.self_attn = NeedleAttention(config)
+        self.self_attn_gate = nn.Parameter(torch.zeros(1))
+        self.encoder_attn_layer_norm = NeedleRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.encoder_attn = NeedleAttention(config)
+        self.cross_attn_gate = nn.Parameter(torch.zeros(1))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        self_mask: torch.Tensor,
+        encoder_mask: torch.Tensor,
+        rope: tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        normed = self.input_layernorm(hidden_states)
+        attn = self.self_attn(normed, normed, self_mask, rope)
+        hidden_states = _add_clipped(hidden_states, torch.sigmoid(self.self_attn_gate).to(dtype=attn.dtype) * attn)
+        attn = self.encoder_attn(self.encoder_attn_layer_norm(hidden_states), encoder_hidden_states, encoder_mask, None)
+        return _add_clipped(hidden_states, torch.sigmoid(self.cross_attn_gate).to(dtype=attn.dtype) * attn)
+
+
+class NeedleEncoder(nn.Module):
+    def __init__(self, config: NeedleConfig) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList([NeedleEncoderLayer(config) for _ in range(config.num_encoder_layers)])
+        self.final_norm = NeedleRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.rope_theta = float(config.rope_theta)
+        self.register_buffer("inv_freq", _build_inv_freq(self.head_dim, self.rope_theta), persistent=False)
+
+    def reset_rope(self) -> None:
+        self.inv_freq = _build_inv_freq(self.head_dim, self.rope_theta).to(device=self.inv_freq.device)
+
+    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        inv_freq = _build_inv_freq(self.head_dim, self.rope_theta).to(device=hidden_states.device)
+        rope = _rotary_tables(
+            inv_freq,
+            hidden_states.shape[0],
+            hidden_states.shape[1],
+            hidden_states.device,
+            hidden_states.dtype,
+        )
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, attention_mask, rope)
+        return self.final_norm(hidden_states)
+
+
+class NeedleDecoder(nn.Module):
+    def __init__(self, config: NeedleConfig) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList([NeedleDecoderLayer(config) for _ in range(config.num_decoder_layers)])
+        self.norm = NeedleRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.rope_theta = float(config.rope_theta)
+        self.register_buffer("inv_freq", _build_inv_freq(self.head_dim, self.rope_theta), persistent=False)
+
+    def reset_rope(self) -> None:
+        self.inv_freq = _build_inv_freq(self.head_dim, self.rope_theta).to(device=self.inv_freq.device)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        self_mask: torch.Tensor,
+        encoder_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        inv_freq = _build_inv_freq(self.head_dim, self.rope_theta).to(device=hidden_states.device)
+        rope = _rotary_tables(
+            inv_freq,
+            hidden_states.shape[0],
+            hidden_states.shape[1],
+            hidden_states.device,
+            hidden_states.dtype,
+        )
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, encoder_hidden_states, self_mask, encoder_mask, rope)
+        return self.norm(hidden_states)
+
+
+class NeedleModel(PreTrainedModel):
+    config_class = NeedleConfig
+    base_model_prefix = "model"
+    main_input_name = "input_ids"
+
+    def __init__(self, config: NeedleConfig) -> None:
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.embed_scale = math.sqrt(float(config.hidden_size))
+        self.encoder = NeedleEncoder(config)
+        self.decoder = NeedleDecoder(config)
+        self.post_init()
+        self.reset_rope()
+
+    def reset_rope(self) -> None:
+        self.encoder.reset_rope()
+        self.decoder.reset_rope()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        decoder_input_ids: torch.Tensor | None = None,
+        **_: Any,
+    ) -> BaseModelOutput:
+        decoder_input_ids = input_ids if decoder_input_ids is None else decoder_input_ids
+        encoder_mask = _padding_mask(input_ids, self.config.pad_token_id)
+        if attention_mask is not None:
+            encoder_mask = encoder_mask & attention_mask[:, None, None, :].to(dtype=torch.bool)
+        encoder_hidden = self.embed_tokens(input_ids) * self.embed_scale
+        encoder_hidden = self.encoder(encoder_hidden, encoder_mask)
+
+        self_mask = _causal_mask(decoder_input_ids.shape[1], decoder_input_ids.device)
+        decoder_hidden = self.embed_tokens(decoder_input_ids) * self.embed_scale
+        decoder_hidden = self.decoder(decoder_hidden, encoder_hidden, self_mask, encoder_mask)
+        return BaseModelOutput(last_hidden_state=decoder_hidden)
+
+
+class NeedleForCausalLM(PreTrainedModel):
+    config_class = NeedleConfig
+    base_model_prefix = "model"
+    main_input_name = "input_ids"
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config: NeedleConfig) -> None:
+        super().__init__(config)
+        self.model = NeedleModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+        self.model.reset_rope()
+        self.tie_weights()
+
+    def get_encoder(self) -> NeedleEncoder:
+        return self.model.encoder
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head
+
+    def set_output_embeddings(self, value: nn.Linear) -> None:
+        self.lm_head = value
+
+    def tie_weights(self, *args: Any, **kwargs: Any) -> None:
+        del args, kwargs
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        decoder_input_ids: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> Seq2SeqLMOutput:
+        hidden_states = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            decoder_input_ids=decoder_input_ids,
+            **kwargs,
+        ).last_hidden_state
+        return Seq2SeqLMOutput(logits=self.lm_head(hidden_states))
+
+    def cactus_source_encode(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        base_mask = _padding_mask(input_ids, self.config.pad_token_id)
+        base_mask = base_mask & attention_mask[:, None, None, :].to(dtype=torch.bool)
+        encoder_mask = base_mask.expand(
+            -1,
+            int(self.config.num_attention_heads),
+            input_ids.shape[1],
+            -1,
+        ).contiguous()
+        encoder_hidden = self.model.embed_tokens(input_ids) * self.model.embed_scale
+        encoder_hidden = self.model.encoder(encoder_hidden, encoder_mask)
+        decoder_mask = base_mask.expand(
+            -1,
+            int(self.config.num_attention_heads),
+            1,
+            -1,
+        ).contiguous()
+        return encoder_hidden, decoder_mask.to(dtype=encoder_hidden.dtype)
+
+    def cactus_decoder_cross_kv(
+        self,
+        encoder_hidden_states: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, ...]:
+        del encoder_attention_mask
+        outputs: list[torch.Tensor] = []
+        for layer in self.model.decoder.layers:
+            k, v = layer.encoder_attn.project_kv(encoder_hidden_states)
+            outputs.extend((k, v))
+        return tuple(outputs)
+
+    def cactus_decoder_step(
+        self,
+        decoder_input_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
+        *cross_kv: torch.Tensor,
+    ) -> torch.Tensor:
+        encoder_attention_mask = encoder_attention_mask != 0
+        hidden_states = self.model.embed_tokens(decoder_input_ids) * self.model.embed_scale
+        inv_freq = _build_inv_freq(
+            self.model.decoder.head_dim,
+            self.model.decoder.rope_theta,
+        ).to(device=hidden_states.device)
+        rope = _rotary_tables_for_position_ids(
+            inv_freq,
+            position_ids.to(dtype=torch.long),
+            hidden_states.dtype,
+        )
+        for layer_index, layer in enumerate(self.model.decoder.layers):
+            normed = layer.input_layernorm(hidden_states)
+            attn = layer.self_attn(normed, normed, None, rope)
+            hidden_states = _add_clipped(hidden_states, torch.sigmoid(layer.self_attn_gate).to(dtype=attn.dtype) * attn)
+
+            cross_attn = layer.encoder_attn.forward_with_kv(
+                layer.encoder_attn_layer_norm(hidden_states),
+                cross_kv[layer_index * 2],
+                cross_kv[layer_index * 2 + 1],
+                encoder_attention_mask,
+                None,
+            )
+            hidden_states = _add_clipped(hidden_states, torch.sigmoid(layer.cross_attn_gate).to(dtype=cross_attn.dtype) * cross_attn)
+        hidden_states = self.model.decoder.norm(hidden_states)
+        return self.lm_head(hidden_states)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, NeedleRMSNorm):
+            nn.init.zeros_(module.weight)

special_tokens_map.json

@@ -0,0 +1,11 @@
+{
+  "additional_special_tokens": [
+    "<tool_call>",
+    "<tools>"
+  ],
+  "bos_token": "<s>",
+  "eos_token": "</s>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}
+

tokenization_needle.py

@@ -0,0 +1,121 @@
+"""Slow SentencePiece tokenizer for Needle."""
+
+from __future__ import annotations
+
+import os
+import shutil
+from typing import Any
+
+import sentencepiece as spm
+from transformers import PreTrainedTokenizer
+
+
+VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"}
+
+
+class NeedleTokenizer(PreTrainedTokenizer):
+    vocab_files_names = VOCAB_FILES_NAMES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file: str,
+        unk_token: str = "<unk>",
+        bos_token: str = "<s>",
+        eos_token: str = "</s>",
+        pad_token: str = "<pad>",
+        tool_call_token: str = "<tool_call>",
+        tools_token: str = "<tools>",
+        **kwargs: Any,
+    ) -> None:
+        self.vocab_file = vocab_file
+        self.sp_model = spm.SentencePieceProcessor()
+        self.sp_model.Load(vocab_file)
+        self.sp = self.sp_model
+        self.tool_call_token = tool_call_token
+        self.tools_token = tools_token
+        additional = list(kwargs.pop("additional_special_tokens", []) or [])
+        for token in (tool_call_token, tools_token):
+            if token not in additional:
+                additional.append(token)
+        super().__init__(
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            additional_special_tokens=additional,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self) -> int:
+        return int(self.sp_model.GetPieceSize())
+
+    @property
+    def tool_call_token_id(self) -> int:
+        return int(self.sp_model.PieceToId(self.tool_call_token))
+
+    @property
+    def tools_token_id(self) -> int:
+        return int(self.sp_model.PieceToId(self.tools_token))
+
+    def get_vocab(self) -> dict[str, int]:
+        vocab = {self.sp_model.IdToPiece(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def _tokenize(self, text: str) -> list[str]:
+        return list(self.sp_model.EncodeAsPieces(text))
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return int(self.sp_model.PieceToId(token))
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return str(self.sp_model.IdToPiece(int(index)))
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        return self.sp_model.DecodePieces(tokens)
+
+    def build_inputs_with_special_tokens(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+    ) -> list[int]:
+        if token_ids_1 is None:
+            return list(token_ids_0)
+        return list(token_ids_0) + list(token_ids_1)
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+        already_has_special_tokens: bool = False,
+    ) -> list[int]:
+        if already_has_special_tokens:
+            all_ids = list(token_ids_0)
+        else:
+            all_ids = self.build_inputs_with_special_tokens(token_ids_0, token_ids_1)
+        special = {
+            self.pad_token_id,
+            self.eos_token_id,
+            self.bos_token_id,
+            self.unk_token_id,
+            self.tool_call_token_id,
+            self.tools_token_id,
+        }
+        return [1 if token_id in special else 0 for token_id in all_ids]
+
+    def create_token_type_ids_from_sequences(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+    ) -> list[int]:
+        return [0] * len(self.build_inputs_with_special_tokens(token_ids_0, token_ids_1))
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
+        os.makedirs(save_directory, exist_ok=True)
+        out_name = "tokenizer.model" if filename_prefix is None else f"{filename_prefix}-tokenizer.model"
+        out_path = os.path.join(save_directory, out_name)
+        if os.path.abspath(self.vocab_file) != os.path.abspath(out_path):
+            shutil.copyfile(self.vocab_file, out_path)
+        return (out_path,)

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0823f5b9133c68a8140addc5d7a425fa9119c4c8cb4a550363b4bffa4ba1c8c7
+size 124960

tokenizer_config.json

@@ -0,0 +1,27 @@
+{
+  "additional_special_tokens": [
+    "<tool_call>",
+    "<tools>"
+  ],
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_needle.NeedleTokenizer",
+      null
+    ]
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 1024,
+  "pad_token": "<pad>",
+  "padding_side": "right",
+  "tokenizer_class": "NeedleTokenizer",
+  "tool_call_token": "<tool_call>",
+  "tools_token": "<tools>",
+  "unk_token": "<unk>"
+}
+