Spaces:
Sleeping
Sleeping
File size: 5,045 Bytes
2e818da | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | import math
import pytest
from app.schemas.visual_lesson import AttentionLessonDraft, EvidenceClaim, EvidenceSource
from app.services.attention_compiler import AttentionCompiler, AttentionCompilerError, _matmul, _softmax_rows
def _draft(tokens=None):
return AttentionLessonDraft(
title="Inside scaled dot-product attention",
interpretation="A small deterministic walkthrough of self-attention.",
requested_variant="standard_self_attention",
lesson_variant="standard_self_attention",
variant_notice="",
tokens=tokens or ["Research", "Mate", "reads", "papers"],
assumptions=["The displayed tensors are illustrative teaching values."],
teaching_steps=[
"Start with tokens and embeddings.",
"Project embeddings into queries, keys, and values.",
"Compare queries with keys.",
"Scale the logits.",
"Optionally apply a causal mask.",
"Normalize each row with softmax.",
"Aggregate values into the output.",
],
claim_ids=["canonical-attention", "illustrative-values"],
)
def _evidence():
sources = [
EvidenceSource(
source_id="builtin:vaswani-2017",
origin="builtin",
title="Attention Is All You Need",
url="https://arxiv.org/abs/1706.03762",
authority="canonical",
)
]
claims = [
EvidenceClaim(
claim_id="canonical-attention",
text="Scaled dot-product attention applies softmax to scaled query-key products before aggregating values.",
claim_type="standard_definition",
support_level="canonical",
source_ids=["builtin:vaswani-2017"],
),
EvidenceClaim(
claim_id="illustrative-values",
text="The small matrices in this lesson are illustrative rather than trained model weights.",
claim_type="illustrative_choice",
support_level="illustrative",
source_ids=[],
),
]
return sources, claims
def test_compiler_produces_verified_unmasked_and_causal_branches():
sources, claims = _evidence()
payload = AttentionCompiler().build_payload(
project_id="project-1",
prompt="visualize attention",
draft=_draft(),
evidence_sources=sources,
evidence_claims=claims,
warnings=[],
)
assert [branch.branch_id for branch in payload.compiled.branches] == ["unmasked", "causal"]
assert len(payload.spec.timeline) == 7
assert payload.compiled.assertions_passed is True
for branch in payload.compiled.branches:
for row in branch.tensors["attentionWeights"]:
assert math.isclose(sum(row), 1.0, abs_tol=1e-6)
causal = payload.compiled.branches[1].tensors["attentionWeights"]
for row_index, row in enumerate(causal):
assert all(abs(value) <= 1e-6 for value in row[row_index + 1 :])
def test_compiler_replaces_unsafe_or_wrong_sized_token_labels():
sources, claims = _evidence()
payload = AttentionCompiler().build_payload(
project_id="project-1",
prompt="visualize attention",
draft=_draft(tokens=["only", "three", "tokens"]),
evidence_sources=sources,
evidence_claims=claims,
warnings=[],
)
assert payload.spec.tokens == ["Research", "Mate", "reads", "papers"]
def test_compiler_rejects_an_unapproved_operation_graph():
compiler = AttentionCompiler()
operations = compiler.operation_graph()
operations[0] = operations[0].model_copy(update={"op": "identity"})
with pytest.raises(AttentionCompilerError, match="approved attention operation graph"):
compiler.compile_operations(operations, causal=False)
def test_matrix_multiplication_and_stable_softmax_are_deterministic():
assert _matmul([[1.0, 2.0]], [[3.0], [4.0]]) == [[11.0]]
probabilities = _softmax_rows([[1000.0, 1001.0, 999.0]])[0]
assert all(math.isfinite(value) for value in probabilities)
assert math.isclose(sum(probabilities), 1.0, abs_tol=1e-12)
assert probabilities[1] > probabilities[0] > probabilities[2]
def test_fixed_fixture_recompiles_to_identical_tensors():
sources, claims = _evidence()
compiler = AttentionCompiler()
payload = compiler.build_payload(
project_id="project-1", prompt="attention", draft=_draft(),
evidence_sources=sources, evidence_claims=claims, warnings=[],
)
recompiled = compiler.compile_spec(payload.spec)
assert recompiled.model_dump() == payload.compiled.model_dump()
def test_compiler_rejects_nonfinite_or_wrong_shaped_persisted_inputs():
compiler = AttentionCompiler()
with pytest.raises(AttentionCompilerError, match="embeddings must have shape"):
compiler.compile_operations(
compiler.operation_graph(), causal=False,
input_tensors={"embeddings": [[1.0]], "Wq": [[1.0, 0.0]] * 3, "Wk": [[1.0, 0.0]] * 3, "Wv": [[1.0, 0.0]] * 3},
)
|