PythiaModel constructor

PythiaModel(PythiaConfig config)

Implementation

PythiaModel(this.config)
  : embedIn = Embedding(
      config.vocabSize,
      config.embedDim,
      device: config.device,
      seed: config.seed,
    ),
    finalLn = LayerNorm(config.embedDim, device: config.device),
    embedOut = Linear(
      config.embedDim,
      config.vocabSize,
      bias: false,
      device: config.device,
      seed: config.seed + 900000,
    ),
    rope = RopeCache(
      maxCtx: config.maxCtx,
      headDim: config.embedDim ~/ config.numHeads,
      rotaryDim: _computeRotaryDim(
        config.embedDim ~/ config.numHeads,
        config.rotaryPct,
      ),
      base: config.ropeBase,
      device: config.device,
    ),
    blocks = <PythiaBlock>[] {
  // Rope must be built before blocks so we can hand it to each MHA.
  for (int i = 0; i < config.numLayers; i++) {
    blocks.add(
      PythiaBlock(
        config.embedDim,
        config.numHeads,
        ffnDim: config.ffnDim,
        rope: rope,
        device: config.device,
        seed: config.seed + 100000 + i * 1000,
      ),
    );
  }
}