model.go 5.17 KB
Newer Older
Patrick Devine's avatar
Patrick Devine committed
1
2
3
package gemma3

import (
Michael Yang's avatar
Michael Yang committed
4
5
	"bytes"
	"image"
Michael Yang's avatar
Michael Yang committed
6
	"math"
7
	"slices"
Patrick Devine's avatar
Patrick Devine committed
8

9
	"github.com/ollama/ollama/fs"
Patrick Devine's avatar
Patrick Devine committed
10
11
	"github.com/ollama/ollama/kvcache"
	"github.com/ollama/ollama/ml"
Michael Yang's avatar
Michael Yang committed
12
	"github.com/ollama/ollama/ml/nn"
Patrick Devine's avatar
Patrick Devine committed
13
14
15
16
17
18
	"github.com/ollama/ollama/model"
	"github.com/ollama/ollama/model/input"
)

type Model struct {
	model.Base
19
	model.TextProcessor
Patrick Devine's avatar
Patrick Devine committed
20

21
	*VisionModel `gguf:"v"`
Patrick Devine's avatar
Patrick Devine committed
22
23
	*TextModel

Michael Yang's avatar
Michael Yang committed
24
	*MultiModalProjector `gguf:"mm"`
Patrick Devine's avatar
Patrick Devine committed
25
26
27
28

	ImageProcessor
}

Michael Yang's avatar
Michael Yang committed
29
30
31
32
33
var _ model.MultimodalProcessor = (*Model)(nil)

type MultiModalProjector struct {
	SoftEmbNorm     *nn.RMSNorm `gguf:"mm_soft_emb_norm"`
	InputProjection *nn.Linear  `gguf:"mm_input_projection"`
Michael Yang's avatar
Michael Yang committed
34
35

	tokensPerImage int
Michael Yang's avatar
Michael Yang committed
36
37
}

Michael Yang's avatar
Michael Yang committed
38
39
40
41
42
43
44
45
46
47
48
func (p *MultiModalProjector) Forward(ctx ml.Context, visionOutputs ml.Tensor, imageSize, patchSize int, eps float32) ml.Tensor {
	l := visionOutputs.Dim(0)

	visionOutputs = visionOutputs.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx)
	patchesPerImage := imageSize / patchSize
	visionOutputs = visionOutputs.Reshape(ctx, patchesPerImage, patchesPerImage, l)

	kernelSize := patchesPerImage / int(math.Sqrt(float64(p.tokensPerImage)))
	visionOutputs = visionOutputs.AvgPool2D(ctx, kernelSize, kernelSize, 0)
	visionOutputs = visionOutputs.Reshape(ctx, visionOutputs.Dim(0)*visionOutputs.Dim(1), l)
	visionOutputs = visionOutputs.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx)
Michael Yang's avatar
Michael Yang committed
49
50
51
52
53
54
55
	visionOutputs = p.SoftEmbNorm.Forward(ctx, visionOutputs, eps)

	// TODO: inputProjection must be transposed since they're incompatible with visionOutputs
	visionOutputs = p.InputProjection.Weight.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx).Mulmat(ctx, visionOutputs)
	return visionOutputs
}

56
func New(c fs.Config) (model.Model, error) {
57
58
59
60
61
62
63
64
65
66
67
	vocabulary := model.Vocabulary{
		Values: c.Strings("tokenizer.ggml.tokens"),
		Scores: c.Floats("tokenizer.ggml.scores"),
		Types:  c.Ints("tokenizer.ggml.token_type"),
		Merges: c.Strings("tokenizer.ggml.merges"),
		AddBOS: c.Bool("tokenizer.ggml.add_bos_token", true),
		BOS:    []int32{int32(c.Uint("tokenizer.ggml.bos_token_id"))},
		AddEOS: c.Bool("tokenizer.ggml.add_eos_token", false),
		EOS: append(
			[]int32{
				int32(c.Uint("tokenizer.ggml.eos_token_id")),
Patrick Devine's avatar
Patrick Devine committed
68
			},
69
			c.Ints("tokenizer.ggml.eos_token_ids")...,
Patrick Devine's avatar
Patrick Devine committed
70
		),
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
	}

	var processor model.TextProcessor
	switch c.String("tokenizer.ggml.model") {
	case "gpt2":
		processor = model.NewBytePairEncoding(&vocabulary)
	default:
		// Previous uploads of Gemma 3 on Ollama did not have token 106
		// (i.e. "<end_of_turn>") so we need to add in case it's not already present
		vocabulary.EOS = append(vocabulary.EOS, int32(c.Uint("tokenizer.ggml.eot_token_id", 106)))
		processor = model.NewSentencePiece(&vocabulary)
	}

	m := Model{
		TextProcessor:  processor,
Patrick Devine's avatar
Patrick Devine committed
86
		ImageProcessor: newImageProcessor(c),
Michael Yang's avatar
Michael Yang committed
87
88
		VisionModel:    newVisionModel(c),
		TextModel:      newTextModel(c),
Michael Yang's avatar
Michael Yang committed
89
90
91
		MultiModalProjector: &MultiModalProjector{
			tokensPerImage: int(c.Uint("mm_tokens_per_image", 256)),
		},
Patrick Devine's avatar
Patrick Devine committed
92
93
	}

Michael Yang's avatar
Michael Yang committed
94
	slidingWindowLen := int32(c.Uint("attention.sliding_window"))
Patrick Devine's avatar
Patrick Devine committed
95
96
97
98
99
	m.Cache = kvcache.NewWrapperCache(kvcache.NewSWACache(slidingWindowLen, m.Shift), kvcache.NewCausalCache(m.Shift))

	return &m, nil
}

