tests/test_decode.py

"""Decoding a direction back to words, against a synthetic vocabulary."""
import numpy as np
import pytest

import ulens
from ulens import decode
from ulens._math import unit

HIDDEN = 24
WORDS = ["clear", "cloudy", "maybe", "certain", "doubtful", "obvious", "ambiguous",
         "plain", "murky", "sure", "unsure", "vague", "a", "xy", "123", "!!", "Clear"]


@pytest.fixture
def voc(monkeypatch, tmp_path):
    """A vocabulary of len(WORDS) tokens whose transported rows we control."""
    rng = np.random.default_rng(0)
    n = len(WORDS)
    W = rng.normal(size=(n, HIDDEN)).astype(np.float32)
    gamma = np.ones(HIDDEN, np.float32)
    J = np.eye(HIDDEN, dtype=np.float32)

    from ulens import lensfile
    monkeypatch.setattr(lensfile, "matrix", lambda lens, block: J)
    spec = ulens.ModelSpec(key="fakevoc", hf="fake/voc", revision="0" * 40, n_blocks=6,
                           hidden_size=HIDDEN, readout_block=3,
                           lens=ulens.models.LensFile(path=str(tmp_path / "n.npy")),
                           think_end="</think>")
    ulens.MODELS["fakevoc"] = spec
    v = decode.Vocabulary(spec, 3, cache_dir=tmp_path)
    monkeypatch.setattr(v, "_open", lambda: (None, "w", gamma))
    monkeypatch.setattr(v, "_rows", lambda chunk: iter([(0, W * gamma[None, :])]))
    monkeypatch.setattr(type(v), "n_tokens", property(lambda self: n))
    monkeypatch.setattr(v, "word", lambda t: WORDS[int(t)])
    v.W = W
    yield v
    del ulens.MODELS["fakevoc"]
    ulens.models._ALIASES.pop("fakevoc", None)


def test_norms_match_the_transported_row_norms(voc):
    np.testing.assert_allclose(voc.norms, np.linalg.norm(voc.W @ voc.J, axis=1), rtol=1e-5)


def test_norms_are_cached_to_disk(voc):
    p = voc.cache_path
    assert not p.exists()
    first = voc.norms
    assert p.exists()
    voc._norms = None
    np.testing.assert_array_equal(voc.norms, first)


def test_scores_are_the_cosine_with_every_transported_row(voc):
    u = unit(np.random.default_rng(1).normal(size=HIDDEN).astype(np.float32))
    D = unit(voc.W @ voc.J)
    np.testing.assert_allclose(voc.scores(u), D @ u, rtol=1e-4, atol=1e-5)


def test_scores_without_cosine_are_the_raw_dot_product(voc):
    u = np.random.default_rng(2).normal(size=HIDDEN).astype(np.float32)
    np.testing.assert_allclose(voc.scores(u, cosine=False), (voc.W @ voc.J) @ u, rtol=1e-4)


def test_a_direction_decodes_to_its_own_word_first(voc):
    """The direction of one token's transported row must rank that token top."""
    for t in (0, 3, 6):
        d = unit(voc.W[t] @ voc.J)
        assert voc.top(d, k=1, min_len=1, alpha_only=False)[0].word == WORDS[t]


def test_the_negative_end_is_the_opposite_ranking(voc):
    d = unit(voc.W[0] @ voc.J)
    pos = voc.top(d, k=3, min_len=1, alpha_only=False)
    neg = voc.top(d, k=3, sign="-", min_len=1, alpha_only=False)
    assert pos[0].score > 0 > neg[0].score
    assert pos[0].word != neg[0].word


def test_filters_drop_fragments_and_dedupe_case(voc):
    d = unit(np.ones(HIDDEN, np.float32))
    hits = voc.top(d, k=20, scan=100)
    got = [h.word for h in hits]
    assert "a" not in got and "xy" not in got and "123" not in got and "!!" not in got
    assert len({w.lower() for w in got}) == len(got)
    assert all(len(w) >= decode.MIN_LEN for w in got)


