File size: 2,867 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 | """Read-only access to the verified official 1017 shards (x86host /opt/dsv41/x86host-data/native-1017).
Tensors come straight from the safetensors bytes through a memory map: weights stay the official bytes (uint8 views
for I8 / F8_E4M3 / F8_E8M0, bf16 / fp32 otherwise), nothing is converted here. ``verify()`` re-hashes tensors against
NATIVE_MANIFEST.json (sha256 of the tensor bytes)."""
from __future__ import annotations
import hashlib
import json
import os
import struct
import numpy as np
import torch
MANIFEST = os.path.join(os.path.dirname(os.path.abspath(__file__)), "NATIVE_MANIFEST.json")
class OfficialShards:
def __init__(self, root: str, manifest: str | None = MANIFEST):
self.root = root
with open(os.path.join(root, "model.safetensors.index.json")) as f:
self.wmap = json.load(f)["weight_map"]
self._hdr: dict = {}
self._mm: dict = {}
self.manifest = None
if manifest and os.path.exists(manifest):
with open(manifest) as f:
self.manifest = {t["name"]: t for t in json.load(f)["tensors"]}
def _file(self, fname: str):
if fname not in self._hdr:
path = os.path.join(self.root, fname)
with open(path, "rb") as f:
n = struct.unpack("<Q", f.read(8))[0]
h = json.loads(f.read(n))
h.pop("__metadata__", None)
self._hdr[fname] = (h, 8 + n)
self._mm[fname] = np.memmap(path, dtype=np.uint8, mode="r")
return self._hdr[fname], self._mm[fname]
def info(self, name: str) -> dict:
(h, off), _ = self._file(self.wmap[name])
return h[name]
def raw(self, name: str) -> np.ndarray:
"""The tensor's bytes (a memory-mapped uint8 view: no copy)."""
fname = self.wmap[name]
(h, off), mm = self._file(fname)
a, b = h[name]["data_offsets"]
return mm[off + a:off + b]
def get(self, name: str) -> torch.Tensor:
"""A CPU tensor (a private copy) in the official dtype: uint8 for I8 / F8_*, bf16, fp32."""
info = self.info(name)
t = torch.from_numpy(np.array(self.raw(name)))
dt = info["dtype"]
if dt == "F32":
t = t.view(torch.float32)
elif dt == "BF16":
t = t.view(torch.bfloat16)
elif dt not in ("I8", "F8_E4M3", "F8_E8M0"):
raise ValueError(f"{name}: unexpected dtype {dt}")
return t.view(info["shape"])
def verify(self, name: str) -> None:
if self.manifest is None:
raise RuntimeError("no NATIVE_MANIFEST.json to verify against")
want = self.manifest[name]["sha256"]
got = hashlib.sha256(memoryview(np.ascontiguousarray(self.raw(name)))).hexdigest()
if got != want:
raise RuntimeError(f"{name}: sha256 {got} != manifest {want}")
|