File size: 1,785 Bytes
87a0bf5 | 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 | #!/usr/bin/env python3
"""Prepare the GSM8K JSONL files used by ADSD inference."""
import argparse
import hashlib
import json
from pathlib import Path
import urllib.request
ROOT = Path(__file__).resolve().parents[1]
def validate(content, split, protocol):
if hashlib.sha256(content).hexdigest() != protocol[split + '_sha256']:
raise ValueError(f'{split}.jsonl does not match the configured GSM8K version')
rows = [json.loads(line) for line in content.splitlines() if line.strip()]
expected = {'train': 7473, 'test': 1319}[split]
if len(rows) != expected:
raise ValueError(f'{split}.jsonl: expected {expected} questions, got {len(rows)}')
return len(rows)
def prepare(data_dir):
protocol = json.loads((ROOT / 'configs/protocol.json').read_text())['dataset']
data_dir = Path(data_dir).expanduser().resolve()
data_dir.mkdir(parents=True, exist_ok=True)
for split in ('train', 'test'):
path = data_dir / f'{split}.jsonl'
if path.exists():
count = validate(path.read_bytes(), split, protocol)
else:
url = (f"https://raw.githubusercontent.com/{protocol['source']}/"
f"{protocol['revision']}/grade_school_math/data/{split}.jsonl")
with urllib.request.urlopen(url, timeout=60) as response:
content = response.read()
count = validate(content, split, protocol)
with path.open('xb') as output:
output.write(content)
print(f'{path}: {count} questions, verified')
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--data-dir', type=Path, required=True)
args = parser.parse_args()
prepare(args.data_dir)
if __name__ == '__main__':
main()
|