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