add needle support

cactus-engine/tests/test_needle.cpp

@@ -4,26 +4,21 @@
 #include <cstring>
 #include <fstream>
 #include <iostream>
-#include <sstream>
 #include <string>
 
 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() {

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:

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 = [

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)

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"}:

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"

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)

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)