| """Learn a byte BPE using training text only; tokenize complete supervised records.""" |
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
| import numpy as np |
| from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders |
|
|
| SPECIAL=['<|pad|>','<|context|>','<|user|>','<|assistant|>','<|end|>'] |
|
|
|
|
| def file_sha256(path): |
| h=hashlib.sha256() |
| with Path(path).open('rb') as stream: |
| for block in iter(lambda:stream.read(4*1024*1024),b''):h.update(block) |
| return h.hexdigest() |
|
|
|
|
| def rows(path): |
| with open(path) as stream: |
| for line in stream: yield json.loads(line) |
|
|
|
|
| def main(): |
| p=argparse.ArgumentParser(); p.add_argument('--data',default='data/tinyquery') |
| p.add_argument('--vocab',type=int,default=24000); p.add_argument('--context',type=int,default=2048) |
| p.add_argument('--reuse-tokenizer',action='store_true') |
| args=p.parse_args(); base=Path(args.data) |
| if args.reuse_tokenizer: |
| tokenizer=Tokenizer.from_file(str(base/'tokenizer.json')) |
| else: |
| tokenizer=Tokenizer(models.BPE()) |
| tokenizer.pre_tokenizer=pre_tokenizers.ByteLevel(add_prefix_space=False) |
| tokenizer.decoder=decoders.ByteLevel() |
| trainer=trainers.BpeTrainer(vocab_size=args.vocab,min_frequency=2,special_tokens=SPECIAL, |
| initial_alphabet=pre_tokenizers.ByteLevel.alphabet()) |
| tokenizer.train_from_iterator((r['prompt']+r['response']+'<|end|>' for r in rows(base/'train.jsonl')),trainer) |
| tokenizer.save(str(base/'tokenizer.json')) |
| stats={'vocab_size':tokenizer.get_vocab_size(),'context':args.context,'special_tokens':{s:tokenizer.token_to_id(s) for s in SPECIAL}} |
| for split in ['train','validation','test']: |
| data=list(rows(base/(split+'.jsonl'))) |
| encoded=[]; boundaries=[]; lengths=[]; actions=[]; kept=[]; rejected=[] |
| for i,r in enumerate(data): |
| prefix=tokenizer.encode(r['prompt']).ids |
| suffix=tokenizer.encode(r['response']+'<|end|>').ids |
| ids=prefix+suffix |
| if len(ids)>args.context+1: |
| rejected.append({'id':r['id'],'tokens':len(ids)}); continue |
| encoded.append(ids); boundaries.append(len(prefix)-1); lengths.append(len(ids)) |
| actions.append({'call':0,'clarify':1,'answer':2}[r['target']['action']]); kept.append(i) |
| width=max(lengths) |
| array=np.zeros((len(encoded),width),dtype=np.uint16) |
| for i,ids in enumerate(encoded): array[i,:len(ids)]=ids |
| np.save(base/(split+'-ids.npy'),array) |
| np.savez(base/(split+'-meta.npz'),boundaries=np.array(boundaries,dtype=np.int32), |
| lengths=np.array(lengths,dtype=np.int32),actions=np.array(actions,dtype=np.int32), |
| kept=np.array(kept,dtype=np.int32),sample_weights=np.array([data[i].get('sample_weight',1.0) for i in kept],dtype=np.float32)) |
| stats[split]={'examples':len(encoded),'tokens':sum(lengths),'response_tokens':sum(l-b-1 for l,b in zip(lengths,boundaries)), |
| 'ids_sha256':file_sha256(base/(split+'-ids.npy')), |
| 'length_p50':float(np.median(lengths)),'length_p95':float(np.quantile(lengths,.95)), |
| 'max_length':max(lengths),'rejected':rejected} |
| print(split,stats[split],flush=True) |
| (base/'tokenization.json').write_text(json.dumps(stats,indent=2)) |
|
|
|
|
| if __name__=='__main__': main() |
|
|
