File size: 3,056 Bytes
b296ad4 | 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 | """Teach exact project-scope and discovery-table copying across identifier formats."""
import argparse
import hashlib
import json
from pathlib import Path
import random
import re
import shutil
import string
from tokenizers import Tokenizer
from sqlglot import Tokenizer as SQLTokenizer
from tinyquery.data import compact,ddls,serialize
def main():
p=argparse.ArgumentParser();p.add_argument('--source',default='data/tinyquery-v4')
p.add_argument('--seed-source',default='data/tinyquery');p.add_argument('--out',default='data/tinyquery-v5')
args=p.parse_args();source=Path(args.source);seeds=Path(args.seed_source);out=Path(args.out);out.mkdir(parents=True,exist_ok=True)
for name in ['validation.jsonl','test.jsonl','manual.jsonl','tokenizer.json']:shutil.copy2(source/name,out/name)
tokenizer=Tokenizer.from_file(str(source/'tokenizer.json'))
chunks=sorted({tokenizer.decode([i]) for i in range(tokenizer.get_vocab_size()) if re.fullmatch('[a-z]{1,8}',tokenizer.decode([i]))})
rng=random.Random(99335);counts={'replay':0,'scope':0,'discovery':0};seen=set()
with (out/'train.jsonl').open('w') as stream:
def emit(row,kind):
key=hashlib.sha256(row['prompt'].encode()).hexdigest()
if key in seen:return
seen.add(key);counts[kind]+=1;stream.write(json.dumps(row,ensure_ascii=False)+'\n')
for line in (source/'train.jsonl').open():emit(json.loads(line),'replay')
for line in (seeds/'train.jsonl').open():
row=json.loads(line)
if row['target']['action']!='call':continue
a=row['target']['arguments'];kind=None
if 'project_id' in a:
style=rng.randrange(3)
project=(''.join(rng.choices(string.ascii_lowercase,k=20)) if style==0 else
''.join(rng.choices(chunks,k=4)) if style==1 else
'-'.join(''.join(rng.choices(string.hexdigits[:16],k=k)) for k in [8,4,4,4,12]))
row['context']['project_id']=project;a['project_id']=project;kind='scope'
if 'table' in a:
old=a['table']
while True:
new=''.join(rng.choices(chunks,k=4))
if len(new)>=6 and new.upper() not in SQLTokenizer.KEYWORDS:break
row['question']=re.sub(r'(?<![A-Za-z0-9_])'+re.escape(old)+r'(?![A-Za-z0-9_])',new,row['question'])
row['slots']['table']=new;a['table']=new
if row['context']['schema']:row['context']['schema']=ddls(row['slots'])
kind='discovery'
if kind:
row['id']+='_scope_99335';row['sample_weight']=4
row['provenance']+='; exact copying of arbitrary project scope and discovery table identifiers'
row['response']=compact(row['target']);row['prompt']=serialize(row['context'],row['question'])
emit(row,kind)
(out/'grounding-stats.json').write_text(json.dumps(counts,indent=2));print(counts)
if __name__=='__main__':main()
|