Spaces:
Running on Zero
Running on Zero
| """Trace geometry, segmentation and persistence, without a model.""" | |
| import numpy as np | |
| import pytest | |
| import ulens | |
| from ulens import runtime | |
| from ulens.runtime import Trace, _segment | |
| OPEN, CLOSE = 900, 901 | |
| def make_trace(n_blocks=5, T=20, k=4, prompt_len=6, eot=14, seed=0, **kw): | |
| rng = np.random.default_rng(seed) | |
| opts: dict = { | |
| "model": "qwen3.5-4b", "prompt": "p", | |
| "tokens": list(range(T)), "token_strs": [f"t{i}" for i in range(T)], | |
| "prompt_len": prompt_len, "eot": eot, | |
| "entropy": np.concatenate( | |
| [[np.nan], rng.uniform(0.1, 2.0, size=T - 1)]).astype(np.float32), | |
| "field": rng.normal(size=(n_blocks, T, k)).astype(np.float32), | |
| "blocks": list(range(10, 10 + n_blocks)), | |
| "names": ulens.terms.categories()[:k], | |
| "readout_block": 12, "generated": "g", "truncated": False, | |
| "config": {"seed": 1}} | |
| opts.update(kw) | |
| return Trace(**opts) | |
| # ------------------------------------------------------------------ segmentation | |
| def test_block_already_open_at_the_prompt_end(): | |
| """The usual case: the chat template opens the channel, the generation closes it.""" | |
| ids = [1, 2, OPEN] + [5, 6, 7, CLOSE, 8, 9] | |
| eot, trunc = _segment(ids, 3, OPEN, CLOSE) | |
| assert eot == 6 and ids[eot] == CLOSE and not trunc | |
| def test_block_opened_inside_the_generation(): | |
| ids = [1, 2, 3] + [OPEN, 5, 6, CLOSE, 8] | |
| eot, trunc = _segment(ids, 3, OPEN, CLOSE) | |
| assert eot == 6 and ids[eot] == CLOSE and not trunc | |
| def test_opened_and_never_closed_is_all_thinking(): | |
| """What a token budget produces. It must not read as a zero-token answer.""" | |
| ids = [1, 2, OPEN] + [5, 6, 7, 8] | |
| eot, trunc = _segment(ids, 3, OPEN, CLOSE) | |
| assert eot == len(ids) - 1 and trunc | |
| tr = make_trace(T=len(ids), prompt_len=3, eot=eot, truncated=trunc) | |
| assert tr.think_len == len(ids) - 3 | |
| assert (tr.segments == "answer").sum() == 0 | |
| def test_no_thinking_block_at_all_is_all_answer(): | |
| ids = [1, 2, 3] + [5, 6, 7] | |
| eot, trunc = _segment(ids, 3, OPEN, CLOSE) | |
| assert eot == 2 and not trunc | |
| tr = make_trace(T=6, prompt_len=3, eot=eot) | |
| assert tr.think_len == 0 | |
| assert list(tr.segments) == ["prompt"] * 3 + ["answer"] * 3 | |
| def test_a_closer_in_the_prompt_does_not_count(): | |
| """A prompt that mentions the closer must not end the block before it starts.""" | |
| ids = [1, CLOSE, 2, OPEN] + [5, 6, CLOSE, 9] | |
| eot, _ = _segment(ids, 4, OPEN, CLOSE) | |
| assert eot == 6 | |
| def test_a_reopened_block_uses_the_last_opener(): | |
| ids = [OPEN, 1, CLOSE, 2, OPEN] + [5, 6, CLOSE] | |
| eot, trunc = _segment(ids, 5, OPEN, CLOSE) | |
| assert eot == 7 and not trunc | |
| def test_without_a_closer_token_everything_is_thinking(): | |
| ids = [1, 2, 3, 4, 5] | |
| eot, trunc = _segment(ids, 2, OPEN, None) | |
| assert eot == len(ids) - 1 and trunc | |
| # ------------------------------------------------------------------ geometry | |
| def test_segments_partition_the_sequence(): | |
| tr = make_trace() | |
| s = tr.segments | |
| assert len(s) == tr.n_tokens | |
| assert (s == "prompt").sum() == tr.prompt_len | |
| assert (s == "think").sum() == tr.think_len | |
| assert (s == "answer").sum() == tr.n_tokens - tr.eot - 1 | |
| def test_a_cone_series_matches_the_readout_at_the_eot(): | |
| tr = make_trace() | |
| assert tr.a_cone == pytest.approx(tr.a_cone_series()[tr.eot]) | |
| z = tr.projections()[tr.eot] | |
| assert tr.a_cone == pytest.approx(float(ulens.a_cone_of(z))) | |
| def test_the_field_agrees_with_a_cone_computed_from_states(): | |
| """a_cone_of on stored projections equals a_cone on the states they came from.""" | |
| rng = np.random.default_rng(1) | |
| U = ulens._math.lowdin(rng.normal(size=(32, 4)).astype(np.float32)) | |
| H = rng.normal(size=(7, 32)).astype(np.float32) | |
| z = ulens.projections(H, U) | |
| np.testing.assert_allclose(ulens.a_cone_of(z), ulens.a_cone(H, U), rtol=1e-6) | |
| def test_surface_is_a_cone_over_depth_and_token(): | |
| tr = make_trace() | |
| S = tr.surface() | |
| assert S.shape == (len(tr.blocks), tr.n_tokens) | |
| np.testing.assert_allclose(S[tr.block_index()], tr.a_cone_series(), rtol=1e-6) | |
| def test_components_are_the_rectified_projections(): | |
| tr = make_trace() | |
| np.testing.assert_allclose(tr.components(), ulens.rectify(tr.projections()), rtol=1e-6) | |
| assert np.all(tr.components() > 0) | |
| def test_think_entropy_spans_the_thinking_block_inclusive_of_the_eot(): | |
| tr = make_trace() | |
| want = tr.entropy[tr.prompt_len: tr.eot + 1].mean() | |
| assert tr.think_entropy == pytest.approx(want) | |
| assert tr.u_lens == pytest.approx(tr.a_cone * tr.think_entropy) | |
| def test_the_first_token_has_no_entropy_and_it_is_excluded(): | |
| """Nothing predicted token 0, so its entropy is NaN — and must not reach the mean.""" | |
| tr = make_trace() | |
| assert np.isnan(tr.entropy[0]) | |
| assert not np.isnan(tr.think_entropy) | |
| def test_an_uncaptured_block_is_refused(): | |
| tr = make_trace() | |
| with pytest.raises(KeyError, match="was not captured"): | |
| tr.projections(999) | |
| def test_summary_and_axis_table_render(): | |
| tr = make_trace() | |
| assert "U-Lens" in tr.summary() and "readout block 12" in tr.summary() | |
| assert "closed" in tr.summary() and tr.readout_valid | |
| assert "argmax" in tr.axis_table() | |
| def test_a_truncated_trace_says_the_readout_is_not_the_published_one(): | |
| tr = make_trace(truncated=True) | |
| assert not tr.readout_valid | |
| assert "TRUNCATED" in tr.summary() and "not the published one" in tr.summary() | |
| assert np.isfinite(tr.u_lens) # still computed; the field is real | |
| def test_a_trace_with_no_thinking_block_is_not_a_valid_readout(): | |
| tr = make_trace(prompt_len=6, eot=5) | |
| assert tr.think_len == 0 and not tr.readout_valid | |
| # ------------------------------------------------------------------ persistence | |
| def test_trace_round_trips_and_needs_no_model(tmp_path): | |
| tr = make_trace() | |
| back = runtime.load(tr.save(tmp_path / "t.npz")) | |
| np.testing.assert_allclose(back.field, tr.field) | |
| np.testing.assert_array_equal(np.isnan(back.entropy), np.isnan(tr.entropy)) | |
| assert back.u_lens == pytest.approx(tr.u_lens) | |
| assert back.tokens == tr.tokens and back.token_strs == tr.token_strs | |
| assert back.blocks == tr.blocks and back.names == tr.names | |
| assert back.eot == tr.eot and back.truncated == tr.truncated | |
| assert back.config == tr.config | |
| def test_the_sidecar_carries_the_scores(tmp_path): | |
| import json | |
| tr = make_trace() | |
| p = tr.save(tmp_path / "t.npz") | |
| side = json.loads(p.with_suffix(".json").read_text()) | |
| assert side["scores"]["u_lens"] == pytest.approx(tr.u_lens) | |
| assert side["readout_block"] == tr.readout_block | |
| def test_the_sidecar_is_strict_json_without_a_thinking_block(tmp_path): | |
| """Thinking off: no thinking-block entropy, no score. The sidecar says null, not NaN.""" | |
| import json | |
| tr = make_trace(prompt_len=6, eot=5) | |
| assert tr.think_len == 0 | |
| p = tr.save(tmp_path / "t.npz") | |
| text = p.with_suffix(".json").read_text() | |
| assert "NaN" not in text | |
| side = json.loads(text, parse_constant=lambda c: pytest.fail(f"non-finite {c} in sidecar")) | |
| assert side["scores"]["think_entropy"] is None and side["scores"]["u_lens"] is None | |
| assert side["scores"]["a_cone"] == pytest.approx(tr.a_cone) | |
| # ------------------------------------------------------------------ cone stacks | |
| def test_cone_stack_round_trips_and_slices(tmp_path): | |
| rng = np.random.default_rng(2) | |
| cats = ulens.terms.categories() | |
| cones = {b: ulens.axes.from_axes( | |
| [ulens.Axis(c, rng.normal(size=12).astype(np.float32)) for c in cats], | |
| model="m", block=b) for b in (3, 5, 9)} | |
| st = ulens.axes.stack(cones, model="m", meta={"terms_version": "v1"}) | |
| assert st.blocks == [3, 5, 9] and st.U.shape == (3, 12, 4) | |
| assert st.nearest(6) == 5 and st.nearest(100) == 9 | |
| np.testing.assert_array_equal(st.at(5).U, cones[5].U) | |
| back = ulens.axes.load_stack(st.save(tmp_path / "s.npz")) | |
| np.testing.assert_allclose(back.U, st.U) | |
| assert back.blocks == st.blocks and back.names == st.names | |
| def test_cone_stack_rejects_mismatched_axis_names(): | |
| v = np.ones(8, np.float32) | |
| a = ulens.axes.from_axes([ulens.Axis("x", v)], block=1) | |
| b = ulens.axes.from_axes([ulens.Axis("y", v)], block=2) | |
| with pytest.raises(ValueError, match="expected"): | |
| ulens.axes.stack({1: a, 2: b}) | |
| with pytest.raises(ValueError, match="at least one cone"): | |
| ulens.axes.stack({}) | |
| def test_stack_index_refuses_an_uncovered_block(): | |
| v = np.ones(8, np.float32) | |
| st = ulens.axes.stack({4: ulens.axes.from_axes([ulens.Axis("x", v)], block=4)}) | |
| with pytest.raises(KeyError, match="not in this stack"): | |
| st.index(5) | |