Add dataset builder script

scripts/build_dataset.py

@@ -0,0 +1,51 @@
+"""Build and push Nomen-AI SFT/DPO datasets."""
+import argparse, random, re
+from datasets import Dataset, DatasetDict
+from nomen_ai.control import ROOT_FAMILIES, THEMES, ControlVector, SYSTEM_PROMPT
+from nomen_ai.synth import make_example
+from nomen_ai.antidup import AntiDuplicationMatrix
+from nomen_ai.phonetics import count_syllables
+GENERIC_SUFFIXES=["ify","ly","ster","hub","ged","io","zy","able"]
+GENERIC_PREFIXES=["Tech","Smart","Get","My","Go","The","Best","Pro"]
+GENERIC_CORES=["brand","name","tube","vlog","media","studio","world","zone"]
+
+def quality_ok(name):
+    return 3<=len(name)<=13 and re.search(r"[aeiouy]",name.lower()) and count_syllables(name)<=4 and not re.search(r"(.)\1\1",name.lower())
+
+def derivative_name(rng):
+    x=rng.random()
+    if x<0.4: return rng.choice(GENERIC_PREFIXES)+rng.choice(GENERIC_CORES).capitalize()
+    if x<0.7: return rng.choice(["Stream","Click","Pixel","Brand","Creat","Snap","Vibe"])+rng.choice(GENERIC_SUFFIXES)
+    return rng.choice(GENERIC_PREFIXES)+rng.choice(GENERIC_SUFFIXES).capitalize()
+
+def sample_cv(rng):
+    k=rng.choices([1,2,3],weights=[5,4,1])[0]
+    return rng.sample(ROOT_FAMILIES,k), rng.choice(THEMES), rng.choice([0.1,0.3,0.5,0.7,0.9])
+
+def build(n_sft,n_dpo,seed=42):
+    rng=random.Random(seed); anti=AntiDuplicationMatrix(min_novelty=0.7); sft=[]; dpo=[]
+    while len(sft)<n_sft:
+        roots,theme,creativity=sample_cv(rng); ex=make_example(roots,theme,creativity,rng,anti)
+        if not ex: continue
+        name,syl,clen=ex
+        if not quality_ok(name): continue
+        cv=ControlVector(roots=roots,theme=theme,syllables=syl,char_len=clen,creativity=creativity)
+        sft.append({'messages':[{'role':'system','content':SYSTEM_PROMPT},{'role':'user','content':cv.to_prompt()},{'role':'assistant','content':name}]})
+    while len(dpo)<n_dpo:
+        roots,theme,creativity=sample_cv(rng); ex=make_example(roots,theme,creativity,rng,anti)
+        if not ex: continue
+        name,syl,clen=ex
+        if not quality_ok(name): continue
+        cv=ControlVector(roots=roots,theme=theme,syllables=syl,char_len=clen,creativity=creativity)
+        dpo.append({'prompt':[{'role':'system','content':SYSTEM_PROMPT},{'role':'user','content':cv.to_prompt()}], 'chosen':[{'role':'assistant','content':name}], 'rejected':[{'role':'assistant','content':derivative_name(rng)}]})
+    return sft,dpo
+
+def split(rows):
+    cut=int(len(rows)*0.97); return DatasetDict(train=Dataset.from_list(rows[:cut]), test=Dataset.from_list(rows[cut:]))
+
+def main():
+    ap=argparse.ArgumentParser(); ap.add_argument('--n_sft',type=int,default=12000); ap.add_argument('--n_dpo',type=int,default=6000); ap.add_argument('--sft_repo',default='krystv/nomen-ai-sft'); ap.add_argument('--dpo_repo',default='krystv/nomen-ai-dpo'); ap.add_argument('--push',action='store_true'); ap.add_argument('--out',default='data'); args=ap.parse_args()
+    sft,dpo=build(args.n_sft,args.n_dpo); print('SFT rows',len(sft),'DPO rows',len(dpo)); sft_ds,dpo_ds=split(sft),split(dpo); sft_ds.save_to_disk(f'{args.out}/sft'); dpo_ds.save_to_disk(f'{args.out}/dpo')
+    print('sample', sft[0]['messages'][1]['content'], '->', sft[0]['messages'][2]['content']); print('dpo', dpo[0]['chosen'][0]['content'], 'vs', dpo[0]['rejected'][0]['content'])
+    if args.push: sft_ds.push_to_hub(args.sft_repo); dpo_ds.push_to_hub(args.dpo_repo)
+if __name__=='__main__': main()