Repository: ollama/ollama
File: model/models/nomicbert/model.go (Go)

Code before:

type Attention struct {
	QKV    *nn.Linear `gguf:"attn_qkv"`
	Output *nn.Linear `gguf:"attn_output"`
}

type MLP struct {
	Gate *nn.Linear `gguf:"ffn_gate"`
	Up   *nn.Linear `gguf:"ffn_up"`
	Down *nn.Linear `gguf:"ffn_down"`
}

// Forward implements model.Model interface
func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
	seqLen := batch.Inputs.Dim(0)

	hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)

	typeEmbed := m.TypeEmbedding.Weight.View(ctx, 0, m.hiddenSize)
	hiddenStates = hiddenStates.Add(ctx, typeEmbed)

	// no position embedding - we use ROPE instead

	hiddenStates = m.TokenEmbeddingNorm.Forward(ctx, hiddenStates, m.eps)

	positions := ctx.Arange(0, float32(seqLen), 1, ml.DTypeI32)

	for _, layer := range m.Layers {
		hiddenStates = layer.Forward(ctx, hiddenStates, positions, &m.Options)
	}

	hiddenStates = m.poolingType.Forward(ctx, hiddenStates)

	if m.normalize {
		hiddenStates = hiddenStates.L2Norm(ctx, 1e-12)
	}

	return hiddenStates, nil
}

func (e *EncoderLayer) Forward(ctx ml.Context, hiddenStates ml.Tensor, positions ml.Tensor, opts *Options) ml.Tensor {
	residual := hiddenStates
	hiddenStates = e.Attention.Forward(ctx, hiddenStates, positions, opts)
	hiddenStates = hiddenStates.Add(ctx, residual)
	hiddenStates = e.AttentionNorm.Forward(ctx, hiddenStates, opts.eps)

	residual = hiddenStates
	hiddenStates = e.MLP.Forward(ctx, hiddenStates)
	hiddenStates = hiddenStates.Add(ctx, residual)
	hiddenStates = e.MLPNorm.Forward(ctx, hiddenStates, opts.eps)


Review comment (performance):
positional ids are already set by the caller

```suggestion
	positions := ctx.Input().FromInts(batch.Positions, len(batch.Positions))
```

Code after:
}

type Attention struct {
	QKV    *nn.Linear `gguf:"attn_qkv"`
	Output *nn.Linear `gguf:"attn_output"`
}

type MLP struct {
	Gate *nn.Linear `gguf:"ffn_gate"`
	Up   *nn.Linear `gguf:"ffn_up"`
	Down *nn.Linear `gguf:"ffn_down"`
}

func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
	hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)

	typeEmbed := m.TypeEmbedding.Weight.Slice(ctx, 1, 0, 1, 1)
	hiddenStates = hiddenStates.Add(ctx, typeEmbed)

	hiddenStates = m.TokenEmbeddingNorm.Forward(ctx, hiddenStates, m.eps)

	positions := ctx.Input().FromInts(batch.Positions, len(batch.Positions))

	for _, layer := range m.Layers {
		hiddenStates = layer.Forward(ctx, hiddenStates, positions, &m.Options)
	}

	hiddenStates = m.poolingType.Forward(ctx, hiddenStates)

	if m.normalize {
		hiddenStates = hiddenStates.L2Norm(ctx, 1e-12)
	}

	return hiddenStates, nil
}

func (e *EncoderLayer) Forward(ctx ml.Context, hiddenStates ml.Tensor, positions ml.Tensor, opts *Options) ml.Tensor {
	residual := hiddenStates
	hiddenStates = e.Attention.Forward(ctx, hiddenStates, positions, opts)
	hiddenStates = hiddenStates.Add(ctx, residual)
	hiddenStates = e.AttentionNorm.Forward(ctx, hiddenStates, opts.eps)

	residual = hiddenStates
	hiddenStates = e.MLP.Forward(ctx, hiddenStates)
	hiddenStates = hiddenStates.Add(ctx, residual)
	hiddenStates = e.MLPNorm.Forward(ctx, hiddenStates, opts.eps)

	return hiddenStates
}

func (a *Attention) Forward(ctx ml.Context, hiddenStates ml.Tensor, positions ml.Tensor, opts *Options) ml.Tensor {