import json, struct, sys, urllib.request, concurrent.futures as cf
S = sys.argv[1]
REPOS = ['DeepSeek-V4.1-Flash', 'DeepSeek-V4.1-Flash-1017']
shards = sorted(set(json.load(open(f'{S}/DeepSeek-V4.1-Flash-1017.index.json'))['weight_map'].values()))
def rng(repo, f, a, b):
    req = urllib.request.Request(f'https://huggingface.co/deepseek-ai/{repo}/resolve/main/{f}',
                                 headers={'Range': f'bytes={a}-{b}', 'User-Agent': 'curl/8.7.1'})
    for _ in range(4):
        try:
            return urllib.request.urlopen(req, timeout=120).read()
        except Exception as e:
            err = e
    raise err
def header(job):
    repo, f = job
    n = struct.unpack('<Q', rng(repo, f, 0, 7))[0]
    h = json.loads(rng(repo, f, 8, 8 + n - 1))
    h.pop('__metadata__', None)
    return repo, f, 8 + n, h
out = {r: {} for r in REPOS}
with cf.ThreadPoolExecutor(16) as ex:
    for repo, f, base, h in ex.map(header, [(r, f) for r in REPOS for f in shards]):
        for k, v in h.items():
            out[repo][k] = dict(file=f, base=base, dtype=v['dtype'], shape=v['shape'], off=v['data_offsets'])
json.dump(out, open(f'{S}/headers.json', 'w'))
for r in REPOS: print(r, len(out[r]))