File size: 4,482 Bytes
be6377b | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """FastPDN-PII FP32 ONNX -> model.safetensors for BertForTokenClassification (plan Task 4a Step 2, Review Gate 14 A).
Named initializers keep their torch names. Each linear weight is the MatMul input that feeds the Add or BiasGelu taking
its named bias, transposed from ONNX's [in, out] to torch's [out, in]. Every weight is used once, the state dict loads
strictly, and no unused float tensor is left in the graph; anything else stops. Weights are copied, never changed.
Usage: python convert.py # writes model.safetensors and convert.json here; create-only
"""
from pathlib import Path
import hashlib, json
import numpy as np
from onnx import numpy_helper
HERE = Path(__file__).resolve().parent
REPO, REVISION = 'ArkadiuszPawlak/fastpdn-ner-polish-pii', '636c57fe77b0afad8f065e7c76243dc051437035'
SNAPSHOT = Path.home()/'.cache/huggingface/hub/models--ArkadiuszPawlak--fastpdn-ner-polish-pii/snapshots'/REVISION
ONNX_SHA256 = '11924bfd20eb478ee614e5e1184536825db4747753d341af7c04b940efa84645' # model.onnx, LFS pin at REVISION
CONFIG_SHA256 = '3cef51cd7d184478adbdec9f68730db279386fe145cb32473a0622edea923578'
NAMED = ('bert.', 'classifier.')
BUFFERS = ('bert.embeddings.position_ids',) # non-persistent in transformers: checked, not saved
def sha(path):
h = hashlib.sha256()
with open(path, 'rb') as f:
for block in iter(lambda: f.read(1 << 20), b''):
h.update(block)
return h.hexdigest()
def weights(graph):
"""({torch name: array}, used initializer names)."""
init = {i.name: numpy_helper.to_array(i) for i in graph.initializer}
producer = {o: n for n in graph.node for o in n.output}
state = {k: v for k, v in init.items() if k.startswith(NAMED)}
used = set(state)
for n in graph.node:
bias = [x for x in n.input if x.endswith('.bias') and 'LayerNorm' not in x and x in init]
if n.op_type not in ('Add', 'BiasGelu') or not bias:
continue
(b,), (other,) = bias, [x for x in n.input if x not in bias]
matmul = producer.get(other)
assert matmul is not None and matmul.op_type == 'MatMul' and matmul.input[1] in init, b
w, key = matmul.input[1], b[:-len('bias')] + 'weight'
assert w not in used and key not in state, (w, key)
state[key] = init[w].T.copy()
used.add(w)
return state, used
def leftovers(init, used):
"""Unused float tensors with more than one value: a weight the mapping missed. Scalars and int shapes pass."""
return sorted(k for k, v in init.items() if k not in used and np.asarray(v).dtype.kind == 'f' and np.size(v) > 1)
def main():
import onnx, torch
from safetensors.torch import save_file
from transformers import BertConfig, BertForTokenClassification
out = HERE/'model.safetensors'
assert not out.exists(), 'create-only'
assert sha(SNAPSHOT/'model.onnx') == ONNX_SHA256 and sha(SNAPSHOT/'config.json') == CONFIG_SHA256
graph = onnx.load(SNAPSHOT/'model.onnx').graph
init = {i.name: numpy_helper.to_array(i) for i in graph.initializer}
state, used = weights(graph)
left = leftovers(init, used)
assert not left, f'{len(left)} unused float tensors'
for k in BUFFERS:
v = state.pop(k)
assert np.array_equal(v, np.arange(v.shape[-1]).reshape(v.shape)), k
tensors = {k: torch.from_numpy(np.ascontiguousarray(v)) for k, v in sorted(state.items())}
model = BertForTokenClassification(BertConfig.from_pretrained(SNAPSHOT))
expected = model.state_dict()
assert {k: tuple(v.shape) for k, v in tensors.items()} == {k: tuple(v.shape) for k, v in expected.items()}, 'keys/shapes'
assert all(t.dtype == torch.float32 for t in tensors.values())
model.load_state_dict(tensors, strict=True)
save_file(tensors, out, metadata=dict(source=f'{REPO}@{REVISION}/model.onnx', onnx_sha256=ONNX_SHA256,
note='converted from ONNX, weights unchanged'))
report = dict(repo=REPO, revision=REVISION, onnx_sha256=ONNX_SHA256, config_sha256=CONFIG_SHA256,
initializers=len(init), used=len(used), tensors=len(tensors),
transposed=len(used) - sum(k.startswith(NAMED) for k in used),
unused_constants=len(init) - len(used), safetensors_sha256=sha(out), convert_sha256=sha(__file__))
(HERE/'convert.json').write_text(json.dumps(report, indent=1) + '\n')
return report
if __name__ == '__main__':
print(json.dumps(main()))
|