77 lines
2.8 KiB
Python
77 lines
2.8 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
# Copyright (C) 2026 David (Corey) Herrera
|
|
"""
|
|
fetch_mtp_tensors.py — fetch ONLY the `mtp.*` tensors of Qwen/Qwen3.5-9B (STEP 4 S2: the
|
|
grafted MTP draft head for Ornith) by HTTP range reads of the safetensors shards, instead of
|
|
downloading ~15 GB of shards for ~0.5 GB of tensors.
|
|
|
|
Uses the logged-in huggingface_hub token. Writes one safetensors file (all mtp tensors,
|
|
original dtypes) plus the index and config into the given directory (gitignored: models/).
|
|
|
|
Usage: fetch_mtp_tensors.py <out_dir> (run with the python that has huggingface_hub)
|
|
"""
|
|
import json
|
|
import os
|
|
import struct
|
|
import sys
|
|
import urllib.request
|
|
|
|
from huggingface_hub import get_token, hf_hub_url
|
|
|
|
REPO = "Qwen/Qwen3.5-9B"
|
|
|
|
|
|
def ranged(url, token, start, end):
|
|
"""Bytes [start, end] of `url` (inclusive), following redirects."""
|
|
req = urllib.request.Request(url, headers={"Range": f"bytes={start}-{end}",
|
|
"Authorization": f"Bearer {token}"})
|
|
with urllib.request.urlopen(req) as resp:
|
|
return resp.read()
|
|
|
|
|
|
def main():
|
|
out = sys.argv[1]
|
|
token = get_token()
|
|
index = json.load(open(os.path.join(out, "model.safetensors.index.json")))
|
|
wanted = {k: v for k, v in index["weight_map"].items() if k.startswith("mtp.")}
|
|
by_shard = {}
|
|
for name, shard in wanted.items():
|
|
by_shard.setdefault(shard, []).append(name)
|
|
|
|
tensors = {}
|
|
for shard, names in sorted(by_shard.items()):
|
|
url = hf_hub_url(REPO, shard)
|
|
head_len = struct.unpack("<Q", ranged(url, token, 0, 7))[0]
|
|
header = json.loads(ranged(url, token, 8, 8 + head_len - 1))
|
|
base = 8 + head_len
|
|
for name in sorted(names):
|
|
meta = header[name]
|
|
begin, end = meta["data_offsets"]
|
|
data = ranged(url, token, base + begin, base + end - 1)
|
|
assert len(data) == end - begin, (name, len(data), end - begin)
|
|
tensors[name] = (meta["dtype"], meta["shape"], data)
|
|
print(f"{name:55s} {meta['dtype']:5s} {meta['shape']} {len(data)} B", flush=True)
|
|
|
|
# One safetensors file, tensors in name order.
|
|
header = {}
|
|
offset = 0
|
|
for name in sorted(tensors):
|
|
dtype, shape, data = tensors[name]
|
|
header[name] = {"dtype": dtype, "shape": shape,
|
|
"data_offsets": [offset, offset + len(data)]}
|
|
offset += len(data)
|
|
blob = json.dumps(header).encode()
|
|
blob += b" " * ((8 - len(blob) % 8) % 8)
|
|
path = os.path.join(out, "mtp.safetensors")
|
|
with open(path, "wb") as f:
|
|
f.write(struct.pack("<Q", len(blob)))
|
|
f.write(blob)
|
|
for name in sorted(tensors):
|
|
f.write(tensors[name][2])
|
|
print(f"wrote {path}: {len(tensors)} tensors, {offset} B of data")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|