70 lines
3.0 KiB
Python
70 lines
3.0 KiB
Python
"""Run on 2033lab04ztr: mirror download with per-file size/SHA256 checks."""
|
|
import concurrent.futures
|
|
import hashlib
|
|
import json
|
|
import pathlib
|
|
import subprocess
|
|
|
|
ROOT = pathlib.Path('/mnt/58462F4C462F29E8/xiaobailing-turn-detection-experiments')
|
|
REPO = 'Soul-AILab/SoulX-Duplug-0.6B'
|
|
REVISION = '61701c4ab8193cc1ee2220d3848872ac6c720142'
|
|
TARGET = ROOT / 'SoulX-Duplug/pretrained_models'
|
|
|
|
|
|
def main():
|
|
raw = subprocess.check_output(['curl', '-fLsS', '--retry', '3', '--max-time', '60',
|
|
f'https://hf-mirror.com/api/models/{REPO}/revision/{REVISION}?blobs=true'])
|
|
metadata = json.loads(raw)
|
|
assert metadata['sha'] == REVISION
|
|
(ROOT / 'results/model-metadata.json').write_bytes(raw)
|
|
TARGET.mkdir(parents=True, exist_ok=True)
|
|
ms_raw = subprocess.check_output(['curl', '-fLsS', '--retry', '3', '--max-time', '60',
|
|
f'https://modelscope.cn/api/v1/models/{REPO}/repo/files?Revision=master&Recursive=true'])
|
|
ms_data = json.loads(ms_raw)
|
|
assert ms_data['Code'] == 200
|
|
ms_files = {x['Path']: x for x in ms_data['Data']['Files'] if x['Type'] != 'tree'}
|
|
(ROOT / 'results/modelscope-metadata.json').write_bytes(ms_raw)
|
|
|
|
def download(item):
|
|
name, size = item['rfilename'], item['size']
|
|
path = TARGET / name
|
|
expected = item.get('lfs', {}).get('sha256')
|
|
def valid():
|
|
if not path.is_file() or path.stat().st_size != size:
|
|
return False
|
|
if expected:
|
|
with path.open('rb') as f:
|
|
actual = hashlib.file_digest(f, 'sha256').hexdigest()
|
|
return actual == expected
|
|
return True
|
|
if not valid():
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
tmp = path.with_name(path.name + '.download')
|
|
url = f'https://hf-mirror.com/{REPO}/resolve/{REVISION}/{name}?download=true'
|
|
if name in ms_files and size > 1000000:
|
|
ms = ms_files[name]
|
|
assert ms['Size'] == size
|
|
if expected:
|
|
assert ms['Sha256'] == expected, f'Mirror mismatch: {name}'
|
|
else:
|
|
expected = ms['Sha256']
|
|
url = f'https://modelscope.cn/api/v1/models/{REPO}/repo?Revision=master&FilePath={name}'
|
|
subprocess.run(['curl', '-fL', '--retry', '5', '--retry-delay', '3',
|
|
'--connect-timeout', '15', '--speed-limit', '1024', '--speed-time', '90',
|
|
'-sS', '-o', str(tmp), url], check=True)
|
|
tmp.replace(path)
|
|
if not valid():
|
|
raise ValueError(f'Invalid downloaded file: {name}')
|
|
print(f'OK {name} {size} bytes', flush=True)
|
|
return dict(file=name, bytes=size, lfs_sha256=expected)
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=3) as pool:
|
|
checked = list(pool.map(download, metadata['siblings']))
|
|
(ROOT / 'results/model-download.json').write_text(json.dumps(dict(
|
|
repo=REPO, revision=REVISION, files=checked), indent=2) + '\n')
|
|
print('DOWNLOAD_COMPLETE', flush=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|