Spaces:
Running on Zero
Running on Zero
File size: 2,915 Bytes
a81897d 5273448 a81897d 5273448 a81897d 5273448 a81897d | 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 | import random
import unittest
import numpy as np
from codec import (
ArithmeticEncoder, ArithmeticDecoder, base91_encode, base91_decode,
logits_to_cdf, pack, unpack, verify_text, profile_id, MAX_TOKENS, MAX_BYTES, MAX_CODE,
)
class CodecTests(unittest.TestCase):
def test_base91_binary_roundtrip(self):
rng = random.Random(721)
for size in [0, 1, 2, 3, 16, 100, 511, 10000]:
raw = rng.randbytes(size)
code = base91_encode(raw)
self.assertTrue(all(33 <= ord(c) <= 126 for c in code))
self.assertEqual(base91_decode(code), raw)
self.assertEqual(base91_encode(b'Hello World!'), '>OwJh>Io0Tv!8PE')
def test_arithmetic_changing_distributions(self):
rng = np.random.default_rng(52)
for length, vocab in [(0, 3), (1, 7), (1000, 256), (15, 262144)]:
distributions = [logits_to_cdf(rng.normal(0, 8, vocab)) for _ in range(length)]
symbols = rng.integers(0, vocab, length).tolist()
enc = ArithmeticEncoder()
for symbol, cdf in zip(symbols, distributions):
enc.write(symbol, cdf)
dec = ArithmeticDecoder(enc.finish())
self.assertEqual([dec.read(cdf) for cdf in distributions], symbols)
def test_distribution_has_full_support(self):
cdf = logits_to_cdf(np.array([-10000.0, 0.0, 10000.0]))
self.assertTrue(np.all(np.diff(cdf) > 0))
self.assertLessEqual(int(cdf[-1]), 1 << 24)
with self.assertRaises(ValueError):
logits_to_cdf([float('nan')])
def test_packet_roundtrip_and_damage(self):
profile = profile_id({'model': 'test'})
for text in ['', 'hello', ' Hello 🌍\n你好\t café ']:
code = pack(b'\x40', text, 0, profile)
packet = unpack(' \n' + code + '\t')
self.assertEqual(packet.profile, profile)
verify_text(text, packet)
with self.assertRaises(ValueError):
verify_text(text + '!', packet)
replacement = 'A' if code[5] != 'A' else 'B'
with self.assertRaises(ValueError):
unpack(code[:5] + replacement + code[6:])
with self.assertRaises(ValueError):
unpack(code[:-2])
def test_limits_and_profile(self):
p = profile_id({'revision': 'a'})
self.assertNotEqual(p, profile_id({'revision': 'b'}))
with self.assertRaises(ValueError):
pack(b'\x40', 'x', MAX_TOKENS + 1, p)
with self.assertRaises(ValueError):
unpack('A' * (MAX_CODE + 1))
text = 'a' * MAX_BYTES
packet = unpack(pack(b'\x40', text, MAX_TOKENS, p))
self.assertEqual(packet.token_count, MAX_TOKENS)
verify_text(text, packet)
with self.assertRaises(ValueError):
pack(b'\x40', text + 'a', 1, p)
if __name__ == '__main__':
unittest.main()
|