File size: 4,130 Bytes
97a739f | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """tools/corpus_split.py — Corpus Factory statistics + template-aware split.
Directive (Phase II): split train/test BY TEMPLATE, not by random rows.
A template is the program source with all integer literals masked to '#',
so two programs differing only in constants share a template. Whole
template groups go to one side; equivalence variants follow their primary.
Outputs:
<out>/train.jsonl
<out>/test.jsonl
and prints corpus statistics (categories, token lengths, uniqueness,
equivalence share, mutant modes, template counts, split sizes).
"""
import argparse
import collections
import json
import os
import random
import re
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import kenchat
def mask_literals(src):
return re.sub(r'\b\d+\b', '#', src)
def main():
ap = argparse.ArgumentParser()
ap.add_argument('manifest')
ap.add_argument('--test-frac', type=float, default=0.1)
ap.add_argument('--seed', type=int, default=13)
ap.add_argument('--out', default='minds/corpus_factory/split')
args = ap.parse_args()
recs = []
with open(args.manifest, encoding='utf-8') as f:
for line in f:
line = line.strip()
if line:
recs.append(json.loads(line))
codec = kenchat.load_codec_vocab('minds/kenga_full.pkl')
# ---- statistics ----
per_cat = collections.Counter(r['category'] for r in recs)
lens = []
for r in recs:
lens.append(len(kenchat.tokenize(r['src'], codec)))
lens_sorted = sorted(lens)
n = len(lens_sorted)
med = lens_sorted[n // 2]
p90 = lens_sorted[int(n * 0.9)]
uniq = len(set(r['src'] for r in recs))
with_var = sum(1 for r in recs if r['variants'])
n_var = sum(len(r['variants']) for r in recs)
mut = collections.Counter(m['mode'] for r in recs for m in r['mutants'])
templates = collections.defaultdict(list)
for r in recs:
templates[(r['category'], mask_literals(r['src']))].append(r)
print(f'manifest: {args.manifest}')
print(f'programs: {len(recs)} unique sources: {uniq}')
print(f'per category: ' + ', '.join(f'{c}={per_cat[c]}' for c in sorted(per_cat)))
print(f'token length: min={lens_sorted[0]} median={med} p90={p90} max={lens_sorted[-1]}')
print(f'equivalence: {with_var} programs with variants ({100*with_var/len(recs):.1f}%), '
f'{n_var} variants total')
print(f'mutants: ' + ', '.join(f'{k}={v}' for k, v in sorted(mut.items())))
print(f'templates: {len(templates)} '
f'(avg {len(recs)/len(templates):.1f} programs per template)')
# ---- template-aware split ----
rng = random.Random(args.seed)
train, test = [], []
by_cat = collections.defaultdict(list)
for key, group in templates.items():
by_cat[key[0]].append(group)
for cat, groups in sorted(by_cat.items()):
rng.shuffle(groups)
k_test = max(1, int(len(groups) * args.test_frac))
for g in groups[:k_test]:
test.extend(g)
for g in groups[k_test:]:
train.extend(g)
# leakage check: no template on both sides
tr_t = set((r['category'], mask_literals(r['src'])) for r in train)
te_t = set((r['category'], mask_literals(r['src'])) for r in test)
overlap = tr_t & te_t
os.makedirs(args.out, exist_ok=True)
for name, part in (('train', train), ('test', test)):
with open(os.path.join(args.out, f'{name}.jsonl'), 'w', encoding='utf-8') as f:
for r in part:
f.write(json.dumps(r) + '\n')
tr_cat = collections.Counter(r['category'] for r in train)
te_cat = collections.Counter(r['category'] for r in test)
print(f'\nsplit: train={len(train)} ({len(tr_t)} templates), '
f'test={len(test)} ({len(te_t)} templates)')
print(f' train: ' + ', '.join(f'{c}={tr_cat[c]}' for c in sorted(tr_cat)))
print(f' test: ' + ', '.join(f'{c}={te_cat[c]}' for c in sorted(te_cat)))
print(f'template overlap between train/test: {len(overlap)} (must be 0)')
return 0 if not overlap else 1
if __name__ == '__main__':
sys.exit(main())
|