| """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' |
| CONFIG_SHA256 = '3cef51cd7d184478adbdec9f68730db279386fe145cb32473a0622edea923578' |
| NAMED = ('bert.', 'classifier.') |
| BUFFERS = ('bert.embeddings.position_ids',) |
|
|
|
|
| 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())) |
|
|
