forward method

List<ValueVector> forward(List<int> idx)

The forward pass for the entire Transformer model.

Takes a list of integer token indices idx and returns logits for the next token.

Implementation

List<ValueVector> forward(List<int> idx) {
  final T = idx.length;
  if (T > blockSize) {
    throw ArgumentError(
        "Input sequence length ($T) exceeds model's block size ($blockSize)");
  }

  // 1. Get token and position embeddings and sum them
  var x = List.generate(T, (t) {
    final tok_emb = tokenEmbeddings[idx[t]];
    final pos_emb = positionEmbeddings[t];
    return tok_emb + pos_emb;
  });

  // 2. Pass sequence through all transformer blocks
  for (final block in blocks) {
    x = block.forward(x);
  }

  // 3. Apply final layer norm
  x = List.generate(T, (t) => finalLayerNorm.forward(x[t]));

  // 4. Language model head to get final logits
  final logits = List.generate(T, (t) => lmHead.forward(x[t]));

  return logits;
}