Hergebruik keys en values zonder voorspellingen te veranderen

import numpy as np

def masked_softmax(scores,allowed):
    scores=np.asarray(scores,float)
    allowed=np.asarray(allowed)
    if scores.ndim<2 or scores.shape[-1]==0 or allowed.dtype.kind!='b' or not np.isfinite(scores).all():
        raise ValueError('Gebruik eindige scores en een booleaans masker')
    allowed=np.broadcast_to(allowed,scores.shape)
    masked=np.where(allowed,scores,-np.inf)
    maximum=np.max(masked,axis=-1,keepdims=True)
    maximum=np.where(np.isfinite(maximum),maximum,0.)
    weights=np.exp(masked-maximum)
    total=weights.sum(axis=-1,keepdims=True)
    return np.divide(weights,total,out=np.zeros_like(weights),where=total>0)

def attention(q,k,v,allowed):
    q,k,v=(np.asarray(value,float) for value in (q,k,v))
    if min(q.ndim,k.ndim,v.ndim)<2 or q.shape[:-2]!=k.shape[:-2] or q.shape[:-2]!=v.shape[:-2] or q.shape[-1]!=k.shape[-1] or k.shape[-2]!=v.shape[-2] or q.shape[-1]==0:
        raise ValueError('Vormen van Q, K en V passen niet')
    scores=q@np.swapaxes(k,-2,-1)/np.sqrt(q.shape[-1])
    weights=masked_softmax(scores,allowed)
    return weights@v,{'q':q,'k':k,'v':v,'weights':weights}

