| """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() |
|
|
