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

Code before:
	batchSize := hiddenStates.Dim(1)

	qkv := a.QKV.Forward(ctx, hiddenStates)

	query := qkv.View(ctx, 0, opts.hiddenSize, qkv.Stride(0), batchSize).Contiguous(ctx)
	key := qkv.View(ctx, opts.hiddenSize*qkv.Stride(0), opts.hiddenSize, qkv.Stride(0), batchSize).Contiguous(ctx)
	value := qkv.View(ctx, 2*opts.hiddenSize*qkv.Stride(0), opts.hiddenSize, qkv.Stride(0), batchSize).Contiguous(ctx)

	query = query.Reshape(ctx, opts.headDim, opts.numHeads, batchSize)
	key = key.Reshape(ctx, opts.headDim, opts.numHeads, batchSize)
	value = value.Reshape(ctx, opts.headDim, opts.numHeads, batchSize)

	query = fast.RoPE(ctx, query, positions, opts.headDim, opts.ropeFreqBase, 1.0)
	key = fast.RoPE(ctx, key, positions, opts.headDim, opts.ropeFreqBase, 1.0)

	attention := nn.Attention(ctx, query, key, value, 1.0/math.Sqrt(float64(opts.headDim)), nil)

	attention = attention.Reshape(ctx, opts.hiddenSize, batchSize)

	return a.Output.Forward(ctx, attention)
}

func (m *MLP) Forward(ctx ml.Context, hiddenStates ml.Tensor) ml.Tensor {
	gate := m.Gate.Forward(ctx, hiddenStates).SILU(ctx)
	up := m.Up.Forward(ctx, hiddenStates)
	hidden := gate.Mul(ctx, up)

	return m.Down.Forward(ctx, hidden)
}

func New(c fs.Config) (model.Model, error) {
	hiddenSize := int(c.Uint("embedding_length"))
	numHeads := int(c.Uint("attention.head_count"))
	headDim := hiddenSize / numHeads

	processor := model.NewWordPiece(
		&model.Vocabulary{
			Values: c.Strings("tokenizer.ggml.tokens"),
			Scores: c.Floats("tokenizer.ggml.scores"),
			Types:  c.Ints("tokenizer.ggml.token_type"),
			AddBOS: c.Bool("tokenizer.ggml.add_bos_token", true),
			BOS: []int32{
				int32(cmp.Or(
					c.Uint("tokenizer.ggml.cls_token_id"),
					c.Uint("tokenizer.ggml.bos_token_id"),
				)),
			},
			AddEOS: c.Bool("tokenizer.ggml.add_eos_token", true),
			EOS: []int32{
				int32(cmp.Or(
					c.Uint("tokenizer.ggml.separator_token_id"),

Review comment (suggestion):
this model doesn't want its inputs lowercased. there's a todo in `WordPiece.Encode(...)` to handle this case. you should make that change as well

Code after:
	qkv = qkv.Reshape(ctx, opts.headDim, opts.numHeads*3, batchSize)
	chunks := qkv.Chunk(ctx, 1, opts.numHeads)
	query, key, value := chunks[0], chunks[1], chunks[2]

	query = fast.RoPE(ctx, query, positions, opts.headDim, opts.ropeFreqBase, 1.0, rope.WithTypeNeoX())
	key = fast.RoPE(ctx, key, positions, opts.headDim, opts.ropeFreqBase, 1.0, rope.WithTypeNeoX())

	attention := nn.Attention(ctx, query, key, value, 1.0/math.Sqrt(float64(opts.headDim)), nil)

	attention = attention.Reshape(ctx, opts.hiddenSize, batchSize)

	return a.Output.Forward(ctx, attention)
}

func (m *MLP) Forward(ctx ml.Context, hiddenStates ml.Tensor) ml.Tensor {
	hidden := m.Gate.Forward(ctx, hiddenStates).SILU(ctx, m.Up.Forward(ctx, hiddenStates))

	return m.Down.Forward(ctx, hidden)
}

func New(c fs.Config) (model.Model, error) {
	hiddenSize := int(c.Uint("embedding_length"))
	numHeads := int(c.Uint("attention.head_count"))
	headDim := hiddenSize / numHeads

	processor := model.NewWordPiece(
		&model.Vocabulary{
			Values: c.Strings("tokenizer.ggml.tokens"),
			Scores: c.Floats("tokenizer.ggml.scores"),
			Types:  c.Ints("tokenizer.ggml.token_type"),
			AddBOS: c.Bool("tokenizer.ggml.add_bos_token", true),
			BOS: []int32{
				int32(cmp.Or(
					c.Uint("tokenizer.ggml.cls_token_id"),
					c.Uint("tokenizer.ggml.bos_token_id"),
				)),
			},
			AddEOS: c.Bool("tokenizer.ggml.add_eos_token", true),
			EOS: []int32{
				int32(cmp.Or(
					c.Uint("tokenizer.ggml.separator_token_id"),
					c.Uint("tokenizer.ggml.eos_token_id"),
				)),
			},
		},
		false,
	)

	return &Model{
		TextProcessor: processor,
		Layers:        make([]EncoderLayer, c.Uint("block_count")),