codec.py

"""Versioned, checksummed arithmetic packets in printable basE91 ASCII."""
import bisect
import hashlib
import json
import struct
import zlib
from dataclasses import dataclass

import numpy as np

ALPHABET = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!#$%&()*+,./:;<=>?@[]^_`{|}~"'
VALUES = {char: index for index, char in enumerate(ALPHABET)}
MAX_TOKENS = 8192
MAX_BYTES = 65536
MAX_CODE = 100000
FULL, HALF, QUARTER = 1 << 32, 1 << 31, 1 << 30
TOTAL_BUDGET = 1 << 24


def base91_encode(data):
    buffer = bits = 0
    out = []
    for byte in data:
        buffer |= byte << bits
        bits += 8
        if bits > 13:
            value = buffer & 8191
            take = 13 if value > 88 else 14
            if take == 14:
                value = buffer & 16383
            buffer >>= take
            bits -= take
            out.extend((ALPHABET[value % 91], ALPHABET[value // 91]))
    if bits:
        out.append(ALPHABET[buffer % 91])
        if bits > 7 or buffer > 90:
            out.append(ALPHABET[buffer // 91])
    return ''.join(out)


def base91_decode(text):
    buffer = bits = 0
    pending = -1
    out = bytearray()
    for char in text:
        if char not in VALUES:
            raise ValueError('Invalid ASCII code. Copy the complete basE91 string.')
        digit = VALUES[char]
        if pending < 0:
            pending = digit
            continue
        value = pending + digit * 91
        buffer |= value << bits
        bits += 13 if value & 8191 > 88 else 14
        while bits >= 8:
            out.append(buffer & 255)
            buffer >>= 8
            bits -= 8
        pending = -1
    if pending >= 0:
        out.append((buffer | pending << bits) & 255)
    if base91_encode(out) != text:
        raise ValueError('Invalid basE91 ending.')
    return bytes(out)


class BitWriter:
    def __init__(self):
        self.data = bytearray()
        self.current = self.used = 0

    def write(self, bit):
        self.current = self.current * 2 + bit
        self.used += 1
        if self.used == 8:
            self.data.append(self.current)
            self.current = self.used = 0

    def finish(self):
        if self.used:
            self.data.append(self.current << (8 - self.used))
        return bytes(self.data)


class BitReader:
    def __init__(self, data):
        self.data, self.position = data, 0

    def read(self):
        pos = self.position
        self.position += 1
        if pos >= len(self.data) * 8 + 32:
            raise ValueError('Truncated arithmetic stream.')
        return (self.data[pos // 8] >> (7 - pos % 8)) & 1 if pos < len(self.data) * 8 else 0


class ArithmeticEncoder:
    def __init__(self):
        self.low, self.high, self.pending = 0, FULL - 1, 0
        self.bits = BitWriter()

    def emit(self, bit):
        self.bits.write(bit)
        for _ in range(self.pending):
            self.bits.write(1 - bit)
        self.pending = 0

    def write(self, symbol, cdf):
        if not 0 <= symbol < len(cdf) - 1:
            raise ValueError('Invalid token.')
        total, width = int(cdf[-1]), self.high - self.low + 1
        self.high = self.low + width * int(cdf[symbol + 1]) // total - 1
        self.low += width * int(cdf[symbol]) // total
        while True:
            if self.high < HALF:
                self.emit(0)
            elif self.low >= HALF:
                self.emit(1)
                self.low -= HALF
                self.high -= HALF
            elif self.low >= QUARTER and self.high < 3 * QUARTER:
                self.pending += 1
                self.low -= QUARTER
                self.high -= QUARTER
            else:
                break
            self.low *= 2
            self.high = self.high * 2 + 1

    def finish(self):
        self.pending += 1
        self.emit(0 if self.low < QUARTER else 1)
        return self.bits.finish()


class ArithmeticDecoder:
    def __init__(self, data):
        self.low, self.high, self.value = 0, FULL - 1, 0
        self.bits = BitReader(data)
        for _ in range(32):
            self.value = self.value * 2 + self.bits.read()

    def read(self, cdf):
        total, width = int(cdf[-1]), self.high - self.low + 1
        scaled = ((self.value - self.low + 1) * total - 1) // width
        symbol = bisect.bisect_right(cdf, scaled) - 1
        if not 0 <= symbol < len(cdf) - 1:
            raise ValueError('Invalid arithmetic interval.')
        self.high = self.low + width * int(cdf[symbol + 1]) // total - 1
        self.low += width * int(cdf[symbol]) // total
        while True:
            if self.high < HALF:
                pass
            elif self.low >= HALF:
                self.low -= HALF
                self.high -= HALF
                self.value -= HALF
            elif self.low >= QUARTER and self.high < 3 * QUARTER:
                self.low -= QUARTER
                self.high -= QUARTER
                self.value -= QUARTER
            else:
                break
            self.low *= 2
            self.high = self.high * 2 + 1
            self.value = self.value * 2 + self.bits.read()
        return symbol


def logits_to_cdf(logits):
    values = np.asarray(logits, dtype=np.float64)
    if values.ndim != 1 or not np.isfinite(values).all() or not 0 < len(values) < TOTAL_BUDGET:
        raise ValueError('Invalid model logits.')
    # Rounding reduces sensitivity; it does not establish cross-hardware determinism.
    values = np.rint(values * 64.0) / 64.0
    weights = np.exp(values - values.max())
    frequencies = 1 + np.floor(weights * ((TOTAL_BUDGET - len(values)) / weights.sum())).astype(np.int64)
    return np.concatenate((np.zeros(1, dtype=np.int64), np.cumsum(frequencies)))


def profile_id(settings):
    return hashlib.blake2s(json.dumps(settings, sort_keys=True).encode(), digest_size=8).digest()


def varint(value):
    out = bytearray()
    while value >= 128:
        out.append((value & 127) | 128)
        value >>= 7
    out.append(value)
    return out


def pack(payload, text, token_count, profile):
    raw = text.encode('utf-8', errors='strict')
    if len(profile) != 8 or token_count > MAX_TOKENS or len(raw) > MAX_BYTES:
        raise ValueError('Input exceeds the format limits.')
    out = bytearray(b'QC\x01') + profile + varint(token_count) + varint(len(raw))
    out += struct.pack('<I', zlib.crc32(raw)) + payload
    out += struct.pack('<I', zlib.crc32(out))
    return base91_encode(out)


@dataclass
class Packet:
    profile: bytes
    token_count: int
    original_bytes: int
    checksum: int
    payload: bytes
    packet_bytes: int


def unpack(code):
    if len(code) > MAX_CODE:
        raise ValueError('Compressed code is too long.')
    compact = ''.join(code.split())
    data = base91_decode(compact)
    if len(data) < 22:
        raise ValueError('Compressed code is incomplete.')
    if zlib.crc32(data[:-4]) != struct.unpack('<I', data[-4:])[0]:
        raise ValueError('Checksum mismatch: code was changed or truncated.')
    if data[:3] != b'QC\x01':
        raise ValueError('This is not a Qwen compressor v1 packet. SmolLM2 codes use the local demo.')
    offset = 11

    def read_varint():
        nonlocal offset
        start, value = offset, 0
        for shift in range(0, 28, 7):
            if offset >= len(data) - 9:
                raise ValueError('Truncated packet header.')
            byte = data[offset]
            offset += 1
            value |= (byte & 127) << shift
            if not byte & 128:
                if data[start:offset] != varint(value):
                    raise ValueError('Invalid packet length encoding.')
                return value
        raise ValueError('Invalid packet length.')

    count, raw_size = read_varint(), read_varint()
    if count > MAX_TOKENS or raw_size > MAX_BYTES:
        raise ValueError('This demo supports at most 8,192 tokens and 64 KiB of text.')
    checksum = struct.unpack('<I', data[offset:offset + 4])[0]
    payload = data[offset + 4:-4]
    if not payload:
        raise ValueError('Missing arithmetic data.')
    return Packet(data[3:11], count, raw_size, checksum, payload, len(data))


def verify_text(text, packet):
    raw = text.encode('utf-8', errors='strict')
    if len(raw) != packet.original_bytes or zlib.crc32(raw) != packet.checksum:
        raise ValueError('Exact decoding failed. Model probabilities differed; no decoded text will be returned.')