File size: 2,073 Bytes
3f98d52 | 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 | """Download released source files with retries and streamed checksum verification."""
import argparse
import hashlib
from pathlib import Path
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
p=argparse.ArgumentParser()
p.add_argument('dataset',choices=['tahoe','sciplex'])
p.add_argument('--output',required=True)
p.add_argument('--shards',type=int,default=2)
p.add_argument('--indices',type=int,nargs='+',help='Explicit official Tahoe shard indices')
a=p.parse_args()
out=Path(a.output);out.mkdir(parents=True,exist_ok=True)
s=requests.Session()
s.mount('https://',HTTPAdapter(max_retries=Retry(total=5,backoff_factor=2,status_forcelist=[429,500,502,503,504],respect_retry_after_header=True)))
base='https://huggingface.co/datasets/tahoebio/Tahoe-100M/resolve/main/'
files=[(base+'metadata/'+n+'.parquet',n+'.parquet',None) for n in ['sample_metadata','gene_metadata','drug_metadata']]
if a.dataset=='tahoe':
files += [(base+f'data/train-{i:05d}-of-03388.parquet',f'train-{i:05d}-of-03388.parquet',None) for i in (a.indices if a.indices is not None else range(a.shards))]
else:
files += [('https://zenodo.org/records/13350497/files/SrivatsanTrapnell2020_sciplex3.h5ad?download=1','sciplex3.h5ad','c9e70629505d98c7ca1a837f62b14e89')]
for url,name,expected in files:
path=out/name
if path.exists():
if expected and hashlib.file_digest(path.open('rb'),'md5').hexdigest()!=expected:
raise ValueError(f'Checksum mismatch for {name}')
print(f'Using {path}',flush=True);continue
temp=path.with_suffix(path.suffix+'.part')
with s.get(url,stream=True,timeout=(30,180)) as r:
r.raise_for_status()
with temp.open('wb') as f:
for block in r.iter_content(8*1024*1024):f.write(block)
if expected:
with temp.open('rb') as f:digest=hashlib.file_digest(f,'md5').hexdigest()
if digest!=expected:raise ValueError(f'Checksum mismatch for {name}')
temp.replace(path)
print(f'Downloaded {path} ({path.stat().st_size:,} bytes)',flush=True)
|