From 85b3881d7e5c0eff95b2ffbf2307b33034222c1c Mon Sep 17 00:00:00 2001 From: jakmro Date: Mon, 13 Jul 2026 17:11:04 +0200 Subject: [PATCH] add needle support --- cactus-engine/tests/test_needle.cpp | 23 +- python/cactus/cli/utils.py | 8 +- .../cactus_adapters/weight_patterns.py | 8 +- python/cactus/convert/cli.py | 9 +- .../cactus/convert/model_adapters/naming.py | 4 + .../cactus/convert/tests/test_naming_qdq.py | 15 + python/cactus/models/__init__.py | 0 python/cactus/models/needle/__init__.py | 25 + .../models/needle/configuration_needle.py | 63 +++ .../cactus/models/needle/modeling_needle.py | 429 ++++++++++++++++++ .../models/needle/tokenization_needle.py | 121 +++++ python/cactus/transpile/hf_model.py | 3 + python/requirements.txt | 1 + python/tests/test_component_plan.py | 16 + python/tests/test_download.py | 9 + 15 files changed, 711 insertions(+), 23 deletions(-) create mode 100644 python/cactus/models/__init__.py create mode 100644 python/cactus/models/needle/__init__.py create mode 100644 python/cactus/models/needle/configuration_needle.py create mode 100644 python/cactus/models/needle/modeling_needle.py create mode 100644 python/cactus/models/needle/tokenization_needle.py diff --git a/cactus-engine/tests/test_needle.cpp b/cactus-engine/tests/test_needle.cpp index 4f114e8b3..54fb1728b 100644 --- a/cactus-engine/tests/test_needle.cpp +++ b/cactus-engine/tests/test_needle.cpp @@ -4,26 +4,21 @@ #include #include #include -#include #include static const char* g_model_path = std::getenv("CACTUS_TEST_MODEL"); static bool model_is_needle(const std::string& bundle_path) { - std::ifstream f(bundle_path + "/config.json"); + std::ifstream f(bundle_path + "/config.txt"); if (!f) return false; - std::stringstream ss; - ss << f.rdbuf(); - const std::string content = ss.str(); - const size_t key = content.find("\"model_type\""); - if (key == std::string::npos) return false; - const size_t colon = content.find(':', key); - if (colon == std::string::npos) return false; - const size_t q1 = content.find('"', colon); - if (q1 == std::string::npos) return false; - const size_t q2 = content.find('"', q1 + 1); - if (q2 == std::string::npos) return false; - return content.substr(q1 + 1, q2 - q1 - 1) == "needle"; + std::string line; + while (std::getline(f, line)) { + while (!line.empty() && (line.back() == '\r' || line.back() == ' ')) line.pop_back(); + if (line.rfind("model_type=", 0) == 0) { + return line.substr(11) == "needle"; + } + } + return false; } bool test_needle_tool_call() { diff --git a/python/cactus/cli/utils.py b/python/cactus/cli/utils.py index 418085700..82d374a92 100644 --- a/python/cactus/cli/utils.py +++ b/python/cactus/cli/utils.py @@ -259,9 +259,13 @@ def validate_extracted_bundle(output_dir: Path) -> None: tokenizer_config[key.strip()] = value.strip() tokenizer_type = tokenizer_config.get("tokenizer_type", "").lower() - optional_required = ["special_tokens.json", "tokenizer.json"] - if tokenizer_type == "bpe": + optional_required = ["special_tokens.json"] + if tokenizer_type == "sentencepiece": optional_required.append("merges.txt") + else: + optional_required.append("tokenizer.json") + if tokenizer_type == "bpe": + optional_required.append("merges.txt") missing_sidecars = [name for name in optional_required if not (output_dir / name).exists()] if missing_sidecars: diff --git a/python/cactus/convert/cactus_adapters/weight_patterns.py b/python/cactus/convert/cactus_adapters/weight_patterns.py index b71869fa3..67add2b28 100644 --- a/python/cactus/convert/cactus_adapters/weight_patterns.py +++ b/python/cactus/convert/cactus_adapters/weight_patterns.py @@ -223,12 +223,8 @@ NEEDLE_GLOBAL_WEIGHTS = [ - ('embedding.embedding', 'token_embeddings.weights', False), - ('embedding.embedding', 'output_weight.weights', False), - ('encoder.final_norm.scale', 'encoder_layer_norm_weight.weights', False), - ('decoder.ZCRMSNorm_0.scale', 'output_norm.weights', False), - ('contrastive_proj.kernel', 'contrastive_proj.weights', True), - ('log_temp', 'log_temp.weights', False), + ('model.encoder.final_norm.weight', 'encoder_layer_norm_weight.weights'), + ('model.decoder.norm.weight', 'output_norm.weights'), ] NEEDLE_ENCODER_LAYER_WEIGHTS = [ diff --git a/python/cactus/convert/cli.py b/python/cactus/convert/cli.py index 128901d09..3a494f961 100644 --- a/python/cactus/convert/cli.py +++ b/python/cactus/convert/cli.py @@ -50,6 +50,9 @@ def _load_hf(model_id_or_path: str, device: str, *, skip_model_load: bool = Fals warnings.filterwarnings("ignore", message=".*You are using a model of type.*") for note in patch_transformers_import_compat(): print(f"note={note}") + from ..models.needle import register_with_transformers + + register_with_transformers() nemo_export = ensure_parakeet_tdt_nemo_source(model_id_or_path, cache_dir=_hf_cache_dir()) if nemo_export is not None: model_id_or_path = nemo_export @@ -434,7 +437,11 @@ def convert(args: argparse.Namespace) -> None: } rows: list[dict[str, Any]] = [] pending: list[tuple[str, Any, Any, Any]] = [] - num_layers = int(model_config.get("num_layers", 0) or 0) or None + num_layers = max( + int(model_config.get("num_layers", 0) or 0), + int(cfg_get(cfg, "num_encoder_layers", 0) or 0), + int(cfg_get(cfg, "num_decoder_layers", 0) or 0), + ) or None target_modules: set[str] = set() for name, tensor in state_dict.items(): match = adapter.name_tensor(name, tensor, num_layers) diff --git a/python/cactus/convert/model_adapters/naming.py b/python/cactus/convert/model_adapters/naming.py index b44a4a072..3b66871af 100644 --- a/python/cactus/convert/model_adapters/naming.py +++ b/python/cactus/convert/model_adapters/naming.py @@ -328,6 +328,10 @@ def _global_match(name: str, family: str) -> str | None: found = _candidate_table_match(name, table) if found: return found + elif family == "needle": + found = _candidate_table_match(name, wp.NEEDLE_GLOBAL_WEIGHTS) + if found: + return found if name in {"wte.weight", "word_embeddings.weight"}: return "token_embeddings.weights" if name in {"wpe.weight", "position_embeddings.weight"}: diff --git a/python/cactus/convert/tests/test_naming_qdq.py b/python/cactus/convert/tests/test_naming_qdq.py index 860c8ffd4..f5d2976e8 100644 --- a/python/cactus/convert/tests/test_naming_qdq.py +++ b/python/cactus/convert/tests/test_naming_qdq.py @@ -21,6 +21,21 @@ def test_gemma4_audio_aux_roundtrip_name(): assert "q_proj_proj" not in restore_hf_key_for_family(match.adapter_name, "gemma4") +def test_needle_encoder_layers_beyond_decoder_depth_are_recognized(): + match = cactus_name_for_tensor("model.encoder.layers.11.self_attn.k_proj.weight", "needle", 12) + assert match.recognized + assert match.output_name == "encoder_layer_11_attn_k.weights" + + +def test_needle_global_norms_map_to_bundle_names(): + encoder_norm = cactus_name_for_tensor("model.encoder.final_norm.weight", "needle") + assert encoder_norm.recognized + assert encoder_norm.output_name == "encoder_layer_norm_weight.weights" + decoder_norm = cactus_name_for_tensor("model.decoder.norm.weight", "needle") + assert decoder_norm.recognized + assert decoder_norm.output_name == "output_norm.weights" + + def test_manifest_qdq_restores_aux_stats_without_suffix(tmp_path): cactus = tmp_path / "cactus" out = tmp_path / "qdq" diff --git a/python/cactus/models/__init__.py b/python/cactus/models/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/cactus/models/needle/__init__.py b/python/cactus/models/needle/__init__.py new file mode 100644 index 000000000..8436b054a --- /dev/null +++ b/python/cactus/models/needle/__init__.py @@ -0,0 +1,25 @@ +"""Needle encoder-decoder model (mirrors the files published on Cactus-Compute/needle).""" + + +def register_with_transformers() -> None: + """Register the needle architecture with transformers' Auto* factories. + + Idempotent; imports torch/sentencepiece lazily so importing cactus stays light. + """ + from transformers import ( + AutoConfig, + AutoModel, + AutoModelForCausalLM, + AutoModelForSeq2SeqLM, + AutoTokenizer, + ) + + from .configuration_needle import NeedleConfig + from .modeling_needle import NeedleForCausalLM, NeedleModel + from .tokenization_needle import NeedleTokenizer + + AutoConfig.register("needle", NeedleConfig, exist_ok=True) + AutoModel.register(NeedleConfig, NeedleModel, exist_ok=True) + AutoModelForCausalLM.register(NeedleConfig, NeedleForCausalLM, exist_ok=True) + AutoModelForSeq2SeqLM.register(NeedleConfig, NeedleForCausalLM, exist_ok=True) + AutoTokenizer.register(NeedleConfig, slow_tokenizer_class=NeedleTokenizer, exist_ok=True) diff --git a/python/cactus/models/needle/configuration_needle.py b/python/cactus/models/needle/configuration_needle.py new file mode 100644 index 000000000..4fd2d3c98 --- /dev/null +++ b/python/cactus/models/needle/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) diff --git a/python/cactus/models/needle/modeling_needle.py b/python/cactus/models/needle/modeling_needle.py new file mode 100644 index 000000000..b8e3cbb0c --- /dev/null +++ b/python/cactus/models/needle/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) diff --git a/python/cactus/models/needle/tokenization_needle.py b/python/cactus/models/needle/tokenization_needle.py new file mode 100644 index 000000000..8fe3450d2 --- /dev/null +++ b/python/cactus/models/needle/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 = "", + bos_token: str = "", + eos_token: str = "", + pad_token: str = "", + tool_call_token: str = "", + tools_token: str = "", + **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,) diff --git a/python/cactus/transpile/hf_model.py b/python/cactus/transpile/hf_model.py index 5a8a2d282..7bf5d2ca3 100644 --- a/python/cactus/transpile/hf_model.py +++ b/python/cactus/transpile/hf_model.py @@ -2791,6 +2791,9 @@ def _lower_preoptimized_ir(ir: IRGraph) -> TranspiledGraph: def main() -> int: + from cactus.models.needle import register_with_transformers + + register_with_transformers() parser = argparse.ArgumentParser( description=( "Load a Hugging Face model, canonicalize it into a generic transpile task, " diff --git a/python/requirements.txt b/python/requirements.txt index 4e509d70b..de9259adb 100644 --- a/python/requirements.txt +++ b/python/requirements.txt @@ -5,6 +5,7 @@ numpy==2.4.6 pillow==11.3.0 huggingface_hub==1.15.0 scipy==1.17.1 +sentencepiece==0.2.2 fastapi>=0.100.0 uvicorn>=0.20.0 python-multipart>=0.0.5 diff --git a/python/tests/test_component_plan.py b/python/tests/test_component_plan.py index 4bd7bbb73..d09cd0e10 100644 --- a/python/tests/test_component_plan.py +++ b/python/tests/test_component_plan.py @@ -34,3 +34,19 @@ def test_gemma4_image_audio_config_keeps_audio_component() -> None: assert plan is not None assert plan.components == ("vision_encoder", "audio_encoder", "lm_encoder", "decoder") assert plan.needs_audio + + +def test_needle_config_uses_component_pipeline_plan() -> None: + config = { + "model_type": "needle", + "architectures": ["NeedleForCausalLM"], + "num_encoder_layers": 12, + "num_decoder_layers": 8, + } + + plan = infer_component_plan_from_config(config, model_id="Cactus-Compute/needle") + + assert plan is not None + assert plan.task == "causal_lm_logits" + assert plan.components == ("source_encoder", "decoder_cross_kv", "decoder_step") + assert plan.force_component_pipeline diff --git a/python/tests/test_download.py b/python/tests/test_download.py index 4591d94a8..23c90df64 100644 --- a/python/tests/test_download.py +++ b/python/tests/test_download.py @@ -178,6 +178,15 @@ def test_validate_extracted_bundle_requires_tokenizer_sidecars(self): with self.assertRaisesRegex(RuntimeError, "missing tokenizer sidecar"): validate_extracted_bundle(package) + def test_validate_extracted_bundle_accepts_sentencepiece_without_tokenizer_json(self): + package = self.root / "pkg" + self._write_minimal_package(package) + (package / "tokenizer.json").unlink() + (package / "tokenizer_config.txt").write_text( + "tokenizer_type=sentencepiece\nsp_model_type=bpe\n", encoding="utf-8" + ) + validate_extracted_bundle(package) + def test_validate_extracted_bundle_accepts_seq2seq_embeddings(self): package = self.root / "pkg" self._write_minimal_package(package)