| """B93Q1 self-contained, JSON-safe mixed-precision model storage. |
| |
| 93-level and 8649-level symmetric weight quantisation. Buffers and quantisation |
| scales are lossless IEEE float32 carried as five Base93 digits. The alphabet is |
| the user's exact Scratch alphabet. This is weight storage, not an integer-only |
| inference engine. |
| """ |
| from pathlib import Path |
| import json |
| import numpy as np |
|
|
| ALPHABET=''.join(chr(i) for i in range(32,127) if i not in (34,92)) |
| LUT=np.frombuffer(ALPHABET.encode('ascii'),dtype=np.uint8) |
| INV=np.full(128,-1,dtype=np.int16) |
| INV[LUT]=np.arange(93,dtype=np.int16) |
| SEP='\\' |
| MAGIC='B93Q1' |
|
|
| def encode_uint(values,digits): |
| values=np.asarray(values,dtype=np.uint64).reshape(-1).copy() |
| if np.any(values>=93**digits):raise ValueError('Integer does not fit') |
| out=np.empty((len(values),digits),dtype=np.uint8) |
| for i in range(digits-1,-1,-1): |
| out[:,i]=LUT[values%93];values//=93 |
| return out.tobytes().decode('ascii') |
|
|
| def decode_uint(text,digits): |
| if digits<1 or len(text)%digits:raise ValueError('Invalid digit count') |
| chars=np.frombuffer(text.encode('ascii'),dtype=np.uint8) |
| if np.any(chars>=128):raise ValueError('Non-ASCII digit') |
| codes=INV[chars] |
| if np.any(codes<0):raise ValueError('Invalid Base93 digit') |
| codes=codes.reshape(-1,digits).astype(np.uint64) |
| values=np.zeros(len(codes),dtype=np.uint64) |
| for i in range(digits):values=values*93+codes[:,i] |
| return values |
|
|
| def pack_f32(values): |
| a=np.asarray(values,dtype='<f4').reshape(-1) |
| if not np.isfinite(a).all():raise ValueError('Non-finite float') |
| return encode_uint(a.view('<u4'),5) |
|
|
| def unpack_f32(text): |
| bits=decode_uint(text,5) |
| if np.any(bits>0xffffffff):raise ValueError('Invalid float32 bits') |
| a=bits.astype('<u4').view('<f4') |
| if not np.isfinite(a).all():raise ValueError('Non-finite float') |
| return a |
|
|
| def row_shape(shape): |
| return (shape[0],int(np.prod(shape[1:]))) if len(shape)>=2 else (1,int(np.prod(shape))) |
|
|
| def quantize(values,digits,group_size): |
| a=np.asarray(values,dtype=np.float32) |
| if digits not in (1,2):raise ValueError('Weight precision must be 1 or 2') |
| rows,width=row_shape(a.shape);groups=(width+group_size-1)//group_size |
| padded=np.zeros((rows,groups*group_size),dtype=np.float32) |
| padded[:,:width]=a.reshape(rows,width) |
| blocks=padded.reshape(rows,groups,group_size) |
| qmax=(93**digits-1)//2 |
| scales=np.max(np.abs(blocks),axis=-1)/qmax |
| scales=np.where(scales>0,scales,1).astype(np.float32) |
| q=np.rint(blocks/scales[...,None]).clip(-qmax,qmax).astype(np.int32) |
| codes=(q.reshape(rows,-1)[:,:width]+qmax).astype(np.uint16).reshape(-1) |
| restored=(q*scales[...,None]).reshape(rows,-1)[:,:width].reshape(a.shape).astype(np.float32) |
| return restored,scales.reshape(-1),codes |
|
|
| def dequantize(codes,scales,shape,digits,group_size): |
| rows,width=row_shape(shape);groups=(width+group_size-1)//group_size |
| if len(codes)!=rows*width or len(scales)!=rows*groups:raise ValueError('Payload shape mismatch') |
| scales=np.asarray(scales).reshape(rows,groups) |
| per_value=scales[:,np.arange(width)//group_size] |
| qmax=(93**digits-1)//2 |
| signed=np.asarray(codes,dtype=np.int32).reshape(rows,width)-qmax |
| return (signed*per_value).astype(np.float32).reshape(shape) |
|
|
| def encode_state(state,parameter_names,config,provenance,plan): |
| fields=[MAGIC,ALPHABET,json.dumps(config,separators=(',',':'),ensure_ascii=True), |
| json.dumps(provenance,separators=(',',':'),ensure_ascii=True),str(len(state))] |
| records=[];position=sum(len(f)+1 for f in fields) |
| for name,t in state.items(): |
| a=t.detach().cpu().numpy() if hasattr(t,'detach') else np.asarray(t) |
| shape=list(a.shape);start=position |
| if name in parameter_names: |
| digits,group_size=plan[name] |
| restored,scales,codes=quantize(a,digits,group_size) |
| kind='q';scale_text=pack_f32(scales);payload=encode_uint(codes,digits) |
| elif a.dtype.kind=='f': |
| kind='f';digits=5;group_size=0;scale_text='';payload=pack_f32(a) |
| elif a.dtype.kind in 'iu': |
| kind='i';digits=0;group_size=0;scale_text='';payload=','.join(map(str,a.reshape(-1))) |
| else:raise ValueError('Unsupported dtype '+str(a.dtype)) |
| head=[name,','.join(map(str,shape)),kind,str(digits),str(group_size)] |
| scale_start=position+sum(len(f)+1 for f in head) |
| payload_start=scale_start+len(scale_text)+1 |
| record=head+[scale_text,payload] |
| if any(SEP in x for x in record):raise ValueError('Separator inside field') |
| fields+=record;position+=sum(len(f)+1 for f in record) |
| records.append({'name':name,'shape':shape,'kind':kind,'chars_per_value':digits, |
| 'group_size':group_size,'num_values':int(a.size),'scale_offset':scale_start, |
| 'scale_count':len(scale_text)//5,'payload_offset':payload_start, |
| 'payload_length':len(payload),'record_offset':start}) |
| text=SEP.join(fields)+SEP |
| manifest={'format':MAGIC,'alphabet':ALPHABET,'offset_base':0,'config':config, |
| 'provenance':provenance,'text_bytes':len(text.encode('ascii')), |
| 'json_string_bytes':len(json.dumps(text,ensure_ascii=True).encode('ascii')), |
| 'tensors':records} |
| return text,manifest |
|
|
| def decode_state(text): |
| fields=text.split(SEP) |
| if len(fields)<6 or fields[0]!=MAGIC or fields[1]!=ALPHABET or fields[-1]!='': |
| raise ValueError('Invalid B93Q1 header or terminator') |
| config=json.loads(fields[2]);provenance=json.loads(fields[3]);n=int(fields[4]) |
| if len(fields)!=6+n*7:raise ValueError('Incorrect field count') |
| result={} |
| for i in range(n): |
| name,shape,kind,digits,group_size,scale_text,payload=fields[5+i*7:12+i*7] |
| if name in result:raise ValueError('Duplicate tensor') |
| shape=tuple(int(x) for x in shape.split(',')) if shape else () |
| digits=int(digits);group_size=int(group_size) |
| if kind=='q': |
| if digits not in (1,2) or group_size<1:raise ValueError('Invalid quantisation') |
| a=dequantize(decode_uint(payload,digits),unpack_f32(scale_text),shape,digits,group_size) |
| elif kind=='f':a=unpack_f32(payload).reshape(shape) |
| elif kind=='i':a=np.array([int(x) for x in payload.split(',')],dtype=np.int64).reshape(shape) |
| else:raise ValueError('Unknown tensor encoding') |
| result[name]=a |
| return result,config,provenance |
|
|
| def load_model(path,device='cpu'): |
| import torch |
| from semantic_model import SemanticColorizer |
| state,config,_=decode_state(Path(path).read_text(encoding='ascii')) |
| model=SemanticColorizer(**{k:config[k] for k in ('head','width','queries')}) |
| model.load_state_dict({k:torch.from_numpy(v.copy()) for k,v in state.items()},strict=True) |
| return model.to(device).eval() |
|
|
| if __name__=='__main__': |
| import argparse |
| p=argparse.ArgumentParser();p.add_argument('text');p.add_argument('--output',default='decoded_model') |
| args=p.parse_args() |
| from safetensors.torch import save_file |
| import torch |
| arrays,config,provenance=decode_state(Path(args.text).read_text(encoding='ascii')) |
| dest=Path(args.output);dest.mkdir(exist_ok=True,parents=True) |
| save_file({k:torch.from_numpy(v.copy()) for k,v in arrays.items()},str(dest/'model.safetensors')) |
| (dest/'config.json').write_text(json.dumps(config,indent=2)) |
| print('Decoded',len(arrays),'tensors to',dest) |
|
|
