--!native
-- SmolLM2-135M-Instruct forward pass. CPU inference is entirely Luau.
local Encoding = game:GetService('EncodingService')
local Http = game:GetService('HttpService')
local Engine = {}
local D,H,KV,HD,FF,LAYERS,VOCAB = 576,9,3,64,1536,30,49152
local EPS,THETA = 1e-5,100000
local rd,wr,ri = buffer.readf32,buffer.writef32,buffer.readi8
local function vec(n) return buffer.create(n*4) end
local lastYield=os.clock()
local cancelCheck=nil
local function checkpoint(force)
 if force or os.clock()-lastYield>0.012 then
  task.wait();lastYield=os.clock()
  if cancelCheck and cancelCheck() then error('CANCELLED',0) end
 end
end
local function matvec(w,x,out)
 local n,m,b=w.rows,w.cols,w.data
 local so=n*m
 for r=0,n-1 do
  local a=r*m;local s=0
  for g=0,m/64-1 do
   local partial=0;local start=g*64
   for c=start,start+63,4 do
    partial+=ri(b,a+c)*rd(x,c*4)+ri(b,a+c+1)*rd(x,(c+1)*4)+ri(b,a+c+2)*rd(x,(c+2)*4)+ri(b,a+c+3)*rd(x,(c+3)*4)
   end
   s+=partial*rd(b,so+(r*m/64+g)*4)
  end
  wr(out,r*4,s)
  if r%128==127 then checkpoint(false) end
 end
end
local function norm(x,w,out)
 local s=0
 for i=0,D-1 do local v=rd(x,i*4);s+=v*v end
 local inv=1/math.sqrt(s/D+EPS)
 for i=0,D-1 do wr(out,i*4,rd(x,i*4)*inv*rd(w.data,i*4)) end
end
local function add(x,b)
 for i=0,D-1 do wr(x,i*4,rd(x,i*4)+rd(b,i*4)) end
