| """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') |
| 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') |
| |
| 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') |
|
|
