File size: 4,503 Bytes
ca1e261 | 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 | """Bodhan decoder-only MLX prototype. Original checkpoint naming and scaling."""
from pathlib import Path
import json, math, time
import mlx.core as mx
class Decoder:
def __init__(self, root):
root=Path(root)
self.config=json.loads((root/'config.json').read_text())
self.w=mx.load(str(root/'decoder.safetensors'))
mx.eval(self.w)
self.heads=self.config['decoder_attention_heads']
self.dim=self.config['d_model']//self.heads
self.layers=self.config['decoder_layers']
self.attn_scale=math.sqrt(math.sqrt(self.dim))
def linear(self,x,prefix):
weight=self.w[prefix+'.weight'];bias=self.w.get(prefix+'.bias')
return x @ weight.T if bias is None else mx.addmm(bias,x,weight.T)
def norm(self,x,prefix):
return mx.fast.layer_norm(x,self.w[prefix+'.weight'],self.w[prefix+'.bias'],1e-5)
def split(self,x):
return x.reshape(x.shape[0],x.shape[1],self.heads,self.dim).transpose(0,2,1,3)
def project(self,x,prefix):
return (self.split(self.linear(x,prefix+'.key_net'))/self.attn_scale,
self.split(self.linear(x,prefix+'.value_net')))
def attention(self,x,kv,prefix,mask=None):
q=self.split(self.linear(x,prefix+'.query_net'))/self.attn_scale
output=mx.fast.scaled_dot_product_attention(q,*kv,scale=1.0,mask=mask)
output=output.transpose(0,2,1,3).reshape(x.shape[0],x.shape[1],-1)
return self.linear(output,prefix+'.out_projection')
def cross_cache(self,encoder):
result=[self.project(encoder,f'layers.{i}.second_sub_layer') for i in range(self.layers)]
mx.eval(result)
return result
def step(self,ids,position,cross,cache):
positions=mx.arange(position,position+ids.shape[1])
x=self.w['embedding.token_embedding.weight'][ids]+self.w['embedding.position_embedding.pos_enc'][positions][None]
x=self.norm(x,'embedding.layer_norm')
# Only populated cache positions participate in attention.
mask='causal' if ids.shape[1]>1 else None
new_cache=[]
for i in range(self.layers):
p=f'layers.{i}'
n=self.norm(x,p+'.layer_norm_1')
k,v=self.project(n,p+'.first_sub_layer')
if cache is not None:
k=mx.concatenate([cache[i][0],k],axis=2)
v=mx.concatenate([cache[i][1],v],axis=2)
new_cache.append((k,v))
x=x+self.attention(n,(k,v),p+'.first_sub_layer',mask)
x=x+self.attention(self.norm(x,p+'.layer_norm_2'),cross[i],p+'.second_sub_layer')
x=x+self.linear(mx.maximum(self.linear(self.norm(x,p+'.layer_norm_3'),p+'.third_sub_layer.dense_in'),0),p+'.third_sub_layer.dense_out')
return self.linear(self.norm(x,'final_layer_norm'),'lm_head'),new_cache
def generate(self,encoder,tokenizer,language='hi',mixed=False,max_tokens=256):
start=time.perf_counter();cross=self.cross_cache(encoder);cross_seconds=time.perf_counter()-start
start=time.perf_counter()
if language=='auto':
prefix=tokenizer['prompts']['hi'][:3]
logits,_=self.step(mx.array([prefix],dtype=mx.int32),0,cross,None)
scores=logits[0,-1];mx.eval(scores)
language=max(tokenizer['prompts'],key=lambda k:float(scores[tokenizer['prompts'][k][3]].item()))
prompts=tokenizer['mixed_prompts'] if mixed else tokenizer['prompts']
ids=mx.array([prompts[language]],dtype=mx.int32)
position=0;cache=None;tokens=[];prefill_seconds=None
for index in range(max_tokens):
logits,cache=self.step(ids,position,cross,cache)
next_token=mx.argmax(logits[0,-1],axis=-1)
mx.eval(next_token,cache)
next_token=int(next_token.item())
if index==0:prefill_seconds=time.perf_counter()-start
position+=ids.shape[1]
if next_token==tokenizer['eos_id']:break
tokens.append(next_token);ids=mx.array([[next_token]],dtype=mx.int32)
else:raise RuntimeError('No EOS within token limit; refusing truncated transcript')
seconds=time.perf_counter()-start
text=''.join(tokenizer['pieces'][t] for t in tokens if t>=tokenizer['special_count']).replace('▁',' ').strip()
return dict(text=text,language=language,tokens=len(tokens),endedWithEOS=True,crossSeconds=cross_seconds,decodeSeconds=seconds,prefillSeconds=prefill_seconds,generationSeconds=seconds-prefill_seconds)
|