def split_heads(x,heads):
    if x.ndim!=3 or type(heads) is not int or heads<=0 or x.shape[-1]%heads:
        raise ValueError('Aantal heads moet de modelbreedte delen')
    b,t,c=x.shape
    return x.reshape(b,t,heads,c//heads).transpose(0,2,1,3)

def merge_heads(x):
    if x.ndim!=4:
        raise ValueError('Verwacht (batch,heads,tijd,headbreedte)')
    b,h,t,d=x.shape
    return x.transpose(0,2,1,3).reshape(b,t,h*d)

def multihead(x,params,heads,allowed):
    x=np.asarray(x,float)
    params={name:np.asarray(value,float) for name,value in params.items()}
    q,k,v=(split_heads(x@params[name],heads) for name in ['wq','wk','wv'])
    out,acache=attention(q,k,v,allowed)
    joined=merge_heads(out)
    return joined@params['wo'],{'x':x,'heads':heads,'attention':acache,'joined':joined}

def layernorm_forward(x,gamma,beta,epsilon=1e-5):
    x=np.asarray(x,float);gamma=np.asarray(gamma,float);beta=np.asarray(beta,float)
    inv=1/np.sqrt(x.var(axis=-1,keepdims=True)+epsilon)
    norm=(x-x.mean(axis=-1,keepdims=True))*inv
    return norm*gamma+beta,{'norm':norm,'inv':inv,'gamma':gamma}

def feedforward(x,params):
    x=np.asarray(x,float)
    params={name:np.asarray(value,float) for name,value in params.items()}
    z=x@params['w1']+params['b1']
    hidden=np.maximum(z,0)
    return hidden@params['w2']+params['b2'],{'x':x,'z':z,'hidden':hidden}

def transformer_block(x,params,heads,allowed):
    x=np.asarray(x,float)
    n1,c1=layernorm_forward(x,params['g1'],params['b1n'])
    a,ca=multihead(n1,params,heads,allowed)
    residual=x+a
    n2,c2=layernorm_forward(residual,params['g2'],params['b2n'])
    f,cf=feedforward(n2,params)
    return residual+f,{'norm1':c1,'attention':ca,'residual':residual,'norm2':c2,'ff':cf}

def init_block(width=4,seed=7):
    rng=np.random.default_rng(seed)
    params={name:rng.normal(0,.2,(width,width)) for name in ['wq','wk','wv','wo']}
    params.update({'g1':np.ones(width),'b1n':np.zeros(width),'g2':np.ones(width),'b2n':np.zeros(width),
                   'w1':rng.normal(0,.2,(width,2*width)),'b1':rng.normal(0,.2,2*width),
                   'w2':rng.normal(0,.2,(2*width,width)),'b2':np.zeros(width)})
    return params

def init_gpt(vocab_size,context=16,width=24,seed=9):
    rng=np.random.default_rng(seed)
    params=init_block(width,seed)
    params.update({'tokens':rng.normal(0,.1,(vocab_size,width)),
                   'positions':rng.normal(0,.1,(context,width)),
                   'gf':np.ones(width),'bf':np.zeros(width),
                   'head':rng.normal(0,.1,(width,vocab_size)),'bias':np.zeros(vocab_size)})
    return params

def gpt_forward(ids,params,heads=4):
    params={name:np.asarray(value,float) for name,value in params.items()}
    ids=np.asarray(ids)
    if ids.ndim!=2 or ids.dtype.kind not in 'iu' or ids.size==0 or ids.shape[1]>len(params['positions']) or ids.min()<0 or ids.max()>=len(params['tokens']):
        raise ValueError('Ongeldige tokenreeks of contextlengte')
    x=params['tokens'][ids]+params['positions'][:ids.shape[1]]
    allowed=np.tril(np.ones((ids.shape[1],ids.shape[1]),bool))
    h,bc=transformer_block(x,params,heads,allowed)
    n,nc=layernorm_forward(h,params['gf'],params['bf'])
    return n@params['head']+params['bias'],{'ids':ids,'block':bc,'finalnorm':nc,'n':n}

def kv_bytes(layers,batch,tokens,kv_heads,head_dim,bytes_per_value):
    values=[layers,batch,tokens,kv_heads,head_dim,bytes_per_value]
    if any(type(v) is not int for v in values) or tokens<0 or min(layers,batch,kv_heads,head_dim,bytes_per_value)<=0:
        raise ValueError('Ongeldige cacheafmetingen')
    return 2*layers*batch*tokens*kv_heads*head_dim*bytes_per_value

def prefill(ids,params,heads=4):
    ids=np.asarray(ids)
    if ids.ndim!=1 or ids.size==0:
        raise ValueError('Gebruik een niet-lege tokenlijst')
    logits,cache=gpt_forward(ids[None,:],params,heads)
    acache=cache['block']['attention']['attention']
    return logits[0,-1].copy(),{'k':acache['k'].copy(),'v':acache['v'].copy(),'length':len(ids)}

def cached_step(token,params,cache=None,heads=4):
    params={name:np.asarray(value,float) for name,value in params.items()}
    width=params['tokens'].shape[1]
    if type(heads) is not int or heads<=0 or width%heads or type(token) is not int or not 0<=token<len(params['tokens']):
        raise ValueError('Ongeldige token of heads')
    if cache is None:
        cache={'k':np.zeros((1,heads,0,width//heads)),'v':np.zeros((1,heads,0,width//heads)),'length':0}
    if not isinstance(cache,dict) or set(cache)!={'k','v','length'} or type(cache['length']) is not int:
        raise ValueError('Ongeldig cachecontract')
    length=cache['length']
    expected=(1,heads,length,width//heads)
    if length<0 or length>=len(params['positions']) or any(not isinstance(cache[k],np.ndarray) or cache[k].shape!=expected or not np.isfinite(cache[k]).all() for k in ['k','v']):
        raise ValueError('Cachevorm of contextlengte ongeldig')
    x=(params['tokens'][token]+params['positions'][length])[None,None,:]
    n1,_=layernorm_forward(x,params['g1'],params['b1n'])
    q,k,v=(split_heads(n1@params[name],heads) for name in ['wq','wk','wv'])
    keys=np.concatenate([cache['k'],k],axis=2);values=np.concatenate([cache['v'],v],axis=2)
    mixed,_=attention(q,keys,values,np.ones((1,length+1),bool))
    residual=x+merge_heads(mixed)@params['wo']
    n2,_=layernorm_forward(residual,params['g2'],params['b2n'])
    transformed=residual+feedforward(n2,params)[0]
    final,_=layernorm_forward(transformed,params['gf'],params['bf'])
    logits=final@params['head']+params['bias']
    return logits[0,0],{'k':keys,'v':values,'length':length+1}

Voer de code uit en verklaar voor jezelf waarom elke stap werkt.