fixed pretokenizer

convert/convert.go

@@ -200,6 +200,8 @@ func ConvertModel(fsys fs.FS, f *os.File) error {
 		conv = &qwen25VLModel{}
 	case "Qwen3VLForConditionalGeneration", "Qwen3VLMoeForConditionalGeneration":
 		conv = &qwen3VLModel{}
+	case "OLMo2ForCausalLM", "Olmo2ForCausalLM", "OLMo3ForCausalLM", "Olmo3ForCausalLM":
+		conv = &olmoModel{}
 	case "BertModel":
 		conv = &bertModel{}
 	case "CohereForCausalLM":

convert/convert_olmo.go

@@ -0,0 +1,94 @@
+package convert
+
+import (
+	"cmp"
+
+	"github.com/ollama/ollama/fs/ggml"
+)
+
+type olmoModel struct {
+	ModelParameters
+
+	HiddenSize            uint32   `json:"hidden_size"`
+	NumHiddenLayers       uint32   `json:"num_hidden_layers"`
+	IntermediateSize      uint32   `json:"intermediate_size"`
+	NumAttentionHeads     uint32   `json:"num_attention_heads"`
+	NumKeyValueHeads      uint32   `json:"num_key_value_heads"`
+	MaxPositionEmbeddings uint32   `json:"max_position_embeddings"`
+	RMSNormEPS            float32  `json:"rms_norm_eps"`
+	RopeTheta             float32  `json:"rope_theta"`
+	ClampKQV              float32  `json:"f_clamp_kqv"`
+	SlidingWindow         uint32   `json:"sliding_window"`
+	LayerTypes            []string `json:"layer_types"`
+}
+
+var _ ModelConverter = (*olmoModel)(nil)
+
+func (p *olmoModel) KV(t *Tokenizer) ggml.KV {
+	kv := p.ModelParameters.KV(t)
+	kv["general.architecture"] = "olmo"
+	kv["olmo.block_count"] = p.NumHiddenLayers
+	kv["olmo.context_length"] = p.MaxPositionEmbeddings
+	kv["olmo.embedding_length"] = p.HiddenSize
+	kv["olmo.feed_forward_length"] = p.IntermediateSize
+	kv["olmo.attention.head_count"] = p.NumAttentionHeads
+	kv["olmo.attention.head_count_kv"] = cmp.Or(p.NumKeyValueHeads, p.NumAttentionHeads)
+
+	if p.RopeTheta > 0 {
+		kv["olmo.rope.freq_base"] = p.RopeTheta
+	} else {
+		kv["olmo.rope.freq_base"] = float32(10000.0)
+	}
+
+	if p.RMSNormEPS > 0 {
+		kv["olmo.attention.layer_norm_rms_epsilon"] = p.RMSNormEPS
+	}
+
+	if p.ClampKQV > 0 {
+		kv["olmo.attention.clamp_kqv"] = p.ClampKQV
+	}
+
+	if p.SlidingWindow > 0 {
+		kv["olmo.attention.sliding_window"] = p.SlidingWindow
+	}
+
+	if len(p.LayerTypes) > 0 {
+		kv["olmo.attention.layer_types"] = p.LayerTypes
+	}
+
+	return kv
+}
+
+func (p *olmoModel) Tensors(ts []Tensor) []*ggml.Tensor {
+	out := make([]*ggml.Tensor, 0, len(ts))
+	for _, t := range ts {
+		out = append(out, &ggml.Tensor{
+			Name:     t.Name(),
+			Kind:     t.Kind(),
+			Shape:    t.Shape(),
+			WriterTo: t,
+		})
+	}
+
+	return out
+}
+
+func (p *olmoModel) Replacements() []string {
+	return []string{
+		"lm_head", "output",
+		"model.embed_tokens", "token_embd",
+		"model.layers", "blk",
+		"model.norm", "output_norm",
+		"self_attn.q_proj", "attn_q",
+		"self_attn.k_proj", "attn_k",
+		"self_attn.v_proj", "attn_v",
+		"self_attn.o_proj", "attn_output",
+		"self_attn.q_norm", "attn_q_norm",
+		"self_attn.k_norm", "attn_k_norm",
+		"post_attention_layernorm", "post_attention_norm",
+		"post_feedforward_layernorm", "post_ffw_norm",
+		"mlp.gate_proj", "ffn_gate",
+		"mlp.down_proj", "ffn_down",
+		"mlp.up_proj", "ffn_up",
+	}
+}

model/models/models.go

@@ -13,6 +13,7 @@ import (
 	_ "github.com/ollama/ollama/model/models/mistral3"
 	_ "github.com/ollama/ollama/model/models/mllama"
 	_ "github.com/ollama/ollama/model/models/nomicbert"
+	_ "github.com/ollama/ollama/model/models/olmo"
 	_ "github.com/ollama/ollama/model/models/qwen2"
 	_ "github.com/ollama/ollama/model/models/qwen25vl"
 	_ "github.com/ollama/ollama/model/models/qwen3"

model/models/olmo/model.go

@@ -0,0 +1,233 @@
+package olmo
+
+import (
+	"cmp"
+	"math"
+
+	"github.com/ollama/ollama/fs"
+	"github.com/ollama/ollama/kvcache"
+	"github.com/ollama/ollama/ml"
+	"github.com/ollama/ollama/ml/nn"
+	"github.com/ollama/ollama/ml/nn/fast"
+	"github.com/ollama/ollama/ml/nn/rope"
+	"github.com/ollama/ollama/model"
+	"github.com/ollama/ollama/model/input"
+)
+
+type Options struct {
+	hiddenSize, numHeads, numKVHeads int
+	headDim, ropeDim                 int
+	eps, ropeBase, ropeScale         float32
+	clampKQV                         float32
+
+	originalContextLength int
+	attnFactor            float32
+}
+
+type Model struct {
+	model.Base
+	model.TextProcessor
+
+	TokenEmbedding *nn.Embedding `gguf:"token_embd"`
+	Layers         []Layer       `gguf:"blk"`
+	OutputNorm     *nn.RMSNorm   `gguf:"output_norm"`
+	Output         *nn.Linear    `gguf:"output,alt:token_embd"`
+
+	Options
+}
+
+func New(c fs.Config) (model.Model, error) {
+	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"))},
+			c.Ints("tokenizer.ggml.eos_token_ids")...,
+		),
+	}
+
+	if c.String("tokenizer.ggml.model") != "gpt2" {
+		return nil, model.ErrUnsupportedTokenizer
+	}
+
+	var pretokenizers []string
+	if c.String("tokenizer.ggml.pre") != "default" {
+		pretokenizers = []string{
+			`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+`,
+		}
+	}
+	processor := model.NewBytePairEncoding(&vocabulary, pretokenizers...)
+
+	m := Model{
+		TextProcessor: processor,
+		Layers:        make([]Layer, c.Uint("block_count")),
+		Options: Options{
+			hiddenSize:            int(c.Uint("embedding_length")),
+			numHeads:              int(c.Uint("attention.head_count")),
+			numKVHeads:            int(c.Uint("attention.head_count_kv")),
+			headDim:               int(c.Uint("attention.key_length")),
+			ropeDim:               int(c.Uint("rope.dimension_count")),
+			eps:                   c.Float("attention.layer_norm_rms_epsilon"),
+			ropeBase:              c.Float("rope.freq_base", 1e4),
+			ropeScale:             c.Float("rope.scaling.factor", 1),
+			clampKQV:              c.Float("attention.clamp_kqv", 0),
+			originalContextLength: int(c.Uint("rope.scaling.original_context_length")),
+			attnFactor:            c.Float("rope.scaling.attn_factor", 1),
+		},
+	}
+
+	m.Cache = kvcache.NewCausalCache(m.Shift)
+
+	return &m, nil
+}
+
+type SelfAttention struct {
+	Query       *nn.Linear  `gguf:"attn_q"`
+	Key         *nn.Linear  `gguf:"attn_k"`
+	Value       *nn.Linear  `gguf:"attn_v"`
+	Output      *nn.Linear  `gguf:"attn_output"`
+	QNorm       *nn.RMSNorm `gguf:"attn_q_norm"`
+	KNorm       *nn.RMSNorm `gguf:"attn_k_norm"`
+	RopeFactors ml.Tensor   `gguf:"rope_freqs.weight"`
+}
+
+func (o *Options) ropeOptions(factors ml.Tensor, isSWA bool) []func(*rope.Options) {
+	opts := []func(*rope.Options){
+		rope.WithFactors(factors),
+	}
+
+	if !isSWA && o.originalContextLength > 0 {
+		opts = append(opts,
+			rope.WithOriginalContextLength(o.originalContextLength),
+			rope.WithExtrapolationFactor(1.),
+			rope.WithAttentionFactor(o.attnFactor),
+		)
+	}
+
+	return opts
+}
+
+func (sa *SelfAttention) Forward(ctx ml.Context, hiddenState, positions ml.Tensor, cache kvcache.Cache, opts *Options, isSWA bool) ml.Tensor {
+	batchSize := hiddenState.Dim(1)
+	headDim := cmp.Or(opts.headDim, opts.hiddenSize/opts.numHeads)
+	ropeDim := cmp.Or(opts.ropeDim, headDim)
+
+	query := sa.Query.Forward(ctx, hiddenState)
+	if sa.QNorm != nil {
+		query = sa.QNorm.Forward(ctx, query, opts.eps)
+	}
+	query = query.Reshape(ctx, headDim, opts.numHeads, batchSize)
+
+	key := sa.Key.Forward(ctx, hiddenState)
+	if sa.KNorm != nil {
+		key = sa.KNorm.Forward(ctx, key, opts.eps)
+	}
+	key = key.Reshape(ctx, headDim, opts.numKVHeads, batchSize)
+
+	value := sa.Value.Forward(ctx, hiddenState)
+	value = value.Reshape(ctx, headDim, opts.numKVHeads, batchSize)
+
+	freqScale := float32(1.0)
+	if !isSWA {
+		freqScale = 1. / opts.ropeScale
+	}
+
+	ropeOpts := opts.ropeOptions(sa.RopeFactors, isSWA)
+	query = fast.RoPE(ctx, query, positions, ropeDim, opts.ropeBase, freqScale, ropeOpts...)
+	key = fast.RoPE(ctx, key, positions, ropeDim, opts.ropeBase, freqScale, ropeOpts...)
+
+	attention := nn.Attention(ctx, query, key, value, 1.0/math.Sqrt(float64(headDim)), cache)
+	attention = attention.Reshape(ctx, headDim*opts.numHeads, batchSize)
+	return sa.Output.Forward(ctx, attention)
+}
+
+func (m *Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
+	ropeDim := cmp.Or(m.ropeDim, m.hiddenSize/m.numHeads)
+	isSWA := isSWALayer(layer)
+
+	freqScale := float32(1.0)
+	if !isSWA {
+		freqScale = 1. / m.ropeScale
+	}
+
+	ropeOpts := m.Options.ropeOptions(m.Layers[layer].SelfAttention.RopeFactors, isSWA)
+	return fast.RoPE(ctx, key, shift, ropeDim, m.ropeBase, freqScale, ropeOpts...), nil
+}
+
+type MLP struct {
+	Up   *nn.Linear `gguf:"ffn_up"`
+	Down *nn.Linear `gguf:"ffn_down"`
+	Gate *nn.Linear `gguf:"ffn_gate"`
+}
+
+func (mlp *MLP) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *Options) ml.Tensor {
+	hiddenState = mlp.Gate.Forward(ctx, hiddenState).SILU(ctx, mlp.Up.Forward(ctx, hiddenState))
+	return mlp.Down.Forward(ctx, hiddenState)
+}
+
+type Layer struct {
+	SelfAttention     *SelfAttention
+	PostAttentionNorm *nn.RMSNorm `gguf:"post_attention_norm"`
+	MLP               *MLP
+	PostFFWNorm       *nn.RMSNorm `gguf:"post_ffw_norm"`
+}
+
+func (l *Layer) Forward(ctx ml.Context, hiddenState, positions, outputs ml.Tensor, cache kvcache.Cache, opts *Options, isSWA bool) ml.Tensor {
+	residual := hiddenState
+
+	hiddenState = l.SelfAttention.Forward(ctx, hiddenState, positions, cache, opts, isSWA)
+
+	if outputs != nil {
+		hiddenState = hiddenState.Rows(ctx, outputs)
+		residual = residual.Rows(ctx, outputs)
+	}
+
+	if l.PostAttentionNorm != nil {
+		hiddenState = l.PostAttentionNorm.Forward(ctx, hiddenState, opts.eps)
+	}
+
+	ffnInput := hiddenState.Add(ctx, residual)
+
+	hiddenState = l.MLP.Forward(ctx, ffnInput, opts)
+
+	if l.PostFFWNorm != nil {
+		hiddenState = l.PostFFWNorm.Forward(ctx, hiddenState, opts.eps)
+	}
+
+	return hiddenState.Add(ctx, ffnInput)
+}
+
+func isSWALayer(layerIdx int) bool {
+	return (layerIdx+1)%4 != 0
+}
+
+func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
+	positions := ctx.Input().FromInts(batch.Positions, len(batch.Positions))
+
+	hiddenState := m.TokenEmbedding.Forward(ctx, batch.Inputs)
+
+	for i, layer := range m.Layers {
+		m.Cache.SetLayer(i)
+
+		isSWA := isSWALayer(i)
+
+		var outputs ml.Tensor
+		if i == len(m.Layers)-1 {
+			outputs = batch.Outputs
+		}
+
+		hiddenState = layer.Forward(ctx, hiddenState, positions, outputs, m.Cache, &m.Options, isSWA)
+	}
+
+	hiddenState = m.OutputNorm.Forward(ctx, hiddenState, m.eps)
+	return m.Output.Forward(ctx, hiddenState), nil
+}
+
+func init() {
+	model.Register("olmo2", New)
+}