| """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) |
|
|
