File size: 3,650 Bytes
e1acc15 | 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 | """PP3 bounded sparse-q codec; importing it makes no CUDA call.
Negotiated sample_mode=2/capability8, STEP flag2, reply flag2. Semantic1 is
dyadic24-floor-first-v1 with top32/tau.7 and fixed depth4/5. Old peers refuse
HELLO. The original frame version and all OFF bytes remain unchanged.
"""
from __future__ import annotations
import struct
import zlib
from typing import Any, TypeAlias
import numpy as np
from numpy.typing import NDArray
SAMPLE_TOP32 = 2
CAP_TOP32 = 8
FLAG = 2
META = struct.Struct('<qI4x') # user seed and target-temperature-positive mask
HEADER = struct.Struct('<4sIII')
MASS = 1 << 24
K = 32
I32: TypeAlias = NDArray[np.int32]
F32: TypeAlias = NDArray[np.float32]
def append_q(body: bytes, mask: I32, ids: I32, q: F32) -> bytes:
n = len(mask)
if ids.shape != q.shape or ids.ndim != 3 or ids.shape[0] != n or ids.shape[2] != K:
raise ValueError('sampled reply tensor shapes')
data = (HEADER.pack(b'PP3Q', 1, n, ids.shape[1]) + np.asarray(mask, dtype='<i4').tobytes()
+ np.asarray(ids, dtype='<i4').tobytes() + np.asarray(q, dtype='<f4').tobytes())
return bytes(body + data + struct.pack('<I', zlib.crc32(body + data)))
def unpack_q(body: Any, offset: int, n: int, depth: int) -> tuple[I32, I32, F32]:
size = HEADER.size + n * 4 + n * depth * K * 8 + 4
if len(body) != offset + size:
raise ValueError('sampled reply size mismatch')
if HEADER.unpack_from(body, offset) != (b'PP3Q', 1, n, depth):
raise ValueError('sampled reply semantic/shape mismatch')
if zlib.crc32(body[:-4]) != struct.unpack_from('<I', body, len(body)-4)[0]:
raise ValueError('sampled reply CRC mismatch')
off = offset + HEADER.size
mask = np.frombuffer(body, dtype='<i4', count=n, offset=off).copy()
off += n * 4
ids = np.frombuffer(body, dtype='<i4', count=n*depth*K, offset=off).reshape(n, depth, K).copy()
off += n * depth * K * 4
q = np.frombuffer(body, dtype='<f4', count=n*depth*K, offset=off).reshape(n, depth, K).copy()
return mask, ids, q
def validate(mask: I32, ids: I32, q: F32, tokens: I32, needed: NDArray[np.bool_],
vocab: int, expected_mask: I32 | None = None) -> None:
n, depth = tokens.shape
if (mask.shape != (n,) or ids.shape != (n, depth, K) or q.shape != ids.shape
or depth not in (4, 5) or not np.isin(mask, (0, 1)).all()):
raise ValueError('invalid sampled q shape/mask')
if expected_mask is not None and not np.array_equal(mask[needed], expected_mask[needed]):
raise ValueError('sampled mask does not match target temperature')
# Every row is scattered on GPU, including zero-q greedy/failed rows.
if np.any(ids < 0) or np.any(ids >= vocab):
raise ValueError('q support ID outside target vocabulary')
for row in range(n):
if not needed[row] or mask[row] == 0:
if np.any(q[row] != 0):
raise ValueError('q present for a greedy/unused row')
continue
rids, rq = ids[row], q[row]
if (np.any(rids < 0) or np.any(rids >= vocab) or np.any(np.diff(rids.astype(np.int64), axis=1) <= 0)
or not np.isfinite(rq).all() or np.any(rq < 0)):
raise ValueError('invalid q support or values')
counts = rq.astype(np.float64) * MASS
if not np.equal(counts, np.floor(counts)).all() or not np.equal(counts.sum(1), MASS).all():
raise ValueError('q must contain dyadic24 counts summing exactly to 2^24')
if not ((rids == tokens[row, :, None]) & (rq > 0)).any(1).all():
raise ValueError('sampled token missing positive proposal mass')
|