def load_model

bend-ml-llama@0.2.1.0 · main.bend:1695

import Base
import bend-ml-tensor-array@0.1.19.0/main.bend as TA

1695# the whole model; the embedding gets one more level of bands than the layers (2^5 with 16 threads)
1696def load_model(+c: Cfg, +dir: String) -> IO(Model<c>):
1697  do IO<Model<c>>:
1698    threads : U32 <- IO.thread_count()
1699    +d : Nat = depth_for(threads)
1700    +st : List<&2, U32> <- read_st(dir, Cfg.layers(c))
1701    emb : TA.Bands<Cfg.vocab(c), Cfg.dim(c)> <- get_bands(dir, st, 0n, Nat.add(d, Nat.min(d, 1n)), Cfg.vocab(c), Cfg.dim(c))
1702    ls : List<&1, LW<c>> <- load_layers(Cfg.layers(c), c, dir, st, d, 0n, Nil{})
1703    lnf : List<&2, F32> <- get_norm(dir, st, t_layer(Cfg.layers(c), 0n), Cfg.dim(c))
1704    return Model{d, emb, ls, lnf}