| """Cached C++ serializers shared by scalar and batched inference paths.""" | |
| from __future__ import annotations | |
| import numpy as np | |
| from .native_backend import native_kernel | |
| class NativeCacheMixin: | |
| """Serialize numerical state while rebuilding process-local C++ caches.""" | |
| def __getstate__(self): | |
| state = vars(self).copy() | |
| for key in ("_compiled_serializer", "_compiled_spatial", "_coverage_metadata"): | |
| state.pop(key, None) | |
| return state | |
| def __setstate__(self, state): | |
| vars(self).update(state) | |
| def compiled_serializer(owner): | |
| """Return a codec for a serializer/vocabulary, or the Python override.""" | |
| spatial = getattr(owner, "spatial", None) | |
| temporal = getattr(owner, "temporal_merges", getattr(owner, "merges", ())) | |
| scheme = getattr(owner, "scheme", "set_bpe" if spatial is not None else "bpe") | |
| constructor = native_kernel("SetBPE" if scheme == "set_bpe" else "PlainBPE") | |
| if constructor is None: | |
| return None | |
| cached = getattr(owner, "_compiled_serializer", None) | |
| if cached is None: | |
| bins = getattr(owner, "bins", spatial.bins if spatial is not None else ()) | |
| if scheme == "set_bpe": | |
| merges = getattr(owner, "spatial_merges", spatial.merges) | |
| cached = constructor(bins, owner.lookahead, merges, temporal) | |
| else: | |
| cached = constructor(bins, owner.merges) | |
| owner._compiled_serializer = cached | |
| return cached | |
| def encode_native(owner, dense: np.ndarray): | |
| codec = compiled_serializer(owner) | |
| options = getattr(owner, "encoding_search", None) | |
| if options is not None and options.enabled: | |
| if codec is None or not hasattr(codec, "encode_search_batch"): | |
| raise RuntimeError( | |
| "packing search requires the rebuilt actioncodec2 native extension; " | |
| "run 'python setup.py build_ext --inplace'" | |
| ) | |
| return codec.encode_search_batch( | |
| np.asarray(dense, dtype=np.int64), | |
| options.spatial_search, | |
| options.temporal_search, | |
| options.search_window, | |
| options.search_beam, | |
| ) | |
| return None if codec is None else codec.encode_batch(np.asarray(dense, dtype=np.int64)) | |
| def decode_native(owner, rows, dims, *, validate=False): | |
| codec = compiled_serializer(owner) | |
| if codec is None: | |
| return None | |
| bins = getattr(owner, "bins", owner.spatial.bins if hasattr(owner, "spatial") else ()) | |
| if validate: | |
| if not hasattr(codec, "decode_checked_batch"): | |
| return None | |
| support = dims if hasattr(codec, "decode_active_batch") else sum(1 << d for d in dims) | |
| return codec.decode_checked_batch(rows, len(bins), support) | |
| if hasattr(codec, "decode_active_batch"): | |
| return codec.decode_active_batch(rows, len(bins), dims) | |
| return codec.decode_batch(rows, len(bins), sum(1 << d for d in dims)) | |