| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
| from transformers import AutoTokenizer |
|
|
| p = argparse.ArgumentParser() |
| p.add_argument('--weights', required=True) |
| p.add_argument('--root', required=True) |
| a = p.parse_args() |
| root = Path(a.root) |
| tokenizer = AutoTokenizer.from_pretrained(a.weights, local_files_only=True) |
| records = [] |
| provenance = {'dataset': 'Salesforce/wikitext', 'config': 'wikitext-2-raw-v1', 'licenses': ['cc-by-sa-3.0', 'gfdl'], 'sequence_length': 256, 'selection': 'First 100 source rows per official split, contiguous non-overlapping token windows; no chat template or added special tokens', 'sources': {}} |
| for source_split, split, count in [('train', 'calibration', 16), ('validation', 'validation', 8), ('test', 'heldout', 8)]: |
| source = root / (source_split + '-rows.json') |
| raw = source.read_bytes() |
| rows = json.loads(raw)['rows'] |
| text = '\n'.join(row['row']['text'] for row in rows) |
| tokens = tokenizer.encode(text, add_special_tokens=False) |
| if len(tokens) < count * 256: |
| raise ValueError(f'{source_split}: insufficient tokens: {len(tokens)}') |
| provenance['sources'][source_split] = {'sha256': hashlib.sha256(raw).hexdigest(), 'rows': len(rows), 'available_tokens': len(tokens), 'used_tokens': count * 256} |
| for index in range(count): |
| records.append({'id': f'{split}-{index:03d}', 'split': split, 'token_ids': tokens[index * 256:(index + 1) * 256]}) |
| (root / 'corpus.jsonl').write_text(''.join(json.dumps(record) + '\n' for record in records)) |
| smoke = dict(records[0], id='smoke', token_ids=records[0]['token_ids'][:32]) |
| (root / 'smoke.jsonl').write_text(json.dumps(smoke) + '\n') |
| provenance['corpus_sha256'] = hashlib.sha256((root / 'corpus.jsonl').read_bytes()).hexdigest() |
| (root / 'corpus-provenance.json').write_text(json.dumps(provenance, indent=2) + '\n') |
| print(json.dumps(provenance, indent=2)) |
|
|
