| """Small decoder-only code model; tied embeddings and causal SDPA attention."""
|
| import json
|
| from dataclasses import dataclass,asdict
|
| import torch
|
| from torch import nn
|
| from torch.nn import functional as F
|
|
|
| @dataclass
|
| class Config:
|
| vocab_size:int=8192
|
| context:int=512
|
| width:int=512
|
| layers:int=8
|
| heads:int=8
|
| dropout:float=.1
|
| class Attention(nn.Module):
|
| def __init__(self,c):
|
| super().__init__();self.heads=c.heads;self.dropout=c.dropout
|
| self.qkv=nn.Linear(c.width,c.width*3);self.out=nn.Linear(c.width,c.width)
|
| def forward(self,x,cache=None,use_cache=False):
|
| b,t,d=x.shape;q,k,v=self.qkv(x).chunk(3,dim=-1)
|
| q,k,v=[z.view(b,t,self.heads,d//self.heads).transpose(1,2) for z in (q,k,v)]
|
| mask=None
|
| if cache is not None:
|
| offset=cache[0].shape[2];k=torch.cat((cache[0],k),dim=2);v=torch.cat((cache[1],v),dim=2)
|
| mask=torch.arange(k.shape[2],device=x.device)[None,:]<=torch.arange(offset,offset+t,device=x.device)[:,None]
|
| a=F.scaled_dot_product_attention(q,k,v,attn_mask=mask,is_causal=cache is None,dropout_p=self.dropout if self.training else 0.)
|
| result=self.out(a.transpose(1,2).contiguous().view(b,t,d))
|
| return (result,(k,v)) if use_cache else result
|
| class Block(nn.Module):
|
| def __init__(self,c):
|
| super().__init__();self.ln1=nn.LayerNorm(c.width);self.attn=Attention(c);self.ln2=nn.LayerNorm(c.width)
|
| self.mlp=nn.Sequential(nn.Linear(c.width,c.width*4),nn.GELU(),nn.Linear(c.width*4,c.width),nn.Dropout(c.dropout))
|
| def forward(self,x,cache=None,use_cache=False):
|
| if use_cache:
|
| a,cache=self.attn(self.ln1(x),cache,use_cache=True);return self.finish(x+a),cache
|
| return self.finish(x+self.attn(self.ln1(x)))
|
| def finish(self,x):return x+self.mlp(self.ln2(x))
|
| class CodeSprout(nn.Module):
|
| def __init__(self,c=Config()):
|
| super().__init__();self.config=c;self.tokens=nn.Embedding(c.vocab_size,c.width);self.positions=nn.Embedding(c.context,c.width)
|
| self.blocks=nn.ModuleList([Block(c) for _ in range(c.layers)]);self.ln=nn.LayerNorm(c.width);self.lm_head=nn.Linear(c.width,c.vocab_size,bias=False);self.lm_head.weight=self.tokens.weight
|
| self.apply(self.init)
|
| def init(self,m):
|
| if isinstance(m,(nn.Linear,nn.Embedding)):
|
| nn.init.normal_(m.weight,std=.02)
|
| if isinstance(m,nn.Linear) and m.bias is not None:nn.init.zeros_(m.bias)
|
| def forward(self,ids,targets=None,past=None,use_cache=False):
|
| offset=past[0][0].shape[2] if past else 0
|
| if not torch.jit.is_tracing() and ids.shape[1]+offset>self.config.context:raise ValueError('Prompt exceeds context limit')
|
| x=self.tokens(ids)+self.positions(torch.arange(offset,offset+ids.shape[1],device=ids.device));caches=[]
|
| for i,block in enumerate(self.blocks):
|
| if use_cache:x,c=block(x,past[i] if past else None,use_cache=True);caches.append(c)
|
| else:x=block(x)
|
| logits=self.lm_head(self.ln(x));loss=F.cross_entropy(logits.reshape(-1,logits.size(-1)),targets.reshape(-1)) if targets is not None else None
|
| return (logits,loss,caches) if use_cache else (logits,loss)
|
| @torch.inference_mode()
|
| def generate(self,ids,steps=120,temperature=.7,top_k=30,eos=0):
|
| self.eval()
|
| steps=min(steps,self.config.context-1)
|
| limit=max(1,self.config.context-steps) |
| prefix=ids if ids.shape[1]<=limit else torch.cat((ids[:,:1],ids[:,-(limit-1):]),dim=1) if limit>1 else ids[:,:1] |
| logits,_,cache=self(prefix,use_cache=True) |
| for i in range(steps):
|
| scores=logits[:,-1,:]/max(.05,temperature) |
| if temperature<=0:nxt=scores.argmax(dim=-1,keepdim=True) |
| else: |
| k=min(top_k,scores.shape[-1]);cut=torch.topk(scores,k).values[:,-1,None];scores=scores.masked_fill(scores<cut,float('-inf'));nxt=torch.multinomial(F.softmax(scores,dim=-1),1) |
| ids=torch.cat((ids,nxt),dim=1) |
| if torch.all(nxt==eos):break
|
| if i<steps-1:logits,_,cache=self(nxt,past=cache,use_cache=True)
|
| return ids
|
|
|
