mirror of
https://github.com/cactus-compute/needle.git
synced 2026-10-23 06:30:56 +08:00
127 lines
4.1 KiB
Python
127 lines
4.1 KiB
Python
import os
|
|
|
|
import pytest
|
|
|
|
|
|
def _engine_available(generation=3):
|
|
"""Whether ``needle._library_path`` would find an engine without downloading.
|
|
|
|
Every location this looks in has to be one that function looks in, or the
|
|
gate skips tests an installed engine could have run.
|
|
"""
|
|
try:
|
|
import needle
|
|
from needle.agent import fetch
|
|
|
|
override = os.environ.get(fetch.lib_path_env(generation))
|
|
if generation == 2 and not override:
|
|
override = os.environ.get("NEEDLE_LIB_PATH")
|
|
if override:
|
|
return os.path.exists(override)
|
|
|
|
here = os.path.dirname(needle.__file__)
|
|
name = fetch.lib_name(generation)
|
|
stem, suffix = os.path.splitext(name)
|
|
local_names = [name if generation == fetch.WHISTLE else f"{stem}{generation}{suffix}"]
|
|
if generation == 2:
|
|
local_names.append(name)
|
|
if any(os.path.exists(os.path.join(here, local)) for local in local_names):
|
|
return True
|
|
|
|
return os.path.exists(os.path.join(fetch.cache_dir(generation), name))
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
requires_engine = pytest.mark.skipif(
|
|
not _engine_available(),
|
|
reason="needle C++ engine not installed (auto-fetched from HF on first real use)")
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def tiny_checkpoint(tmp_path_factory):
|
|
import numpy as np
|
|
import jax
|
|
import jax.numpy as jnp
|
|
from needle.model.architecture import SimpleAttentionNetwork, TransformerConfig
|
|
from needle.model.checkpoints import write_checkpoint
|
|
|
|
config = TransformerConfig(
|
|
vocab_size=8192, out_vocab=8192, d_model=64, num_heads=4, num_kv_heads=2,
|
|
num_layers=4, qk_head_dim=16, v_head_dim=16, max_seq_len=128,
|
|
engram_layers=(1, 3), engram_slots=64, global_layers=(3,), sliding_window=32,
|
|
mhc_lanes=2, qkv_conv_taps=3, flash=False,
|
|
)
|
|
model = SimpleAttentionNetwork(config)
|
|
params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 8), jnp.int32))["params"]
|
|
params = jax.tree_util.tree_map(lambda x: np.asarray(x), params)
|
|
|
|
path = tmp_path_factory.mktemp("ckpt") / "tiny.safetensors"
|
|
write_checkpoint(path, {"format_version": 2, "params": params, "config": dict(vars(config))})
|
|
return str(path)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def tiny_checkpoint_safetensors(tiny_checkpoint):
|
|
return tiny_checkpoint
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def tiny_base_archive(tiny_checkpoint, tmp_path_factory):
|
|
from needle.model.export import write_export
|
|
from needle.model.run import load_checkpoint
|
|
from needle.model.tokenizer import get_tokenizer
|
|
|
|
params, config = load_checkpoint(tiny_checkpoint)
|
|
path = tmp_path_factory.mktemp("base") / "needle3.cact"
|
|
write_export(params, config, str(path), bits=4, tokenizer=get_tokenizer(config.vocab_size))
|
|
return str(path)
|
|
|
|
|
|
@pytest.fixture
|
|
def published_base(tiny_base_archive, monkeypatch):
|
|
from needle.agent import fetch
|
|
|
|
calls = []
|
|
|
|
def fake_fetch(generation=2, dest_dir=None, force=False):
|
|
calls.append((generation, force))
|
|
return tiny_base_archive
|
|
|
|
monkeypatch.setattr(fetch, "fetch_weights", fake_fetch)
|
|
return calls
|
|
|
|
|
|
@pytest.fixture
|
|
def fake_cact():
|
|
return _fake_cact
|
|
|
|
|
|
def _fake_cact(path, heads):
|
|
import struct
|
|
|
|
header_fmt, record_fmt = "<48If", "<BBHIIIIQQII"
|
|
body = bytearray()
|
|
records = []
|
|
|
|
def add(dtype, shape, blob):
|
|
records.append((dtype, len(shape), 0, *(list(shape) + [0] * (4 - len(shape))), len(body), len(blob), 0, 0))
|
|
body.extend(blob)
|
|
|
|
add(1, (4,), struct.pack("<4e", 0, 0, 0, 0))
|
|
if heads:
|
|
add(1, (len(heads),), struct.pack(f"<{len(heads)}e", *heads))
|
|
for _ in heads:
|
|
for _ in range(6):
|
|
add(1, (2,), struct.pack("<2e", 0, 0))
|
|
add(4, (3,), b"tok")
|
|
head = struct.calcsize(header_fmt) + struct.calcsize(record_fmt) * len(records)
|
|
fields = [0x05E12A84, len(records), 0] + [0] * 45
|
|
out = bytearray(struct.pack(header_fmt, *fields, 0.0))
|
|
for rec in records:
|
|
rec = list(rec)
|
|
rec[7] += head
|
|
out += struct.pack(record_fmt, *rec)
|
|
out += body
|
|
path.write_bytes(bytes(out))
|