Spaces:
Running on Zero
Running on Zero
| """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"] | |
| 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]) | |