Buckets:
| """Compact, backward-compatible generated BSQ rollout artifacts.""" | |
| from __future__ import annotations | |
| from pathlib import Path | |
| import math | |
| import os | |
| import uuid | |
| import torch | |
| from .bsq_codes import pack_visual_codes, unpack_visual_codes | |
| SCHEMA = "zebra_generated_codes_packed_v1" | |
| def pack_generated_codes(codes: list[torch.Tensor | None]) -> dict: | |
| entries = [] | |
| for value in codes: | |
| if value is None: | |
| entries.append(None) | |
| continue | |
| if value.numel() and not bool(((value == 0) | (value == 1)).all()): | |
| raise ValueError("generated BSQ codes must be binary") | |
| packed, shape = pack_visual_codes(value) | |
| entries.append( | |
| { | |
| "packed": packed, | |
| "shape": list(shape), | |
| "dtype": str(value.dtype).removeprefix("torch."), | |
| } | |
| ) | |
| return {"schema": SCHEMA, "entries": entries} | |
| def unpack_generated_codes(payload) -> list[torch.Tensor | None]: | |
| if not isinstance(payload, dict): | |
| # Historical artifacts stored the raw list directly. | |
| if not isinstance(payload, list): | |
| raise ValueError("unsupported generated-code artifact payload") | |
| return payload | |
| if payload.get("schema") != SCHEMA: | |
| raise ValueError(f"unsupported generated-code schema: {payload.get('schema')!r}") | |
| result = [] | |
| for entry in payload["entries"]: | |
| if entry is None: | |
| result.append(None) | |
| continue | |
| if set(entry) != {"packed", "shape", "dtype"}: | |
| raise ValueError("generated-code entry has an invalid schema") | |
| shape = entry["shape"] | |
| if ( | |
| not isinstance(shape, list) | |
| or len(shape) < 2 | |
| or any(type(value) is not int or value < 0 for value in shape) | |
| ): | |
| raise ValueError("generated-code entry has an invalid shape") | |
| if entry["dtype"] not in {"bool", "uint8", "int8", "int16", "int32", "int64"}: | |
| raise ValueError("generated-code entry has an unsupported dtype") | |
| packed = entry["packed"] | |
| expected_packed_shape = ( | |
| shape[0], | |
| math.ceil(math.prod(shape[1:]) / 8), | |
| ) | |
| if not torch.is_tensor(packed) or tuple(packed.shape) != expected_packed_shape: | |
| raise ValueError("generated-code packed tensor has an invalid shape") | |
| dtype = getattr(torch, entry["dtype"]) | |
| result.append( | |
| unpack_visual_codes(packed, shape).to(dtype=dtype) | |
| ) | |
| return result | |
| def save_generated_codes(path: Path, codes: list[torch.Tensor | None]) -> None: | |
| temporary = path.with_name(f".{path.name}.tmp-{uuid.uuid4().hex}") | |
| try: | |
| torch.save(pack_generated_codes(codes), temporary) | |
| os.replace(temporary, path) | |
| finally: | |
| temporary.unlink(missing_ok=True) | |
| def load_generated_codes(path: Path) -> list[torch.Tensor | None]: | |
| try: | |
| payload = torch.load( | |
| path, map_location="cpu", weights_only=False, mmap=True | |
| ) | |
| except RuntimeError as error: | |
| if "mmap" not in str(error).lower(): | |
| raise | |
| payload = torch.load(path, map_location="cpu", weights_only=False) | |
| return unpack_generated_codes(payload) | |
| __all__ = [ | |
| "load_generated_codes", | |
| "pack_generated_codes", | |
| "save_generated_codes", | |
| "unpack_generated_codes", | |
| ] | |
Xet Storage Details
- Size:
- 3.38 kB
- Xet hash:
- 2399ca8d0243c5294b8c84a960de7fca94e036bf2ace641386d8938a5d8266cc
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.