100
func (m *Model) EncodeMultimodal(ctx ml.Context, multimodalData []byte) ([]input.Multimodal, error) {
101
102
103
104
	if len(m.VisionModel.Layers) == 0 {
		return nil, model.ErrNoVisionModel
	}

Michael Yang's avatar
Michael Yang committed
105
106
107
108
109
110
111
112
113
114
	image, _, err := image.Decode(bytes.NewReader(multimodalData))
	if err != nil {
		return nil, err
	}

	f32s, err := m.ImageProcessor.ProcessImage(image)
	if err != nil {
		return nil, err
	}

Michael Yang's avatar
Michael Yang committed
115
	pixelValues := ctx.Input().FromFloats(f32s,
Michael Yang's avatar
Michael Yang committed
116
117
118
119
120
		m.ImageProcessor.imageSize,
		m.ImageProcessor.imageSize,
		m.ImageProcessor.numChannels,
	)

Michael Yang's avatar
Michael Yang committed
121
	visionOutputs := m.VisionModel.Forward(ctx, pixelValues)
Michael Yang's avatar
Michael Yang committed
122
	visionOutputs = m.MultiModalProjector.Forward(ctx, visionOutputs, m.imageSize, m.patchSize, m.VisionModel.eps)
123
	return []input.Multimodal{{Tensor: visionOutputs}}, nil
Michael Yang's avatar
Michael Yang committed
124
125
}

126
127
func (m *Model) PostTokenize(inputs []*input.Input) ([]*input.Input, error) {
	var result []*input.Input
Michael Yang's avatar
Michael Yang committed
128

129
	for _, inp := range inputs {
130
		if len(inp.Multimodal) == 0 {
131
			result = append(result, inp)
Michael Yang's avatar
Michael Yang committed
132
		} else {
133
			inputMultimodal := inp.Multimodal[0].Tensor
134

135
			result = append(result,
136
137
138
				&input.Input{Token: 108, SameBatch: inputMultimodal.Dim(1) + 3}, // "\n\n"
				&input.Input{Token: 255999},                                     // "<start_of_image>""
				&input.Input{Multimodal: []input.Multimodal{{Tensor: inputMultimodal}}, MultimodalHash: inp.MultimodalHash}, // image data is on the first placeholder
139
			)
Michael Yang's avatar
Michael Yang committed
140

141
			// add image token placeholders
142
			result = append(result, slices.Repeat([]*input.Input{{Token: 0}}, inputMultimodal.Dim(1)-1)...)
143

144
			result = append(result,
145
146
				&input.Input{Token: 256000}, // <end_of_image>
				&input.Input{Token: 108},    // "\n\n"
147
			)
Michael Yang's avatar
Michael Yang committed
148
149
		}
	}
Michael Yang's avatar
Michael Yang committed
150

151
	return result, nil
Michael Yang's avatar
Michael Yang committed
152
153
}

Jesse Gross's avatar
Jesse Gross committed
154
func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
155
156
157
158
159
160
161
162
163
164
	hiddenState := m.TextModel.Forward(ctx, batch, m.Cache)
	hiddenState = m.Output.Forward(ctx, hiddenState)

	if m.TextConfig.finalLogitSoftcap > 0.0 {
		hiddenState = hiddenState.Scale(ctx, 1.0/float64(m.TextConfig.finalLogitSoftcap))
		hiddenState = hiddenState.Tanh(ctx)
		hiddenState = hiddenState.Scale(ctx, float64(m.TextConfig.finalLogitSoftcap))
	}

	return hiddenState, nil
Patrick Devine's avatar
Patrick Devine committed
165
166
167
168
}

func init() {
	model.Register("gemma3", New)
Michael Yang's avatar
Michael Yang committed
169
	model.Register("gemma3_embed", newEmbedModel)
Patrick Devine's avatar
Patrick Devine committed
170
}