83 lines
3.2 KiB
Python
83 lines
3.2 KiB
Python
"""Read safetensors BF16 directly; never instantiate a model or use pickle."""
|
|
import json
|
|
import math
|
|
import mmap
|
|
import re
|
|
import struct
|
|
|
|
import numpy as np
|
|
|
|
DTYPES = {'BF16': ('<u2', 2), 'F16': ('<f2', 2), 'F32': ('<f4', 4), 'F64': ('<f8', 8)}
|
|
|
|
|
|
def decode(raw, dtype):
|
|
if dtype == 'BF16':
|
|
return (np.asarray(raw).astype(np.uint32) << 16).view(np.float32)
|
|
return np.asarray(raw, dtype=np.float32)
|
|
|
|
|
|
class Tensor:
|
|
def __init__(self, file, base, name, meta):
|
|
self.name, self.shape, self.dtype = name, tuple(meta['shape']), meta['dtype']
|
|
self.numel = math.prod(self.shape)
|
|
self.bytes = meta['data_offsets'][1] - meta['data_offsets'][0]
|
|
code, itemsize = DTYPES[self.dtype]
|
|
if self.bytes != self.numel * itemsize:
|
|
raise ValueError(f'Invalid byte length: {name}')
|
|
self.raw = np.memmap(file, mode='r', dtype=code,
|
|
offset=base + meta['data_offsets'][0], shape=(self.numel,))
|
|
self.role = re.sub(r'\.(?:layers|blocks)\.\d+', '.layers.*', name)
|
|
self.role = re.sub(r'\.experts\.\d+', '.experts.*', self.role)
|
|
m = re.search(r'\.(?:layers|blocks)\.(\d+)', name)
|
|
self.layer = int(m[1]) if m else -1
|
|
|
|
def read(self, selector=slice(None)):
|
|
return decode(self.raw[selector], self.dtype)
|
|
|
|
def chunked(self, size=262144):
|
|
for start in range(0, self.numel, size):
|
|
yield self.read(slice(start, min(start + size, self.numel)))
|
|
|
|
def evict(self):
|
|
if hasattr(self.raw._mmap, 'madvise'):
|
|
self.raw._mmap.madvise(mmap.MADV_DONTNEED)
|
|
|
|
|
|
def load_tensors(directory):
|
|
tensors = {}
|
|
for file in sorted(directory.glob('*.safetensors')):
|
|
with file.open('rb') as f:
|
|
size = struct.unpack('<Q', f.read(8))[0]
|
|
if size > 100_000_000:
|
|
raise ValueError('Oversized safetensors header')
|
|
header = json.loads(f.read(size))
|
|
for name, meta in header.items():
|
|
if name == '__metadata__':
|
|
continue
|
|
if name in tensors:
|
|
raise ValueError(f'Duplicate tensor name: {name}')
|
|
tensors[name] = Tensor(file, 8 + size, name, meta)
|
|
if not tensors:
|
|
raise ValueError('No safetensors found')
|
|
return tensors
|
|
|
|
|
|
def metrics(x, y):
|
|
"""x is member; y is representative. Float64 accumulation."""
|
|
x, y = x.astype(np.float64), y.astype(np.float64)
|
|
xx, yy, xy = float(x @ x), float(y @ y), float(x @ y)
|
|
delta = float((x - y) @ (x - y))
|
|
return finish(xx, yy, xy, delta, float(x.sum()), float(y.sum()), x.size)
|
|
|
|
|
|
def finish(xx, yy, xy, delta, sx, sy, n):
|
|
tiny = 1e-30
|
|
vx, vy = max(0., xx - sx * sx / n), max(0., yy - sy * sy / n)
|
|
return {'cosine': xy / np.sqrt(xx * yy) if xx * yy > tiny else None,
|
|
'pearson': (xy - sx * sy / n) / np.sqrt(vx * vy) if vx * vy > tiny else None,
|
|
'error_member_to_rep': np.sqrt(delta / max(xx, tiny)),
|
|
'error_rep_to_member': np.sqrt(delta / max(yy, tiny)),
|
|
'symmetric_error': np.sqrt(2 * delta / max(xx + yy, tiny)),
|
|
'scaled_error': np.sqrt(max(0., xx - xy * xy / max(yy, tiny)) / max(xx, tiny)),
|
|
'scale': xy / max(yy, tiny)}
|