Spaces:
Sleeping
Sleeping
| """Inference-only Morton and spatial BPE codec.""" | |
| from __future__ import annotations | |
| from collections import defaultdict | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Any, Dict, List, Optional, Sequence, Tuple | |
| import json | |
| import re | |
| import numpy as np | |
| _DIR_TO_VEC = {0: (1, 0, 0), 1: (0, 1, 0), 2: (0, 0, 1)} | |
| MORTON_GRID_SIZE = 32 | |
| MORTON_BITS = 5 | |
| MORTON_MAX_CODE = (1 << 15) - 1 | |
| MORTON_TOKEN_RE = re.compile(r"<morton_(\d+)>") | |
| MESH_TOKEN_RE = re.compile(r"<mesh_(\d+)>") | |
| MORTON_MESH_PAIR_RE = re.compile(r"<morton_(\d+)>\s*<mesh_(\d+)>") | |
| PairKey = Tuple[int, int, int, int, int] | |
| class MergeEntry: | |
| new_id: int | |
| pair: Tuple[int, int] | |
| direction: int | |
| rel_offset: Tuple[int, int, int] | |
| leaf_tokens: List[int] | |
| leaf_offsets: List[Tuple[int, int, int]] | |
| class MortonMeshParseResult: | |
| ids: np.ndarray | |
| anchors: np.ndarray | |
| dropped_count: int = 0 | |
| def _as_int64_array(x: Any, name: str, ndim: int) -> np.ndarray: | |
| arr = np.asarray(x, dtype=np.int64) | |
| if arr.ndim != ndim: | |
| raise ValueError(f'{name} must be {ndim}D, got shape={arr.shape}') | |
| return arr | |
| def morton3d_encode_xyz32(coords: Any) -> Any: | |
| """Encode block coordinates ``[0, 31]^3`` with fixed x/y/z bit interleave. | |
| Bit order is fixed for the whole project: | |
| ``x_i -> bit 3*i``, ``y_i -> bit 3*i+1``, ``z_i -> bit 3*i+2``. | |
| """ | |
| arr = np.asarray(coords, dtype=np.int64) | |
| single = arr.ndim == 1 | |
| if single: | |
| if arr.shape[0] != 3: | |
| raise ValueError(f'coords must be [3] or [N,3], got shape={arr.shape}') | |
| arr2 = arr.reshape(1, 3) | |
| elif arr.ndim == 2 and arr.shape[1] == 3: | |
| arr2 = arr | |
| else: | |
| raise ValueError(f'coords must be [3] or [N,3], got shape={arr.shape}') | |
| if arr2.size and ((arr2 < 0).any() or (arr2 >= MORTON_GRID_SIZE).any()): | |
| raise ValueError(f'Morton coordinates must be in [0, {MORTON_GRID_SIZE - 1}], got min={arr2.min()} max={arr2.max()}') | |
| x = arr2[:, 0] | |
| y = arr2[:, 1] | |
| z = arr2[:, 2] | |
| code = np.zeros(arr2.shape[0], dtype=np.int64) | |
| for bit in range(MORTON_BITS): | |
| code |= (x >> bit & 1) << 3 * bit | |
| code |= (y >> bit & 1) << 3 * bit + 1 | |
| code |= (z >> bit & 1) << 3 * bit + 2 | |
| return int(code[0]) if single else code | |
| def morton3d_decode_xyz32(codes: Any) -> Any: | |
| """Decode fixed 15-bit Morton code(s) back to block coordinates.""" | |
| arr = np.asarray(codes, dtype=np.int64) | |
| single = arr.ndim == 0 | |
| arr1 = arr.reshape(-1) | |
| if arr1.size and ((arr1 < 0).any() or (arr1 > MORTON_MAX_CODE).any()): | |
| raise ValueError(f'Morton codes must be in [0, {MORTON_MAX_CODE}], got min={arr1.min()} max={arr1.max()}') | |
| coords = np.zeros((arr1.shape[0], 3), dtype=np.int64) | |
| for bit in range(MORTON_BITS): | |
| coords[:, 0] |= (arr1 >> 3 * bit & 1) << bit | |
| coords[:, 1] |= (arr1 >> 3 * bit + 1 & 1) << bit | |
| coords[:, 2] |= (arr1 >> 3 * bit + 2 & 1) << bit | |
| return coords[0] if single else coords | |
| def serialize_morton_mesh_pairs(ids: Sequence[int], anchors: Any) -> str: | |
| """Serialize BPE/VQ ids with explicit block-anchor Morton tokens.""" | |
| ids_np = np.asarray(ids, dtype=np.int64).reshape(-1) | |
| anchors_np = _as_int64_array(anchors, 'anchors', 2) | |
| if anchors_np.shape[1] != 3: | |
| raise ValueError(f'anchors must be [K,3], got {anchors_np.shape}') | |
| if ids_np.shape[0] != anchors_np.shape[0]: | |
| raise ValueError('ids and anchors length mismatch') | |
| if ids_np.size == 0: | |
| return '<mesh_start><mesh_end>' | |
| morton_codes = morton3d_encode_xyz32(anchors_np) | |
| parts = [f'<morton_{int(m)}><mesh_{int(t)}>' for (m, t) in zip(morton_codes.tolist(), ids_np.tolist())] | |
| return '<mesh_start>' + ''.join(parts) + '<mesh_end>' | |
| def parse_morton_mesh_pairs(text: str, max_mesh_id: Optional[int]=None) -> MortonMeshParseResult: | |
| """Parse valid ``<morton_*><mesh_*>`` pairs, dropping damaged fragments. | |
| Only adjacent Morton/Mesh pairs are accepted. Lone Morton tokens, lone mesh | |
| tokens, out-of-range Morton codes, and mesh ids >= ``max_mesh_id`` are | |
| skipped and counted in ``dropped_count``. | |
| """ | |
| ids: List[int] = [] | |
| morton_codes: List[int] = [] | |
| used_morton_spans = set() | |
| used_mesh_spans = set() | |
| dropped = 0 | |
| for match in MORTON_MESH_PAIR_RE.finditer(text): | |
| morton = int(match.group(1)) | |
| mesh_id = int(match.group(2)) | |
| used_morton_spans.add((match.start(1), match.end(1))) | |
| used_mesh_spans.add((match.start(2), match.end(2))) | |
| if morton < 0 or morton > MORTON_MAX_CODE: | |
| dropped += 1 | |
| continue | |
| if max_mesh_id is not None and mesh_id >= int(max_mesh_id): | |
| dropped += 1 | |
| continue | |
| morton_codes.append(morton) | |
| ids.append(mesh_id) | |
| total_morton = len(list(MORTON_TOKEN_RE.finditer(text))) | |
| total_mesh = len(list(MESH_TOKEN_RE.finditer(text))) | |
| dropped += max(0, total_morton - len(used_morton_spans)) | |
| dropped += max(0, total_mesh - len(used_mesh_spans)) | |
| if ids: | |
| anchors = morton3d_decode_xyz32(np.asarray(morton_codes, dtype=np.int64)) | |
| ids_np = np.asarray(ids, dtype=np.int64) | |
| else: | |
| anchors = np.zeros((0, 3), dtype=np.int64) | |
| ids_np = np.zeros((0,), dtype=np.int64) | |
| return MortonMeshParseResult(ids=ids_np, anchors=anchors, dropped_count=dropped) | |
| def _morton_order(coords: np.ndarray) -> np.ndarray: | |
| if coords.size == 0: | |
| return np.array([], dtype=np.int64) | |
| cmax = int(np.max(coords)) | |
| return _morton_sort_indices(coords.astype(np.int64), coord_max=max(cmax + 1, 1)) | |
| def _morton_sort_indices(coords_xyz: np.ndarray, coord_max: int) -> np.ndarray: | |
| n = coords_xyz.shape[0] | |
| if n == 0: | |
| return np.array([], dtype=np.int64) | |
| bits = max(1, int(np.ceil(np.log2(coord_max + 1)))) | |
| x = coords_xyz[:, 0].astype(np.int64) | |
| y = coords_xyz[:, 1].astype(np.int64) | |
| z = coords_xyz[:, 2].astype(np.int64) | |
| bit_idx = np.arange(bits, dtype=np.int64) | |
| x_bits = (x[:, None] >> bit_idx & 1) << 3 * bit_idx | |
| y_bits = (y[:, None] >> bit_idx & 1) << 3 * bit_idx + 1 | |
| z_bits = (z[:, None] >> bit_idx & 1) << 3 * bit_idx + 2 | |
| code = x_bits.sum(axis=1) + y_bits.sum(axis=1) + z_bits.sum(axis=1) | |
| return np.argsort(code, kind='stable') | |
| def _canonical_min_offset(offsets: Sequence[Tuple[int, int, int]]) -> Tuple[int, int, int]: | |
| return min(((int(c[0]), int(c[1]), int(c[2])) for c in offsets)) | |
| def _add3(a: Tuple[int, int, int], b: Tuple[int, int, int]) -> Tuple[int, int, int]: | |
| return (a[0] + b[0], a[1] + b[1], a[2] + b[2]) | |
| def _sub3(a: Tuple[int, int, int], b: Tuple[int, int, int]) -> Tuple[int, int, int]: | |
| return (a[0] - b[0], a[1] - b[1], a[2] - b[2]) | |
| def _make_pair_key(a_tok: int, b_tok: int, a_anchor: Sequence[int], b_anchor: Sequence[int]) -> PairKey: | |
| return (int(a_tok), int(b_tok), int(b_anchor[0]) - int(a_anchor[0]), int(b_anchor[1]) - int(a_anchor[1]), int(b_anchor[2]) - int(a_anchor[2])) | |
| class _IncEncodeState: | |
| """Per-sample slot graph for fast BPE encode (replays ``merge_table`` in order).""" | |
| __slots__ = ('tok', 'M', 'cell_xyz', 'directions', 'coord_to_idx', 'slot_of_cell', 'token_of_slot', 'slot_anchor', 'cells_by_slot', 'pair_slot_pairs', 'incidence') | |
| def __init__(self, tok: 'BPE3DTokenizer', tokens: np.ndarray, coords: np.ndarray) -> None: | |
| self.tok = tok | |
| self.M = int(tokens.shape[0]) | |
| self.cell_xyz = coords.astype(np.int64, copy=False) | |
| self.directions = tok.directions | |
| self.coord_to_idx: Dict[Tuple[int, int, int], int] = {} | |
| for i in range(self.M): | |
| t = (int(self.cell_xyz[i, 0]), int(self.cell_xyz[i, 1]), int(self.cell_xyz[i, 2])) | |
| self.coord_to_idx[t] = i | |
| self.slot_of_cell = np.arange(self.M, dtype=np.int64) | |
| self.token_of_slot = tokens.astype(np.int64, copy=False).copy() | |
| self.slot_anchor = self.cell_xyz.copy() | |
| self.cells_by_slot: Dict[int, np.ndarray] = {i: np.array([i], dtype=np.int64) for i in range(self.M)} | |
| self.pair_slot_pairs: Dict[PairKey, set] = {} | |
| self.incidence: Dict[int, Dict[PairKey, Dict[int, int]]] = defaultdict(dict) | |
| if self.M > 0: | |
| cmax = int(np.max(self.cell_xyz)) | |
| K = int(cmax) + 2 | |
| xyz = self.cell_xyz | |
| codes = ((xyz[:, 0] * K + xyz[:, 1]) * K + xyz[:, 2]).astype(np.int64) | |
| order = np.argsort(codes, kind='stable') | |
| sorted_codes = codes[order] | |
| tokv = self.token_of_slot | |
| for d in self.directions: | |
| dv = np.asarray(_DIR_TO_VEC[int(d)], dtype=np.int64) | |
| nbr = xyz + dv | |
| nc = ((nbr[:, 0] * K + nbr[:, 1]) * K + nbr[:, 2]).astype(np.int64) | |
| idx = np.searchsorted(sorted_codes, nc) | |
| hit = (idx < self.M) & (sorted_codes[np.minimum(idx, self.M - 1)] == nc) | |
| a_cell = np.nonzero(hit)[0] | |
| if a_cell.size == 0: | |
| continue | |
| b_cell = order[idx[hit]] | |
| mask = b_cell != a_cell | |
| a_cell = a_cell[mask] | |
| b_cell = b_cell[mask] | |
| for k in range(int(a_cell.shape[0])): | |
| i = int(a_cell[k]) | |
| j = int(b_cell[k]) | |
| self.register_directed(_make_pair_key(int(tokv[i]), int(tokv[j]), self.slot_anchor[i], self.slot_anchor[j]), int(i), int(j)) | |
| def register_directed(self, key: PairKey, a_slot: int, b_slot: int) -> None: | |
| bucket = self.pair_slot_pairs.setdefault(key, set()) | |
| if (a_slot, b_slot) in bucket: | |
| return | |
| bucket.add((a_slot, b_slot)) | |
| self.incidence[a_slot].setdefault(key, {})[b_slot] = 0 | |
| self.incidence[b_slot].setdefault(key, {})[a_slot] = 1 | |
| def remove_directed(self, key: PairKey, a_slot: int, b_slot: int) -> None: | |
| bucket = self.pair_slot_pairs.get(key) | |
| if bucket is None or (a_slot, b_slot) not in bucket: | |
| return | |
| bucket.discard((a_slot, b_slot)) | |
| if not bucket: | |
| del self.pair_slot_pairs[key] | |
| if key in self.incidence[a_slot]: | |
| self.incidence[a_slot][key].pop(b_slot, None) | |
| if not self.incidence[a_slot][key]: | |
| del self.incidence[a_slot][key] | |
| if key in self.incidence[b_slot]: | |
| self.incidence[b_slot][key].pop(a_slot, None) | |
| if not self.incidence[b_slot][key]: | |
| del self.incidence[b_slot][key] | |
| def remove_all_incident(self, slot: int) -> None: | |
| for (key, others) in list(self.incidence[slot].items()): | |
| for (other, role) in list(others.items()): | |
| if int(role) == 0: | |
| self.remove_directed(key, int(slot), int(other)) | |
| else: | |
| self.remove_directed(key, int(other), int(slot)) | |
| def merge_slots(self, sa: int, sb: int, new_id: int) -> None: | |
| if sa == sb: | |
| return | |
| ca = self.cells_by_slot[sa] | |
| cb = self.cells_by_slot[sb] | |
| merged = np.concatenate([ca, cb], axis=0) | |
| self.cells_by_slot[sa] = merged | |
| del self.cells_by_slot[sb] | |
| for idx in cb.tolist(): | |
| self.slot_of_cell[int(idx)] = sa | |
| self.token_of_slot[sa] = int(new_id) | |
| self.token_of_slot[sb] = -1 | |
| if tuple((int(v) for v in self.slot_anchor[sb].tolist())) < tuple((int(v) for v in self.slot_anchor[sa].tolist())): | |
| self.slot_anchor[sa] = self.slot_anchor[sb] | |
| def add_edges_for_slot(self, slot: int) -> None: | |
| if int(self.token_of_slot[slot]) < 0: | |
| return | |
| ta = int(self.token_of_slot[slot]) | |
| cells = self.cells_by_slot[slot] | |
| for ci in cells.tolist(): | |
| c = (int(self.cell_xyz[int(ci), 0]), int(self.cell_xyz[int(ci), 1]), int(self.cell_xyz[int(ci), 2])) | |
| for d in self.directions: | |
| dv = _DIR_TO_VEC[int(d)] | |
| n = (c[0] + dv[0], c[1] + dv[1], c[2] + dv[2]) | |
| j = self.coord_to_idx.get(n) | |
| if j is not None: | |
| sj = int(self.slot_of_cell[int(j)]) | |
| if sj != int(slot): | |
| tb = int(self.token_of_slot[sj]) | |
| if tb >= 0: | |
| self.register_directed(_make_pair_key(ta, tb, self.slot_anchor[slot], self.slot_anchor[sj]), int(slot), sj) | |
| n_sub = (c[0] - dv[0], c[1] - dv[1], c[2] - dv[2]) | |
| k = self.coord_to_idx.get(n_sub) | |
| if k is not None: | |
| sk = int(self.slot_of_cell[int(k)]) | |
| if sk != int(slot): | |
| t_left = int(self.token_of_slot[sk]) | |
| if t_left >= 0: | |
| self.register_directed(_make_pair_key(t_left, ta, self.slot_anchor[sk], self.slot_anchor[slot]), sk, int(slot)) | |
| def apply_merge_table(self) -> None: | |
| for entry in self.tok.merge_table: | |
| (a_tok, b_tok) = (int(entry.pair[0]), int(entry.pair[1])) | |
| rel_offset = tuple((int(v) for v in entry.rel_offset)) | |
| new_id = int(entry.new_id) | |
| key = (a_tok, b_tok, rel_offset[0], rel_offset[1], rel_offset[2]) | |
| if key not in self.pair_slot_pairs: | |
| continue | |
| snapshot = list(self.pair_slot_pairs[key]) | |
| pairs = [(int(x), int(y)) for (x, y) in snapshot if int(self.token_of_slot[int(x)]) == a_tok and int(self.token_of_slot[int(y)]) == b_tok] | |
| if not pairs: | |
| continue | |
| if len(pairs) > 1: | |
| anchors = np.stack([self.slot_anchor[pa] for (pa, _) in pairs]) | |
| order = _morton_order(anchors) | |
| pairs = [pairs[int(i)] for i in order.tolist()] | |
| used: set = set() | |
| for (sa, sb) in pairs: | |
| if sa in used or sb in used: | |
| continue | |
| if int(self.token_of_slot[int(sa)]) < 0 or int(self.token_of_slot[int(sb)]) < 0: | |
| continue | |
| if int(self.token_of_slot[int(sa)]) != a_tok or int(self.token_of_slot[int(sb)]) != b_tok: | |
| continue | |
| self.remove_all_incident(int(sa)) | |
| self.remove_all_incident(int(sb)) | |
| self.merge_slots(int(sa), int(sb), new_id) | |
| self.add_edges_for_slot(int(sa)) | |
| used.add(sa) | |
| used.add(sb) | |
| def final_ids_anchors(self) -> Tuple[np.ndarray, np.ndarray]: | |
| if self.M == 0: | |
| return (np.zeros((0,), dtype=np.int64), np.zeros((0, 3), dtype=np.int64)) | |
| alive_slots = [s for s in self.cells_by_slot if int(self.token_of_slot[s]) >= 0] | |
| if not alive_slots: | |
| return (np.zeros((0,), dtype=np.int64), np.zeros((0, 3), dtype=np.int64)) | |
| ids = np.asarray([int(self.token_of_slot[s]) for s in alive_slots], dtype=np.int64) | |
| anchors = np.stack([self.slot_anchor[s] for s in alive_slots]).astype(np.int64, copy=False) | |
| order = _morton_order(anchors) | |
| return (ids[order], anchors[order]) | |
| class BPE3DTokenizer: | |
| def __init__(self, base_vocab_size: int=8192, directions: Tuple[int, ...]=(0, 1, 2)) -> None: | |
| self.base_vocab_size = int(base_vocab_size) | |
| self.directions = tuple((int(d) for d in directions)) | |
| for d in self.directions: | |
| if d not in _DIR_TO_VEC: | |
| raise ValueError(f'unsupported direction={d}, expected one of {tuple(_DIR_TO_VEC.keys())}') | |
| self.merge_table: List[MergeEntry] = [] | |
| def vocab_size(self) -> int: | |
| return self.base_vocab_size + len(self.merge_table) | |
| def _build_flat_leaf_templates(self) -> None: | |
| flat: Dict[int, Tuple[List[int], List[Tuple[int, int, int]]]] = {i: ([i], [(0, 0, 0)]) for i in range(self.base_vocab_size)} | |
| for entry in self.merge_table: | |
| (a, b) = entry.pair | |
| shift = tuple((int(v) for v in entry.rel_offset)) | |
| (a_leaf_ids, a_leaf_off) = flat[a] | |
| (b_leaf_ids, b_leaf_off) = flat[b] | |
| leaf_ids = list(a_leaf_ids) | |
| leaf_offs = list(a_leaf_off) | |
| for (t, off) in zip(b_leaf_ids, b_leaf_off): | |
| leaf_ids.append(int(t)) | |
| leaf_offs.append(_add3(off, shift)) | |
| anchor_shift = _canonical_min_offset(leaf_offs) | |
| leaf_offs = [_sub3(off, anchor_shift) for off in leaf_offs] | |
| entry.leaf_tokens = leaf_ids | |
| entry.leaf_offsets = leaf_offs | |
| flat[entry.new_id] = (leaf_ids, leaf_offs) | |
| def _encode_single_incremental(self, tokens: np.ndarray, coords: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: | |
| """Replay ``merge_table`` with local edge updates (fast encode path).""" | |
| tokens = _as_int64_array(tokens, 'tokens', 1) | |
| coords = _as_int64_array(coords, 'coords', 2) | |
| if tokens.shape[0] == 0: | |
| return (np.zeros((0,), dtype=np.int64), np.zeros((0, 3), dtype=np.int64)) | |
| st = _IncEncodeState(self, tokens, coords) | |
| st.apply_merge_table() | |
| return st.final_ids_anchors() | |
| def _encode_single(self, tokens: np.ndarray, coords: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: | |
| return self._encode_single_incremental(tokens, coords) | |
| def encode_sparse_numpy(self, tokens: np.ndarray, coords: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: | |
| """Encode one mesh sample on CPU only (no per-point ``torch.tensor``). | |
| Returns ``(macro_ids [K], macro_anchors [K,3])`` in Morton order of anchors. | |
| """ | |
| return self._encode_single(tokens, coords) | |
| def decode_to_leaf_batches(self, batches: Sequence[Dict[str, np.ndarray]], *, coordinate_min: int=0, coordinate_max: int=MORTON_GRID_SIZE - 1, duplicate_policy: str='first') -> List[Dict[str, Any]]: | |
| """Expand BPE macros and sanitize leaf VQ tokens before sparse decode. | |
| ``duplicate_policy='first'`` is deliberately deterministic: the first | |
| generated leaf at a coordinate is retained and later leaves at that | |
| coordinate are counted and discarded. Out-of-range coordinates and VQ | |
| ids never reach ``SparseTensor``. | |
| """ | |
| if duplicate_policy != 'first': | |
| raise ValueError(f"unsupported duplicate_policy={duplicate_policy!r}; expected 'first'") | |
| lo = int(coordinate_min) | |
| hi = int(coordinate_max) | |
| if lo > hi: | |
| raise ValueError(f'invalid coordinate range [{lo}, {hi}]') | |
| outputs: List[Dict[str, Any]] = [] | |
| for rec in batches: | |
| ids = _as_int64_array(rec['ids'], 'ids', 1) | |
| anchors = _as_int64_array(rec['anchors'], 'anchors', 2) | |
| if anchors.shape[1] != 3: | |
| raise ValueError(f'anchors must be [K,3], got {anchors.shape}') | |
| if ids.shape[0] != anchors.shape[0]: | |
| raise ValueError('ids and anchors length mismatch') | |
| leaf_ids: List[int] = [] | |
| leaf_coords: List[Tuple[int, int, int]] = [] | |
| seen_coords: set[Tuple[int, int, int]] = set() | |
| seen_anchors: set[Tuple[int, int, int]] = set() | |
| stats = {'macro_pair_count': int(ids.shape[0]), 'duplicate_macro_anchors': 0, 'invalid_bpe_ids': 0, 'expanded_leaf_token_count': 0, 'invalid_vq_ids': 0, 'out_of_range_leaf_coords': 0, 'duplicate_leaf_coords': 0, 'valid_leaf_token_count': 0, 'duplicate_policy': duplicate_policy, 'coordinate_min': lo, 'coordinate_max': hi} | |
| for (tid, anc_raw) in zip(ids.tolist(), anchors.tolist()): | |
| aid = int(tid) | |
| anchor = tuple((int(v) for v in anc_raw)) | |
| if anchor in seen_anchors: | |
| stats['duplicate_macro_anchors'] += 1 | |
| else: | |
| seen_anchors.add(anchor) | |
| if aid < 0 or aid >= self.vocab_size: | |
| stats['invalid_bpe_ids'] += 1 | |
| continue | |
| if aid < self.base_vocab_size: | |
| expanded = [(aid, (0, 0, 0))] | |
| else: | |
| entry = self.merge_table[aid - self.base_vocab_size] | |
| expanded = list(zip(entry.leaf_tokens, entry.leaf_offsets)) | |
| for (leaf_id_raw, offset_raw) in expanded: | |
| stats['expanded_leaf_token_count'] += 1 | |
| leaf_id = int(leaf_id_raw) | |
| if leaf_id < 0 or leaf_id >= self.base_vocab_size: | |
| stats['invalid_vq_ids'] += 1 | |
| continue | |
| coord = _add3(anchor, tuple((int(v) for v in offset_raw))) | |
| if any((v < lo or v > hi for v in coord)): | |
| stats['out_of_range_leaf_coords'] += 1 | |
| continue | |
| if coord in seen_coords: | |
| stats['duplicate_leaf_coords'] += 1 | |
| continue | |
| seen_coords.add(coord) | |
| leaf_ids.append(leaf_id) | |
| leaf_coords.append(coord) | |
| stats['valid_leaf_token_count'] = len(leaf_ids) | |
| outputs.append({'ids': np.asarray(leaf_ids, dtype=np.int64), 'coords': np.asarray(leaf_coords, dtype=np.int64).reshape(-1, 3), 'stats': stats}) | |
| return outputs | |
| def load(cls, path: str) -> 'BPE3DTokenizer': | |
| p = Path(path) | |
| data = json.loads(p.read_text(encoding='utf-8')) | |
| tok = cls(base_vocab_size=int(data['base_vocab_size']), directions=tuple((int(x) for x in data.get('directions', [0, 1, 2])))) | |
| table: List[MergeEntry] = [] | |
| for r in data['merge_table']: | |
| if 'rel_offset' not in r: | |
| raise ValueError('This merge_table.json was created by an older non-reversible 3D BPE format without rel_offset. Please retrain the BPE merge table before using decode_to_sparse.') | |
| table.append(MergeEntry(new_id=int(r['new_id']), pair=(int(r['pair'][0]), int(r['pair'][1])), direction=int(r['direction']), rel_offset=tuple((int(v) for v in r['rel_offset'])), leaf_tokens=[int(x) for x in r.get('leaf_tokens', [])], leaf_offsets=[tuple((int(v) for v in t)) for t in r.get('leaf_offsets', [])])) | |
| tok.merge_table = table | |
| if len(tok.merge_table) > 0 and len(tok.merge_table[0].leaf_tokens) == 0: | |
| tok._build_flat_leaf_templates() | |
| return tok | |