def test_rejects_a_direction_of_the_wrong_width(voc):
    with pytest.raises(ValueError, match="wide"):
        voc.scores(np.ones(HIDDEN + 1, np.float32))


def test_decode_handles_an_axis_a_cone_and_a_bare_direction(voc):
    ax = ulens.Axis("one", unit(voc.W[0] @ voc.J))
    cone = ulens.axes.from_axes([ax, ulens.Axis("two", unit(voc.W[3] @ voc.J))])
    assert [d.name for d in voc.decode(ax, k=2, min_len=1)] == ["one"]
    assert [d.name for d in voc.decode(cone, k=2, min_len=1)] == ["one", "two"]
    assert [d.name for d in voc.decode(np.ones(HIDDEN), k=2, min_len=1)] == ["direction"]


def test_report_shows_both_ends_of_every_axis(voc):
    cone = ulens.axes.from_axes([ulens.Axis(n, unit(voc.W[i] @ voc.J))
                                 for i, n in ((0, "alpha"), (3, "beta"))])
    txt = voc.report(cone, k=2, min_len=1)
    assert "alpha" in txt and "beta" in txt
    assert txt.count("+") >= 2 and txt.count("-") >= 2


def test_purify_removes_the_shared_direction():
    """The two removals are sequential, not jointly orthogonalised — the same order the
    reference implementation uses — so a trace of `mu` survives removing the components
    after it. The principal components themselves come out exactly."""
    rng = np.random.default_rng(3)
    shared = unit(rng.normal(size=HIDDEN).astype(np.float32))
    D = unit(shared[None, :] * 3.0 + rng.normal(size=(500, HIDDEN)).astype(np.float32) * 0.2)
    mu, V = decode.common_component(D, k=2)
    assert abs(float(unit(mu) @ shared)) > 0.9
    v = unit(rng.normal(size=HIDDEN).astype(np.float32))
    p = decode.purify(v, mu, V)
    assert abs(float(unit(p) @ mu)) < 1e-3
    for row in V:
        assert abs(float(p @ row)) < 1e-5


def test_purify_with_no_components_removes_the_mean_exactly():
    rng = np.random.default_rng(31)
    D = unit(rng.normal(size=(200, HIDDEN)).astype(np.float32))
    mu, V = decode.common_component(D, k=0)
    v = unit(rng.normal(size=HIDDEN).astype(np.float32))
    assert abs(float(decode.purify(v, mu, V) @ mu)) < 1e-6


def test_purify_with_k_zero_only_removes_the_mean():
    rng = np.random.default_rng(4)
    D = unit(rng.normal(size=(100, HIDDEN)).astype(np.float32))
    mu, V = decode.common_component(D, k=0)
    assert V.shape == (0, HIDDEN)
    v = rng.normal(size=HIDDEN).astype(np.float32)
    assert abs(float(unit(decode.purify(v, mu, V)) @ mu)) < 1e-4


def test_purifying_a_direction_changes_what_it_decodes_to(voc):
    d = unit(np.ones(HIDDEN, np.float32))
    pur = decode.common_component(unit(voc.W @ voc.J), k=1)
    plain = [h.word for h in voc.top(d, k=4, min_len=1, alpha_only=False)]
    cleaned = [h.word for h in voc.top(d, k=4, min_len=1, alpha_only=False,
                                       purify_with=pur)]
    assert plain != cleaned


def test_sample_directions_are_unit_and_the_right_shape(voc):
    S = voc.sample_directions(n=8, seed=0)
    assert S.shape == (8, HIDDEN)
    np.testing.assert_allclose(np.linalg.norm(S, axis=1), 1.0, atol=1e-5)


def test_save_decode_writes_both_ends(voc, tmp_path):
    ds = voc.decode(unit(voc.W[0] @ voc.J), k=2, min_len=1)
    import json
    blob = json.loads(decode.save_decode(ds, tmp_path / "d.json").read_text())
    assert blob[0]["name"] == "direction"
    assert len(blob[0]["positive"]) == 2 and len(blob[0]["negative"]) == 2
    assert {"word", "score", "id"} == set(blob[0]["positive"][0])