Spaces:
Running on Zero
Running on Zero
File size: 4,354 Bytes
aa7bfed | 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 | import unittest
import numpy as np
from src.absolute_pe import add_positional_encoding, sinusoidal_pe
from src.extract import random_qk, select_head
from src.rope import (
apply_rope,
attention_scores,
l2_norms,
pair_frequencies,
pair_xy,
rotate_pair,
row_cosine,
)
class RopeTests(unittest.TestCase):
def test_even_dim_required(self):
with self.assertRaises(ValueError):
apply_rope(np.ones((4, 5)))
def test_position_zero_is_identity(self):
x = np.random.default_rng(0).standard_normal((8, 16))
y = apply_rope(x, base=10000.0, style="interleaved")
np.testing.assert_allclose(y[0], x[0], atol=1e-12)
def test_preserves_l2(self):
x = np.random.default_rng(1).standard_normal((12, 32))
for style in ("interleaved", "llama"):
y = apply_rope(x, style=style)
np.testing.assert_allclose(l2_norms(x), l2_norms(y), atol=1e-10)
def test_matches_scratch_even_odd(self):
x = np.random.default_rng(2).standard_normal((6, 10))
seq, dim = x.shape
positions = np.arange(seq)[:, None]
pair_indices = np.arange(0, dim, 2)
inv_freq = 1 / (10000 ** (pair_indices / dim))
angles = positions * inv_freq
cos, sin = np.cos(angles), np.sin(angles)
expected = np.empty_like(x)
expected[:, 0::2] = x[:, 0::2] * cos - x[:, 1::2] * sin
expected[:, 1::2] = x[:, 0::2] * sin + x[:, 1::2] * cos
np.testing.assert_allclose(apply_rope(x, style="interleaved"), expected)
def test_relative_angle_depends_on_offset(self):
omega = pair_frequencies(8, base=10000.0)[0]
x = np.zeros((5, 8))
x[:, 0] = 1.0
y = apply_rope(x, style="interleaved")
xm, ym = pair_xy(y, 3, 0)
xn, yn = pair_xy(y, 1, 0)
a_m = np.arctan2(ym, xm)
a_n = np.arctan2(yn, xn)
self.assertAlmostEqual(a_m - a_n, (3 - 1) * omega, places=10)
def test_rotate_pair_and_batched_heads(self):
even, odd = np.array([1.0]), np.array([0.0])
re, ro = rotate_pair(even, odd, np.array([np.pi / 2]))
np.testing.assert_allclose(re, 0.0, atol=1e-12)
np.testing.assert_allclose(ro, 1.0, atol=1e-12)
x = np.random.default_rng(3).standard_normal((4, 7, 16))
y = apply_rope(x, style="llama")
self.assertEqual(y.shape, x.shape)
np.testing.assert_allclose(l2_norms(x), l2_norms(y), atol=1e-10)
def test_cosine_identity_at_zero(self):
x = np.random.default_rng(4).standard_normal((5, 12))
y = apply_rope(x)
np.testing.assert_allclose(row_cosine(x, y)[0], 1.0, atol=1e-10)
def test_attention_scores_shape(self):
q = np.ones((3, 4))
k = np.ones((3, 4))
s = attention_scores(q, k)
self.assertEqual(s.shape, (3, 3))
class AbsolutePeTests(unittest.TestCase):
def test_matches_loop_formula(self):
seq, dim, base = 4, 5, 10000.0
pe = sinusoidal_pe(seq, dim, base=base)
for k in range(seq):
for i in range(dim):
if i % 2 == 0:
expected = np.sin(k / base ** (i / dim))
else:
expected = np.cos(k / base ** ((i - 1) / dim))
self.assertAlmostEqual(pe[k, i], expected, places=12)
def test_add_changes_norm(self):
emb = np.random.default_rng(0).standard_normal((6, 8))
pe, combined = add_positional_encoding(emb)
self.assertEqual(pe.shape, emb.shape)
self.assertFalse(np.allclose(l2_norms(emb), l2_norms(combined)))
class ExtractHelpersTests(unittest.TestCase):
def test_random_qk(self):
data = random_qk(seq_len=8, dim=16, seed=0, base=10000.0)
self.assertEqual(data["q_before"].shape, (8, 16))
np.testing.assert_allclose(l2_norms(data["q_before"]), l2_norms(data["q_after"]), atol=1e-10)
slice_q = select_head(data["q_before"], 99)
self.assertEqual(slice_q.shape, (8, 16))
from src.plots import bulk_before_after_delta, attention_heatmaps
fig = bulk_before_after_delta(data["q_before"], data["q_after"])
self.assertTrue(len(fig.data) >= 1)
s = attention_scores(data["q_before"], data["k_before"])
attention_heatmaps(s, s)
if __name__ == "__main__":
unittest.main()
|