File size: 3,237 Bytes
02a39fb | 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 | #!/usr/bin/env python3
"""Fetch pinned training files from the JSON index; dry-run unless --execute."""
import argparse,hashlib,json,shutil
from pathlib import Path,PurePosixPath
def safe(value):
p=PurePosixPath(value)
if p.is_absolute() or '..' in p.parts or not p.parts:raise ValueError('Unsafe relative data path')
return Path(*p.parts)
def sha(path):
h=hashlib.sha256()
with Path(path).open('rb') as f:
for b in iter(lambda:f.read(8*1024*1024),b''):h.update(b)
return h.hexdigest()
def main():
p=argparse.ArgumentParser(description=__doc__)
directory=Path(__file__).absolute().parent
default_index=directory/'TRAINING_DATA_INDEX.json'
if not default_index.exists():default_index=directory.parent/'TRAINING_DATA_INDEX.json'
p.add_argument('--index',type=Path,default=default_index)
p.add_argument('--dataset',default='legacy_mix6_sft')
p.add_argument('--output',type=Path,default=Path('training_data_download'))
p.add_argument('--execute',action='store_true');a=p.parse_args()
index=json.loads(a.index.read_text());datasets={d['dataset_id']:d for d in index['datasets']}
if a.dataset=='all_private':ids={r['artifact_id'] for r in index['private_artifacts']}
elif a.dataset in datasets:ids=set(datasets[a.dataset].get('artifact_ids',[]))
else:raise ValueError('Unknown dataset ID; inspect index.datasets')
rows=[r for r in index['private_artifacts'] if r['artifact_id'] in ids]
if not rows:raise ValueError('No private snapshot artifacts for this dataset; inspect public download location')
source=index['private_snapshot']
if not source['repo_id'].startswith('diffusion-gif/training-data-vault-') or source['repo_type']!='dataset' or not source['private']:raise ValueError('Unexpected private data provider')
print(json.dumps({'dataset':a.dataset,'files':len(rows),'bytes':sum(r['bytes'] for r in rows),'source_repo':source['repo_id'],'revision':source['revision'],'login_required':True,'execute':a.execute},indent=2))
if not a.execute:return
from huggingface_hub import get_token,hf_hub_download
token=get_token()
if not token:raise RuntimeError('Run hf auth login locally using an account authorized for the private vault. Never put tokens in this file or chat.')
a.output.mkdir(parents=True,exist_ok=True)
for r in rows:
target=a.output/safe(r['path'])
if target.exists():
if not target.is_file() or target.stat().st_size!=r['bytes'] or sha(target)!=r['sha256']:raise ValueError('Existing user file differs; will not overwrite')
continue
cached=Path(hf_hub_download(repo_id=source['repo_id'],repo_type='dataset',revision=source['revision'],filename=r['path'],token=token,endpoint='https://huggingface.co'))
if cached.stat().st_size!=r['bytes'] or sha(cached)!=r['sha256']:raise ValueError('Downloaded file SHA256 mismatch')
target.parent.mkdir(parents=True,exist_ok=True);temporary=target.with_name(target.name+'.copy-partial')
shutil.copyfile(cached,temporary)
if sha(temporary)!=r['sha256']:raise ValueError('Copied file SHA256 mismatch')
temporary.replace(target);print('VERIFIED',r['path'],flush=True)
print('Done: all selected frozen files downloaded and SHA256 verified. See restore_legacy_mix6.py for a relocatable SFT working copy.')
if __name__=='__main__':main()
|