end
function Engine.load(folder,onProgress)
 local manifest=Http:JSONDecode(folder.Manifest.Value)
 local weights={};local names={}
 for name in manifest.tensors do table.insert(names,name) end
 table.sort(names)
 for index,name in names do
  local meta=manifest.tensors[name];local node=folder[meta.key]
  local chunks=node:GetChildren();table.sort(chunks,function(a,b) return a.Name<b.Name end)
  local data=buffer.create(meta.bytes);local off=0
  for _,chunk in chunks do
   local decoded=Encoding:Base64Decode(buffer.fromstring(chunk.Value))
   buffer.copy(data,off,decoded);off+=buffer.len(decoded)
  end
  assert(off==meta.bytes,'Corrupt tensor '..name)
  if meta.kind=='q4g64' then
   local count=meta.shape[1]*meta.shape[2]
   local expanded=buffer.create(count+count/64*4)
   for j=0,count/2-1 do
    local packed=buffer.readu8(data,j)
    buffer.writei8(expanded,j*2,bit32.band(packed,15)-8)
    buffer.writei8(expanded,j*2+1,bit32.rshift(packed,4)-8)
    if j%131072==131071 then checkpoint(false) end
   end
   buffer.copy(expanded,count,data,count/2,count/64*4)
   data=expanded
  end
  weights[name]={data=data,rows=meta.shape[1],cols=meta.shape[2]}
  node:Destroy()
  if onProgress then onProgress(index/#names) end
  checkpoint(true)
 end
 Engine.weights=weights;Engine.manifest=manifest
 Engine.scratch={x=vec(D),n=vec(D),q=vec(D),k=vec(KV*HD),v=vec(KV*HD),a=vec(D),o=vec(D),gate=vec(FF),up=vec(FF),logits=vec(VOCAB)}
 return Engine
end
function Engine.newSession(maxContext)
 local s={position=0,maxContext=maxContext or 256,keys={},values={}}
 for l=1,LAYERS do s.keys[l]=vec(s.maxContext*KV*HD);s.values[l]=vec(s.maxContext*KV*HD) end
 return s
end
function Engine.forward(session,token,project,shouldCancel)
 cancelCheck=shouldCancel;checkpoint(false)
 local pos=session.position
 assert(pos<session.maxContext,'Context limit reached')
 local w=Engine.weights;local t=Engine.scratch;local x=t.x
 local emb=w['model.embed_tokens.weight']
 for i=0,D-1 do
  local scale=rd(emb.data,emb.rows*D+(token*D/64+math.floor(i/64))*4)
  wr(x,i*4,ri(emb.data,token*D+i)*scale)
 end
 for layer=0,LAYERS-1 do
  local p='model.layers.'..layer..'.';local att=p..'self_attn.'
  norm(x,w[p..'input_layernorm.weight'],t.n)
  matvec(w[att..'q_proj.weight'],t.n,t.q)
  matvec(w[att..'k_proj.weight'],t.n,t.k)
  matvec(w[att..'v_proj.weight'],t.n,t.v)
  for h=0,H-1 do
   for j=0,HD/2-1 do
    local angle=pos/(THETA^(2*j/HD));local co,si=math.cos(angle),math.sin(angle)
    local a=(h*HD+j)*4;local b=a+HD/2*4;local v1,v2=rd(t.q,a),rd(t.q,b)
    wr(t.q,a,v1*co-v2*si);wr(t.q,b,v2*co+v1*si)
    if h<KV then
     v1,v2=rd(t.k,a),rd(t.k,b);wr(t.k,a,v1*co-v2*si);wr(t.k,b,v2*co+v1*si)
    end
   end
  end
  local kc,vc=session.keys[layer+1],session.values[layer+1]
  buffer.copy(kc,pos*KV*HD*4,t.k);buffer.copy(vc,pos*KV*HD*4,t.v)
  for h=0,H-1 do
   local kh=math.floor(h/(H/KV));local scores=table.create(pos+1);local peak=-math.huge
   for step=0,pos do
    local base=(step*KV*HD+kh*HD)*4;local qbase=h*HD*4;local dot=0
    for j=0,HD-1 do dot+=rd(t.q,qbase+j*4)*rd(kc,base+j*4) end
    dot/=8;scores[step+1]=dot;peak=math.max(peak,dot)
   end
   local total=0
   for step=1,pos+1 do local e=math.exp(scores[step]-peak);scores[step]=e;total+=e end
   for j=0,HD-1 do
    local sum=0
    for step=0,pos do sum+=scores[step+1]/total*rd(vc,(step*KV*HD+kh*HD+j)*4) end
    wr(t.a,(h*HD+j)*4,sum)
   end
  end
  matvec(w[att..'o_proj.weight'],t.a,t.o);add(x,t.o)
  norm(x,w[p..'post_attention_layernorm.weight'],t.n)
  matvec(w[p..'mlp.gate_proj.weight'],t.n,t.gate)
  matvec(w[p..'mlp.up_proj.weight'],t.n,t.up)
  for i=0,FF-1 do local g=rd(t.gate,i*4);wr(t.gate,i*4,g/(1+math.exp(-g))*rd(t.up,i*4)) end
  matvec(w[p..'mlp.down_proj.weight'],t.gate,t.o);add(x,t.o)
  checkpoint(false)
 end
 session.position+=1
 if not project then return nil end
 norm(x,w['model.norm.weight'],t.n);matvec(emb,t.n,t.logits)
 return t.logits
end
function Engine.sample(logits,temperature,rng,previous)
 local best={};local repetition={}
 for _,id in previous or {} do repetition[id]=true end
 for id=0,VOCAB-1 do
  local score=rd(logits,id*4)
  if repetition[id] then score=if score>0 then score/1.08 else score*1.08 end
  if id==0 or id==1 or (id>=3 and id<=16) then score=-math.huge end
  if #best<32 or score>best[#best][2] then
   local k=#best+1
   while k>1 and best[k-1][2]<score do best[k]=best[k-1];k-=1 end
   best[k]={id,score};if #best>32 then best[33]=nil end
  end
 end
 if temperature==0 then return best[1][1] end
 local sum=0
 for _,v in best do v[3]=math.exp((v[2]-best[1][2])/temperature);sum+=v[3] end
 local z=rng:NextNumber()*sum
 for _,v in best do z-=v[3];if z<=0 then return v[1] end end
 return best[1][1]
end
return Engine