| """Labelled typed requests and leakage-aware deterministic data splitting.""" |
| import argparse |
| import hashlib |
| import json |
| import random |
| from pathlib import Path |
|
|
| from .decisions import validate_request |
|
|
|
|
| def canonical(value): |
| return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(',', ':'), allow_nan=False) |
|
|
|
|
| def state_hash(row): |
| return hashlib.sha256(canonical(row['state']).encode()).hexdigest() |
|
|
|
|
| def request_of(row): |
| return validate_request({'state': row['state'], 'questions': { |
| key: {k: v for k, v in q.items() if k != 'label'} for key, q in row['questions'].items()}}) |
|
|
|
|
| def option_keys(question): |
| if question['type'] == 'choice': |
| return list(question['criteria']) |
| if question['type'] == 'noul': |
| return ['false', 'true'] |
| return [str(i) for i in range(len(question['criteria']))] |
|
|
|
|
| def target_index(question): |
| label = question['label'] |
| if question['type'] == 'noul': |
| if type(label) is not bool: |
| raise ValueError('Noul label must be boolean') |
| return int(label) |
| if question['type'] == 'score': |
| if type(label) is not int or not 0 <= label < len(question['criteria']): |
| raise ValueError('Score label must be a zero-based level index') |
| return label |
| if not isinstance(label, str) or label not in question['criteria']: |
| raise ValueError('Choice label must name an option') |
| return list(question['criteria']).index(label) |
|
|
|
|
| def load_records(path): |
| records = [] |
| seen = set() |
| for line in Path(path).read_text().splitlines(): |
| if not line.strip(): |
| continue |
| row = json.loads(line) |
| if (not isinstance(row, dict) or set(row) - {'id', 'group', 'domain', 'source', 'state', 'questions'} |
| or not isinstance(row.get('id'), str) or not row['id'] or row['id'] in seen): |
| raise ValueError('Records require unique nonempty IDs and supported fields') |
| for name in ('group', 'domain', 'source'): |
| if name in row and (not isinstance(row[name], str) or not row[name]): |
| raise ValueError('Record metadata must be nonempty strings') |
| request_of(row) |
| for question in row['questions'].values(): |
| if set(question) - {'type', 'instructions', 'criteria', 'label'} or 'label' not in question: |
| raise ValueError('Every question needs a label') |
| target_index(question) |
| seen.add(row['id']) |
| records.append(row) |
| if not records: |
| raise ValueError('Dataset is empty') |
| return records |
|
|
|
|
| def require_disjoint(*partitions): |
| prior_ids, prior_states, prior_groups = set(), set(), set() |
| for records in partitions: |
| ids = {r['id'] for r in records} |
| states = {state_hash(r) for r in records} |
| groups = {r['group'] for r in records if 'group' in r} |
| if ids & prior_ids or states & prior_states or groups & prior_groups: |
| raise ValueError('Partitions overlap by ID, exact state, or supplied group') |
| prior_ids |= ids |
| prior_states |= states |
| prior_groups |= groups |
|
|
|
|
| def split_records(records, seed=42): |
| |
| parent = list(range(len(records))) |
| def find(i): |
| while parent[i] != i: |
| parent[i] = parent[parent[i]] |
| i = parent[i] |
| return i |
| owners = {} |
| for i, row in enumerate(records): |
| keys = [('state', state_hash(row))] |
| if 'group' in row: |
| keys.append(('group', row['group'])) |
| for key in keys: |
| if key in owners: |
| parent[find(i)] = find(owners[key]) |
| owners[key] = i |
| groups = {} |
| for i, row in enumerate(records): |
| groups.setdefault(find(i), []).append(row) |
| values = list(groups.values()) |
| if len(values) < 4: |
| raise ValueError('Need at least four independent groups for train/dev/calibration/test') |
| random.Random(seed).shuffle(values) |
| n = max(1, len(values) // 10) |
| boundaries = (len(values) - 3*n, len(values) - 2*n, len(values) - n, len(values)) |
| result = {} |
| start = 0 |
| for name, stop in zip(('train', 'development', 'calibration', 'test'), boundaries): |
| result[name] = [row for group in values[start:stop] for row in group] |
| start = stop |
| require_disjoint(*result.values()) |
| return result |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--data', type=Path, required=True) |
| parser.add_argument('--out', type=Path, required=True) |
| parser.add_argument('--seed', type=int, default=42) |
| args = parser.parse_args() |
| records = load_records(args.data) |
| partitions = split_records(records, args.seed) |
| args.out.mkdir(parents=True, exist_ok=False) |
| manifest = {'seed': args.seed, 'sourceSha256': hashlib.sha256(args.data.read_bytes()).hexdigest(), |
| 'grouping': 'connected exact states and supplied family groups; not fuzzy decontamination', 'files': {}} |
| for name, rows in partitions.items(): |
| path = args.out / (name + '.jsonl') |
| path.write_text(''.join(canonical(row) + '\n' for row in rows)) |
| manifest['files'][path.name] = {'records': len(rows), 'sha256': hashlib.sha256(path.read_bytes()).hexdigest()} |
| (args.out / 'manifest.json').write_text(json.dumps(manifest, indent=2) + '\n') |
| print(json.dumps(manifest)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
