File size: 6,259 Bytes
c83c617 | 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 113 114 115 116 117 118 119 120 121 122 123 124 125 | """Frozen V4.1 virtual union; deterministic maps can be rebuilt on AutoDL."""
import json
from pathlib import Path
import numpy as np
from vimeml.training.data import file_sha, write_json
from vimeml.training.data_v2 import PrefixCropWindowDataset
from vimeml.training.data_mixture import JoinedPrefixWindows
from vimeml.posttraining.streams import ConversionWindows
ROOT=Path(__file__).resolve().parents[3]
def window_sentences(dataset, ids):
position=np.searchsorted(dataset.end,ids,side='right')
previous=np.zeros(len(ids),dtype=np.int64)
present=position>0
previous[present]=dataset.end[position[present]-1]-dataset.sentences[position[present]-1]-1
sentences=ids-previous
candidates=np.flatnonzero(position<len(dataset.end))
extra=candidates[ids[candidates]>=dataset.first[position[candidates]]]
sentences[extra]=dataset.sentences[position[extra]]
return sentences
def prepare_maps(directory, root=ROOT):
directory,root=Path(directory),Path(root)
request=json.loads((directory/'request.json').read_text('utf-8'))
expected=json.loads((directory/'manifest.json').read_text('utf-8')) if (directory/'manifest.json').exists() else None
if expected and all((directory/r['file']).exists() and file_sha(directory/r['file'])==r['sha256'] for r in expected['epochs']):
return expected
keep_path=root/request['old_keep']
if file_sha(keep_path)!=request['old_keep_sha256']:
raise ValueError('Old training exclusion mask changed')
keep=np.load(keep_path,mmap_mode='r',allow_pickle=False)
bases=[]
for record in request['joined_stores']:
if file_sha(root/record['token_dir']/'manifest.json')!=record['token_manifest_sha256'] or file_sha(root/record['index_dir']/'manifest.json')!=record['window_manifest_sha256']:
raise ValueError('Joined bank identity changed')
bases.append(PrefixCropWindowDataset(root/record['token_dir'],root/record['index_dir'],
seed=request['seed'],probability=request['crop_probability'],min_remaining_tokens=request['crop_minimum']))
joined=JoinedPrefixWindows(bases,0,[{int(k):v for k,v in r['source_map'].items()} for r in request['joined_stores']])
try:
chunks=[]
for start in range(0,len(bases[0]),262144):
ids=np.arange(start,min(len(bases[0]),start+262144),dtype=np.int64)
chunks.append(ids[keep[window_sentences(bases[0],ids)]])
web=np.concatenate(chunks)
chat=np.arange(len(bases[0]),len(joined),dtype=np.int64)
records=[]
for epoch in range(request['epochs']):
joined.set_epoch(epoch)
web_tokens=0
for start in range(0,len(web),262144):
web_tokens+=int(joined.window_lengths(web[start:start+262144]).sum())
chat_lengths=joined.window_lengths(chat)
target=round(web_tokens*request['chat_fraction']/(1-request['chat_fraction']))
rng=np.random.default_rng(request['seed']+epoch)
chosen,used=[],0
while used<target:
order=rng.permutation(len(chat))
cumulative=np.cumsum(chat_lengths[order],dtype=np.int64)
take=int(np.searchsorted(cumulative,target-used,side='right'))
if take:
chosen.append(chat[order[:take]])
used+=int(cumulative[take-1])
if take<len(order):
break
chosen=np.concatenate(chosen)
mapping=np.concatenate((web,chosen)).astype('<u4')
rng.shuffle(mapping)
path=directory/f'epoch-{epoch}.npy'
temp=path.with_suffix('.tmp.npy')
np.save(temp,mapping,allow_pickle=False)
temp.replace(path)
record={'epoch':epoch,'file':path.name,'bytes':path.stat().st_size,'sha256':file_sha(path),
'windows':len(mapping),'web_windows':len(web),'chat_windows':len(chosen),
'web_tokens':web_tokens,'chat_tokens':used,'effective_tokens':web_tokens+used,
'chat_fraction':used/(web_tokens+used),'chat_effective_passes':used/int(chat_lengths.sum())}
if expected and record!=expected['epochs'][epoch]:
raise ValueError(f'Rebuilt map differs from frozen epoch {epoch}')
records.append(record)
print(json.dumps({'map_epoch':epoch+1,'windows':len(mapping),'chat_fraction':record['chat_fraction']}),flush=True)
result={'format':'vimeml_token_mixture_v1','status':'complete','request':request,'epochs':records,
'policy':'Each retained V2.1/V4 LM window once per epoch; chat-only repeated shuffled cycles at 5% actual cropped LM targets; maps SHA256-pinned'}
if expected and result!=expected:
raise ValueError('Rebuilt mixture manifest differs')
# Preserve the frozen bytes across Windows/Linux newline conventions.
# Rebuilding maps must not rewrite an existing manifest with a new SHA.
if expected is None:
write_json(directory/'manifest.json',result)
return result
finally:
joined.close()
def conversion_quota(cursor, updates, examples):
if not 0<=cursor<updates or updates<1 or examples<updates:
raise ValueError('Conversion epoch cannot be spread across these updates')
return ((cursor+1)*examples)//updates-(cursor*examples)//updates
class IndexedConversionWindows(ConversionWindows):
def __init__(self,data,index_path):
super().__init__(data)
self.index_path=Path(index_path)
self.indices=np.load(self.index_path,mmap_mode='r',allow_pickle=False)
if self.indices.dtype!=np.dtype('<u4') or self.indices.ndim!=1 or not len(self.indices):
raise ValueError('Invalid frozen conversion subset')
def __len__(self):
return len(self.indices)
def __getitem__(self,index):
return self.data[int(self.indices[int(index)])]
def window_lengths(self,indices):
return super().window_lengths(self.indices[np.asarray(indices,dtype=np.int64)])
def __getstate__(self):
return {'data':self.data,'index_path':self.index_path}
def __setstate__(self,state):
self.__init__(**state)
|