Mirror_Repos/ollama_ollama
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

New models (#15861)

Browse Source 
* mlx: add laguna model support

* convert: support fp8 safetensors import

Decode HF F8_E4M3 safetensors with block scale companions into GGUF-supported tensor types, and record which output tensors came from FP8 source weights.

Use that source-precision metadata during create quantization: default FP8-sourced GGUFs to Q8_0, keep non-FP8 tensors at their original precision for Q8_0, and promote non-FP8 quantizable tensors to Q8_0 for Q4_K requests.

* ggml: add laguna model support

* server: preserve generate logprobs with builtin parsers

Generate requests were dropping logprob-only chunks whenever a builtin parser buffered visible content. Chat already handled this case, but generate only forwarded chunks with visible response, thinking, or tool-call output.

Keep generate chunks that carry logprobs even when the builtin parser has not flushed visible content yet, and add a regression test that exercises the behavior with a generic thinking parser.

* review comments - perf improvements

* ggml: implement nemotron 3 nano omni

* add poolside integration

* update poolside doc

* adapt to new cache setup

* fix test

* fix test

---------

Co-authored-by: Eva Ho <hoyyeva@gmail.com>
This commit is contained in:
Daniel Hiltgen
2026-04-28 11:50:12 -07:00
committed by GitHub
parent 2bbe2405fe
commit 87288ced4f
62 changed files with 11284 additions and 633 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/models/laguna/model.go
+444
View File
@@ -0,0 +1,444 @@
package laguna
import (
"fmt"
"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/rope"
"github.com/ollama/ollama/model"
"github.com/ollama/ollama/model/input"
"github.com/ollama/ollama/tokenizer"
)
const (
cacheTypeSWA = iota
cacheTypeCausal
)
type Options struct {
hiddenSize int
headDim int
numHeads []int
numKVHeads int
eps float32
slidingWindow int
slidingWindowPattern []bool
fullRopeDim int
fullRopeBase, fullRopeScale float32
fullRopeOriginalContextLength int
fullRopeAttentionFactor float32
fullRopeBetaFast float32
fullRopeBetaSlow float32
swaRopeDim int
swaRopeBase, swaRopeScale float32
numExperts, numExpertsUsed int
normTopKProb bool
routedScalingFactor float32
decoderSparseStep int
denseLayers map[int]bool
}
func (o *Options) numHeadsForLayer(layer int) int {
if layer < len(o.numHeads) && o.numHeads[layer] > 0 {
return o.numHeads[layer]
}
if len(o.numHeads) > 0 && o.numHeads[0] > 0 {
return o.numHeads[0]
}
return 1
}
func (o *Options) layerIsSliding(layer int) bool {
return layer < len(o.slidingWindowPattern) && o.slidingWindowPattern[layer]
}
func (o *Options) layerUsesMoE(layer int) bool {
if o.numExperts == 0 || o.denseLayers[layer] {
return false
}
step := o.decoderSparseStep
if step <= 0 {
step = 1
}
return (layer+1)%step == 0
}
func (o *Options) applyRotaryPositionEmbeddings(ctx ml.Context, layer int, states, positions ml.Tensor) ml.Tensor {
opts := []func(*rope.Options){rope.WithTypeNeoX()}
if o.layerIsSliding(layer) {
return nn.RoPE(ctx, states, positions, o.swaRopeDim, o.swaRopeBase, 1./o.swaRopeScale, opts...)
}
opts = append(opts,
rope.WithOriginalContextLength(o.fullRopeOriginalContextLength),
rope.WithExtrapolationFactor(1),
rope.WithAttentionFactor(o.fullRopeAttentionFactor),
rope.WithBetaFast(o.fullRopeBetaFast),
rope.WithBetaSlow(o.fullRopeBetaSlow),
)
return nn.RoPE(ctx, states, positions, o.fullRopeDim, o.fullRopeBase, 1./o.fullRopeScale, opts...)
}
type Attention struct {
Query *nn.Linear `gguf:"attn_q"`
QueryNorm *nn.RMSNorm `gguf:"attn_q_norm"`
Key *nn.Linear `gguf:"attn_k"`
KeyNorm *nn.RMSNorm `gguf:"attn_k_norm"`
Value *nn.Linear `gguf:"attn_v"`
Gate *nn.Linear `gguf:"attn_g"`
Output *nn.Linear `gguf:"attn_output"`
}
func (sa *Attention) Forward(ctx ml.Context, layer int, hiddenStates, positions ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor {
batchSize := hiddenStates.Dim(1)
numHeads := opts.numHeadsForLayer(layer)
query := sa.Query.Forward(ctx, hiddenStates)
key := sa.Key.Forward(ctx, hiddenStates)
value := sa.Value.Forward(ctx, hiddenStates)
gate := sa.Gate.Forward(ctx, hiddenStates)
query = query.Reshape(ctx, opts.headDim, numHeads, batchSize)
key = key.Reshape(ctx, opts.headDim, opts.numKVHeads, batchSize)
value = value.Reshape(ctx, opts.headDim, opts.numKVHeads, batchSize)
query = sa.QueryNorm.Forward(ctx, query, opts.eps)
key = sa.KeyNorm.Forward(ctx, key, opts.eps)
query = opts.applyRotaryPositionEmbeddings(ctx, layer, query, positions)
key = opts.applyRotaryPositionEmbeddings(ctx, layer, key, positions)
attention := nn.Attention(ctx, query, key, value, 1./math.Sqrt(float64(opts.headDim)), cache)
// Laguna applies the per-head gate softplus in float32, then casts back.
gate = gate.Cast(ctx, ml.DTypeF32).Softplus(ctx).Cast(ctx, attention.DType())
attention = attention.Mul(ctx, gate.Reshape(ctx, 1, numHeads, batchSize))
attention = attention.Reshape(ctx, opts.headDim*numHeads, batchSize)
return sa.Output.Forward(ctx, attention)
}
type MLP interface {
Forward(ml.Context, ml.Tensor, *Options) ml.Tensor
}
type dense struct {
Gate *nn.Linear `gguf:"ffn_gate"`
Up *nn.Linear `gguf:"ffn_up"`
Down *nn.Linear `gguf:"ffn_down"`
}
func (mlp *dense) Forward(ctx ml.Context, hiddenStates ml.Tensor, _ *Options) ml.Tensor {
hiddenStates = mlp.Gate.Forward(ctx, hiddenStates).SILU(ctx, mlp.Up.Forward(ctx, hiddenStates))
return mlp.Down.Forward(ctx, hiddenStates)
}
type sparse struct {
Router *nn.Linear `gguf:"ffn_gate_inp"`
Gate *nn.LinearBatch `gguf:"ffn_gate_exps"`
Up *nn.LinearBatch `gguf:"ffn_up_exps"`
Down *nn.LinearBatch `gguf:"ffn_down_exps"`
SharedExpert *dense `gguf:",suf:_shexp"`
ExpProbsBias ml.Tensor `gguf:"exp_probs_b.bias,alt:exp_probs_b"`
}
func (moe *sparse) topKIndices(ctx ml.Context, scores ml.Tensor, opts *Options) ml.Tensor {
if moe.ExpProbsBias != nil {
scores = scores.Add(ctx, moe.ExpProbsBias)
}
return scores.TopK(ctx, opts.numExpertsUsed)
}
func (moe *sparse) Forward(ctx ml.Context, hiddenStates ml.Tensor, opts *Options) ml.Tensor {
residual := hiddenStates
scores := moe.Router.Forward(ctx, hiddenStates).Cast(ctx, ml.DTypeF32).Sigmoid(ctx)
selectedExperts := moe.topKIndices(ctx, scores, opts)
routingWeights := scores.Reshape(ctx, 1, opts.numExperts, hiddenStates.Dim(1)).Rows(ctx, selectedExperts)
if opts.normTopKProb {
routingWeights = routingWeights.Reshape(ctx, opts.numExpertsUsed, hiddenStates.Dim(1))
routingWeights = routingWeights.Div(ctx, routingWeights.SumRows(ctx))
routingWeights = routingWeights.Reshape(ctx, 1, opts.numExpertsUsed, hiddenStates.Dim(1))
}
routingWeights = routingWeights.Scale(ctx, float64(opts.routedScalingFactor))
hiddenStates = hiddenStates.Reshape(ctx, hiddenStates.Dim(0), 1, hiddenStates.Dim(1))
upStates := moe.Up.Forward(ctx, hiddenStates, selectedExperts)
hiddenStates = moe.Gate.Forward(ctx, hiddenStates, selectedExperts).SILU(ctx, upStates)
experts := moe.Down.Forward(ctx, hiddenStates, selectedExperts)
experts = experts.Mul(ctx, routingWeights)
nextStates := experts.View(ctx, 0, experts.Dim(0), experts.Stride(2), experts.Dim(2))
for i := 1; i < opts.numExpertsUsed; i++ {
nextStates = nextStates.Add(ctx, experts.View(ctx, i*experts.Stride(1), experts.Dim(0), experts.Stride(2), experts.Dim(2)))
}
return nextStates.Add(ctx, moe.SharedExpert.Forward(ctx, residual, opts))
}
type Layer struct {
AttentionNorm *nn.RMSNorm `gguf:"attn_norm"`
*Attention
MLPNorm *nn.RMSNorm `gguf:"ffn_norm"`
MLP
}
func (l *Layer) Forward(ctx ml.Context, layer int, hiddenStates, positions, outputs ml.Tensor, cache kvcache.Cache, opts *Options) ml.Tensor {
residual := hiddenStates
hiddenStates = l.AttentionNorm.Forward(ctx, hiddenStates, opts.eps)
hiddenStates = l.Attention.Forward(ctx, layer, hiddenStates, positions, cache, opts)
if outputs != nil {
hiddenStates = hiddenStates.Rows(ctx, outputs)
residual = residual.Rows(ctx, outputs)
}
hiddenStates = hiddenStates.Add(ctx, residual)
residual = hiddenStates
hiddenStates = l.MLPNorm.Forward(ctx, hiddenStates, opts.eps)
hiddenStates = l.MLP.Forward(ctx, hiddenStates, opts)
return hiddenStates.Add(ctx, residual)
}
type Model struct {
model.Base
tokenizer.Tokenizer
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) {
if c.Bool("attention.sink_enabled") {
return nil, fmt.Errorf("laguna: SWA attention sinks are not supported")
}
if c.Uint("attention.gating_type") != 1 {
return nil, fmt.Errorf("laguna: unsupported attention gating type %d", c.Uint("attention.gating_type"))
}
if !c.Bool("attention.qk_norm") {
return nil, fmt.Errorf("laguna: Q/K RMSNorm is required")
}
if gating := c.Uint("expert_gating_func"); gating != 2 {
return nil, fmt.Errorf("laguna: unsupported expert gating function %d", gating)
}
numLayers := int(c.Uint("block_count"))
opts := newOptions(c, numLayers)
layers := make([]Layer, numLayers)
for i := range layers {
if opts.layerUsesMoE(i) {
layers[i].MLP = &sparse{}
} else {
layers[i].MLP = &dense{}
}
}
var pre []string
switch c.String("tokenizer.ggml.pre") {
case "laguna":
pre = []string{
`(?:\r?\n)+(?!\r?\n)`,
`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
}
default:
return nil, model.ErrUnsupportedTokenizer
}
m := Model{
Tokenizer: tokenizer.NewBytePairEncoding(
&tokenizer.Vocabulary{
Values: c.Strings("tokenizer.ggml.tokens"),
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")...,
),
},
pre...,
),
Layers: layers,
Options: opts,
}
m.Cache = kvcache.NewWrapperCache(
kvcache.NewSWACache(int32(opts.slidingWindow), m.Shift),
kvcache.NewCausalCache(m.Shift),
)
return &m, nil
}
func newOptions(c fs.Config, numLayers int) *Options {
denseLayers := make(map[int]bool)
for _, layer := range configUints(c, "dense_layers") {
denseLayers[int(layer)] = true
}
for i := range c.Uint("leading_dense_block_count") {
denseLayers[int(i)] = true
}
fullRopeScale := c.Float("rope.scaling.factor", 1)
if fullRopeScale == 0 {
fullRopeScale = 1
}
swaRopeScale := c.Float("rope.swa.scaling.factor", 1)
if swaRopeScale == 0 {
swaRopeScale = 1
}
fullRopeType := c.String("rope.scaling.type")
fullRopeAttentionFactor := lagunaAttentionFactor(fullRopeType, fullRopeScale, c.Float("rope.scaling.attn_factor"))
return &Options{
hiddenSize: int(c.Uint("embedding_length")),
headDim: int(c.Uint("attention.key_length")),
numHeads: expandIntArray(configUints(c, "attention.head_count"), numLayers, c.Uint("attention.head_count", 1)),
numKVHeads: int(c.Uint("attention.head_count_kv")),
eps: c.Float("attention.layer_norm_rms_epsilon", 1e-6),
slidingWindow: int(c.Uint("attention.sliding_window", 512)),
slidingWindowPattern: slidingWindowPattern(c, numLayers),
fullRopeDim: int(c.Uint("rope.dimension_count", c.Uint("attention.key_length"))),
fullRopeBase: c.Float("rope.freq_base", 500000),
fullRopeScale: fullRopeScale,
fullRopeOriginalContextLength: int(c.Uint("rope.scaling.original_context_length", 4096)),
fullRopeAttentionFactor: fullRopeAttentionFactor,
fullRopeBetaFast: c.Float("rope.scaling.beta_fast", 64),
fullRopeBetaSlow: c.Float("rope.scaling.beta_slow", 1),
swaRopeDim: int(c.Uint("rope.swa.dimension_count", c.Uint("attention.key_length"))),
swaRopeBase: c.Float("rope.swa.freq_base", 10000),
swaRopeScale: swaRopeScale,
numExperts: int(c.Uint("expert_count")),
numExpertsUsed: int(c.Uint("expert_used_count")),
normTopKProb: c.Bool("expert_weights_norm", true),
routedScalingFactor: c.Float("expert_weights_scale", 1),
decoderSparseStep: int(c.Uint("decoder_sparse_step", 1)),
denseLayers: denseLayers,
}
}
func lagunaAttentionFactor(ropeType string, scaleFactor, attentionFactor float32) float32 {
if attentionFactor != 0 {
return attentionFactor
}
if ropeType == "yarn" && scaleFactor > 1 {
return float32(0.1*math.Log(float64(scaleFactor)) + 1)
}
return 1
}
func slidingWindowPattern(c fs.Config, numLayers int) []bool {
pattern := c.Bools("attention.sliding_window_pattern")
if len(pattern) == numLayers {
return pattern
}
layerTypes := configUints(c, "attention.layer_types")
if len(layerTypes) == numLayers {
pattern = make([]bool, numLayers)
for i, layerType := range layerTypes {
pattern[i] = layerType == 1
}
return pattern
}
return make([]bool, numLayers)
}
func configUints(c fs.Config, key string) []uint32 {
keyExists := c.Value(c.Architecture()+"."+key) != nil || c.Value(key) != nil
if cc, ok := c.(interface {
Uints(string, ...[]uint32) []uint32
}); ok {
if values := cc.Uints(key); len(values) > 0 && (keyExists || !(len(values) == 1 && values[0] == 0)) {
return values
}
}
ints := c.Ints(key)
if len(ints) > 0 && (keyExists || !(len(ints) == 1 && ints[0] == 0)) {
values := make([]uint32, len(ints))
for i, v := range ints {
values[i] = uint32(v)
}
return values
}
if scalar := c.Uint(key); scalar != 0 {
return []uint32{scalar}
}
return nil
}
func expandIntArray(values []uint32, n int, fallback uint32) []int {
if len(values) == 0 {
values = []uint32{fallback}
}
defaultValue := values[0]
if len(values) == 1 {
defaultValue = values[0]
}
out := make([]int, n)
for i := range out {
if i < len(values) {
out[i] = int(values[i])
} else {
out[i] = int(defaultValue)
}
}
return out
}
func (m *Model) Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
return m.Options.applyRotaryPositionEmbeddings(ctx, layer, key, shift), nil
}
func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
positions := ctx.Input().FromInts(batch.Positions, len(batch.Positions))
hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)
for i, layer := range m.Layers {
if m.Cache != nil {
m.Cache.SetLayer(i)
if wrapper, ok := m.Cache.(*kvcache.WrapperCache); ok {
cacheType := cacheTypeCausal
if m.Options.layerIsSliding(i) {
cacheType = cacheTypeSWA
}
wrapper.SetLayerType(cacheType)
}
}
var outputs ml.Tensor
if i == len(m.Layers)-1 {
outputs = batch.Outputs
}
hiddenStates = layer.Forward(ctx, i, hiddenStates, positions, outputs, m.Cache, m.Options)
}
hiddenStates = m.OutputNorm.Forward(ctx, hiddenStates, m.eps)
return m.Output.Forward(ctx, hiddenStates), nil
}
func init() {
model.Register("laguna", New)
}
var _ model.Model = (*Model)(nil)
model/models/laguna/model_test.go
+237
View File
@@ -0,0 +1,237 @@
package laguna
import (
"iter"
"math"
"testing"
)
type testConfig map[string]any
func (c testConfig) Architecture() string { return "laguna" }
func (c testConfig) key(key string) string {
switch {
case len(key) >= len("tokenizer.") && key[:len("tokenizer.")] == "tokenizer.":
return key
case len(key) >= len("general.") && key[:len("general.")] == "general.":
return key
default:
return "laguna." + key
}
}
func (c testConfig) String(key string, defaultValue ...string) string {
if v, ok := c[c.key(key)].(string); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return ""
}
func (c testConfig) Uint(key string, defaultValue ...uint32) uint32 {
switch v := c[c.key(key)].(type) {
case uint32:
return v
case int:
return uint32(v)
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return 0
}
func (c testConfig) Float(key string, defaultValue ...float32) float32 {
if v, ok := c[c.key(key)].(float32); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return 0
}
func (c testConfig) Bool(key string, defaultValue ...bool) bool {
if v, ok := c[c.key(key)].(bool); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return false
}
func (c testConfig) Strings(key string, defaultValue ...[]string) []string {
if v, ok := c[c.key(key)].([]string); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return nil
}
func (c testConfig) Ints(key string, defaultValue ...[]int32) []int32 {
if v, ok := c[c.key(key)].([]int32); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return nil
}
func (c testConfig) Uints(key string, defaultValue ...[]uint32) []uint32 {
if v, ok := c[c.key(key)].([]uint32); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return nil
}
func (c testConfig) Floats(key string, defaultValue ...[]float32) []float32 {
if v, ok := c[c.key(key)].([]float32); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return nil
}
func (c testConfig) Bools(key string, defaultValue ...[]bool) []bool {
if v, ok := c[c.key(key)].([]bool); ok {
return v
}
if len(defaultValue) > 0 {
return defaultValue[0]
}
return nil
}
func (c testConfig) Len() int { return len(c) }
func (c testConfig) Keys() iter.Seq[string] {
return func(yield func(string) bool) {
for key := range c {
if !yield(key) {
return
}
}
}
}
func (c testConfig) Value(key string) any { return c[key] }
func TestNewOptionsLayerConfig(t *testing.T) {
cfg := testConfig{
"laguna.block_count": uint32(4),
"laguna.embedding_length": uint32(128),
"laguna.attention.key_length": uint32(16),
"laguna.attention.head_count": []uint32{8, 16, 16, 16},
"laguna.attention.head_count_kv": uint32(4),
"laguna.attention.layer_norm_rms_epsilon": float32(1e-6),
"laguna.attention.sliding_window": uint32(512),
"laguna.attention.sliding_window_pattern": []bool{false, true, true, true},
"laguna.rope.dimension_count": uint32(8),
"laguna.rope.freq_base": float32(500000),
"laguna.rope.scaling.factor": float32(32),
"laguna.rope.scaling.original_context_length": uint32(4096),
"laguna.rope.swa.dimension_count": uint32(16),
"laguna.rope.swa.freq_base": float32(10000),
"laguna.expert_count": uint32(32),
"laguna.expert_used_count": uint32(4),
"laguna.expert_weights_norm": true,
"laguna.expert_weights_scale": float32(2.5),
"laguna.decoder_sparse_step": uint32(1),
"laguna.dense_layers": []uint32{0},
}
opts := newOptions(cfg, 4)
if got := opts.numHeadsForLayer(0); got != 8 {
t.Fatalf("layer 0 heads = %d, want 8", got)
}
if got := opts.numHeadsForLayer(1); got != 16 {
t.Fatalf("layer 1 heads = %d, want 16", got)
}
if opts.layerIsSliding(0) {
t.Fatal("layer 0 should be full attention")
}
if !opts.layerIsSliding(1) {
t.Fatal("layer 1 should be sliding attention")
}
if opts.layerUsesMoE(0) {
t.Fatal("layer 0 should be dense")
}
if !opts.layerUsesMoE(1) {
t.Fatal("layer 1 should use MoE")
}
if opts.fullRopeDim != 8 || opts.swaRopeDim != 16 {
t.Fatalf("rope dims = full %d swa %d, want 8/16", opts.fullRopeDim, opts.swaRopeDim)
}
}
func TestNewOptionsYarnAttentionFactorFallback(t *testing.T) {
cfg := testConfig{
"laguna.block_count": uint32(1),
"laguna.embedding_length": uint32(128),
"laguna.attention.key_length": uint32(16),
"laguna.attention.head_count": uint32(8),
"laguna.attention.head_count_kv": uint32(4),
"laguna.rope.scaling.type": "yarn",
"laguna.rope.scaling.factor": float32(32),
}
opts := newOptions(cfg, 1)
want := float32(0.1*math.Log(32) + 1)
if got := opts.fullRopeAttentionFactor; math.Abs(float64(got-want)) > 1e-6 {
t.Fatalf("fullRopeAttentionFactor = %v, want %v", got, want)
}
}
func TestNewRejectsUnsupportedLagunaVariants(t *testing.T) {
tests := []struct {
name string
cfg testConfig
}{
{
name: "attention sinks",
cfg: testConfig{
"laguna.attention.sink_enabled": true,
},
},
{
name: "non per-head gate",
cfg: testConfig{
"laguna.attention.gating_type": uint32(0),
},
},
{
name: "missing qk norm",
cfg: testConfig{
"laguna.attention.gating_type": uint32(1),
},
},
{
name: "non sigmoid experts",
cfg: testConfig{
"laguna.attention.gating_type": uint32(1),
"laguna.attention.qk_norm": true,
"laguna.expert_gating_func": uint32(1),
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if _, err := New(tt.cfg); err == nil {
t.Fatal("expected unsupported variant error")
}
})
}
}
model/models/models.go
+1
View File
@@ -11,6 +11,7 @@ import (
_ "github.com/ollama/ollama/model/models/glm4moelite"
_ "github.com/ollama/ollama/model/models/glmocr"
_ "github.com/ollama/ollama/model/models/gptoss"
_ "github.com/ollama/ollama/model/models/laguna"
_ "github.com/ollama/ollama/model/models/lfm2"
_ "github.com/ollama/ollama/model/models/llama"
_ "github.com/ollama/ollama/model/models/llama4"
model/models/nemotronh/imageproc.go
+355
View File
@@ -0,0 +1,355 @@
package nemotronh
import (
"errors"
"image"
"math"
"slices"
"github.com/ollama/ollama/fs"
"github.com/ollama/ollama/model/imageproc"
)
type ImageProcessor struct {
imageSize int
patchSize int
numChannels int
maxTiles int
minNumPatches int
maxNumPatches int
useThumbnail bool
projectorScale int
imageMean [3]float32
imageStd [3]float32
}
type processedVisionTile struct {
data []float32
size image.Point
}
func newImageProcessor(c fs.Config) ImageProcessor {
mean := c.Floats("vision.image_mean")
std := c.Floats("vision.image_std")
processor := ImageProcessor{
imageSize: int(c.Uint("vision.image_size", 512)),
patchSize: int(c.Uint("vision.patch_size", 16)),
numChannels: int(c.Uint("vision.num_channels", 3)),
maxTiles: int(c.Uint("vision.max_tiles", 12)),
minNumPatches: int(c.Uint("vision.min_num_patches")),
maxNumPatches: int(c.Uint("vision.max_num_patches")),
useThumbnail: c.Bool("vision.use_thumbnail", true),
projectorScale: int(c.Uint("vision.projector.scale_factor", 2)),
imageMean: imageproc.ClipDefaultMean,
imageStd: imageproc.ClipDefaultSTD,
}
if len(mean) >= 3 {
processor.imageMean = [3]float32{mean[0], mean[1], mean[2]}
}
if len(std) >= 3 {
processor.imageStd = [3]float32{std[0], std[1], std[2]}
}
if processor.imageSize <= 0 {
processor.imageSize = 512
}
if processor.patchSize <= 0 {
processor.patchSize = 16
}
if processor.numChannels <= 0 {
processor.numChannels = 3
}
if processor.maxTiles <= 0 {
processor.maxTiles = 12
}
if processor.projectorScale <= 0 {
processor.projectorScale = 2
}
return processor
}
func (p ImageProcessor) ProcessImage(img image.Image) ([]processedVisionTile, error) {
img = imageproc.Composite(img)
if p.useDynamicResolution() {
return p.processDynamicImage(img)
}
return p.processTiledImage(img), nil
}
func (p ImageProcessor) useDynamicResolution() bool {
return p.minNumPatches > 0 || p.maxNumPatches > 0
}
func (p ImageProcessor) processTiledImage(img image.Image) []processedVisionTile {
bounds := img.Bounds()
origWidth := bounds.Dx()
origHeight := bounds.Dy()
targetRatios := nemotronTargetRatios(p.maxTiles)
gridWidth, gridHeight := findClosestAspectRatio(float64(origWidth)/float64(origHeight), targetRatios, origWidth, origHeight, p.imageSize)
targetWidth := p.imageSize * gridWidth
targetHeight := p.imageSize * gridHeight
resized := resizeImageBicubicCHW(img, targetWidth, targetHeight)
tiles := make([]processedVisionTile, 0, gridWidth*gridHeight+1)
for row := range gridHeight {
for col := range gridWidth {
tile := cropCHWRegion(
resized,
targetWidth,
targetHeight,
p.numChannels,
col*p.imageSize,
row*p.imageSize,
p.imageSize,
p.imageSize,
)
tiles = append(tiles, processedVisionTile{
data: normalizeVisionCHW(tile, p.imageMean, p.imageStd),
size: image.Point{X: p.imageSize, Y: p.imageSize},
})
}
}
if p.useThumbnail && len(tiles) > 1 {
thumbnail := resizeImageBicubicCHW(img, p.imageSize, p.imageSize)
tiles = append(tiles, processedVisionTile{
data: normalizeVisionCHW(thumbnail, p.imageMean, p.imageStd),
size: image.Point{X: p.imageSize, Y: p.imageSize},
})
}
return tiles
}
func (p ImageProcessor) processDynamicImage(img image.Image) ([]processedVisionTile, error) {
bounds := img.Bounds()
origWidth := bounds.Dx()
origHeight := bounds.Dy()
patchesWidth, patchesHeight := p.dynamicPatchGrid(origWidth, origHeight)
if patchesWidth <= 0 || patchesHeight <= 0 {
return nil, errors.New("nemotron_h_omni: invalid dynamic image patch grid")
}
targetWidth := patchesWidth * p.patchSize
targetHeight := patchesHeight * p.patchSize
resized := resizeImageBicubicCHW(img, targetWidth, targetHeight)
return []processedVisionTile{{
data: normalizeVisionCHW(resized, p.imageMean, p.imageStd),
size: image.Point{X: targetWidth, Y: targetHeight},
}}, nil
}
func (p ImageProcessor) dynamicPatchGrid(origWidth, origHeight int) (int, int) {
patchesHeight := max(1, int(math.Round(float64(origHeight)/float64(p.patchSize)+0.5)))
patchesWidth := max(1, int(math.Round(float64(origWidth)/float64(p.patchSize)+0.5)))
patches := patchesHeight * patchesWidth
currentNumPatchesAvailable := p.maxNumPatches
if currentNumPatchesAvailable <= 0 {
currentNumPatchesAvailable = max(patches, p.minNumPatches)
}
factor := math.Min(math.Sqrt(float64(currentNumPatchesAvailable)/float64(patches)), 1.0)
targetPatchesHeight := max(1, int(math.Floor(factor*float64(patchesHeight))))
targetPatchesWidth := max(1, int(math.Floor(factor*float64(patchesWidth))))
if currentNumPatchesAvailable > p.minNumPatches && targetPatchesHeight*targetPatchesWidth < p.minNumPatches {
upFactor := math.Sqrt(float64(p.minNumPatches) / float64(targetPatchesHeight*targetPatchesWidth))
targetPatchesHeight = int(math.Ceil(upFactor * float64(targetPatchesHeight)))
targetPatchesWidth = int(math.Ceil(upFactor * float64(targetPatchesWidth)))
}
targetPatchesHeight = roundPatchGridForPixelShuffle(targetPatchesHeight, targetPatchesWidth, currentNumPatchesAvailable, p.projectorScale)
targetPatchesWidth = roundPatchGridForPixelShuffle(targetPatchesWidth, targetPatchesHeight, currentNumPatchesAvailable, p.projectorScale)
return targetPatchesWidth, targetPatchesHeight
}
func roundPatchGridForPixelShuffle(v, other, maxPatches, divisor int) int {
if divisor <= 1 {
return v
}
rem := v % divisor
if rem == 0 {
return v
}
inc := divisor - rem
if (v+inc)*other <= maxPatches {
return v + inc
}
return max(divisor, v-rem)
}
type nemotronImageRatio struct {
width int
height int
}
func nemotronTargetRatios(maxTiles int) []nemotronImageRatio {
targetRatios := make([]nemotronImageRatio, 0, maxTiles*maxTiles)
for n := 1; n <= maxTiles; n++ {
for w := 1; w <= n; w++ {
for h := 1; h <= n; h++ {
if w*h > maxTiles {
continue
}
targetRatios = append(targetRatios, nemotronImageRatio{width: w, height: h})
}
}
}
unique := targetRatios[:0]
for _, ratio := range targetRatios {
if slices.Contains(unique, ratio) {
continue
}
unique = append(unique, ratio)
}
slices.SortFunc(unique, func(a, b nemotronImageRatio) int {
return a.width*a.height - b.width*b.height
})
return unique
}
func findClosestAspectRatio(aspectRatio float64, targetRatios []nemotronImageRatio, width, height, imageSize int) (int, int) {
bestRatio := nemotronImageRatio{width: 1, height: 1}
bestRatioDiff := math.MaxFloat64
area := width * height
for _, ratio := range targetRatios {
targetAspectRatio := float64(ratio.width) / float64(ratio.height)
ratioDiff := math.Abs(aspectRatio - targetAspectRatio)
if ratioDiff < bestRatioDiff {
bestRatioDiff = ratioDiff
bestRatio = ratio
continue
}
if ratioDiff == bestRatioDiff && area > int(0.5*float64(imageSize*imageSize*ratio.width*ratio.height)) {
bestRatio = ratio
}
}
return bestRatio.width, bestRatio.height
}
func resizeImageBicubicCHW(img image.Image, outW, outH int) []float32 {
bounds := img.Bounds()
inW := bounds.Dx()
inH := bounds.Dy()
src := make([]float32, 3*inW*inH)
for y := range inH {
for x := range inW {
r, g, b, _ := img.At(bounds.Min.X+x, bounds.Min.Y+y).RGBA()
src[y*inW+x] = float32(r>>8) / 255.0
src[inW*inH+y*inW+x] = float32(g>>8) / 255.0
src[2*inW*inH+y*inW+x] = float32(b>>8) / 255.0
}
}
dst := make([]float32, 3*outW*outH)
scaleX := float64(inW) / float64(outW)
scaleY := float64(inH) / float64(outH)
for oy := range outH {
srcY := scaleY*(float64(oy)+0.5) - 0.5
yBase := int(math.Floor(srcY))
yFrac := clampUnit(srcY - float64(yBase))
wy := torchBicubicWeights(yFrac)
for ox := range outW {
srcX := scaleX*(float64(ox)+0.5) - 0.5
xBase := int(math.Floor(srcX))
xFrac := clampUnit(srcX - float64(xBase))
wx := torchBicubicWeights(xFrac)
for c := range 3 {
var sum float64
channelOffset := c * inW * inH
for ky := range 4 {
iy := clampIndex(yBase-1+ky, 0, inH-1)
for kx := range 4 {
ix := clampIndex(xBase-1+kx, 0, inW-1)
sum += float64(src[channelOffset+iy*inW+ix]) * wy[ky] * wx[kx]
}
}
dst[c*outW*outH+oy*outW+ox] = float32(sum)
}
}
}
return dst
}
func cropCHWRegion(values []float32, width, height, channels, left, top, cropW, cropH int) []float32 {
out := make([]float32, channels*cropW*cropH)
channelSize := width * height
cropSize := cropW * cropH
for c := range channels {
srcBase := c * channelSize
dstBase := c * cropSize
for y := range cropH {
copy(out[dstBase+y*cropW:dstBase+(y+1)*cropW], values[srcBase+(top+y)*width+left:srcBase+(top+y)*width+left+cropW])
}
}
return out
}
func normalizeVisionCHW(values []float32, mean, std [3]float32) []float32 {
out := make([]float32, len(values))
channelSize := len(values) / 3
for c := range 3 {
base := c * channelSize
for i := range channelSize {
out[base+i] = (values[base+i] - mean[c]) / std[c]
}
}
return out
}
func torchBicubicWeights(t float64) [4]float64 {
const a = -0.75
return [4]float64{
bicubicConvolution2(t+1.0, a),
bicubicConvolution1(t, a),
bicubicConvolution1(1.0-t, a),
bicubicConvolution2(2.0-t, a),
}
}
func bicubicConvolution1(x, a float64) float64 {
return ((a+2)*x-(a+3))*x*x + 1
}
func bicubicConvolution2(x, a float64) float64 {
return ((a*x-5*a)*x+8*a)*x - 4*a
}
func clampUnit(v float64) float64 {
if v < 0 {
return 0
}
if v > 1 {
return 1
}
return v
}
func clampIndex(v, lo, hi int) int {
if v < lo {
return lo
}
if v > hi {
return hi
}
return v
}
model/models/nemotronh/model.go
+20 -5
View File
@@ -117,9 +117,7 @@ func Shift(ctx ml.Context, layer int, key, shift ml.Tensor) (ml.Tensor, error) {
return key, nil
}
func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)
func (m *Model) forwardHiddenStates(ctx ml.Context, batch input.Batch, hiddenStates ml.Tensor) (ml.Tensor, error) {
cache := m.Cache.(*HybridCache)
for i, layer := range m.Layers {
@@ -137,11 +135,24 @@ func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
}
}
hiddenStates = m.OutputNorm.Forward(ctx, hiddenStates, m.eps)
return m.OutputNorm.Forward(ctx, hiddenStates, m.eps), nil
}
func (m *Model) forwardLogits(ctx ml.Context, batch input.Batch, hiddenStates ml.Tensor) (ml.Tensor, error) {
hiddenStates, err := m.forwardHiddenStates(ctx, batch, hiddenStates)
if err != nil {
return nil, err
}
return m.Output.Forward(ctx, hiddenStates), nil
}
func New(c fs.Config) (model.Model, error) {
func (m *Model) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)
return m.forwardLogits(ctx, batch, hiddenStates)
}
func newTextModel(c fs.Config) (*Model, error) {
numLayers := int(c.Uint("block_count"))
layers := make([]Layer, numLayers)
@@ -306,6 +317,10 @@ func New(c fs.Config) (model.Model, error) {
return &m, nil
}
func New(c fs.Config) (model.Model, error) {
return newTextModel(c)
}
func init() {
model.Register("nemotron_h", New)
model.Register("nemotron_h_moe", New)
model/models/nemotronh/model_audio.go
+511
View File
@@ -0,0 +1,511 @@
package nemotronh
import (
"math"
"sync"
"github.com/ollama/ollama/fs"
"github.com/ollama/ollama/ml"
"github.com/ollama/ollama/ml/nn"
)
type AudioOptions struct {
hiddenSize int
numHeads int
headDim int
intermediateSize int
convKernelSize int
melBins int
sampleRate int
subsamplingKernel int
subsamplingStride int
scaleInput bool
eps float32
}
type AudioFeatureExtractor struct {
FB ml.Tensor `gguf:"fb"`
Window ml.Tensor `gguf:"window"`
mu sync.Mutex
fb []float32
window []float32
fbShape [2]int
}
func (f *AudioFeatureExtractor) windowAndFilters(melBins, freqBins, sampleRate int) ([]float32, []float32) {
if f == nil {
return defaultParakeetWindow(), buildSlaneyMelFilterBank(freqBins, melBins, sampleRate)
}
f.mu.Lock()
defer f.mu.Unlock()
if f.window == nil {
if f.Window != nil {
if values := f.Window.BackendGet(); len(values) == parakeetWinLength {
f.window = values
}
}
if f.window == nil {
f.window = defaultParakeetWindow()
}
}
if f.fb == nil {
if f.FB != nil {
if values := f.FB.BackendGet(); len(values) == melBins*freqBins {
f.fb = values
f.fbShape = [2]int{melBins, freqBins}
}
}
if f.fb == nil {
f.fb = buildSlaneyMelFilterBank(freqBins, melBins, sampleRate)
f.fbShape = [2]int{melBins, freqBins}
}
}
return f.window, f.fb
}
type AudioSubsampling struct {
Conv0 *nn.Conv2D `gguf:"conv0"`
DW1 *AudioDepthwiseConv2D `gguf:"dw1"`
PW1 *nn.Conv2D `gguf:"pw1"`
DW2 *AudioDepthwiseConv2D `gguf:"dw2"`
PW2 *nn.Conv2D `gguf:"pw2"`
Linear *nn.Linear `gguf:"linear"`
}
type AudioDepthwiseConv2D struct {
Weight ml.Tensor `gguf:"weight"`
Bias ml.Tensor `gguf:"bias"`
}
type AudioFeedForward struct {
Up *nn.Linear `gguf:"up"`
Down *nn.Linear `gguf:"down"`
}
type AudioSelfAttention struct {
Query *nn.Linear `gguf:"attn_q"`
Key *nn.Linear `gguf:"attn_k"`
Value *nn.Linear `gguf:"attn_v"`
Output *nn.Linear `gguf:"attn_out"`
RelativeKey *nn.Linear `gguf:"attn_rel_k"`
BiasU ml.Tensor `gguf:"attn_bias_u"`
BiasV ml.Tensor `gguf:"attn_bias_v"`
}
type AudioConvolutionModule struct {
Pointwise1 *nn.Linear `gguf:"conv_pw1"`
Depthwise ml.Tensor `gguf:"conv_dw.weight"`
BatchNorm *AudioBatchNorm1D `gguf:"conv_bn"`
Pointwise2 *nn.Linear `gguf:"conv_pw2"`
}
type AudioBatchNorm1D struct {
Weight ml.Tensor `gguf:"weight"`
Bias ml.Tensor `gguf:"bias"`
RunningMean ml.Tensor `gguf:"running_mean"`
RunningVar ml.Tensor `gguf:"running_var"`
}
type AudioLayer struct {
FFN1Norm *nn.LayerNorm `gguf:"ffn1_norm"`
FFN1Up *nn.Linear `gguf:"ffn1_up"`
FFN1Down *nn.Linear `gguf:"ffn1_down"`
AttentionNorm *nn.LayerNorm `gguf:"attn_norm"`
Attention *AudioSelfAttention
ConvNorm *nn.LayerNorm `gguf:"conv_norm"`
Conv *AudioConvolutionModule
FFN2Norm *nn.LayerNorm `gguf:"ffn2_norm"`
FFN2Up *nn.Linear `gguf:"ffn2_up"`
FFN2Down *nn.Linear `gguf:"ffn2_down"`
OutputNorm *nn.LayerNorm `gguf:"out_norm"`
}
type AudioModel struct {
FeatureExtractor *AudioFeatureExtractor `gguf:"feature_extractor"`
Subsampling *AudioSubsampling `gguf:"subsampling"`
Layers []AudioLayer `gguf:"blk"`
*AudioOptions
}
type AudioProjector struct {
Norm *nn.RMSNorm `gguf:"norm"`
Linear1 *nn.Linear `gguf:"1"`
Linear2 *nn.Linear `gguf:"2"`
}
func (p *AudioProjector) Forward(ctx ml.Context, x ml.Tensor, eps float32) ml.Tensor {
x = p.Norm.Forward(ctx, x, eps)
x = audioF32(ctx, p.Linear1.Forward(ctx, x))
x = x.RELU(ctx)
x = x.Mul(ctx, x)
return audioF32(ctx, p.Linear2.Forward(ctx, x))
}
func (m *AudioModel) ForwardAudio(ctx ml.Context, melFeatures ml.Tensor, validFrames int, projector *AudioProjector) ml.Tensor {
x := melFeatures.Reshape(ctx, melFeatures.Dim(0), melFeatures.Dim(1), 1, 1)
validLen := validFrames
x = forwardAudioConv2D(ctx, m.Subsampling.Conv0, x, m.subsamplingStride, m.subsamplingStride, audioConvPadding(m.subsamplingKernel), audioConvPadding(m.subsamplingKernel), 1, 1)
x = x.RELU(ctx)
validLen = convOutputLength(validLen, m.subsamplingKernel, m.subsamplingStride, audioConvPadding(m.subsamplingKernel))
x = applyAudioTimeMask(ctx, x, validLen)
x = forwardAudioDepthwiseConv2D(ctx, m.Subsampling.DW1, x, m.subsamplingStride, m.subsamplingStride, audioConvPadding(m.subsamplingKernel), audioConvPadding(m.subsamplingKernel), 1, 1)
x = forwardAudioConv2D(ctx, m.Subsampling.PW1, x, 1, 1, 0, 0, 1, 1)
x = x.RELU(ctx)
validLen = convOutputLength(validLen, m.subsamplingKernel, m.subsamplingStride, audioConvPadding(m.subsamplingKernel))
x = applyAudioTimeMask(ctx, x, validLen)
x = forwardAudioDepthwiseConv2D(ctx, m.Subsampling.DW2, x, m.subsamplingStride, m.subsamplingStride, audioConvPadding(m.subsamplingKernel), audioConvPadding(m.subsamplingKernel), 1, 1)
x = forwardAudioConv2D(ctx, m.Subsampling.PW2, x, 1, 1, 0, 0, 1, 1)
x = x.RELU(ctx)
validLen = convOutputLength(validLen, m.subsamplingKernel, m.subsamplingStride, audioConvPadding(m.subsamplingKernel))
x = applyAudioTimeMask(ctx, x, validLen)
x = flattenAudioSubsamplingOutput(ctx, x)
x = m.Subsampling.Linear.Forward(ctx, x)
if m.scaleInput {
x = x.Scale(ctx, math.Sqrt(float64(m.hiddenSize)))
}
if validLen > 0 && validLen < x.Dim(1) {
x = x.Slice(ctx, 1, 0, validLen, 1).Contiguous(ctx)
}
for i := range m.Layers {
x = m.Layers[i].Forward(ctx, x, validLen, m.AudioOptions)
}
if projector != nil {
x = projector.Forward(ctx, x, m.eps)
}
return x
}
func flattenAudioSubsamplingOutput(ctx ml.Context, x ml.Tensor) ml.Tensor {
fOut, tOut, cOut := x.Dim(0), x.Dim(1), x.Dim(2)
// PyTorch flattens the subsampling output after [B, C, T, F] ->
// [B, T, C, F], so F must remain the fastest dimension inside each
// channel block before the linear projection.
x = x.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
return x.Reshape(ctx, cOut*fOut, tOut)
}
func (l *AudioLayer) Forward(ctx ml.Context, x ml.Tensor, validLen int, opts *AudioOptions) ml.Tensor {
residual := x
x = audioFeedForward(ctx, l.FFN1Up, l.FFN1Down, l.FFN1Norm.Forward(ctx, x, opts.eps)).Scale(ctx, 0.5)
x = residual.Add(ctx, x)
residual = x
x = l.Attention.Forward(ctx, l.AttentionNorm.Forward(ctx, x, opts.eps), validLen, opts)
x = residual.Add(ctx, x)
residual = x
x = l.Conv.Forward(ctx, l.ConvNorm.Forward(ctx, x, opts.eps), opts)
x = residual.Add(ctx, x)
residual = x
x = audioFeedForward(ctx, l.FFN2Up, l.FFN2Down, l.FFN2Norm.Forward(ctx, x, opts.eps)).Scale(ctx, 0.5)
x = residual.Add(ctx, x)
return l.OutputNorm.Forward(ctx, x, opts.eps)
}
func audioFeedForward(ctx ml.Context, up, down *nn.Linear, x ml.Tensor) ml.Tensor {
x = audioF32(ctx, up.Forward(ctx, x))
x = x.SILU(ctx)
return audioF32(ctx, down.Forward(ctx, x))
}
func (a *AudioSelfAttention) Forward(ctx ml.Context, x ml.Tensor, validLen int, opts *AudioOptions) ml.Tensor {
seqLen := x.Dim(1)
headDim := opts.headDim
numHeads := opts.numHeads
q := audioF32(ctx, a.Query.Forward(ctx, x)).Reshape(ctx, headDim, numHeads, seqLen)
k := audioF32(ctx, a.Key.Forward(ctx, x)).Reshape(ctx, headDim, numHeads, seqLen)
v := audioF32(ctx, a.Value.Forward(ctx, x)).Reshape(ctx, headDim, numHeads, seqLen)
qU := q
if a.BiasU != nil {
qU = qU.Add(ctx, audioF32(ctx, a.BiasU).Reshape(ctx, headDim, numHeads, 1))
}
qV := q
if a.BiasV != nil {
qV = qV.Add(ctx, audioF32(ctx, a.BiasV).Reshape(ctx, headDim, numHeads, 1))
}
qP := qU.Permute(ctx, 0, 2, 1, 3)
kP := k.Permute(ctx, 0, 2, 1, 3)
logits := kP.MulmatFullPrec(ctx, qP)
positionEmbeddings := parakeetPositionEmbeddings(ctx, seqLen, opts.hiddenSize)
relKey := audioF32(ctx, a.RelativeKey.Forward(ctx, positionEmbeddings)).Reshape(ctx, headDim, numHeads, 2*seqLen-1)
pP := relKey.Permute(ctx, 0, 2, 1, 3)
qVP := qV.Permute(ctx, 0, 2, 1, 3)
relLogits := pP.MulmatFullPrec(ctx, qVP)
relLogits = relativeShiftParakeet(ctx, relLogits, seqLen, numHeads)
logits = logits.Add(ctx, relLogits)
logits = logits.Scale(ctx, math.Pow(float64(headDim), -0.5))
if validLen > 0 && validLen < seqLen {
logits = logits.Add(ctx, audioAttentionMask(ctx, seqLen, validLen))
}
logits = logits.Softmax(ctx)
vP := v.Permute(ctx, 0, 2, 1, 3)
vPT := vP.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx)
out := vPT.Mulmat(ctx, logits)
out = out.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
out = out.Reshape(ctx, opts.hiddenSize, seqLen)
return audioF32(ctx, a.Output.Forward(ctx, out))
}
func (c *AudioConvolutionModule) Forward(ctx ml.Context, x ml.Tensor, opts *AudioOptions) ml.Tensor {
x = audioF32(ctx, c.Pointwise1.Forward(ctx, x))
hidden := x.Dim(0) / 2
value := x.Slice(ctx, 0, 0, hidden, 1).Contiguous(ctx)
gate := x.Slice(ctx, 0, hidden, 2*hidden, 1).Contiguous(ctx).Sigmoid(ctx)
x = value.Mul(ctx, gate)
x = audioDepthwiseConv1DSame(ctx, x, c.Depthwise, audioConvPadding(opts.convKernelSize))
x = c.BatchNorm.Forward(ctx, x, opts.eps)
x = x.SILU(ctx)
return audioF32(ctx, c.Pointwise2.Forward(ctx, x))
}
func audioF32(ctx ml.Context, x ml.Tensor) ml.Tensor {
if x.DType() == ml.DTypeF32 {
return x
}
// Metal binary kernels used by the audio graph require F32 operands here.
// This likely slows audio and should be revisited once the precision vs.
// speed tradeoff is validated against BF16-native elementwise paths.
return x.Cast(ctx, ml.DTypeF32)
}
func (b *AudioBatchNorm1D) Forward(ctx ml.Context, x ml.Tensor, eps float32) ml.Tensor {
if b == nil || b.RunningMean == nil || b.RunningVar == nil {
return x
}
hidden := x.Dim(0)
epsValues := make([]float32, hidden)
for i := range epsValues {
epsValues[i] = eps
}
variance := b.RunningVar.Add(ctx, ctx.Input().FromFloats(epsValues, hidden))
x = x.Sub(ctx, b.RunningMean)
x = x.Div(ctx, variance.Sqrt(ctx))
if b.Weight != nil {
x = x.Mul(ctx, b.Weight)
}
if b.Bias != nil {
x = x.Add(ctx, b.Bias)
}
return x
}
func forwardAudioConv2D(ctx ml.Context, conv *nn.Conv2D, x ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
weight := conv.Weight.Contiguous(ctx)
x = weight.Conv2D(ctx, x, s0, s1, p0, p1, d0, d1)
if conv.Bias != nil {
x = x.Add(ctx, conv.Bias.Reshape(ctx, 1, 1, -1))
}
return x
}
func forwardAudioDepthwiseConv2D(ctx ml.Context, conv *AudioDepthwiseConv2D, x ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
x = audioDepthwiseConv2D(ctx, x, conv.Weight, s0, s1, p0, p1, d0, d1)
if conv.Bias != nil {
x = x.Add(ctx, conv.Bias.Reshape(ctx, 1, 1, -1))
}
return x
}
func applyAudioTimeMask(ctx ml.Context, x ml.Tensor, validLen int) ml.Tensor {
if validLen <= 0 || validLen >= x.Dim(1) {
return x
}
mask := make([]float32, x.Dim(1))
for i := range validLen {
mask[i] = 1
}
return x.Mul(ctx, ctx.Input().FromFloats(mask, 1, x.Dim(1), 1, 1))
}
func audioDepthwiseConv1DSame(ctx ml.Context, x, kernel ml.Tensor, padding int) ml.Tensor {
kernelSize := kernel.Dim(0)
seqLen := x.Dim(1)
kernelT := kernel.Permute(ctx, 1, 0, 2, 3).Contiguous(ctx)
var out ml.Tensor
for k := range kernelSize {
offset := k - padding
shifted := x
switch {
case offset > 0:
shifted = x.Slice(ctx, 1, offset, seqLen, 1).Contiguous(ctx)
shifted = shifted.PadExt(ctx, 0, 0, 0, offset, 0, 0, 0, 0)
case offset < 0:
shift := -offset
shifted = x.Slice(ctx, 1, 0, seqLen-shift, 1).Contiguous(ctx)
shifted = shifted.PadExt(ctx, 0, 0, shift, 0, 0, 0, 0, 0)
}
wk := kernelT.Slice(ctx, 1, k, k+1, 1).Contiguous(ctx)
term := shifted.Mul(ctx, wk)
if out == nil {
out = term
} else {
out = out.Add(ctx, term)
}
}
return out
}
func audioDepthwiseConv2D(ctx ml.Context, x, kernel ml.Tensor, s0, s1, p0, p1, d0, d1 int) ml.Tensor {
if d0 != 1 || d1 != 1 {
panic("audio depthwise conv2d only supports dilation 1")
}
kernel = kernel.Contiguous(ctx)
kernelW, kernelH := kernel.Dim(0), kernel.Dim(1)
outW := convOutputLength(x.Dim(0), kernelW, s0, p0)
outH := convOutputLength(x.Dim(1), kernelH, s1, p1)
padded := x.PadExt(ctx, p0, p0, p1, p1, 0, 0, 0, 0)
var out ml.Tensor
for ky := range kernelH {
for kx := range kernelW {
patch := padded.Slice(ctx, 0, kx, kx+s0*(outW-1)+1, s0).Contiguous(ctx)
patch = patch.Slice(ctx, 1, ky, ky+s1*(outH-1)+1, s1).Contiguous(ctx)
wk := kernel.Slice(ctx, 0, kx, kx+1, 1).Slice(ctx, 1, ky, ky+1, 1).Contiguous(ctx)
if wk.Dim(2) == 1 {
wk = wk.Permute(ctx, 0, 1, 3, 2).Contiguous(ctx)
} else {
wk = wk.Reshape(ctx, 1, 1, wk.Dim(2), wk.Dim(3))
}
term := patch.Mul(ctx, wk)
if out == nil {
out = term
} else {
out = out.Add(ctx, term)
}
}
}
return out
}
func convOutputLength(inputLength, kernel, stride, padding int) int {
if inputLength <= 0 {
return 0
}
return (inputLength+2*padding-kernel)/stride + 1
}
func audioConvPadding(kernel int) int {
return (kernel - 1) / 2
}
func parakeetPositionEmbeddings(ctx ml.Context, seqLen, hiddenSize int) ml.Tensor {
half := hiddenSize / 2
values := make([]float32, hiddenSize*(2*seqLen-1))
for posIdx, pos := 0, seqLen-1; posIdx < 2*seqLen-1; posIdx, pos = posIdx+1, pos-1 {
for i := range half {
invFreq := math.Pow(10000, -float64(2*i)/float64(hiddenSize))
angle := float64(pos) * invFreq
values[posIdx*hiddenSize+2*i] = float32(math.Sin(angle))
values[posIdx*hiddenSize+2*i+1] = float32(math.Cos(angle))
}
}
return ctx.Input().FromFloats(values, hiddenSize, 2*seqLen-1)
}
func relativeShiftParakeet(ctx ml.Context, x ml.Tensor, seqLen, numHeads int) ml.Tensor {
positionLen := 2*seqLen - 1
x = x.PadExt(ctx, 1, 0, 0, 0, 0, 0, 0, 0)
x = x.Reshape(ctx, seqLen, positionLen+1, numHeads)
x = x.Slice(ctx, 1, 1, positionLen+1, 1).Contiguous(ctx)
x = x.Reshape(ctx, positionLen, seqLen, numHeads)
return x.Slice(ctx, 0, 0, seqLen, 1).Contiguous(ctx)
}
func audioAttentionMask(ctx ml.Context, seqLen, validLen int) ml.Tensor {
values := make([]float32, seqLen*seqLen)
for q := range seqLen {
for k := range seqLen {
if q >= validLen || k >= validLen {
values[q*seqLen+k] = -1e9
}
}
}
return ctx.Input().FromFloats(values, seqLen, seqLen, 1)
}
func newAudioModel(c fs.Config) *AudioModel {
numLayers := int(c.Uint("audio.block_count", 0))
if numLayers == 0 {
return nil
}
return &AudioModel{
Layers: make([]AudioLayer, numLayers),
AudioOptions: newAudioOptions(c),
}
}
func newAudioProjector(c fs.Config) *AudioProjector {
if c.Uint("audio.block_count", 0) == 0 {
return nil
}
return &AudioProjector{}
}
func newAudioOptions(c fs.Config) *AudioOptions {
hiddenSize := int(c.Uint("audio.embedding_length", 1024))
numHeads := int(c.Uint("audio.attention.head_count", 8))
headDim := hiddenSize / max(1, numHeads)
return &AudioOptions{
hiddenSize: hiddenSize,
numHeads: numHeads,
headDim: headDim,
intermediateSize: int(c.Uint("audio.feed_forward_length", uint32(hiddenSize*4))),
convKernelSize: int(c.Uint("audio.conv_kernel_size", 9)),
melBins: int(c.Uint("audio.num_mel_bins", 128)),
sampleRate: int(c.Uint("audio.sample_rate", 16000)),
subsamplingKernel: int(c.Uint("audio.subsampling_conv_kernel_size", 3)),
subsamplingStride: int(c.Uint("audio.subsampling_conv_stride", 2)),
scaleInput: c.Bool("audio.scale_input", false),
eps: c.Float("audio.attention.layer_norm_epsilon", 1e-5),
}
}
func defaultAudioOptions() *AudioOptions {
return &AudioOptions{
hiddenSize: 1024,
numHeads: 8,
headDim: 128,
intermediateSize: 4096,
convKernelSize: 9,
melBins: 128,
sampleRate: 16000,
subsamplingKernel: 3,
subsamplingStride: 2,
eps: 1e-5,
}
}
model/models/nemotronh/model_omni.go
+239
View File
@@ -0,0 +1,239 @@
package nemotronh
import (
"bytes"
"errors"
"image"
"slices"
"github.com/ollama/ollama/fs"
"github.com/ollama/ollama/ml"
"github.com/ollama/ollama/model"
"github.com/ollama/ollama/model/input"
)
type OmniModel struct {
*Model
*VisionModel `gguf:"v"`
*AudioModel `gguf:"a"`
*MultiModalProjector `gguf:"mm"`
*AudioProjector `gguf:"mm.a"`
ImageProcessor
imageTokenID int32
imageStartToken int32
imageEndToken int32
audioTokenID int32
}
var _ model.MultimodalProcessor = (*OmniModel)(nil)
func NewOmni(c fs.Config) (model.Model, error) {
textModel, err := newTextModel(c)
if err != nil {
return nil, err
}
imageTokenID := int32(c.Uint("vision.image_token_id", 18))
imageStartToken := int32(c.Uint("vision.image_start_token_id", 19))
imageEndToken := int32(c.Uint("vision.image_end_token_id", 20))
audioTokenID := int32(c.Uint("audio.sound_token_id", 27))
return &OmniModel{
Model: textModel,
VisionModel: newVisionModel(c),
AudioModel: newAudioModel(c),
MultiModalProjector: newMultiModalProjector(c),
AudioProjector: newAudioProjector(c),
ImageProcessor: newImageProcessor(c),
imageTokenID: imageTokenID,
imageStartToken: imageStartToken,
imageEndToken: imageEndToken,
audioTokenID: audioTokenID,
}, nil
}
func (m *OmniModel) EncodeMultimodal(ctx ml.Context, multimodalData []byte) ([]input.Multimodal, error) {
if isAudioData(multimodalData) {
return m.encodeAudioMultimodal(ctx, multimodalData)
}
if m.VisionModel == nil || m.MultiModalProjector == nil || len(m.VisionModel.Layers) == 0 {
return nil, model.ErrNoVisionModel
}
img, _, err := image.Decode(bytes.NewReader(multimodalData))
if err != nil {
return nil, err
}
tiles, err := m.ImageProcessor.ProcessImage(img)
if err != nil {
return nil, err
}
mm := make([]input.Multimodal, 0, len(tiles))
for _, tile := range tiles {
patches := visionPatchGrid{
Width: tile.size.X / m.ImageProcessor.patchSize,
Height: tile.size.Y / m.ImageProcessor.patchSize,
}
if patches.Width == 0 || patches.Height == 0 {
return nil, errors.New("nemotron_h_omni: invalid resized image dimensions")
}
patchInput := packVisionPatchesCHW(tile.data, tile.size.X, tile.size.Y, m.ImageProcessor.numChannels, m.ImageProcessor.patchSize)
visionOutputs := m.VisionModel.ForwardPacked(ctx, patchInput, patches)
projected := m.MultiModalProjector.Forward(ctx, visionOutputs, patches)
mm = append(mm, input.Multimodal{Tensor: projected})
}
return mm, nil
}
type audioTag struct{}
func (m *OmniModel) encodeAudioMultimodal(ctx ml.Context, data []byte) ([]input.Multimodal, error) {
if m.AudioModel == nil || m.AudioProjector == nil || len(m.AudioModel.Layers) == 0 {
return nil, model.ErrNoVisionModel
}
samples, err := decodeWAV(data, m.AudioModel.sampleRate)
if err != nil {
return nil, err
}
melData, frames, validFrames, err := computeParakeetMelSpectrogram(samples, m.AudioModel.FeatureExtractor, m.AudioModel.AudioOptions)
if err != nil {
return nil, err
}
melTensor := ctx.Input().FromFloats(melData, m.AudioModel.melBins, frames)
audioOutputs := m.AudioModel.ForwardAudio(ctx, melTensor, validFrames, m.AudioProjector)
return []input.Multimodal{{Tensor: audioOutputs, Data: audioTag{}}}, nil
}
func (m *OmniModel) PostLoad() error {
return nil
}
func (m *OmniModel) PostTokenize(inputs []*input.Input) ([]*input.Input, error) {
var result []*input.Input
imageToken := m.imageTokenID
if imageToken == 0 {
imageToken = 18
}
for _, inp := range inputs {
if len(inp.Multimodal) == 0 {
result = append(result, inp)
continue
}
totalTokens := 0
for _, mm := range inp.Multimodal {
if mm.Tensor == nil {
continue
}
totalTokens += mm.Tensor.Dim(1)
}
if totalTokens <= 0 {
return nil, errors.New("nemotron_h_omni: multimodal input has no tokens")
}
if _, ok := inp.Multimodal[0].Data.(audioTag); ok {
audioToken := m.audioTokenID
if audioToken == 0 {
audioToken = 27
}
for i, mm := range inp.Multimodal {
tokenCount := 0
if mm.Tensor != nil {
tokenCount = mm.Tensor.Dim(1)
}
if tokenCount <= 0 {
return nil, errors.New("nemotron_h_omni: multimodal input has no tokens")
}
first := &input.Input{Token: audioToken, SameBatch: tokenCount - 1}
if i == 0 {
first.MultimodalHash = inp.MultimodalHash
}
first.Multimodal = []input.Multimodal{mm}
result = append(result, first)
if tokenCount > 1 {
result = append(result, slices.Repeat([]*input.Input{{Token: audioToken}}, tokenCount-1)...)
}
}
continue
}
if m.imageStartToken > 0 {
result = append(result, &input.Input{
Token: m.imageStartToken,
SameBatch: totalTokens + btoi(m.imageEndToken > 0),
})
}
for _, mm := range inp.Multimodal {
tokenCount := 0
if mm.Tensor != nil {
tokenCount = mm.Tensor.Dim(1)
}
if tokenCount <= 0 {
return nil, errors.New("nemotron_h_omni: multimodal input has no tokens")
}
result = append(result, &input.Input{
Token: imageToken,
Multimodal: []input.Multimodal{mm},
MultimodalHash: inp.MultimodalHash,
})
if tokenCount > 1 {
result = append(result, slices.Repeat([]*input.Input{{Token: imageToken}}, tokenCount-1)...)
}
}
if m.imageEndToken > 0 {
result = append(result, &input.Input{Token: m.imageEndToken})
}
}
return result, nil
}
func btoi(v bool) int {
if v {
return 1
}
return 0
}
func (m *OmniModel) Forward(ctx ml.Context, batch input.Batch) (ml.Tensor, error) {
hiddenStates := m.TokenEmbedding.Forward(ctx, batch.Inputs)
if len(batch.Multimodal) > 0 {
hiddenStates = hiddenStates.Duplicate(ctx)
}
for _, mm := range batch.Multimodal {
offset := mm.Index
for _, multimodal := range mm.Multimodal {
if multimodal.Tensor == nil {
continue
}
tensor := multimodal.Tensor
ctx.Forward(tensor.Copy(ctx, hiddenStates.View(ctx, offset*hiddenStates.Stride(1), tensor.Dim(0)*tensor.Dim(1))))
offset += tensor.Dim(1)
}
}
return m.forwardLogits(ctx, batch, hiddenStates)
}
func init() {
model.Register("nemotron_h_omni", NewOmni)
}
model/models/nemotronh/model_omni_test.go
+606
View File
@@ -0,0 +1,606 @@
package nemotronh
import (
"bytes"
"encoding/base64"
"encoding/binary"
"image"
"image/color"
"math"
"os"
"path/filepath"
"slices"
"strings"
"testing"
fsggml "github.com/ollama/ollama/fs/ggml"
"github.com/ollama/ollama/ml"
backendggml "github.com/ollama/ollama/ml/backend/ggml"
"github.com/ollama/ollama/ml/nn"
"github.com/ollama/ollama/model/input"
)
type fakeTensor struct {
*backendggml.Tensor
dims []int
}
func (t *fakeTensor) Dim(i int) int {
return t.dims[i]
}
func setupTestContext(t *testing.T) ml.Context {
t.Helper()
f, err := os.CreateTemp(t.TempDir(), "*.gguf")
if err != nil {
t.Fatal(err)
}
defer f.Close()
if err := fsggml.WriteGGUF(f, fsggml.KV{"general.architecture": "test"}, nil); err != nil {
t.Fatal(err)
}
b, err := ml.NewBackend(f.Name(), ml.BackendParams{AllocMemory: true})
if err != nil {
t.Fatal(err)
}
ctx := b.NewContext().Input()
t.Cleanup(func() {
ctx.Close()
b.Close()
})
return ctx
}
func TestPostTokenizeImageSpans(t *testing.T) {
m := &OmniModel{
imageTokenID: 18,
imageStartToken: 19,
imageEndToken: 20,
}
makeChunk := func() input.Multimodal {
return input.Multimodal{Tensor: &fakeTensor{dims: []int{2688, 256, 1, 1}}}
}
in := []*input.Input{
{Token: 7},
{
Multimodal: []input.Multimodal{makeChunk(), makeChunk()},
MultimodalHash: 99,
},
{Token: 8},
}
out, err := m.PostTokenize(in)
if err != nil {
t.Fatalf("PostTokenize() error = %v", err)
}
if len(out) != 516 {
t.Fatalf("len(out) = %d, want 516", len(out))
}
if out[0].Token != 7 {
t.Fatalf("out[0].Token = %d, want 7", out[0].Token)
}
if out[1].Token != 19 {
t.Fatalf("out[1].Token = %d, want 19", out[1].Token)
}
if out[1].SameBatch != 513 {
t.Fatalf("out[1].SameBatch = %d, want 513", out[1].SameBatch)
}
if out[2].Token != 18 || len(out[2].Multimodal) != 1 || out[2].MultimodalHash != 99 || out[2].SameBatch != 0 {
t.Fatalf("unexpected first image token: %+v", *out[2])
}
if out[258].Token != 18 || len(out[258].Multimodal) != 1 || out[258].MultimodalHash != 99 || out[258].SameBatch != 0 {
t.Fatalf("unexpected second image token: %+v", *out[258])
}
if out[514].Token != 20 {
t.Fatalf("out[514].Token = %d, want 20", out[514].Token)
}
if out[515].Token != 8 {
t.Fatalf("out[515].Token = %d, want 8", out[515].Token)
}
}
func TestProjectorPixelShuffleMatchesReferenceV2Order(t *testing.T) {
ctx := setupTestContext(t)
hidden := 2
width := 4
height := 2
values := make([]float32, 0, hidden*width*height)
for y := range height {
for x := range width {
for c := range hidden {
values = append(values, float32(100*y+10*x+c))
}
}
}
got := pixelShuffleVisionOutputs(ctx, ctx.FromFloats(values, hidden, width*height), visionPatchGrid{
Width: width,
Height: height,
}, 2)
ctx.Forward(got).Compute(got)
want := []float32{
0, 1, 10, 11, 100, 101, 110, 111,
20, 21, 30, 31, 120, 121, 130, 131,
}
if got.Shape()[0] != 8 || got.Shape()[1] != 2 {
t.Fatalf("shape = %v, want [8 2 1]", got.Shape())
}
gotValues := got.BackendGet()
if len(gotValues) != len(want) {
t.Fatalf("len(got) = %d, want %d", len(gotValues), len(want))
}
for i := range want {
if gotValues[i] != want[i] {
t.Fatalf("got[%d] = %v, want %v", i, gotValues[i], want[i])
}
}
}
func TestPostTokenizeAudioSpans(t *testing.T) {
m := &OmniModel{
audioTokenID: 27,
}
in := []*input.Input{
{Token: 7},
{
Multimodal: []input.Multimodal{{
Tensor: &fakeTensor{dims: []int{2688, 13, 1, 1}},
Data: audioTag{},
}},
MultimodalHash: 99,
},
{Token: 8},
}
out, err := m.PostTokenize(in)
if err != nil {
t.Fatalf("PostTokenize() error = %v", err)
}
if len(out) != 15 {
t.Fatalf("len(out) = %d, want 15", len(out))
}
if out[0].Token != 7 || out[14].Token != 8 {
t.Fatalf("unexpected surrounding tokens: first=%d last=%d", out[0].Token, out[14].Token)
}
for i := 1; i <= 13; i++ {
if out[i].Token != 27 {
t.Fatalf("out[%d].Token = %d, want 27", i, out[i].Token)
}
}
if len(out[1].Multimodal) != 1 || out[1].MultimodalHash != 99 {
t.Fatalf("first audio token did not carry multimodal payload: %+v", *out[1])
}
if out[1].SameBatch != 12 {
t.Fatalf("first audio token SameBatch = %d, want 12", out[1].SameBatch)
}
if len(out[2].Multimodal) != 0 {
t.Fatalf("only the first audio token should carry multimodal payload: %+v", *out[2])
}
}
func TestParakeetAudioPreprocessShapes(t *testing.T) {
data := sineWAV(t, 16000, 440, 1.0)
samples, err := decodeWAV(data, 16000)
if err != nil {
t.Fatal(err)
}
if got, want := len(samples), 16000; got != want {
t.Fatalf("sample count = %d, want %d", got, want)
}
mel, frames, validFrames, err := computeParakeetMelSpectrogram(samples, nil, defaultAudioOptions())
if err != nil {
t.Fatal(err)
}
if frames != 101 {
t.Fatalf("frames = %d, want 101", frames)
}
if validFrames != 100 {
t.Fatalf("validFrames = %d, want 100", validFrames)
}
if len(mel) != 101*128 {
t.Fatalf("len(mel) = %d, want %d", len(mel), 101*128)
}
lastFrame := mel[100*128 : 101*128]
if !slices.Equal(lastFrame, make([]float32, 128)) {
t.Fatal("expected masked final frame to be zero")
}
}
func TestParakeetAudioPreprocessMatchesIntegrationWAVReference(t *testing.T) {
data := integrationAudioWAV(t)
samples, err := decodeWAV(data, 16000)
if err != nil {
t.Fatal(err)
}
if got, want := len(samples), 42083; got != want {
t.Fatalf("sample count = %d, want %d", got, want)
}
mel, frames, validFrames, err := computeParakeetMelSpectrogram(samples, nil, defaultAudioOptions())
if err != nil {
t.Fatal(err)
}
if frames != 264 {
t.Fatalf("frames = %d, want 264", frames)
}
if validFrames != 263 {
t.Fatalf("validFrames = %d, want 263", validFrames)
}
if len(mel) != 264*128 {
t.Fatalf("len(mel) = %d, want %d", len(mel), 264*128)
}
lastFrame := mel[263*128 : 264*128]
if !slices.Equal(lastFrame, make([]float32, 128)) {
t.Fatal("expected masked final frame to be zero")
}
// Reference values come from the ParakeetExtractor path used by vLLM:
// pre-emphasis, torch.stft(center=True, pad_mode="constant"), Slaney mel
// filters, log guard 2^-24, and per-mel normalization over valid frames.
checks := map[[2]int]float32{
{0, 0}: -1.0855197,
{0, 50}: -0.93212974,
{1, 10}: -0.9735168,
{2, 100}: -0.6533053,
{50, 0}: 2.2483668,
{50, 127}: -0.3828735,
{100, 50}: 2.9742377,
{262, 0}: -0.9521758,
{262, 127}: -0.4602786,
{263, 50}: 0,
}
for pos, want := range checks {
got := mel[pos[0]*128+pos[1]]
if math.Abs(float64(got-want)) > 1e-4 {
t.Errorf("mel[%d,%d] = %v, want %v", pos[0], pos[1], got, want)
}
}
}
func integrationAudioWAV(t *testing.T) []byte {
t.Helper()
path := filepath.Join("..", "..", "..", "integration", "audio_test_data_test.go")
b, err := os.ReadFile(path)
if err != nil {
t.Fatal(err)
}
const marker = "const audioEncodingPrompt = `"
s := string(b)
start := strings.Index(s, marker)
if start < 0 {
t.Fatal("audioEncodingPrompt marker not found")
}
start += len(marker)
end := strings.Index(s[start:], "`")
if end < 0 {
t.Fatal("audioEncodingPrompt terminator not found")
}
data, err := base64.StdEncoding.DecodeString(strings.TrimSpace(s[start : start+end]))
if err != nil {
t.Fatal(err)
}
return data
}
func TestRelativeShiftParakeetMatchesReference(t *testing.T) {
ctx := setupTestContext(t)
seqLen := 3
positionLen := 2*seqLen - 1
values := make([]float32, seqLen*positionLen)
for q := range seqLen {
for p := range positionLen {
values[q*positionLen+p] = float32(q*10 + p)
}
}
x := ctx.FromFloats(values, positionLen, seqLen, 1)
got := relativeShiftParakeet(ctx, x, seqLen, 1)
ctx.Forward(got).Compute(got)
want := []float32{
2, 3, 4,
11, 12, 13,
20, 21, 22,
}
if !slices.Equal(got.BackendGet(), want) {
t.Fatalf("relative shift mismatch:\n got %v\nwant %v", got.BackendGet(), want)
}
}
func TestAudioDepthwiseConv2DMatchesReference(t *testing.T) {
ctx := setupTestContext(t)
freq, frames, channels := 4, 5, 2
xValues := make([]float32, freq*frames*channels)
for i := range xValues {
xValues[i] = float32(i)/10 - 1
}
kernelValues := make([]float32, 3*3*channels)
for i := range kernelValues {
kernelValues[i] = float32(i)/7 - 1
}
x := ctx.FromFloats(xValues, freq, frames, channels, 1)
kernel := ctx.FromFloats(kernelValues, 3, 3, 1, channels)
bias := ctx.FromFloats([]float32{0.25, -0.5}, channels)
got := audioDepthwiseConv2D(ctx, x, kernel, 2, 2, 1, 1, 1, 1).Add(ctx, bias.Reshape(ctx, 1, 1, -1))
ctx.Forward(got).Compute(got)
want := []float32{
0.86428565, 1.3357141,
1.2785715, 1.3642857,
-0.5928571, -1.7499999,
5.4000001, 8.8142853,
10.514286, 16.042856,
6.6857138, 9.8428574,
}
assertCloseSlice(t, got.BackendGet(), want, 1e-5)
}
func TestFlattenAudioSubsamplingOutputMatchesReference(t *testing.T) {
ctx := setupTestContext(t)
const (
freq = 2
frames = 3
channels = 2
)
values := make([]float32, freq*frames*channels)
for c := range channels {
for t := range frames {
for f := range freq {
values[f+freq*(t+frames*c)] = float32(100*c + 10*t + f)
}
}
}
got := flattenAudioSubsamplingOutput(ctx, ctx.FromFloats(values, freq, frames, channels, 1))
ctx.Forward(got).Compute(got)
want := []float32{
0, 1, 100, 101,
10, 11, 110, 111,
20, 21, 120, 121,
}
assertCloseSlice(t, got.BackendGet(), want, 0)
}
func TestAudioDepthwiseConv1DMatchesReference(t *testing.T) {
ctx := setupTestContext(t)
xValues := make([]float32, 2*5)
for i := range xValues {
xValues[i] = float32(i)/5 - 0.7
}
kernelValues := make([]float32, 3*2)
for i := range kernelValues {
kernelValues[i] = float32(i)/3 - 0.5
}
x := ctx.FromFloats(xValues, 2, 5)
kernel := ctx.FromFloats(kernelValues, 3, 2)
got := audioDepthwiseConv1DSame(ctx, x, kernel, 1)
ctx.Forward(got).Compute(got)
want := []float32{
0.066666655, -0.5333333,
0.41666666, 0.016666688,
0.21666668, 1.0166667,
0.01666667, 2.0166664,
-0.40000004, 1.2666667,
}
assertCloseSlice(t, got.BackendGet(), want, 1e-5)
}
func TestAudioSelfAttentionMatchesReference(t *testing.T) {
ctx := setupTestContext(t)
const (
hiddenSize = 4
numHeads = 2
headDim = 2
seqLen = 3
)
xValues := make([]float32, hiddenSize*seqLen)
for i := range xValues {
xValues[i] = float32(i)/10 - 0.5
}
identity := make([]float32, hiddenSize*hiddenSize)
for i := range hiddenSize {
identity[i*hiddenSize+i] = 1
}
linear := func() *nn.Linear {
return &nn.Linear{Weight: ctx.FromFloats(identity, hiddenSize, hiddenSize)}
}
attn := &AudioSelfAttention{
Query: linear(),
Key: linear(),
Value: linear(),
Output: linear(),
RelativeKey: linear(),
BiasU: ctx.FromFloats([]float32{0.1, -0.2, 0.3, -0.4}, headDim, numHeads),
BiasV: ctx.FromFloats([]float32{-0.05, 0.07, 0.11, -0.13}, headDim, numHeads),
}
got := attn.Forward(ctx, ctx.FromFloats(xValues, hiddenSize, seqLen), seqLen, &AudioOptions{
hiddenSize: hiddenSize,
numHeads: numHeads,
headDim: headDim,
})
ctx.Forward(got).Compute(got)
want := []float32{
-0.08471569, 0.015284289, 0.05532019, 0.1553202,
-0.09135241, 0.008647568, 0.11468154, 0.21468155,
-0.019152153, 0.08084783, 0.1733382, 0.2733382,
}
assertCloseSlice(t, got.BackendGet(), want, 1e-5)
}
func assertCloseSlice(t *testing.T, got, want []float32, tolerance float64) {
t.Helper()
if len(got) != len(want) {
t.Fatalf("len(got) = %d, want %d", len(got), len(want))
}
for i := range want {
if math.Abs(float64(got[i]-want[i])) > tolerance {
t.Fatalf("got[%d] = %v, want %v\nall got: %v", i, got[i], want[i], got)
}
}
}
func TestPackPatchesCHW(t *testing.T) {
values := []float32{
0, 1, 2, 3,
4, 5, 6, 7,
8, 9, 10, 11,
12, 13, 14, 15,
100, 101, 102, 103,
104, 105, 106, 107,
108, 109, 110, 111,
112, 113, 114, 115,
}
got := packVisionPatchesCHW(values, 4, 4, 2, 2)
want := []float32{
0, 1, 4, 5, 100, 101, 104, 105,
2, 3, 6, 7, 102, 103, 106, 107,
8, 9, 12, 13, 108, 109, 112, 113,
10, 11, 14, 15, 110, 111, 114, 115,
}
if len(got) != len(want) {
t.Fatalf("len(got) = %d, want %d", len(got), len(want))
}
for i := range want {
if got[i] != want[i] {
t.Fatalf("got[%d] = %v, want %v", i, got[i], want[i])
}
}
}
func TestResizePositionEmbeddingMatchesReferenceInterpolation(t *testing.T) {
values := []float32{
0, 10,
20, 30,
}
got := resizePositionEmbedding(values, 1, 2, 2, 3, 3)
want := []float32{
0, 5, 10,
10, 15, 20,
20, 25, 30,
}
if len(got) != len(want) {
t.Fatalf("len(got) = %d, want %d", len(got), len(want))
}
for i := range want {
if got[i] != want[i] {
t.Fatalf("got[%d] = %v, want %v", i, got[i], want[i])
}
}
}
func TestDynamicImageProcessorMatchesReferencePatchBudget(t *testing.T) {
p := ImageProcessor{
imageSize: 512,
patchSize: 16,
numChannels: 3,
minNumPatches: 1024,
maxNumPatches: 13312,
projectorScale: 2,
imageMean: [3]float32{0.48145466, 0.4578275, 0.40821073},
imageStd: [3]float32{0.26862954, 0.26130258, 0.27577711},
}
img := image.NewRGBA(image.Rect(0, 0, 400, 250))
bounds := img.Bounds()
width, height := bounds.Dx(), bounds.Dy()
for y := range height {
for x := range width {
img.SetRGBA(x, y, color.RGBA{R: uint8(x), G: uint8(y), B: 128, A: 255})
}
}
tiles, err := p.ProcessImage(img)
if err != nil {
t.Fatalf("ProcessImage() error = %v", err)
}
if got, want := len(tiles), 1; got != want {
t.Fatalf("len(tiles) = %d, want %d", got, want)
}
if got, want := tiles[0].size, (image.Point{X: 672, Y: 416}); got != want {
t.Fatalf("tile size = %v, want %v", got, want)
}
if got, want := len(tiles[0].data), 3*672*416; got != want {
t.Fatalf("tile data len = %d, want %d", got, want)
}
}
func sineWAV(t *testing.T, sampleRate int, frequency float64, seconds float64) []byte {
t.Helper()
samples := int(float64(sampleRate) * seconds)
var pcm bytes.Buffer
for i := range samples {
v := int16(math.Sin(2*math.Pi*frequency*float64(i)/float64(sampleRate)) * 32767)
if err := binary.Write(&pcm, binary.LittleEndian, v); err != nil {
t.Fatal(err)
}
}
var out bytes.Buffer
out.WriteString("RIFF")
if err := binary.Write(&out, binary.LittleEndian, uint32(36+pcm.Len())); err != nil {
t.Fatal(err)
}
out.WriteString("WAVE")
out.WriteString("fmt ")
if err := binary.Write(&out, binary.LittleEndian, uint32(16)); err != nil {
t.Fatal(err)
}
if err := binary.Write(&out, binary.LittleEndian, uint16(1)); err != nil {
t.Fatal(err)
}
if err := binary.Write(&out, binary.LittleEndian, uint16(1)); err != nil {
t.Fatal(err)
}
if err := binary.Write(&out, binary.LittleEndian, uint32(sampleRate)); err != nil {
t.Fatal(err)
}
if err := binary.Write(&out, binary.LittleEndian, uint32(sampleRate*2)); err != nil {
t.Fatal(err)
}
if err := binary.Write(&out, binary.LittleEndian, uint16(2)); err != nil {
t.Fatal(err)
}
if err := binary.Write(&out, binary.LittleEndian, uint16(16)); err != nil {
t.Fatal(err)
}
out.WriteString("data")
if err := binary.Write(&out, binary.LittleEndian, uint32(pcm.Len())); err != nil {
t.Fatal(err)
}
out.Write(pcm.Bytes())
return out.Bytes()
}
model/models/nemotronh/model_vision.go
+348
View File
@@ -0,0 +1,348 @@
package nemotronh
import (
"math"
"sync"
"github.com/ollama/ollama/fs"
"github.com/ollama/ollama/ml"
"github.com/ollama/ollama/ml/nn"
)
const nemotronVisionBatchSize = 1
type visionPatchGrid struct {
Width int
Height int
}
type VisionPatchEmbedding struct {
*nn.Linear
}
func packVisionPatchesCHW(values []float32, width, height, channels, patchSize int) []float32 {
patchesX, patchesY := width/patchSize, height/patchSize
patchDim := channels * patchSize * patchSize
plane := width * height
patches := make([]float32, patchDim*patchesX*patchesY)
offset := 0
for py := range patchesY {
for px := range patchesX {
for c := range channels {
channelBase := c * plane
for yy := range patchSize {
rowBase := (py*patchSize + yy) * width
for xx := range patchSize {
patches[offset] = values[channelBase+rowBase+px*patchSize+xx]
offset++
}
}
}
}
}
return patches
}
func (p *VisionPatchEmbedding) ForwardPacked(ctx ml.Context, patches []float32, patchDim, numPatches int) ml.Tensor {
hiddenState := ctx.Input().FromFloats(patches, patchDim, numPatches)
hiddenState = hiddenState.Duplicate(ctx)
return p.Linear.Forward(ctx, hiddenState)
}
func (p *VisionPatchEmbedding) Forward(ctx ml.Context, pixelValues ml.Tensor, patchSize int) ml.Tensor {
// Match the RADIO patch generator's exact flattening order: patches are laid
// out token-major with each token packed as channel, then patch-row, then
// patch-col. This is more explicit than the prior IM2Col path and likely
// slower, but it avoids backend-specific packing differences that caused the
// converted patch embedder to diverge badly from the reference model.
width, height, channels := pixelValues.Dim(0), pixelValues.Dim(1), pixelValues.Dim(2)
patchesX, patchesY := width/patchSize, height/patchSize
patchDim := channels * patchSize * patchSize
values := pixelValues.BackendGet()
return p.ForwardPacked(ctx, packVisionPatchesCHW(values, width, height, channels, patchSize), patchDim, patchesX*patchesY)
}
type VisionSelfAttention struct {
Query *nn.Linear `gguf:"attn_q"`
Key *nn.Linear `gguf:"attn_k"`
Value *nn.Linear `gguf:"attn_v"`
Output *nn.Linear `gguf:"attn_out"`
}
func (sa *VisionSelfAttention) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *VisionOptions) ml.Tensor {
headDim := opts.hiddenSize / opts.numHeads
query := sa.Query.Forward(ctx, hiddenState)
key := sa.Key.Forward(ctx, hiddenState)
value := sa.Value.Forward(ctx, hiddenState)
query = query.Reshape(ctx, headDim, opts.numHeads, query.Dim(1), nemotronVisionBatchSize)
key = key.Reshape(ctx, headDim, opts.numHeads, key.Dim(1), nemotronVisionBatchSize)
value = value.Reshape(ctx, headDim, opts.numHeads, value.Dim(1), nemotronVisionBatchSize)
attention := nn.Attention(ctx, query, key, value, 1.0/math.Sqrt(float64(headDim)), nil)
attention = attention.Reshape(ctx, opts.hiddenSize, attention.Dim(2), nemotronVisionBatchSize)
return sa.Output.Forward(ctx, attention)
}
type VisionMLP struct {
Up *nn.Linear `gguf:"ffn_up"`
Down *nn.Linear `gguf:"ffn_down"`
}
func (mlp *VisionMLP) Forward(ctx ml.Context, hiddenState ml.Tensor) ml.Tensor {
return mlp.Down.Forward(ctx, mlp.Up.Forward(ctx, hiddenState).GELU(ctx))
}
type VisionEncoderLayer struct {
LayerNorm1 *nn.LayerNorm `gguf:"ln1"`
SelfAttention *VisionSelfAttention
LayerNorm2 *nn.LayerNorm `gguf:"ln2"`
MLP *VisionMLP
}
func (l *VisionEncoderLayer) Forward(ctx ml.Context, hiddenState ml.Tensor, opts *VisionOptions) ml.Tensor {
residual := hiddenState
hiddenState = l.LayerNorm1.Forward(ctx, hiddenState, opts.eps)
hiddenState = l.SelfAttention.Forward(ctx, hiddenState, opts)
hiddenState = hiddenState.Add(ctx, residual)
residual = hiddenState
hiddenState = l.LayerNorm2.Forward(ctx, hiddenState, opts.eps)
hiddenState = l.MLP.Forward(ctx, hiddenState)
return hiddenState.Add(ctx, residual)
}
type VisionOptions struct {
hiddenSize int
numHeads int
imageSize int
patchSize int
eps float32
}
type VisionModel struct {
PatchEmbedding *VisionPatchEmbedding `gguf:"patch_embd"`
PositionEmbedding ml.Tensor `gguf:"position_embd"`
ClassEmbedding ml.Tensor `gguf:"cls_embd"`
Layers []VisionEncoderLayer `gguf:"blk"`
*VisionOptions
resizedPositionEmbeddingsMu sync.Mutex
resizedPositionEmbeddings map[visionPatchGrid][]float32
}
func (m *VisionModel) Forward(ctx ml.Context, pixelValues ml.Tensor, patches visionPatchGrid) ml.Tensor {
numPatches := patches.Width * patches.Height
hiddenState := m.PatchEmbedding.Forward(ctx, pixelValues, m.patchSize)
return m.forwardPatchEmbeddings(ctx, hiddenState, patches, numPatches)
}
func (m *VisionModel) ForwardPacked(ctx ml.Context, patchValues []float32, patches visionPatchGrid) ml.Tensor {
numPatches := patches.Width * patches.Height
patchDim := 0
if numPatches > 0 {
patchDim = len(patchValues) / numPatches
}
hiddenState := m.PatchEmbedding.ForwardPacked(ctx, patchValues, patchDim, numPatches)
return m.forwardPatchEmbeddings(ctx, hiddenState, patches, numPatches)
}
func (m *VisionModel) forwardPatchEmbeddings(ctx ml.Context, hiddenState ml.Tensor, patches visionPatchGrid, numPatches int) ml.Tensor {
if m.PositionEmbedding != nil {
positionEmbeddings := m.positionEmbeddings(ctx, hiddenState, patches, numPatches)
hiddenState = hiddenState.Add(ctx, positionEmbeddings)
}
if m.ClassEmbedding != nil {
numPrefixTokens := m.ClassEmbedding.Dim(1)
classEmbeddings := m.ClassEmbedding.Cast(ctx, hiddenState.DType())
classEmbeddings = classEmbeddings.Reshape(ctx, classEmbeddings.Dim(0), numPrefixTokens, 1)
hiddenState = classEmbeddings.Concat(ctx, hiddenState, 1)
}
for _, layer := range m.Layers {
hiddenState = layer.Forward(ctx, hiddenState, m.VisionOptions)
}
if m.ClassEmbedding != nil {
hiddenState = hiddenState.Slice(ctx, 1, m.ClassEmbedding.Dim(1), hiddenState.Dim(1), 1)
}
return hiddenState.Reshape(ctx, hiddenState.Dim(0), hiddenState.Dim(1))
}
func (m *VisionModel) positionEmbeddings(ctx ml.Context, hiddenState ml.Tensor, patches visionPatchGrid, numPatches int) ml.Tensor {
posTokens := m.PositionEmbedding.Dim(1)
source := int(math.Sqrt(float64(posTokens)))
positionEmbeddings := m.PositionEmbedding.Cast(ctx, hiddenState.DType())
if !(source > 0 && source*source == posTokens && (source != patches.Width || source != patches.Height)) {
if positionEmbeddings.Dim(1) > numPatches {
positionEmbeddings = positionEmbeddings.Slice(ctx, 1, 0, numPatches, 1)
}
return positionEmbeddings
}
if cached, ok := m.cachePositionEmbeddings(ctx, hiddenState.Dim(0), patches); ok {
return ctx.Input().FromFloats(cached, hiddenState.Dim(0), numPatches)
}
// Runner fit/reserve builds worst-case multimodal graphs before weights are
// loaded, so the align-corners CPU cache path cannot materialize source
// values there. Fall back to a graph-only bilinear resize for reservation;
// the loaded inference path above still uses the cached align-corners data.
positionEmbeddings = positionEmbeddings.Reshape(ctx, -1, source, source)
positionEmbeddings = positionEmbeddings.Permute(ctx, 2, 0, 1, 3).Contiguous(ctx)
positionEmbeddings = positionEmbeddings.Interpolate(ctx, [4]int{
patches.Width,
patches.Height,
hiddenState.Dim(0),
1,
}, ml.SamplingModeBilinear)
positionEmbeddings = positionEmbeddings.Permute(ctx, 1, 2, 0, 3)
return positionEmbeddings.Contiguous(ctx, -1, patches.Width*patches.Height)
}
func (m *VisionModel) cachePositionEmbeddings(ctx ml.Context, hidden int, patches visionPatchGrid) ([]float32, bool) {
m.resizedPositionEmbeddingsMu.Lock()
cached := m.resizedPositionEmbeddings[patches]
m.resizedPositionEmbeddingsMu.Unlock()
if cached != nil {
return cached, true
}
if len(m.PositionEmbedding.Bytes()) == 0 {
return nil, false
}
posTokens := m.PositionEmbedding.Dim(1)
source := int(math.Sqrt(float64(posTokens)))
positionEmbeddingsF32 := m.PositionEmbedding.Cast(ctx, ml.DTypeF32)
ctx.Forward(positionEmbeddingsF32).Compute(positionEmbeddingsF32)
// RADIO eval-time CPE uses bilinear interpolation with align_corners=false.
// Cache a CPU-resized token-major embedding here for correctness first. This
// is likely slower than a native graph path and should be revisited if this
// precision vs speed tradeoff is not worthwhile.
cached = resizePositionEmbedding(positionEmbeddingsF32.Floats(), hidden, source, source, patches.Width, patches.Height)
m.resizedPositionEmbeddingsMu.Lock()
if m.resizedPositionEmbeddings == nil {
m.resizedPositionEmbeddings = make(map[visionPatchGrid][]float32)
}
if existing := m.resizedPositionEmbeddings[patches]; existing != nil {
cached = existing
} else {
m.resizedPositionEmbeddings[patches] = cached
}
m.resizedPositionEmbeddingsMu.Unlock()
return cached, true
}
func resizePositionEmbedding(values []float32, hidden, sourceWidth, sourceHeight, targetWidth, targetHeight int) []float32 {
out := make([]float32, hidden*targetWidth*targetHeight)
scaleX := float64(sourceWidth) / float64(targetWidth)
scaleY := float64(sourceHeight) / float64(targetHeight)
for oy := range targetHeight {
srcY := scaleY*(float64(oy)+0.5) - 0.5
y0 := int(math.Floor(srcY))
y1 := min(y0+1, sourceHeight-1)
wy := float32(srcY - float64(y0))
y0 = max(y0, 0)
for ox := range targetWidth {
srcX := scaleX*(float64(ox)+0.5) - 0.5
x0 := int(math.Floor(srcX))
x1 := min(x0+1, sourceWidth-1)
wx := float32(srcX - float64(x0))
x0 = max(x0, 0)
t00 := (y0*sourceWidth + x0) * hidden
t01 := (y0*sourceWidth + x1) * hidden
t10 := (y1*sourceWidth + x0) * hidden
t11 := (y1*sourceWidth + x1) * hidden
dst := (oy*targetWidth + ox) * hidden
for h := range hidden {
v00 := values[t00+h]
v01 := values[t01+h]
v10 := values[t10+h]
v11 := values[t11+h]
top := v00 + (v01-v00)*wx
bot := v10 + (v11-v10)*wx
out[dst+h] = top + (bot-top)*wy
}
}
}
return out
}
func newVisionModel(c fs.Config) *VisionModel {
return &VisionModel{
Layers: make([]VisionEncoderLayer, c.Uint("vision.block_count", 32)),
VisionOptions: &VisionOptions{
hiddenSize: int(c.Uint("vision.embedding_length", 1280)),
numHeads: int(c.Uint("vision.attention.head_count", 16)),
imageSize: int(c.Uint("vision.image_size", 512)),
patchSize: int(c.Uint("vision.patch_size", 16)),
eps: c.Float("vision.attention.layer_norm_epsilon", 1e-6),
},
}
}
type MultiModalProjector struct {
Norm *nn.RMSNorm `gguf:"norm"`
Linear1 *nn.Linear `gguf:"1"`
Linear2 *nn.Linear `gguf:"2"`
scaleFactor int
}
func (p *MultiModalProjector) Forward(ctx ml.Context, visionOutputs ml.Tensor, patches visionPatchGrid) ml.Tensor {
scaleFactor := max(p.scaleFactor, 1)
// The reference projector first pixel-shuffles the vision grid with
// downsample_ratio=0.5 before applying the RMSNorm/MLP. Preserve that exact
// v2 packing order here rather than flattening 2x2 neighborhoods via IM2Col.
merged := pixelShuffleVisionOutputs(ctx, visionOutputs, patches, scaleFactor)
merged = p.Norm.Forward(ctx, merged, 1e-5)
merged = p.Linear1.Forward(ctx, merged)
merged = merged.RELU(ctx)
merged = merged.Mul(ctx, merged)
return p.Linear2.Forward(ctx, merged)
}
func pixelShuffleVisionOutputs(ctx ml.Context, visionOutputs ml.Tensor, patches visionPatchGrid, scaleFactor int) ml.Tensor {
hiddenSize := visionOutputs.Dim(0)
scaleFactor = max(scaleFactor, 1)
merged := visionOutputs.Reshape(ctx, hiddenSize, patches.Width, patches.Height, 1)
width := patches.Width / scaleFactor
height := patches.Height / scaleFactor
channels := hiddenSize * scaleFactor
merged = merged.Reshape(ctx, channels, width, patches.Height, 1)
merged = merged.Reshape(ctx, channels, width, scaleFactor, height)
merged = merged.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
return merged.Reshape(ctx, channels*scaleFactor, width*height, 1)
}
func newMultiModalProjector(c fs.Config) *MultiModalProjector {
return &MultiModalProjector{
scaleFactor: int(c.Uint("vision.projector.scale_factor", 2)),
}
}
model/models/nemotronh/process_audio.go
+328
View File
@@ -0,0 +1,328 @@
package nemotronh
import (
"encoding/binary"
"fmt"
"math"
"math/cmplx"
)
const (
parakeetHopLength = 160
parakeetNFFT = 512
parakeetWinLength = 400
parakeetPreemphasis = 0.97
parakeetLogZeroGuardValue = 1.0 / (1 << 24)
parakeetNormalizeEps = 1e-5
)
func isAudioData(data []byte) bool {
return len(data) >= 12 && string(data[:4]) == "RIFF" && string(data[8:12]) == "WAVE"
}
func decodeWAV(data []byte, targetSampleRate int) ([]float32, error) {
if len(data) < 12 {
return nil, fmt.Errorf("WAV file too short")
}
if !isAudioData(data) {
return nil, fmt.Errorf("not a WAV file")
}
var audioFormat uint16
var numChannels, sampleRate, bitsPerSample int
var audioData []byte
foundFmt := false
offset := 12
for offset+8 <= len(data) {
chunkID := string(data[offset : offset+4])
chunkSize := int(binary.LittleEndian.Uint32(data[offset+4 : offset+8]))
chunkEnd := min(offset+8+chunkSize, len(data))
chunkData := data[offset+8 : chunkEnd]
switch chunkID {
case "fmt ":
if len(chunkData) < 16 {
return nil, fmt.Errorf("fmt chunk too short")
}
audioFormat = binary.LittleEndian.Uint16(chunkData[0:2])
numChannels = int(binary.LittleEndian.Uint16(chunkData[2:4]))
sampleRate = int(binary.LittleEndian.Uint32(chunkData[4:8]))
bitsPerSample = int(binary.LittleEndian.Uint16(chunkData[14:16]))
if audioFormat == 0xfffe && len(chunkData) >= 26 {
audioFormat = binary.LittleEndian.Uint16(chunkData[24:26])
}
foundFmt = true
case "data":
audioData = chunkData
}
offset += 8 + chunkSize
if chunkSize%2 != 0 {
offset++
}
}
if !foundFmt {
return nil, fmt.Errorf("no fmt chunk found in WAV file")
}
if audioFormat != 1 && audioFormat != 3 {
return nil, fmt.Errorf("unsupported WAV format: %d (need PCM=1 or float=3)", audioFormat)
}
if audioData == nil {
return nil, fmt.Errorf("no data chunk found in WAV file")
}
if numChannels <= 0 {
return nil, fmt.Errorf("invalid WAV channel count: %d", numChannels)
}
samples := decodeWAVSamples(audioData, audioFormat, bitsPerSample, numChannels)
if sampleRate != targetSampleRate {
samples = resampleLinear(samples, sampleRate, targetSampleRate)
}
return samples, nil
}
func decodeWAVSamples(data []byte, format uint16, bits, channels int) []float32 {
bytesPerSample := bits / 8
if bytesPerSample <= 0 || channels <= 0 {
return nil
}
totalSamples := len(data) / (bytesPerSample * channels)
mono := make([]float32, totalSamples)
for i := range totalSamples {
var sum float64
for ch := range channels {
off := (i*channels + ch) * bytesPerSample
if off+bytesPerSample > len(data) {
break
}
switch {
case format == 1 && bits == 16:
v := int16(binary.LittleEndian.Uint16(data[off : off+2]))
sum += float64(v) / 32768.0
case format == 1 && bits == 32:
v := int32(binary.LittleEndian.Uint32(data[off : off+4]))
sum += float64(v) / 2147483648.0
case format == 1 && bits == 24:
v := int32(data[off]) | int32(data[off+1])<<8 | int32(data[off+2])<<16
if v&0x800000 != 0 {
v |= ^0xffffff
}
sum += float64(v) / 8388608.0
case format == 3 && bits == 32:
sum += float64(math.Float32frombits(binary.LittleEndian.Uint32(data[off : off+4])))
case format == 1 && bits == 8:
sum += (float64(data[off]) - 128.0) / 128.0
}
}
mono[i] = float32(sum / float64(channels))
}
return mono
}
func resampleLinear(samples []float32, fromRate, toRate int) []float32 {
if fromRate <= 0 || toRate <= 0 || len(samples) == 0 {
return samples
}
n := int(float64(len(samples)) / float64(fromRate) * float64(toRate))
if n <= 1 {
return slicesCloneOne(samples)
}
out := make([]float32, n)
for i := range n {
pos := float64(i) * float64(len(samples)-1) / float64(n-1)
idx := int(pos)
frac := float32(pos - float64(idx))
if idx+1 < len(samples) {
out[i] = samples[idx]*(1-frac) + samples[idx+1]*frac
} else {
out[i] = samples[idx]
}
}
return out
}
func slicesCloneOne(samples []float32) []float32 {
if len(samples) == 0 {
return nil
}
return []float32{samples[0]}
}
func computeParakeetMelSpectrogram(samples []float32, extractor *AudioFeatureExtractor, opts *AudioOptions) ([]float32, int, int, error) {
if len(samples) == 0 {
return nil, 0, 0, fmt.Errorf("audio too short to encode")
}
if opts == nil {
opts = defaultAudioOptions()
}
melBins := opts.melBins
freqBins := parakeetNFFT/2 + 1
window, melFilters := extractor.windowAndFilters(melBins, freqBins, opts.sampleRate)
if len(window) != parakeetWinLength {
return nil, 0, 0, fmt.Errorf("invalid Parakeet window length: %d", len(window))
}
if len(melFilters) != melBins*freqBins {
return nil, 0, 0, fmt.Errorf("invalid Parakeet mel filter shape: %d", len(melFilters))
}
emphasized := make([]float32, len(samples))
emphasized[0] = samples[0]
for i := 1; i < len(samples); i++ {
emphasized[i] = samples[i] - parakeetPreemphasis*samples[i-1]
}
frames := len(samples)/parakeetHopLength + 1
validFrames := max(1, len(samples)/parakeetHopLength)
if validFrames > frames {
validFrames = frames
}
result := make([]float32, frames*melBins)
fftInput := make([]complex128, parakeetNFFT)
winOffset := (parakeetNFFT - parakeetWinLength) / 2
centerPad := parakeetNFFT / 2
for frame := range frames {
for i := range parakeetNFFT {
fftInput[i] = 0
}
for i := range parakeetWinLength {
src := frame*parakeetHopLength + i + winOffset - centerPad
if src >= 0 && src < len(emphasized) {
fftInput[i+winOffset] = complex(float64(emphasized[src])*float64(window[i]), 0)
}
}
fft(fftInput)
for mel := range melBins {
var v float64
filterOffset := mel * freqBins
for freq := range freqBins {
mag := cmplx.Abs(fftInput[freq])
v += float64(melFilters[filterOffset+freq]) * mag * mag
}
result[frame*melBins+mel] = float32(math.Log(v + parakeetLogZeroGuardValue))
}
}
for mel := range melBins {
var sum float64
for frame := range validFrames {
sum += float64(result[frame*melBins+mel])
}
mean := sum / float64(validFrames)
var variance float64
for frame := range validFrames {
d := float64(result[frame*melBins+mel]) - mean
variance += d * d
}
denom := max(1, validFrames-1)
std := math.Sqrt(variance / float64(denom))
for frame := range frames {
idx := frame*melBins + mel
if frame >= validFrames {
result[idx] = 0
continue
}
result[idx] = float32((float64(result[idx]) - mean) / (std + parakeetNormalizeEps))
}
}
return result, frames, validFrames, nil
}
func defaultParakeetWindow() []float32 {
window := make([]float32, parakeetWinLength)
for i := range window {
window[i] = float32(0.5 - 0.5*math.Cos(2*math.Pi*float64(i)/float64(parakeetWinLength-1)))
}
return window
}
func buildSlaneyMelFilterBank(numFreqBins, numMels int, sampleRate int) []float32 {
hzToMel := func(f float64) float64 {
if f < 1000 {
return 3 * f / 200
}
return 15 + math.Log(f/1000)*27/math.Log(6.4)
}
melToHz := func(m float64) float64 {
if m < 15 {
return 200 * m / 3
}
return 1000 * math.Exp(math.Log(6.4)*(m-15)/27)
}
minMel := hzToMel(0)
maxMel := hzToMel(float64(sampleRate) / 2)
mels := make([]float64, numMels+2)
freqs := make([]float64, numMels+2)
for i := range mels {
mels[i] = minMel + (maxMel-minMel)*float64(i)/float64(numMels+1)
freqs[i] = melToHz(mels[i])
}
fftFreqs := make([]float64, numFreqBins)
for i := range fftFreqs {
fftFreqs[i] = float64(i) * float64(sampleRate) / float64(parakeetNFFT)
}
filters := make([]float32, numMels*numFreqBins)
for mel := range numMels {
left, center, right := freqs[mel], freqs[mel+1], freqs[mel+2]
enorm := 2.0 / (right - left)
for freq, fftFreq := range fftFreqs {
var lower, upper float64
if center > left {
lower = (fftFreq - left) / (center - left)
}
if right > center {
upper = (right - fftFreq) / (right - center)
}
v := math.Max(0, math.Min(lower, upper))
filters[mel*numFreqBins+freq] = float32(v * enorm)
}
}
return filters
}
func fft(x []complex128) {
n := len(x)
if n <= 1 {
return
}
j := 0
for i := 1; i < n; i++ {
bit := n >> 1
for j&bit != 0 {
j ^= bit
bit >>= 1
}
j ^= bit
if i < j {
x[i], x[j] = x[j], x[i]
}
}
for size := 2; size <= n; size <<= 1 {
halfSize := size / 2
w := complex(math.Cos(2*math.Pi/float64(size)), -math.Sin(2*math.Pi/float64(size)))
for start := 0; start < n; start += size {
wn := complex(1, 0)
for k := range halfSize {
t := wn * x[start+k+halfSize]
x[start+k+halfSize] = x[start+k] - t
x[start+k] = x[start+k] + t
wn *= w
}
}
}
}
model/parsers/laguna.go
+498
View File
@@ -0,0 +1,498 @@
package parsers
import (
"encoding/json"
"fmt"
"regexp"
"strings"
"unicode"
"github.com/ollama/ollama/api"
)
const (
lagunaThinkingOpenTag = "<think>"
lagunaThinkingCloseTag = "</think>"
lagunaToolCallOpenTag = "<tool_call>"
lagunaToolCallCloseTag = "</tool_call>"
lagunaUserOpenTag = "<user>"
lagunaUserCloseTag = "</user>"
)
type lagunaParserState int
const (
lagunaParserStateThinking lagunaParserState = iota
lagunaParserStateContent
lagunaParserStateTool
)
type LagunaParser struct {
state lagunaParserState
buffer strings.Builder
tools []api.Tool
callIndex int
thinkingEnabled bool
thinkingSuppressed bool
allowLeadingThinkOpen bool
}
func (p *LagunaParser) HasToolSupport() bool {
return true
}
func (p *LagunaParser) HasThinkingSupport() bool {
return true
}
func (p *LagunaParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool {
p.tools = tools
p.callIndex = 0
p.buffer.Reset()
p.thinkingEnabled = thinkValue == nil || thinkValue.Bool()
p.thinkingSuppressed = thinkValue != nil && !thinkValue.Bool()
p.state = lagunaParserStateContent
p.allowLeadingThinkOpen = false
return tools
}
func (p *LagunaParser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) {
p.buffer.WriteString(s)
var contentSB, thinkingSB strings.Builder
for {
progress := false
switch p.state {
case lagunaParserStateThinking:
progress, thinking = p.consumeThinking(done)
if p.thinkingEnabled {
thinkingSB.WriteString(thinking)
}
case lagunaParserStateContent:
var parsedCalls []api.ToolCall
progress, content, parsedCalls, err = p.consumeContent(done)
if err != nil {
return "", "", nil, err
}
contentSB.WriteString(content)
calls = append(calls, parsedCalls...)
case lagunaParserStateTool:
var call api.ToolCall
progress, call, err = p.consumeTool(done)
if err != nil {
return "", "", nil, err
}
if progress {
calls = append(calls, call)
}
}
if !progress {
break
}
}
return contentSB.String(), thinkingSB.String(), calls, nil
}
func (p *LagunaParser) consumeThinking(done bool) (bool, string) {
acc := p.buffer.String()
if p.allowLeadingThinkOpen {
trimmed := strings.TrimLeftFunc(acc, unicode.IsSpace)
if strings.HasPrefix(trimmed, lagunaThinkingOpenTag) {
p.buffer.Reset()
p.buffer.WriteString(strings.TrimLeftFunc(strings.TrimPrefix(trimmed, lagunaThinkingOpenTag), unicode.IsSpace))
p.allowLeadingThinkOpen = false
return true, ""
}
if strings.HasPrefix(lagunaThinkingOpenTag, trimmed) && !done {
return false, ""
}
p.allowLeadingThinkOpen = false
}
if idx := strings.Index(acc, lagunaThinkingCloseTag); idx != -1 {
thinking := acc[:idx]
after := strings.TrimLeftFunc(acc[idx+len(lagunaThinkingCloseTag):], unicode.IsSpace)
p.buffer.Reset()
p.buffer.WriteString(after)
p.state = lagunaParserStateContent
return true, thinking
}
if idx := strings.Index(acc, lagunaToolCallOpenTag); idx != -1 {
thinking := strings.TrimRightFunc(acc[:idx], unicode.IsSpace)
after := acc[idx+len(lagunaToolCallOpenTag):]
p.buffer.Reset()
p.buffer.WriteString(after)
p.state = lagunaParserStateTool
return true, thinking
}
if done {
p.buffer.Reset()
p.state = lagunaParserStateContent
return acc != "", acc
}
overlapLen := max(overlap(acc, lagunaThinkingCloseTag), overlap(acc, lagunaToolCallOpenTag))
trailingLen := trailingWhitespaceLen(acc)
keep := max(overlapLen, trailingLen)
if keep > 0 && keep < len(acc) {
emit := acc[:len(acc)-keep]
p.buffer.Reset()
p.buffer.WriteString(acc[len(acc)-keep:])
return emit != "", emit
}
return false, ""
}
func (p *LagunaParser) consumeContent(done bool) (bool, string, []api.ToolCall, error) {
acc := p.buffer.String()
if p.thinkingEnabled || p.thinkingSuppressed {
if idx := strings.Index(acc, lagunaThinkingOpenTag); idx != -1 {
content := acc[:idx]
after := strings.TrimLeftFunc(acc[idx+len(lagunaThinkingOpenTag):], unicode.IsSpace)
p.buffer.Reset()
p.buffer.WriteString(after)
p.state = lagunaParserStateThinking
p.allowLeadingThinkOpen = false
return true, content, nil, nil
}
if !done {
overlapLen := overlap(acc, lagunaThinkingOpenTag)
if overlapLen > 0 && overlapLen < len(acc) {
content := acc[:len(acc)-overlapLen]
p.buffer.Reset()
p.buffer.WriteString(acc[len(acc)-overlapLen:])
return content != "", content, nil, nil
}
}
}
if p.thinkingEnabled {
trimmed := strings.TrimLeftFunc(acc, unicode.IsSpace)
if strings.HasPrefix(trimmed, lagunaThinkingCloseTag) {
p.buffer.Reset()
p.buffer.WriteString(strings.TrimLeftFunc(strings.TrimPrefix(trimmed, lagunaThinkingCloseTag), unicode.IsSpace))
return true, "", nil, nil
}
if strings.HasPrefix(lagunaThinkingCloseTag, trimmed) && !done {
return false, "", nil, nil
}
}
if p.thinkingSuppressed {
trimmed := strings.TrimLeftFunc(acc, unicode.IsSpace)
if strings.HasPrefix(trimmed, lagunaThinkingCloseTag) {
p.buffer.Reset()
p.buffer.WriteString(strings.TrimLeftFunc(strings.TrimPrefix(trimmed, lagunaThinkingCloseTag), unicode.IsSpace))
return true, "", nil, nil
}
if strings.HasPrefix(lagunaThinkingCloseTag, trimmed) && !done {
return false, "", nil, nil
}
}
if idx := strings.Index(acc, lagunaToolCallOpenTag); idx != -1 {
content := strings.TrimRightFunc(acc[:idx], unicode.IsSpace)
after := acc[idx+len(lagunaToolCallOpenTag):]
p.buffer.Reset()
p.buffer.WriteString(after)
p.state = lagunaParserStateTool
return true, content, nil, nil
}
if idx := strings.Index(acc, lagunaUserOpenTag); idx != -1 && len(p.tools) > 0 {
before := strings.TrimRightFunc(acc[:idx], unicode.IsSpace)
afterOpen := acc[idx+len(lagunaUserOpenTag):]
if closeIdx := strings.Index(afterOpen, lagunaUserCloseTag); closeIdx != -1 {
raw := afterOpen[:closeIdx]
if call, ok := p.parseToolAlias(raw); ok {
after := strings.TrimLeftFunc(afterOpen[closeIdx+len(lagunaUserCloseTag):], unicode.IsSpace)
p.buffer.Reset()
p.buffer.WriteString(after)
return true, before, []api.ToolCall{call}, nil
}
} else if !done {
if idx > 0 {
p.buffer.Reset()
p.buffer.WriteString(acc[idx:])
return true, before, nil, nil
}
return false, "", nil, nil
}
}
if len(p.tools) > 0 {
if progress, content, call, ok, err := p.consumeStandaloneJSONTool(done); ok || err != nil {
if err != nil {
return false, "", nil, err
}
if progress {
return true, content, []api.ToolCall{call}, nil
}
return false, "", nil, nil
}
}
if done {
p.buffer.Reset()
return acc != "", acc, nil, nil
}
overlapLen := max(overlap(acc, lagunaToolCallOpenTag), overlap(acc, lagunaUserOpenTag))
if p.thinkingEnabled || p.thinkingSuppressed {
overlapLen = max(overlapLen, overlap(acc, lagunaThinkingOpenTag))
}
if p.thinkingSuppressed {
overlapLen = max(overlapLen, overlap(acc, lagunaThinkingCloseTag))
}
trailingLen := trailingWhitespaceLen(acc)
keep := max(overlapLen, trailingLen)
if keep > 0 && keep < len(acc) {
emit := acc[:len(acc)-keep]
p.buffer.Reset()
p.buffer.WriteString(acc[len(acc)-keep:])
return emit != "", emit, nil, nil
}
if keep == 0 && acc != "" {
p.buffer.Reset()
return true, acc, nil, nil
}
return false, "", nil, nil
}
func (p *LagunaParser) consumeStandaloneJSONTool(done bool) (progress bool, content string, call api.ToolCall, ok bool, err error) {
acc := p.buffer.String()
jsonIdx := strings.Index(acc, "{")
if jsonIdx == -1 {
return false, "", api.ToolCall{}, false, nil
}
before := strings.TrimRightFunc(acc[:jsonIdx], unicode.IsSpace)
raw := strings.TrimLeftFunc(acc[jsonIdx:], unicode.IsSpace)
if !lagunaLooksLikeJSONToolCall(raw, done) {
return false, "", api.ToolCall{}, false, nil
}
if !done && !json.Valid([]byte(strings.TrimSpace(raw))) {
if before != "" {
p.buffer.Reset()
p.buffer.WriteString(acc[jsonIdx:])
return true, before, api.ToolCall{}, true, nil
}
return false, "", api.ToolCall{}, true, nil
}
call, err = parseLagunaToolCall(raw, p.tools)
if err != nil {
return false, "", api.ToolCall{}, true, err
}
call.Function.Index = p.callIndex
p.callIndex++
p.buffer.Reset()
p.state = lagunaParserStateContent
return true, before, call, true, nil
}
func lagunaLooksLikeJSONToolCall(raw string, done bool) bool {
trimmed := strings.TrimLeftFunc(raw, unicode.IsSpace)
if !strings.HasPrefix(trimmed, "{") {
return false
}
if strings.Contains(trimmed, `"name"`) || strings.Contains(trimmed, `"arguments"`) {
return true
}
if done {
return false
}
return strings.HasPrefix(trimmed, `{"`) || strings.HasPrefix(trimmed, "{\n") || strings.HasPrefix(trimmed, "{\r\n")
}
func (p *LagunaParser) parseToolAlias(raw string) (api.ToolCall, bool) {
raw = cleanLagunaToolCallRaw(raw)
name, ok := lagunaToolCallName(raw)
if !ok {
return api.ToolCall{}, false
}
if _, ok := lagunaResolveToolName(name, p.tools); !ok {
return api.ToolCall{}, false
}
call, err := parseLagunaToolCall(raw, p.tools)
if err != nil {
return api.ToolCall{}, false
}
call.Function.Index = p.callIndex
p.callIndex++
return call, true
}
func lagunaResolveToolName(name string, tools []api.Tool) (string, bool) {
for i := range tools {
if tools[i].Function.Name == name {
return name, true
}
}
aliases := map[string]string{
"read_file": "read",
"write_file": "write",
"edit_file": "edit",
"web_fetch": "webfetch",
}
if alias, ok := aliases[name]; ok {
for i := range tools {
if tools[i].Function.Name == alias {
return alias, true
}
}
}
return name, false
}
func cleanLagunaToolCallRaw(raw string) string {
raw = strings.TrimSpace(raw)
for strings.HasPrefix(raw, lagunaToolCallOpenTag) {
raw = strings.TrimSpace(strings.TrimPrefix(raw, lagunaToolCallOpenTag))
}
if idx := strings.Index(raw, lagunaToolCallCloseTag); idx != -1 {
raw = strings.TrimSpace(raw[:idx])
}
if idx := strings.Index(raw, lagunaToolCallOpenTag); idx != -1 {
before := strings.TrimSpace(raw[:idx])
if before != "" {
return before
}
raw = strings.TrimSpace(raw[idx+len(lagunaToolCallOpenTag):])
}
return raw
}
func lagunaToolCallName(raw string) (string, bool) {
raw = cleanLagunaToolCallRaw(raw)
if strings.HasPrefix(raw, "{") {
var parsed struct {
Name string `json:"name"`
}
if err := json.Unmarshal([]byte(raw), &parsed); err != nil {
return "", false
}
name := strings.TrimSpace(parsed.Name)
return name, name != ""
}
nameEnd := strings.Index(raw, "<arg_key>")
if nameEnd < 0 {
nameEnd = strings.Index(raw, "{")
}
if nameEnd < 0 {
nameEnd = strings.IndexAny(raw, "\r\n")
}
if nameEnd < 0 {
nameEnd = len(raw)
}
name := strings.TrimSpace(raw[:nameEnd])
return name, name != ""
}
func (p *LagunaParser) consumeTool(done bool) (bool, api.ToolCall, error) {
acc := p.buffer.String()
if idx := strings.Index(acc, lagunaToolCallCloseTag); idx != -1 {
raw := acc[:idx]
after := strings.TrimLeftFunc(acc[idx+len(lagunaToolCallCloseTag):], unicode.IsSpace)
p.buffer.Reset()
p.buffer.WriteString(after)
p.state = lagunaParserStateContent
call, err := parseLagunaToolCall(raw, p.tools)
if err != nil {
return false, api.ToolCall{}, err
}
call.Function.Index = p.callIndex
p.callIndex++
return true, call, nil
}
if done && strings.TrimSpace(acc) != "" {
p.buffer.Reset()
p.state = lagunaParserStateContent
call, err := parseLagunaToolCall(acc, p.tools)
if err != nil {
return false, api.ToolCall{}, err
}
call.Function.Index = p.callIndex
p.callIndex++
return true, call, nil
}
return false, api.ToolCall{}, nil
}
var lagunaArgRE = regexp.MustCompile(`(?s)<arg_key>(.*?)</arg_key>\s*<arg_value>(.*?)</arg_value>`)
func parseLagunaToolCall(raw string, tools []api.Tool) (api.ToolCall, error) {
raw = cleanLagunaToolCallRaw(raw)
if strings.HasPrefix(raw, "{") {
var parsed struct {
Name string `json:"name"`
Arguments api.ToolCallFunctionArguments `json:"arguments"`
}
if err := json.Unmarshal([]byte(raw), &parsed); err != nil {
return api.ToolCall{}, fmt.Errorf("failed to parse Laguna JSON tool call: %w", err)
}
if parsed.Name == "" {
return api.ToolCall{}, fmt.Errorf("empty Laguna tool call name")
}
if name, ok := lagunaResolveToolName(parsed.Name, tools); ok {
parsed.Name = name
}
return api.ToolCall{
Function: api.ToolCallFunction{
Name: parsed.Name,
Arguments: parsed.Arguments,
},
}, nil
}
nameEnd := strings.Index(raw, "<arg_key>")
name := raw
argsText := ""
if nameEnd >= 0 {
name = raw[:nameEnd]
argsText = raw[nameEnd:]
} else if jsonStart := strings.Index(raw, "{"); jsonStart >= 0 {
name = raw[:jsonStart]
argsText = raw[jsonStart:]
}
name = strings.TrimSpace(name)
if resolved, ok := lagunaResolveToolName(name, tools); ok {
name = resolved
}
var matchedTool *api.Tool
for i := range tools {
if tools[i].Function.Name == name {
matchedTool = &tools[i]
break
}
}
call := api.ToolCall{
Function: api.ToolCallFunction{
Name: name,
Arguments: api.NewToolCallFunctionArguments(),
},
}
if strings.HasPrefix(strings.TrimSpace(argsText), "{") {
if err := json.Unmarshal([]byte(strings.TrimSpace(argsText)), &call.Function.Arguments); err != nil {
return api.ToolCall{}, fmt.Errorf("failed to parse Laguna JSON tool call arguments: %w", err)
}
return call, nil
}
for _, match := range lagunaArgRE.FindAllStringSubmatch(argsText, -1) {
key := strings.TrimSpace(match[1])
value := match[2]
var paramType api.PropertyType
if matchedTool != nil && matchedTool.Function.Parameters.Properties != nil {
if prop, ok := matchedTool.Function.Parameters.Properties.Get(key); ok {
if len(prop.AnyOf) > 0 {
for _, anyOfProp := range prop.AnyOf {
paramType = append(paramType, anyOfProp.Type...)
}
} else {
paramType = prop.Type
}
}
}
call.Function.Arguments.Set(key, parseValue(value, paramType))
}
return call, nil
}
model/parsers/laguna_test.go
+484
View File
@@ -0,0 +1,484 @@
package parsers
import (
"testing"
"github.com/ollama/ollama/api"
)
func lagunaTestTools() []api.Tool {
props := api.NewToolPropertiesMap()
props.Set("location", api.ToolProperty{Type: api.PropertyType{"string"}})
props.Set("days", api.ToolProperty{Type: api.PropertyType{"integer"}})
return []api.Tool{{
Function: api.ToolFunction{
Name: "get_weather",
Parameters: api.ToolFunctionParameters{
Properties: props,
},
},
}}
}
func TestLagunaParserToolCall(t *testing.T) {
parser := ParserForName("laguna")
if parser == nil {
t.Fatal("expected laguna parser")
}
if !parser.HasToolSupport() || !parser.HasThinkingSupport() {
t.Fatal("laguna parser should advertise tools and thinking")
}
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("<tool_call>get_weather\n<arg_key>location</arg_key>\n<arg_value>Paris</arg_value>\n<arg_key>days</arg_key>\n<arg_value>3</arg_value>\n</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" {
t.Fatalf("content=%q thinking=%q, want empty", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("location"); got != "Paris" {
t.Fatalf("location=%v, want Paris", got)
}
if got, _ := calls[0].Function.Arguments.Get("days"); got != 3 {
t.Fatalf("days=%v, want 3", got)
}
}
func TestLagunaParserJSONToolCall(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
_, _, calls, err := parser.Add("<tool_call>\n{\"name\":\"get_weather\",\"arguments\":{\"location\":\"Paris\",\"days\":3}}\n</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("location"); got != "Paris" {
t.Fatalf("location=%v, want Paris", got)
}
if got, _ := calls[0].Function.Arguments.Get("days"); got != float64(3) {
t.Fatalf("days=%v, want 3", got)
}
}
func TestLagunaParserStandaloneJSONToolCall(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("{\"name\":\"get_weather\",\"arguments\":{\"location\":\"Paris\",\"days\":3}}", true)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" {
t.Fatalf("content=%q thinking=%q", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
}
func TestLagunaParserStandaloneJSONToolCallAfterLeadingContent(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("Let me call the weather tool.\n{\"name\":\"get_weather\",\"arguments\":{\"location\":\"Paris\"}}", true)
if err != nil {
t.Fatal(err)
}
if content != "Let me call the weather tool." || thinking != "" {
t.Fatalf("content=%q thinking=%q", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
}
func TestLagunaParserStreamingStandaloneJSONToolCall(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("{\"name\":\"get_weather\",\"arguments\":{\"location\":\"San Francisco,", false)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" || len(calls) != 0 {
t.Fatalf("first chunk content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
content, thinking, calls, err = parser.Add(" CA\"}}", true)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" || len(calls) != 1 {
t.Fatalf("second chunk content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("location"); got != "San Francisco, CA" {
t.Fatalf("location=%v, want San Francisco, CA", got)
}
}
func TestLagunaParserNameLineJSONToolCall(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
_, _, calls, err := parser.Add("<tool_call>get_weather\n{\"location\":\"San Francisco\"}</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("location"); got != "San Francisco" {
t.Fatalf("location=%v, want San Francisco", got)
}
}
func TestLagunaParserNormalizesCommonToolAliases(t *testing.T) {
props := api.NewToolPropertiesMap()
props.Set("path", api.ToolProperty{Type: api.PropertyType{"string"}})
tools := []api.Tool{{
Function: api.ToolFunction{
Name: "read",
Parameters: api.ToolFunctionParameters{
Properties: props,
},
},
}}
parser := ParserForName("laguna")
parser.Init(tools, nil, nil)
_, _, calls, err := parser.Add("<tool_call>\n{\"name\":\"read_file\",\"arguments\":{\"path\":\"./go.mod\"}}\n</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "read" {
t.Fatalf("name=%q, want read", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("path"); got != "./go.mod" {
t.Fatalf("path=%v, want ./go.mod", got)
}
}
func TestLagunaParserIgnoresDuplicatedNestedToolCall(t *testing.T) {
props := api.NewToolPropertiesMap()
props.Set("name", api.ToolProperty{Type: api.PropertyType{"string"}})
tools := []api.Tool{{
Function: api.ToolFunction{
Name: "skill",
Parameters: api.ToolFunctionParameters{
Properties: props,
},
},
}}
parser := ParserForName("laguna")
parser.Init(tools, nil, nil)
_, _, calls, err := parser.Add("<tool_call>skill\n{\"name\":\"git-diff-review\"}\n<tool_call>skill\n{\"name\":\"git-diff-review\"}</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "skill" {
t.Fatalf("name=%q, want skill", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("name"); got != "git-diff-review" {
t.Fatalf("name arg=%v, want git-diff-review", got)
}
}
func TestLagunaParserThinkingThenTool(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, &api.ThinkValue{Value: true})
content, thinking, calls, err := parser.Add("<think>Need current weather.</think>\n<tool_call>get_weather\n<arg_key>location</arg_key>\n<arg_value>SF</arg_value>\n</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if content != "" {
t.Fatalf("content=%q, want empty", content)
}
if thinking != "Need current weather." {
t.Fatalf("thinking=%q, want reasoning", thinking)
}
if len(calls) != 1 || calls[0].Function.Name != "get_weather" {
t.Fatalf("unexpected calls: %#v", calls)
}
}
func TestLagunaParserUserTaggedToolAlias(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("<user>get_weather\n<arg_key>location</arg_key>\n<arg_value>San Francisco, CA</arg_value>\n</user>", true)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" {
t.Fatalf("content=%q thinking=%q, want empty", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "get_weather" {
t.Fatalf("name=%q, want get_weather", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("location"); got != "San Francisco, CA" {
t.Fatalf("location=%v, want San Francisco, CA", got)
}
}
func TestLagunaParserUserTaggedToolAliasWithLeadingContent(t *testing.T) {
parser := ParserForName("laguna")
props := api.NewToolPropertiesMap()
props.Set("path", api.ToolProperty{Type: api.PropertyType{"string"}})
tools := []api.Tool{{
Function: api.ToolFunction{
Name: "read",
Parameters: api.ToolFunctionParameters{
Properties: props,
},
},
}}
parser.Init(tools, nil, nil)
content, thinking, calls, err := parser.Add("I'll read the file for you.\n<user>read\n<arg_key>path</arg_key>\n<arg_value>/Users/test/code/myproject/go.mod</arg_value>\n</user>", true)
if err != nil {
t.Fatal(err)
}
if content != "I'll read the file for you." || thinking != "" {
t.Fatalf("content=%q thinking=%q", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "read" {
t.Fatalf("name=%q, want read", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("path"); got != "/Users/test/code/myproject/go.mod" {
t.Fatalf("path=%v, want /Users/test/code/myproject/go.mod", got)
}
}
func TestLagunaParserUserTaggedJSONToolCallWithLeadingContent(t *testing.T) {
parser := ParserForName("laguna")
props := api.NewToolPropertiesMap()
props.Set("command", api.ToolProperty{Type: api.PropertyType{"string"}})
tools := []api.Tool{{
Function: api.ToolFunction{
Name: "bash",
Parameters: api.ToolFunctionParameters{
Properties: props,
},
},
}}
parser.Init(tools, nil, nil)
content, thinking, calls, err := parser.Add("I'll run git diff for you.<user>\n{\"name\":\"bash\",\"arguments\":{\"command\":\"git diff main\"}}\n</user>", true)
if err != nil {
t.Fatal(err)
}
if content != "I'll run git diff for you." || thinking != "" {
t.Fatalf("content=%q thinking=%q", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "bash" {
t.Fatalf("name=%q, want bash", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("command"); got != "git diff main" {
t.Fatalf("command=%v, want git diff main", got)
}
}
func TestLagunaParserStreamingUserTaggedToolAliasAfterContent(t *testing.T) {
parser := ParserForName("laguna")
props := api.NewToolPropertiesMap()
props.Set("path", api.ToolProperty{Type: api.PropertyType{"string"}})
tools := []api.Tool{{
Function: api.ToolFunction{
Name: "read",
Parameters: api.ToolFunctionParameters{
Properties: props,
},
},
}}
parser.Init(tools, nil, nil)
content, thinking, calls, err := parser.Add("I'll read the file for you.<us", false)
if err != nil {
t.Fatal(err)
}
if content != "I'll read the file for you." || thinking != "" || len(calls) != 0 {
t.Fatalf("first chunk content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
content, thinking, calls, err = parser.Add("er>read\n<arg_key>path</arg_key>\n<arg_value>/Users/test/code/myproject/go.mod</arg_value>\n</user>", true)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" {
t.Fatalf("second chunk content=%q thinking=%q", content, thinking)
}
if len(calls) != 1 {
t.Fatalf("calls=%d, want 1", len(calls))
}
if calls[0].Function.Name != "read" {
t.Fatalf("name=%q, want read", calls[0].Function.Name)
}
if got, _ := calls[0].Function.Arguments.Get("path"); got != "/Users/test/code/myproject/go.mod" {
t.Fatalf("path=%v, want /Users/test/code/myproject/go.mod", got)
}
}
func TestLagunaParserUserTaggedNonToolContent(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("<user>hello</user>", true)
if err != nil {
t.Fatal(err)
}
if content != "<user>hello</user>" || thinking != "" || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserThinkingDefaultsOn(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(nil, nil, nil)
content, thinking, calls, err := parser.Add("<think>Need to reason.</think>\nDirect answer.", true)
if err != nil {
t.Fatal(err)
}
if content != "Direct answer." || thinking != "Need to reason." || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserThinkingDefaultsOnWhenToolsPresent(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, nil)
content, thinking, calls, err := parser.Add("<think>Need to reason.</think>\n<tool_call>get_weather\n<arg_key>location</arg_key>\n<arg_value>Paris</arg_value>\n</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if thinking != "Need to reason." || len(calls) != 1 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
if content != "" {
t.Fatalf("content=%q, want thinking block suppressed from content when default thinking is enabled", content)
}
}
func TestLagunaParserThinkingExplicitlyDisabled(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(nil, nil, &api.ThinkValue{Value: false})
content, thinking, calls, err := parser.Add("<think>Hidden?</think>\nDirect answer.", true)
if err != nil {
t.Fatal(err)
}
if content != "Direct answer." || thinking != "" || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserThinkingExplicitlyDisabledDropsLeadingCloseTag(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(nil, nil, &api.ThinkValue{Value: false})
content, thinking, calls, err := parser.Add("</think>\nTokyo\n", true)
if err != nil {
t.Fatal(err)
}
if content != "Tokyo\n" || thinking != "" || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserThinkingEnabledDropsLeadingCloseTag(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(nil, nil, &api.ThinkValue{Value: true})
content, thinking, calls, err := parser.Add("</think>\nTokyo\n", true)
if err != nil {
t.Fatal(err)
}
if content != "Tokyo\n" || thinking != "" || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserThinkingDefaultOnDropsLeadingCloseTag(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(nil, nil, nil)
content, thinking, calls, err := parser.Add("</think>\nTokyo\n", true)
if err != nil {
t.Fatal(err)
}
if content != "Tokyo\n" || thinking != "" || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserThinkingEnabledUntaggedAnswerIsContent(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(nil, nil, &api.ThinkValue{Value: true})
content, thinking, calls, err := parser.Add("Direct answer.", true)
if err != nil {
t.Fatal(err)
}
if content != "Direct answer." || thinking != "" || len(calls) != 0 {
t.Fatalf("content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
func TestLagunaParserSplitToolTag(t *testing.T) {
parser := ParserForName("laguna")
parser.Init(lagunaTestTools(), nil, &api.ThinkValue{Value: true})
content, thinking, calls, err := parser.Add("<think>Need lookup<tool_c", false)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "Need lookup" || len(calls) != 0 {
t.Fatalf("first chunk content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
content, thinking, calls, err = parser.Add("all>get_weather\n<arg_key>location</arg_key>\n<arg_value>SF</arg_value>\n</tool_call>", true)
if err != nil {
t.Fatal(err)
}
if content != "" || thinking != "" || len(calls) != 1 {
t.Fatalf("second chunk content=%q thinking=%q calls=%d", content, thinking, len(calls))
}
}
model/parsers/nemotron3nano.go
+66 -4
View File
@@ -16,14 +16,17 @@ const (
)
const (
nemotronThinkOpen = "<think>"
nemotronThinkClose = "</think>"
nemotronToolCallOpen = "<tool_call>"
)
type Nemotron3NanoParser struct {
state Nemotron3NanoParserState
buffer strings.Builder
toolParser *Qwen3CoderParser
state Nemotron3NanoParserState
buffer strings.Builder
toolParser *Qwen3CoderParser
maybeThinkingOpenAtBOL bool
skipThinkingLeadingWS bool
}
func (p *Nemotron3NanoParser) HasToolSupport() bool { return true }
@@ -32,14 +35,18 @@ func (p *Nemotron3NanoParser) HasThinkingSupport() bool { return true }
func (p *Nemotron3NanoParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool {
p.toolParser = &Qwen3CoderParser{}
p.toolParser.Init(tools, nil, nil)
p.buffer.Reset()
p.maybeThinkingOpenAtBOL = false
p.skipThinkingLeadingWS = false
thinkingEnabled := thinkValue != nil && thinkValue.Bool()
thinkingEnabled := thinkValue == nil || thinkValue.Bool()
prefill := lastMessage != nil && lastMessage.Role == "assistant"
if !thinkingEnabled || (prefill && lastMessage.Content != "") {
p.state = Nemotron3NanoCollectingContent
} else {
p.state = Nemotron3NanoCollectingThinking
p.maybeThinkingOpenAtBOL = true
}
return tools
@@ -61,6 +68,29 @@ func (p *Nemotron3NanoParser) Add(s string, done bool) (content string, thinking
// Nemotron3NanoCollectingThinking - buffer and look for end markers
p.buffer.WriteString(s)
if p.skipThinkingLeadingWS {
trimmed := strings.TrimLeftFunc(p.buffer.String(), unicode.IsSpace)
p.buffer.Reset()
p.buffer.WriteString(trimmed)
if trimmed == "" {
return "", "", nil, nil
}
p.skipThinkingLeadingWS = false
}
if p.stripOpeningThinkTag() {
return p.Add("", done)
}
if p.maybeThinkingOpenAtBOL {
bufStr := p.buffer.String()
trimmed := strings.TrimLeftFunc(bufStr, unicode.IsSpace)
if trimmed == "" || overlap(trimmed, nemotronThinkOpen) == len(trimmed) {
if len(trimmed) != len(bufStr) {
p.buffer.Reset()
p.buffer.WriteString(trimmed)
}
return "", "", nil, nil
}
}
bufStr := p.buffer.String()
// Look for end of thinking: </think> or <tool_call> (model may skip </think>)
@@ -124,3 +154,35 @@ func (p *Nemotron3NanoParser) emitThinking(bufStr string) string {
p.buffer.Reset()
return bufStr
}
func (p *Nemotron3NanoParser) stripOpeningThinkTag() bool {
if !p.maybeThinkingOpenAtBOL {
return false
}
bufStr := p.buffer.String()
trimmed := strings.TrimLeftFunc(bufStr, unicode.IsSpace)
if trimmed == "" {
p.buffer.Reset()
return false
}
if strings.HasPrefix(trimmed, nemotronThinkOpen) {
p.buffer.Reset()
p.buffer.WriteString(strings.TrimLeftFunc(trimmed[len(nemotronThinkOpen):], unicode.IsSpace))
p.maybeThinkingOpenAtBOL = false
p.skipThinkingLeadingWS = true
return true
}
if overlap(trimmed, nemotronThinkOpen) == len(trimmed) {
if len(trimmed) != len(bufStr) {
p.buffer.Reset()
p.buffer.WriteString(trimmed)
}
return false
}
p.maybeThinkingOpenAtBOL = false
return false
}
model/parsers/nemotron3nano_test.go
+47 -3
View File
@@ -82,6 +82,20 @@ func TestNemotron3NanoParser(t *testing.T) {
expectedThinking: "My thoughts...",
expectedContent: "Content here.",
},
{
name: "leading open think tag is ignored",
input: "<think>\nLet me think about this...</think>\nHere is my answer.",
thinkValue: &api.ThinkValue{Value: true},
expectedThinking: "Let me think about this...",
expectedContent: "Here is my answer.",
},
{
name: "empty explicit think block is ignored",
input: "<think></think>\nHere is my answer.",
thinkValue: &api.ThinkValue{Value: true},
expectedThinking: "",
expectedContent: "Here is my answer.",
},
}
for _, tt := range tests {
@@ -191,6 +205,13 @@ func TestNemotron3NanoParser_Streaming(t *testing.T) {
},
},
},
{
name: "leading open think tag split across chunks",
chunks: []string{"<th", "ink>", "\nThink first", "</think>", "\nDone."},
thinkValue: &api.ThinkValue{Value: true},
expectedThinking: "Think first",
expectedContent: "Done.",
},
}
for _, tt := range tests {
@@ -265,11 +286,11 @@ func TestNemotron3NanoParser_Init(t *testing.T) {
}
})
t.Run("starts in content state when nil thinkValue", func(t *testing.T) {
t.Run("starts in thinking state when nil thinkValue", func(t *testing.T) {
p := &Nemotron3NanoParser{}
p.Init(nil, nil, nil)
if p.state != Nemotron3NanoCollectingContent {
t.Errorf("expected state Nemotron3NanoCollectingContent, got %v", p.state)
if p.state != Nemotron3NanoCollectingThinking {
t.Errorf("expected state Nemotron3NanoCollectingThinking, got %v", p.state)
}
})
@@ -281,6 +302,29 @@ func TestNemotron3NanoParser_Init(t *testing.T) {
t.Errorf("expected state Nemotron3NanoCollectingContent, got %v", p.state)
}
})
t.Run("reinit clears buffered state", func(t *testing.T) {
p := &Nemotron3NanoParser{}
p.Init(nil, nil, &api.ThinkValue{Value: true})
if _, _, _, err := p.Add("thinking in progress", false); err != nil {
t.Fatalf("unexpected error: %v", err)
}
p.Init(nil, nil, &api.ThinkValue{Value: false})
content, thinking, calls, err := p.Add("content only", true)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if content != "content only" {
t.Fatalf("expected content after reinit, got %q", content)
}
if thinking != "" {
t.Fatalf("expected no thinking after reinit, got %q", thinking)
}
if len(calls) != 0 {
t.Fatalf("expected no tool calls after reinit, got %v", calls)
}
})
}
func TestNemotron3NanoParser_WithTools(t *testing.T) {
model/parsers/parsers.go
+2
View File
@@ -87,6 +87,8 @@ func ParserForName(name string) Parser {
return &LFM2Parser{hasThinkingSupport: false}
case "lfm2-thinking":
return &LFM2Parser{hasThinkingSupport: true}
case "laguna":
return &LagunaParser{}
default:
return nil
}
model/renderers/laguna.go
+111
View File
@@ -0,0 +1,111 @@
package renderers
import (
"strings"
"github.com/ollama/ollama/api"
)
const (
lagunaBOS = "〈|EOS|〉"
lagunaThoughtOpen = "<think>"
lagunaThoughtClose = "</think>"
)
type LagunaRenderer struct{}
func (r *LagunaRenderer) Render(messages []api.Message, tools []api.Tool, think *api.ThinkValue) (string, error) {
var sb strings.Builder
sb.WriteString(lagunaBOS)
thinkingEnabled := think == nil || think.Bool()
systemMessage := ""
firstMessageIsSystem := len(messages) > 0 && messages[0].Role == "system"
if firstMessageIsSystem {
systemMessage = strings.TrimRight(messages[0].Content, "\n")
}
sb.WriteString("<system>\n")
if thinkingEnabled {
sb.WriteString("You should use chain-of-thought reasoning. Put your reasoning inside <think> </think> tags before your response.")
} else {
sb.WriteString("You should respond directly without using chain-of-thought reasoning tags.")
}
if strings.TrimSpace(systemMessage) != "" {
sb.WriteByte('\n')
sb.WriteString(systemMessage)
}
if len(tools) > 0 {
sb.WriteString("\n\n### Tools\n\n")
sb.WriteString("You may call functions to assist with the user query.\n")
sb.WriteString("All available function signatures are listed below:\n")
sb.WriteString("<available_tools>\n")
for _, tool := range tools {
if b, err := marshalWithSpaces(tool); err == nil {
sb.Write(b)
sb.WriteByte('\n')
}
}
sb.WriteString("</available_tools>\n\n")
sb.WriteString("For each function call, return a json object with function name and arguments within '<tool_call>' and '</tool_call>' tags:\n")
sb.WriteString("<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call>")
}
sb.WriteString("\n</system>\n")
for i, message := range messages {
if i == 0 && firstMessageIsSystem {
continue
}
content := message.Content
switch message.Role {
case "user":
sb.WriteString("<user>\n")
sb.WriteString(content)
sb.WriteString("\n</user>\n")
case "assistant":
lastMessage := i == len(messages)-1
prefill := lastMessage && (content != "" || message.Thinking != "" || len(message.ToolCalls) > 0)
sb.WriteString("<assistant>\n")
if thinkingEnabled && message.Thinking != "" {
sb.WriteString(lagunaThoughtOpen)
sb.WriteString(message.Thinking)
sb.WriteString(lagunaThoughtClose)
sb.WriteByte('\n')
}
if strings.Trim(content, "\n") != "" {
sb.WriteString(strings.Trim(content, "\n"))
sb.WriteByte('\n')
}
for _, toolCall := range message.ToolCalls {
sb.WriteString("<tool_call>")
sb.WriteString(toolCall.Function.Name)
sb.WriteByte('\n')
for name, value := range toolCall.Function.Arguments.All() {
sb.WriteString("<arg_key>")
sb.WriteString(name)
sb.WriteString("</arg_key>\n")
sb.WriteString("<arg_value>")
sb.WriteString(formatToolCallArgument(value))
sb.WriteString("</arg_value>\n")
}
sb.WriteString("</tool_call>\n")
}
if !prefill {
sb.WriteString("</assistant>\n")
}
case "tool":
sb.WriteString("<tool_response>\n")
sb.WriteString(content)
sb.WriteString("\n</tool_response>\n")
case "system":
sb.WriteString("<system>\n")
sb.WriteString(content)
sb.WriteString("\n</system>\n")
}
}
if len(messages) == 0 || messages[len(messages)-1].Role != "assistant" {
sb.WriteString("<assistant>\n")
}
return sb.String(), nil
}
model/renderers/laguna_test.go
+339
View File
@@ -0,0 +1,339 @@
package renderers
import (
"encoding/json"
"os"
"os/exec"
"strings"
"testing"
"github.com/google/go-cmp/cmp"
"github.com/ollama/ollama/api"
)
const (
lagunaDirectDirective = "You should respond directly without using chain-of-thought reasoning tags."
lagunaThinkDirective = "You should use chain-of-thought reasoning. Put your reasoning inside <think> </think> tags before your response."
)
func TestLagunaRendererReferenceFlowCoverage(t *testing.T) {
weather := lagunaWeatherTool()
tests := []struct {
name string
messages []api.Message
tools []api.Tool
think *api.ThinkValue
want string
}{
{
name: "user_only_thinking_default_on",
messages: []api.Message{{Role: "user", Content: "Hello"}},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\n</system>\n" +
"<user>\nHello\n</user>\n" +
"<assistant>\n",
},
{
name: "user_only_thinking_enabled",
messages: []api.Message{{Role: "user", Content: "Hello"}},
think: &api.ThinkValue{Value: true},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\n</system>\n" +
"<user>\nHello\n</user>\n" +
"<assistant>\n",
},
{
name: "user_only_thinking_disabled",
messages: []api.Message{{Role: "user", Content: "Hello"}},
think: &api.ThinkValue{Value: false},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaDirectDirective +
"\n</system>\n" +
"<user>\nHello\n</user>\n" +
"<assistant>\n",
},
{
name: "first_system_is_header",
messages: []api.Message{
{Role: "system", Content: "Stay concise.\n\n"},
{Role: "user", Content: "Hi"},
},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\nStay concise." +
"\n</system>\n" +
"<user>\nHi\n</user>\n" +
"<assistant>\n",
},
{
name: "additional_system_message_renders_in_loop",
messages: []api.Message{
{Role: "system", Content: "Primary."},
{Role: "user", Content: "Hi"},
{Role: "system", Content: "Secondary."},
},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\nPrimary." +
"\n</system>\n" +
"<user>\nHi\n</user>\n" +
"<system>\nSecondary.\n</system>\n" +
"<assistant>\n",
},
{
name: "tools_in_header",
messages: []api.Message{
{Role: "system", Content: "Stay concise."},
{Role: "user", Content: "Weather?"},
},
tools: weather,
think: &api.ThinkValue{Value: true},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\nStay concise." +
"\n\n### Tools\n\n" +
"You may call functions to assist with the user query.\n" +
"All available function signatures are listed below:\n" +
"<available_tools>\n" +
`{"type": "function", "function": {"name": "get_weather", "description": "Get weather", "parameters": {"type": "object", "required": ["location"], "properties": {"location": {"type": "string", "description": "City"}}}}}` + "\n" +
"</available_tools>\n\n" +
"For each function call, return a json object with function name and arguments within '<tool_call>' and '</tool_call>' tags:\n" +
"<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call>" +
"\n</system>\n" +
"<user>\nWeather?\n</user>\n" +
"<assistant>\n",
},
{
name: "tools_default_thinking_on_when_unspecified",
messages: []api.Message{
{Role: "user", Content: "Weather?"},
},
tools: weather,
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\n\n### Tools\n\n" +
"You may call functions to assist with the user query.\n" +
"All available function signatures are listed below:\n" +
"<available_tools>\n" +
`{"type": "function", "function": {"name": "get_weather", "description": "Get weather", "parameters": {"type": "object", "required": ["location"], "properties": {"location": {"type": "string", "description": "City"}}}}}` + "\n" +
"</available_tools>\n\n" +
"For each function call, return a json object with function name and arguments within '<tool_call>' and '</tool_call>' tags:\n" +
"<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call>" +
"\n</system>\n" +
"<user>\nWeather?\n</user>\n" +
"<assistant>\n",
},
{
name: "assistant_history_with_thinking_content_tool_and_response",
messages: []api.Message{
{Role: "user", Content: "Add these."},
{
Role: "assistant",
Content: "\nCalling the tool.\n",
Thinking: "Need addition.",
ToolCalls: []api.ToolCall{{
Function: api.ToolCallFunction{
Name: "add",
Arguments: testArgsOrdered([]orderedArg{
{Key: "a", Value: 2},
{Key: "b", Value: 3},
}),
},
}},
},
{Role: "tool", Content: "5"},
{Role: "user", Content: "Thanks"},
},
think: &api.ThinkValue{Value: true},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\n</system>\n" +
"<user>\nAdd these.\n</user>\n" +
"<assistant>\n" +
"<think>Need addition.</think>\n" +
"Calling the tool.\n" +
"<tool_call>add\n" +
"<arg_key>a</arg_key>\n<arg_value>2</arg_value>\n" +
"<arg_key>b</arg_key>\n<arg_value>3</arg_value>\n" +
"</tool_call>\n" +
"</assistant>\n" +
"<tool_response>\n5\n</tool_response>\n" +
"<user>\nThanks\n</user>\n" +
"<assistant>\n",
},
{
name: "final_assistant_prefill_is_continued",
messages: []api.Message{
{Role: "user", Content: "Complete this"},
{Role: "assistant", Content: "Partial"},
},
want: "" +
"〈|EOS|〉<system>\n" +
lagunaThinkDirective +
"\n</system>\n" +
"<user>\nComplete this\n</user>\n" +
"<assistant>\nPartial\n",
},
}
renderer := &LagunaRenderer{}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := renderer.Render(tt.messages, tt.tools, tt.think)
if err != nil {
t.Fatal(err)
}
if diff := cmp.Diff(tt.want, got); diff != "" {
t.Fatalf("renderer output mismatch (-want +got):\n%s", diff)
}
})
}
}
func TestLagunaRendererMatchesLocalJinjaControlFlow(t *testing.T) {
if os.Getenv("VERIFY_LAGUNA_JINJA2") == "" {
t.Skip("set VERIFY_LAGUNA_JINJA2=1 to compare against the local Laguna chat_template.jinja")
}
python := "/Users/daniel/.codex/worktrees/7038/ollama/.venv/bin/python3"
if _, err := os.Stat(python); err != nil {
t.Fatalf("VERIFY_LAGUNA_JINJA2 requires %s with jinja2 installed", python)
}
tests := []struct {
name string
messages []api.Message
think *api.ThinkValue
}{
{
name: "user_only",
messages: []api.Message{{Role: "user", Content: "Hello"}},
},
{
name: "system_user",
messages: []api.Message{
{Role: "system", Content: "Stay concise.\n"},
{Role: "user", Content: "Hello"},
},
},
{
name: "additional_system_and_tool_response",
messages: []api.Message{
{Role: "system", Content: "Primary."},
{Role: "user", Content: "Weather?"},
{Role: "assistant", Content: "Calling."},
{Role: "tool", Content: "Sunny"},
{Role: "system", Content: "Secondary."},
},
},
{
name: "thinking_enabled",
messages: []api.Message{{Role: "user", Content: "Think briefly."}},
think: &api.ThinkValue{Value: true},
},
{
name: "thinking_disabled",
messages: []api.Message{{Role: "user", Content: "Answer directly."}},
think: &api.ThinkValue{Value: false},
},
}
renderer := &LagunaRenderer{}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := renderer.Render(tt.messages, nil, tt.think)
if err != nil {
t.Fatal(err)
}
for _, modelDir := range []string{
"/Users/daniel/Models/poolside/laguna-xs-23-04-2026",
} {
want := renderLagunaChatTemplate(t, python, modelDir, tt.messages, tt.think)
if diff := cmp.Diff(want, got); diff != "" {
t.Fatalf("%s mismatch (-chat_template +renderer):\n%s", modelDir, diff)
}
}
})
}
}
func renderLagunaChatTemplate(t *testing.T, python, modelDir string, messages []api.Message, think *api.ThinkValue) string {
t.Helper()
type templateMessage struct {
Role string `json:"role"`
Content string `json:"content"`
}
templateMessages := make([]templateMessage, 0, len(messages))
for _, msg := range messages {
templateMessages = append(templateMessages, templateMessage{
Role: msg.Role,
Content: msg.Content,
})
}
messagesJSON, err := json.Marshal(templateMessages)
if err != nil {
t.Fatalf("failed to marshal messages: %v", err)
}
enableThinking := "True"
if think != nil && !think.Bool() {
enableThinking = "False"
}
script := `
import json
import sys
from transformers import AutoTokenizer
model_dir = sys.argv[1]
messages = json.loads(sys.argv[2])
enable_thinking = sys.argv[3] == "True"
tok = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
print(tok.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
enable_thinking=enable_thinking,
), end="")
`
cmd := exec.Command(python, "-c", script, modelDir, string(messagesJSON), enableThinking)
var stdout, stderr strings.Builder
cmd.Stdout = &stdout
cmd.Stderr = &stderr
if err := cmd.Run(); err != nil {
t.Fatalf("chat_template render failed: %v\nstderr: %s", err, stderr.String())
}
return stdout.String()
}
func lagunaWeatherTool() []api.Tool {
return []api.Tool{{
Type: "function",
Function: api.ToolFunction{
Name: "get_weather",
Description: "Get weather",
Parameters: api.ToolFunctionParameters{
Type: "object",
Required: []string{"location"},
Properties: testPropsOrdered([]orderedProp{{
Key: "location",
Value: api.ToolProperty{
Type: api.PropertyType{"string"},
Description: "City",
},
}}),
},
},
}}
}
model/renderers/nemotron3nano.go
+263 -34
View File
@@ -3,6 +3,9 @@ package renderers
import (
"encoding/json"
"fmt"
"reflect"
"sort"
"strconv"
"strings"
"github.com/ollama/ollama/api"
@@ -12,15 +15,15 @@ type Nemotron3NanoRenderer struct{}
func (r *Nemotron3NanoRenderer) Render(messages []api.Message, tools []api.Tool, thinkValue *api.ThinkValue) (string, error) {
var sb strings.Builder
imageOffset := 0
// thinking is enabled if user requests it
enableThinking := thinkValue != nil && thinkValue.Bool()
enableThinking := r.resolveThinking(messages, thinkValue)
// Extract system message if present
var systemMessage string
var loopMessages []api.Message
if len(messages) > 0 && messages[0].Role == "system" {
systemMessage = messages[0].Content
systemMessage = r.sanitizeSystemMessage(messages[0].Content)
loopMessages = messages[1:]
} else {
loopMessages = messages
@@ -34,6 +37,7 @@ func (r *Nemotron3NanoRenderer) Render(messages []api.Message, tools []api.Tool,
}
}
sb.WriteString("\n\n\n")
sb.WriteString("<|im_start|>system\n")
if systemMessage != "" {
sb.WriteString(systemMessage)
@@ -45,28 +49,30 @@ func (r *Nemotron3NanoRenderer) Render(messages []api.Message, tools []api.Tool,
}
sb.WriteString(r.renderTools(tools))
}
sb.WriteString("<|im_end|>\n")
sb.WriteString("<|im_end|>\n\n")
for i, message := range loopMessages {
switch message.Role {
case "assistant":
// Build content with thinking tags
content := r.buildContent(message)
shouldTruncate := i < lastUserIdx
if len(message.ToolCalls) > 0 {
sb.WriteString("<|im_start|>assistant\n")
sb.WriteString(r.formatContent(content, shouldTruncate, true))
sb.WriteString(r.formatToolCallContent(content, shouldTruncate))
r.writeToolCalls(&sb, message.ToolCalls)
sb.WriteString("<|im_end|>\n")
} else {
formatted := r.formatContent(content, shouldTruncate, false)
sb.WriteString("<|im_start|>assistant\n" + formatted + "<|im_end|>\n")
formatted := r.formatAssistantContent(content, shouldTruncate)
sb.WriteString("<|im_start|>assistant\n")
sb.WriteString(formatted)
sb.WriteString("<|im_end|>\n")
}
case "user", "system":
sb.WriteString("<|im_start|>" + message.Role + "\n")
sb.WriteString(message.Content)
sb.WriteString(r.renderMessageContent(message, imageOffset))
imageOffset += len(message.Images)
sb.WriteString("<|im_end|>\n")
case "tool":
@@ -90,6 +96,8 @@ func (r *Nemotron3NanoRenderer) Render(messages []api.Message, tools []api.Tool,
}
}
sb.WriteString("\n")
// Add generation prompt
if enableThinking {
sb.WriteString("<|im_start|>assistant\n<think>\n")
@@ -119,7 +127,7 @@ func (r *Nemotron3NanoRenderer) renderTools(tools []api.Tool) string {
sb.WriteString("\n<name>" + paramName + "</name>")
if len(paramFields.Type) > 0 {
sb.WriteString("\n<type>" + strings.Join(paramFields.Type, ", ") + "</type>")
sb.WriteString("\n<type>" + r.formatPropertyType(paramFields.Type) + "</type>")
}
if paramFields.Description != "" {
@@ -127,17 +135,17 @@ func (r *Nemotron3NanoRenderer) renderTools(tools []api.Tool) string {
}
if len(paramFields.Enum) > 0 {
enumJSON, _ := json.Marshal(paramFields.Enum)
sb.WriteString("\n<enum>" + string(enumJSON) + "</enum>")
sb.WriteString("\n<enum>" + r.pythonJSON(paramFields.Enum) + "</enum>")
}
r.renderToolPropertyExtraKeys(&sb, paramFields)
sb.WriteString("\n</parameter>")
}
}
r.renderToolParameterExtraKeys(&sb, fn.Parameters)
if len(fn.Parameters.Required) > 0 {
reqJSON, _ := json.Marshal(fn.Parameters.Required)
sb.WriteString("\n<required>" + string(reqJSON) + "</required>")
sb.WriteString("\n<required>" + r.pythonJSON(fn.Parameters.Required) + "</required>")
}
sb.WriteString("\n</parameters>")
@@ -159,27 +167,38 @@ func (r *Nemotron3NanoRenderer) renderTools(tools []api.Tool) string {
}
func (r *Nemotron3NanoRenderer) buildContent(message api.Message) string {
// The parser always extracts thinking into the Thinking field,
// so Content will never have <think> tags embedded
content := nemotron3NanoRenderContent(message.Content)
if message.Thinking != "" {
return "<think>\n" + message.Thinking + "\n</think>\n" + message.Content
return "<think>\n" + message.Thinking + "\n</think>\n" + content
}
return "<think></think>" + message.Content
if !strings.Contains(content, "<think>") && !strings.Contains(content, "</think>") {
return "<think></think>" + content
}
return content
}
func (r *Nemotron3NanoRenderer) formatContent(content string, truncate bool, addNewline bool) string {
if content == "" {
func (r *Nemotron3NanoRenderer) formatAssistantContent(content string, truncate bool) string {
if !truncate {
return strings.TrimSpace(content)
}
c := content
if strings.Contains(c, "<think>") && strings.Contains(c, "</think>") {
parts := strings.Split(c, "</think>")
c = "<think></think>" + parts[len(parts)-1]
}
return strings.TrimSpace(c)
}
func (r *Nemotron3NanoRenderer) formatToolCallContent(content string, truncate bool) string {
if strings.TrimSpace(content) == "" {
return "<think></think>"
}
if !truncate {
if addNewline {
return strings.TrimSpace(content) + "\n"
}
return strings.TrimSpace(content)
return strings.TrimSpace(content) + "\n"
}
// Truncate thinking - keep only content after </think>
c := content
if strings.Contains(c, "</think>") {
parts := strings.Split(c, "</think>")
@@ -190,13 +209,7 @@ func (r *Nemotron3NanoRenderer) formatContent(content string, truncate bool, add
}
c = "<think></think>" + strings.TrimSpace(c)
if addNewline && len(c) > len("<think></think>") {
return c + "\n"
}
if c == "<think></think>" {
return c
}
return strings.TrimSpace(c)
return strings.TrimSpace(c) + "\n"
}
func (r *Nemotron3NanoRenderer) writeToolCalls(sb *strings.Builder, toolCalls []api.ToolCall) {
@@ -212,9 +225,225 @@ func (r *Nemotron3NanoRenderer) writeToolCalls(sb *strings.Builder, toolCalls []
func (r *Nemotron3NanoRenderer) formatArgValue(value any) string {
switch v := value.(type) {
case map[string]any, []any:
jsonBytes, _ := json.Marshal(v)
return string(jsonBytes)
return r.pythonJSON(v)
default:
return fmt.Sprintf("%v", v)
}
}
func (r *Nemotron3NanoRenderer) renderMessageContent(message api.Message, imageOffset int) string {
content := nemotron3NanoRenderContent(message.Content)
if len(message.Images) == 0 {
return content
}
if strings.Contains(content, "[img-") {
return content
}
if strings.Contains(content, "[img]") {
for i := range message.Images {
content = strings.Replace(content, "[img]", fmt.Sprintf("[img-%d]", imageOffset+i), 1)
}
return content
}
var sb strings.Builder
for i := range message.Images {
sb.WriteString(fmt.Sprintf("[img-%d]", imageOffset+i))
}
sb.WriteString(content)
return sb.String()
}
func nemotron3NanoRenderContent(content any) string {
switch v := content.(type) {
case string:
return v
case []any:
var sb strings.Builder
for _, item := range v {
obj, ok := item.(map[string]any)
if !ok {
bts, _ := json.Marshal(item)
sb.Write(bts)
continue
}
switch obj["type"] {
case "image":
sb.WriteString("<image>")
case "text":
if text, ok := obj["text"].(string); ok {
sb.WriteString(text)
}
default:
bts, _ := json.Marshal(item)
sb.Write(bts)
}
}
return sb.String()
default:
bts, _ := json.Marshal(v)
return string(bts)
}
}
func (r *Nemotron3NanoRenderer) resolveThinking(messages []api.Message, thinkValue *api.ThinkValue) bool {
enableThinking := thinkValue == nil || thinkValue.Bool()
for _, message := range messages {
if message.Role != "user" && message.Role != "system" {
continue
}
content := message.Content
if strings.Contains(strings.ReplaceAll(content, "</think>", ""), "/think") {
enableThinking = true
} else if strings.Contains(content, "/no_think") {
enableThinking = false
}
}
return enableThinking
}
func (r *Nemotron3NanoRenderer) sanitizeSystemMessage(content string) string {
system := nemotron3NanoRenderContent(content)
system = strings.ReplaceAll(system, "</think>", "<_end_think>")
system = strings.ReplaceAll(system, "/think", "")
system = strings.ReplaceAll(system, "/no_think", "")
system = strings.ReplaceAll(system, "<_end_think>", "</think>")
return system
}
func (r *Nemotron3NanoRenderer) formatPropertyType(propertyType api.PropertyType) string {
if len(propertyType) == 1 {
return propertyType[0]
}
quoted := make([]string, 0, len(propertyType))
for _, v := range propertyType {
quoted = append(quoted, "'"+v+"'")
}
return "[" + strings.Join(quoted, ", ") + "]"
}
func (r *Nemotron3NanoRenderer) renderToolPropertyExtraKeys(sb *strings.Builder, prop api.ToolProperty) {
if len(prop.AnyOf) > 0 {
sb.WriteString("\n<anyOf>" + r.pythonJSON(prop.AnyOf) + "</anyOf>")
}
if prop.Items != nil {
sb.WriteString("\n<items>" + r.pythonJSON(prop.Items) + "</items>")
}
if prop.Properties != nil {
sb.WriteString("\n<properties>" + r.pythonJSON(prop.Properties) + "</properties>")
}
if len(prop.Required) > 0 {
sb.WriteString("\n<required>" + r.pythonJSON(prop.Required) + "</required>")
}
}
func (r *Nemotron3NanoRenderer) renderToolParameterExtraKeys(sb *strings.Builder, params api.ToolFunctionParameters) {
if params.Defs != nil {
sb.WriteString("\n<$defs>" + r.pythonJSON(params.Defs) + "</$defs>")
}
if params.Items != nil {
sb.WriteString("\n<items>" + r.pythonJSON(params.Items) + "</items>")
}
}
func (r *Nemotron3NanoRenderer) pythonJSON(v any) string {
switch value := v.(type) {
case nil:
return "null"
case string:
return strconv.Quote(value)
case bool:
if value {
return "true"
}
return "false"
case int, int8, int16, int32, int64:
return fmt.Sprintf("%d", reflect.ValueOf(value).Int())
case uint, uint8, uint16, uint32, uint64:
return fmt.Sprintf("%d", reflect.ValueOf(value).Uint())
case float32, float64:
b, _ := json.Marshal(value)
return string(b)
case api.PropertyType:
return r.pythonJSON([]string(value))
case []string:
parts := make([]string, 0, len(value))
for _, item := range value {
parts = append(parts, r.pythonJSON(item))
}
return "[" + strings.Join(parts, ", ") + "]"
case []any:
parts := make([]string, 0, len(value))
for _, item := range value {
parts = append(parts, r.pythonJSON(item))
}
return "[" + strings.Join(parts, ", ") + "]"
case []api.ToolProperty:
parts := make([]string, 0, len(value))
for _, item := range value {
parts = append(parts, r.pythonJSON(item))
}
return "[" + strings.Join(parts, ", ") + "]"
case map[string]any:
keys := make([]string, 0, len(value))
for key := range value {
keys = append(keys, key)
}
sort.Strings(keys)
parts := make([]string, 0, len(keys))
for _, key := range keys {
parts = append(parts, strconv.Quote(key)+": "+r.pythonJSON(value[key]))
}
return "{" + strings.Join(parts, ", ") + "}"
case *api.ToolPropertiesMap:
if value == nil {
return "null"
}
parts := make([]string, 0, value.Len())
for key, prop := range value.All() {
parts = append(parts, strconv.Quote(key)+": "+r.pythonJSON(prop))
}
return "{" + strings.Join(parts, ", ") + "}"
case api.ToolProperty:
parts := make([]string, 0, 6)
if len(value.AnyOf) > 0 {
parts = append(parts, `"anyOf": `+r.pythonJSON(value.AnyOf))
}
if len(value.Type) > 0 {
if len(value.Type) == 1 {
parts = append(parts, `"type": `+r.pythonJSON(value.Type[0]))
} else {
parts = append(parts, `"type": `+r.pythonJSON([]string(value.Type)))
}
}
if value.Items != nil {
parts = append(parts, `"items": `+r.pythonJSON(value.Items))
}
if value.Description != "" {
parts = append(parts, `"description": `+r.pythonJSON(value.Description))
}
if len(value.Enum) > 0 {
parts = append(parts, `"enum": `+r.pythonJSON(value.Enum))
}
if value.Properties != nil {
parts = append(parts, `"properties": `+r.pythonJSON(value.Properties))
}
if len(value.Required) > 0 {
parts = append(parts, `"required": `+r.pythonJSON(value.Required))
}
return "{" + strings.Join(parts, ", ") + "}"
default:
b, err := json.Marshal(value)
if err != nil {
return "null"
}
var generic any
if err := json.Unmarshal(b, &generic); err != nil {
return string(b)
}
return r.pythonJSON(generic)
}
}
model/renderers/nemotron3nano_reference_test.go
+614
View File
@@ -0,0 +1,614 @@
package renderers
import (
"encoding/json"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
"testing"
"github.com/google/go-cmp/cmp"
"github.com/ollama/ollama/api"
)
const nemotron3NanoTemplate = "testdata/nemotron3nano_chat_template.jinja2"
func TestNemotron3NanoRendererMatchesReference(t *testing.T) {
toolText := `<|im_start|>system
# Tools
You have access to the following functions:
<tools>
<function>
<name>search_docs</name>
<description>Search docs</description>
<parameters>
<parameter>
<name>query</name>
<type>string</type>
<description>Search query</description>
<enum>["api", "cli"]</enum>
</parameter>
<parameter>
<name>mode</name>
<type>['string', 'null']</type>
<description>Mode</description>
<anyOf>[{"type": "string"}, {"type": "number"}]</anyOf>
</parameter>
<parameter>
<name>payload</name>
<type>object</type>
<description>Payload</description>
<properties>{"enabled": {"type": "boolean"}}</properties>
<required>["enabled"]</required>
</parameter>
<parameter>
<name>tags</name>
<type>array</type>
<description>Tags</description>
<items>{"type": "string"}</items>
</parameter>
<$defs>{"shared": {"type": "string"}}</$defs>
<required>["query"]</required>
</parameters>
</function>
</tools>
If you choose to call a function ONLY reply in the following format with NO suffix:
<tool_call>
<function=example_function_name>
<parameter=example_parameter_1>
value_1
</parameter>
<parameter=example_parameter_2>
This is the value for the second parameter
that can span
multiple lines
</parameter>
</function>
</tool_call>
<IMPORTANT>
Reminder:
- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags
- Required parameters MUST be specified
- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after
- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls
</IMPORTANT><|im_end|>
`
toolTextWithSystem := `<|im_start|>system
Follow policy.
# Tools
You have access to the following functions:
<tools>
<function>
<name>search_docs</name>
<description>Search docs</description>
<parameters>
<parameter>
<name>query</name>
<type>string</type>
<description>Search query</description>
<enum>["api", "cli"]</enum>
</parameter>
<parameter>
<name>mode</name>
<type>['string', 'null']</type>
<description>Mode</description>
<anyOf>[{"type": "string"}, {"type": "number"}]</anyOf>
</parameter>
<parameter>
<name>payload</name>
<type>object</type>
<description>Payload</description>
<properties>{"enabled": {"type": "boolean"}}</properties>
<required>["enabled"]</required>
</parameter>
<parameter>
<name>tags</name>
<type>array</type>
<description>Tags</description>
<items>{"type": "string"}</items>
</parameter>
<$defs>{"shared": {"type": "string"}}</$defs>
<required>["query"]</required>
</parameters>
</function>
</tools>
If you choose to call a function ONLY reply in the following format with NO suffix:
<tool_call>
<function=example_function_name>
<parameter=example_parameter_1>
value_1
</parameter>
<parameter=example_parameter_2>
This is the value for the second parameter
that can span
multiple lines
</parameter>
</function>
</tool_call>
<IMPORTANT>
Reminder:
- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags
- Required parameters MUST be specified
- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after
- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls
</IMPORTANT><|im_end|>
`
tests := []struct {
name string
messages []api.Message
tools []api.Tool
think *api.ThinkValue
expected string
}{
{
name: "no system default thinking on",
messages: []api.Message{
{Role: "user", Content: "Hello"},
},
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHello<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "no system explicit thinking off",
messages: []api.Message{
{Role: "user", Content: "Hello"},
},
think: thinkFalse(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHello<|im_end|>\n\n<|im_start|>assistant\n<think></think>",
},
{
name: "literal endthink does not enable thinking",
messages: []api.Message{
{Role: "user", Content: "literal </think> only"},
},
think: thinkFalse(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nliteral </think> only<|im_end|>\n\n<|im_start|>assistant\n<think></think>",
},
{
name: "user no think toggle",
messages: []api.Message{
{Role: "user", Content: "Hello /no_think"},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHello /no_think<|im_end|>\n\n<|im_start|>assistant\n<think></think>",
},
{
name: "system think toggle overrides false",
messages: []api.Message{
{Role: "system", Content: "Policy /think"},
{Role: "user", Content: "Hello"},
},
think: thinkFalse(),
expected: "\n\n\n<|im_start|>system\nPolicy <|im_end|>\n\n<|im_start|>user\nHello<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "later toggle wins",
messages: []api.Message{
{Role: "system", Content: "Policy /no_think"},
{Role: "user", Content: "Actually /think"},
},
think: thinkFalse(),
expected: "\n\n\n<|im_start|>system\nPolicy <|im_end|>\n\n<|im_start|>user\nActually /think<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "system sanitizes toggles but preserves closing tag",
messages: []api.Message{
{Role: "system", Content: "A /think B /no_think C </think>"},
{Role: "user", Content: "Hello"},
},
think: thinkFalse(),
expected: "\n\n\n<|im_start|>system\nA B C </think><|im_end|>\n\n<|im_start|>user\nHello<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "assistant plain content adds empty think block",
messages: []api.Message{
{Role: "user", Content: "Hi"},
{Role: "assistant", Content: "Hello there"},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think></think>Hello there<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "assistant reasoning content",
messages: []api.Message{
{Role: "user", Content: "Hi"},
{Role: "assistant", Content: "Answer", Thinking: "Need to think"},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think>\nNeed to think\n</think>\nAnswer<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "assistant preserves existing think tags",
messages: []api.Message{
{Role: "user", Content: "Hi"},
{Role: "assistant", Content: "<think>kept</think>Answer"},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think>kept</think>Answer<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "tools without system",
messages: []api.Message{
{Role: "user", Content: "Use a tool"},
},
tools: nemotron3NanoReferenceTools(),
think: thinkTrue(),
expected: "\n\n\n" + toolText + "\n<|im_start|>user\nUse a tool<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "system with tools",
messages: []api.Message{
{Role: "system", Content: "Follow policy."},
{Role: "user", Content: "Use a tool"},
},
tools: nemotron3NanoReferenceTools(),
think: thinkTrue(),
expected: "\n\n\n" + toolTextWithSystem + "\n<|im_start|>user\nUse a tool<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "assistant tool call with content",
messages: []api.Message{
{Role: "user", Content: "Weather?"},
{
Role: "assistant",
Content: "Checking now.",
ToolCalls: []api.ToolCall{{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: testArgs(map[string]any{"city": "Paris"}),
},
}},
},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nWeather?<|im_end|>\n<|im_start|>assistant\n<think></think>Checking now.\n<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "assistant tool call with structured arguments",
messages: []api.Message{
{Role: "user", Content: "Create data"},
{
Role: "assistant",
ToolCalls: []api.ToolCall{{
Function: api.ToolCallFunction{
Name: "create",
Arguments: testArgsOrdered([]orderedArg{
{Key: "payload", Value: map[string]any{"count": 42, "nested": map[string]any{"value": "ok"}}},
{Key: "tags", Value: []any{"a", "b"}},
}),
},
}},
},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nCreate data<|im_end|>\n<|im_start|>assistant\n<think></think>\n<tool_call>\n<function=create>\n<parameter=payload>\n{\"count\": 42, \"nested\": {\"value\": \"ok\"}}\n</parameter>\n<parameter=tags>\n[\"a\", \"b\"]\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "assistant tool call truncated with reasoning",
messages: []api.Message{
{Role: "user", Content: "Weather?"},
{
Role: "assistant",
Content: "Checking now.",
Thinking: "Need weather",
ToolCalls: []api.ToolCall{{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: testArgs(map[string]any{"city": "Paris"}),
},
}},
},
{Role: "user", Content: "And tomorrow?"},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nWeather?<|im_end|>\n<|im_start|>assistant\n<think></think>Checking now.\n<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n<|im_start|>user\nAnd tomorrow?<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "assistant tool call truncated open think only",
messages: []api.Message{
{Role: "user", Content: "Weather?"},
{
Role: "assistant",
Content: "<think>draft",
ToolCalls: []api.ToolCall{{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: testArgs(map[string]any{"city": "Paris"}),
},
}},
},
{Role: "user", Content: "And tomorrow?"},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nWeather?<|im_end|>\n<|im_start|>assistant\n<think></think>\n<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n<|im_start|>user\nAnd tomorrow?<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "assistant tool call empty content",
messages: []api.Message{
{Role: "user", Content: "Weather?"},
{
Role: "assistant",
Content: "",
ToolCalls: []api.ToolCall{{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: testArgs(map[string]any{"city": "Paris"}),
},
}},
},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nWeather?<|im_end|>\n<|im_start|>assistant\n<think></think>\n<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "assistant truncated with think pair",
messages: []api.Message{
{Role: "user", Content: "Hi"},
{Role: "assistant", Content: "<think>hidden</think>Visible"},
{Role: "user", Content: "Next"},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think></think>Visible<|im_end|>\n<|im_start|>user\nNext<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "assistant truncated reasoning content",
messages: []api.Message{
{Role: "user", Content: "Hi"},
{Role: "assistant", Thinking: "hidden", Content: "Visible"},
{Role: "user", Content: "Next"},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think></think>\nVisible<|im_end|>\n<|im_start|>user\nNext<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "assistant truncated plain content",
messages: []api.Message{
{Role: "user", Content: "Hi"},
{Role: "assistant", Content: "Visible"},
{Role: "user", Content: "Next"},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think></think>Visible<|im_end|>\n<|im_start|>user\nNext<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "assistant truncated empty content",
messages: []api.Message{
{Role: "user", Content: "Hi"},
{Role: "assistant", Content: ""},
{Role: "user", Content: "Next"},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think></think><|im_end|>\n<|im_start|>user\nNext<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "consecutive tool messages grouped",
messages: []api.Message{
{Role: "user", Content: "Do work"},
{
Role: "assistant",
ToolCalls: []api.ToolCall{{
Function: api.ToolCallFunction{
Name: "step",
Arguments: testArgs(map[string]any{"value": 1}),
},
}},
},
{Role: "tool", Content: "one"},
{Role: "tool", Content: "two"},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nDo work<|im_end|>\n<|im_start|>assistant\n<think></think>\n<tool_call>\n<function=step>\n<parameter=value>\n1\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n<|im_start|>user\n<tool_response>\none\n</tool_response>\n<tool_response>\ntwo\n</tool_response>\n<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "fallback role",
messages: []api.Message{
{Role: "developer", Content: "Custom role content"},
},
think: thinkTrue(),
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>developer\nCustom role content<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
}
verifyJinja2 := os.Getenv("VERIFY_JINJA2") != ""
if verifyJinja2 {
if _, err := os.Stat(filepath.Join(nemotron3NanoRepoRoot(t), ".venv", "bin", "python3")); err != nil {
t.Fatal("VERIFY_JINJA2=1 requires .venv/bin/python3")
}
}
renderer := &Nemotron3NanoRenderer{}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := renderer.Render(tt.messages, tt.tools, tt.think)
if err != nil {
t.Fatalf("Render() error = %v", err)
}
if diff := cmp.Diff(tt.expected, got); diff != "" {
t.Fatalf("renderer mismatch (-want +got):\n%s", diff)
}
if verifyJinja2 {
jinja2Output := renderNemotron3NanoWithJinja2(t, tt.messages, tt.tools, tt.think)
if diff := cmp.Diff(tt.expected, jinja2Output); diff != "" {
t.Fatalf("reference template mismatch (-want +got):\n%s", diff)
}
}
})
}
}
func nemotron3NanoReferenceTools() []api.Tool {
return []api.Tool{{
Type: "function",
Function: api.ToolFunction{
Name: "search_docs",
Description: "Search docs",
Parameters: api.ToolFunctionParameters{
Type: "object",
Defs: map[string]any{"shared": map[string]any{"type": "string"}},
Required: []string{"query"},
Properties: testPropsOrdered([]orderedProp{
{
Key: "query",
Value: api.ToolProperty{
Type: api.PropertyType{"string"},
Description: "Search query",
Enum: []any{"api", "cli"},
},
},
{
Key: "mode",
Value: api.ToolProperty{
Type: api.PropertyType{"string", "null"},
Description: "Mode",
AnyOf: []api.ToolProperty{
{Type: api.PropertyType{"string"}},
{Type: api.PropertyType{"number"}},
},
},
},
{
Key: "payload",
Value: api.ToolProperty{
Type: api.PropertyType{"object"},
Description: "Payload",
Properties: testPropsOrdered([]orderedProp{{Key: "enabled", Value: api.ToolProperty{Type: api.PropertyType{"boolean"}}}}),
Required: []string{"enabled"},
},
},
{
Key: "tags",
Value: api.ToolProperty{
Type: api.PropertyType{"array"},
Description: "Tags",
Items: map[string]any{"type": "string"},
},
},
}),
},
},
}}
}
func renderNemotron3NanoWithJinja2(t *testing.T, messages []api.Message, tools []api.Tool, think *api.ThinkValue) string {
t.Helper()
type jinja2ToolCall struct {
ID string `json:"id,omitempty"`
Function struct {
Name string `json:"name"`
Arguments any `json:"arguments"`
} `json:"function"`
}
type jinja2Message struct {
Role string `json:"role"`
Content string `json:"content"`
ReasoningContent string `json:"reasoning_content,omitempty"`
ToolCalls []jinja2ToolCall `json:"tool_calls,omitempty"`
Name string `json:"name,omitempty"`
ToolCallID string `json:"tool_call_id,omitempty"`
}
var jMsgs []jinja2Message
for _, m := range messages {
jm := jinja2Message{
Role: m.Role,
Content: m.Content,
ReasoningContent: m.Thinking,
Name: m.ToolName,
ToolCallID: m.ToolCallID,
}
for _, tc := range m.ToolCalls {
jtc := jinja2ToolCall{ID: tc.ID}
jtc.Function.Name = tc.Function.Name
var args map[string]any
raw, _ := tc.Function.Arguments.MarshalJSON()
if err := json.Unmarshal(raw, &args); err != nil {
t.Fatalf("failed to unmarshal tool args: %v", err)
}
jtc.Function.Arguments = args
jm.ToolCalls = append(jm.ToolCalls, jtc)
}
jMsgs = append(jMsgs, jm)
}
msgsJSON, err := json.Marshal(jMsgs)
if err != nil {
t.Fatalf("failed to marshal messages: %v", err)
}
toolsJSON := "None"
if len(tools) > 0 {
b, err := json.Marshal(tools)
if err != nil {
t.Fatalf("failed to marshal tools: %v", err)
}
toolsJSON = string(b)
}
thinking := "unset"
if think != nil {
if think.Bool() {
thinking = "true"
} else {
thinking = "false"
}
}
repoRoot := nemotron3NanoRepoRoot(t)
templatePath := filepath.Join(repoRoot, "model", "renderers", nemotron3NanoTemplate)
pythonPath := filepath.Join(repoRoot, ".venv", "bin", "python3")
script := `
import json
import sys
from pathlib import Path
from transformers.utils.chat_template_utils import _compile_jinja_template
template_path, messages_json, tools_json, thinking = sys.argv[1:5]
tmpl = _compile_jinja_template(Path(template_path).read_text())
kwargs = {
"messages": json.loads(messages_json),
"add_generation_prompt": True,
}
if tools_json != "None":
kwargs["tools"] = json.loads(tools_json)
if thinking == "true":
kwargs["enable_thinking"] = True
elif thinking == "false":
kwargs["enable_thinking"] = False
print(tmpl.render(**kwargs), end="")
`
cmd := exec.Command(pythonPath, "-c", script, templatePath, string(msgsJSON), toolsJSON, thinking)
var stdout, stderr strings.Builder
cmd.Stdout = &stdout
cmd.Stderr = &stderr
if err := cmd.Run(); err != nil {
t.Fatalf("python render failed: %v\nstderr: %s", err, stderr.String())
}
return stdout.String()
}
func nemotron3NanoRepoRoot(t *testing.T) string {
t.Helper()
_, filename, _, ok := runtime.Caller(0)
if !ok {
t.Fatal("failed to locate test file")
}
return filepath.Dir(filepath.Dir(filepath.Dir(filename)))
}
model/renderers/nemotron3nano_test.go
+18 -533
View File
@@ -8,561 +8,46 @@ import (
"github.com/ollama/ollama/api"
)
func TestNemotron3NanoRenderer(t *testing.T) {
func TestNemotron3NanoRenderer_Images(t *testing.T) {
tests := []struct {
name string
msgs []api.Message
tools []api.Tool
thinkValue *api.ThinkValue
expected string
name string
msgs []api.Message
expected string
}{
{
name: "basic user message - thinking mode",
name: "single image inserts placeholder",
msgs: []api.Message{
{Role: "user", Content: "Hello!"},
{Role: "user", Content: "Describe this image.", Images: []api.ImageData{api.ImageData("img1")}},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\n<|im_end|>\n" +
"<|im_start|>user\nHello!<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\n[img-0]Describe this image.<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "basic user message - no thinking",
name: "generic image placeholder is rewritten",
msgs: []api.Message{
{Role: "user", Content: "Hello!"},
{Role: "user", Content: "[img]Describe this image.", Images: []api.ImageData{api.ImageData("img1")}},
},
thinkValue: nil,
expected: "<|im_start|>system\n<|im_end|>\n" +
"<|im_start|>user\nHello!<|im_end|>\n" +
"<|im_start|>assistant\n<think></think>",
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\n[img-0]Describe this image.<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
{
name: "with system message",
name: "image offsets increment across messages",
msgs: []api.Message{
{Role: "system", Content: "You are a helpful assistant."},
{Role: "user", Content: "Hello!"},
{Role: "user", Content: "Describe the first image.", Images: []api.ImageData{api.ImageData("img1")}},
{Role: "assistant", Content: "It shows something."},
{Role: "user", Content: "Compare these.", Images: []api.ImageData{api.ImageData("img2"), api.ImageData("img3")}},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" +
"<|im_start|>user\nHello!<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
},
{
name: "multi-turn conversation",
msgs: []api.Message{
{Role: "user", Content: "Hi"},
{Role: "assistant", Content: "Hello! How can I help?"},
{Role: "user", Content: "Tell me a joke"},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\n<|im_end|>\n" +
"<|im_start|>user\nHi<|im_end|>\n" +
"<|im_start|>assistant\n<think></think>Hello! How can I help?<|im_end|>\n" +
"<|im_start|>user\nTell me a joke<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
},
{
name: "with tools",
msgs: []api.Message{
{Role: "user", Content: "What's the weather in Paris?"},
},
tools: []api.Tool{
{
Type: "function",
Function: api.ToolFunction{
Name: "get_weather",
Description: "Get the current weather",
Parameters: api.ToolFunctionParameters{
Type: "object",
Required: []string{"city"},
Properties: testPropsMap(map[string]api.ToolProperty{
"city": {Type: api.PropertyType{"string"}, Description: "The city name"},
}),
},
},
},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\n" +
"# Tools\n\nYou have access to the following functions:\n\n<tools>\n" +
"<function>\n<name>get_weather</name>\n" +
"<description>Get the current weather</description>\n" +
"<parameters>\n" +
"<parameter>\n<name>city</name>\n<type>string</type>\n<description>The city name</description>\n</parameter>\n" +
"<required>[\"city\"]</required>\n" +
"</parameters>\n</function>\n</tools>\n\n" +
"If you choose to call a function ONLY reply in the following format with NO suffix:\n\n" +
"<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n" +
"<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n" +
"</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n" +
"- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n" +
"- Required parameters MUST be specified\n" +
"- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n" +
"- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n" +
"</IMPORTANT><|im_end|>\n" +
"<|im_start|>user\nWhat's the weather in Paris?<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
},
{
name: "tool call with response",
msgs: []api.Message{
{Role: "user", Content: "What's the weather in Paris?"},
{
Role: "assistant",
ToolCalls: []api.ToolCall{
{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: testArgs(map[string]any{"city": "Paris"}),
},
},
},
},
{Role: "tool", Content: "Sunny, 72F"},
},
tools: []api.Tool{
{
Type: "function",
Function: api.ToolFunction{
Name: "get_weather",
Description: "Get the current weather",
Parameters: api.ToolFunctionParameters{
Type: "object",
Required: []string{"city"},
Properties: testPropsMap(map[string]api.ToolProperty{
"city": {Type: api.PropertyType{"string"}, Description: "The city name"},
}),
},
},
},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\n" +
"# Tools\n\nYou have access to the following functions:\n\n<tools>\n" +
"<function>\n<name>get_weather</name>\n" +
"<description>Get the current weather</description>\n" +
"<parameters>\n" +
"<parameter>\n<name>city</name>\n<type>string</type>\n<description>The city name</description>\n</parameter>\n" +
"<required>[\"city\"]</required>\n" +
"</parameters>\n</function>\n</tools>\n\n" +
"If you choose to call a function ONLY reply in the following format with NO suffix:\n\n" +
"<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n" +
"<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n" +
"</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n" +
"- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n" +
"- Required parameters MUST be specified\n" +
"- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n" +
"- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n" +
"</IMPORTANT><|im_end|>\n" +
"<|im_start|>user\nWhat's the weather in Paris?<|im_end|>\n" +
"<|im_start|>assistant\n<think></think>\n" +
"<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n" +
"<|im_start|>user\n<tool_response>\nSunny, 72F\n</tool_response>\n<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
},
{
name: "assistant with content and tool call",
msgs: []api.Message{
{Role: "user", Content: "What's the weather?"},
{
Role: "assistant",
Content: "Let me check that for you.",
ToolCalls: []api.ToolCall{
{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: testArgs(map[string]any{"city": "Paris"}),
},
},
},
},
{Role: "tool", Content: "Sunny"},
},
tools: []api.Tool{
{
Type: "function",
Function: api.ToolFunction{
Name: "get_weather",
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: testPropsMap(map[string]api.ToolProperty{
"city": {Type: api.PropertyType{"string"}},
}),
},
},
},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\n" +
"# Tools\n\nYou have access to the following functions:\n\n<tools>\n" +
"<function>\n<name>get_weather</name>\n" +
"<parameters>\n" +
"<parameter>\n<name>city</name>\n<type>string</type>\n</parameter>\n" +
"</parameters>\n</function>\n</tools>\n\n" +
"If you choose to call a function ONLY reply in the following format with NO suffix:\n\n" +
"<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n" +
"<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n" +
"</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n" +
"- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n" +
"- Required parameters MUST be specified\n" +
"- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n" +
"- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n" +
"</IMPORTANT><|im_end|>\n" +
"<|im_start|>user\nWhat's the weather?<|im_end|>\n" +
"<|im_start|>assistant\n<think></think>Let me check that for you.\n" +
"<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n" +
"<|im_start|>user\n<tool_response>\nSunny\n</tool_response>\n<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
},
{
name: "thinking in history is truncated",
msgs: []api.Message{
{Role: "user", Content: "Hi"},
{Role: "assistant", Content: "Hello!", Thinking: "Let me think about this..."},
{Role: "user", Content: "How are you?"},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\n<|im_end|>\n" +
"<|im_start|>user\nHi<|im_end|>\n" +
"<|im_start|>assistant\n<think></think>Hello!<|im_end|>\n" +
"<|im_start|>user\nHow are you?<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
},
{
name: "parallel tool calls",
msgs: []api.Message{
{Role: "user", Content: "Weather in Paris and London?"},
{
Role: "assistant",
ToolCalls: []api.ToolCall{
{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: testArgs(map[string]any{"city": "Paris"}),
},
},
{
Function: api.ToolCallFunction{
Name: "get_weather",
Arguments: testArgs(map[string]any{"city": "London"}),
},
},
},
},
{Role: "tool", Content: "Sunny"},
{Role: "tool", Content: "Rainy"},
},
tools: []api.Tool{
{
Type: "function",
Function: api.ToolFunction{
Name: "get_weather",
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: testPropsMap(map[string]api.ToolProperty{
"city": {Type: api.PropertyType{"string"}},
}),
},
},
},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\n" +
"# Tools\n\nYou have access to the following functions:\n\n<tools>\n" +
"<function>\n<name>get_weather</name>\n" +
"<parameters>\n" +
"<parameter>\n<name>city</name>\n<type>string</type>\n</parameter>\n" +
"</parameters>\n</function>\n</tools>\n\n" +
"If you choose to call a function ONLY reply in the following format with NO suffix:\n\n" +
"<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n" +
"<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n" +
"</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n" +
"- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n" +
"- Required parameters MUST be specified\n" +
"- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n" +
"- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n" +
"</IMPORTANT><|im_end|>\n" +
"<|im_start|>user\nWeather in Paris and London?<|im_end|>\n" +
"<|im_start|>assistant\n<think></think>\n" +
"<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>\n" +
"<tool_call>\n<function=get_weather>\n<parameter=city>\nLondon\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n" +
"<|im_start|>user\n<tool_response>\nSunny\n</tool_response>\n<tool_response>\nRainy\n</tool_response>\n<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
},
{
name: "thinking disabled when user doesn't request it",
msgs: []api.Message{
{Role: "user", Content: "Hello!"},
},
thinkValue: nil,
expected: "<|im_start|>system\n<|im_end|>\n" +
"<|im_start|>user\nHello!<|im_end|>\n" +
"<|im_start|>assistant\n<think></think>",
},
{
name: "complex message history with thinking, tools, tool calls, tool results and content",
msgs: []api.Message{
{Role: "user", Content: "What's the weather in Paris and London? Also, what's 2+2?"},
{Role: "assistant", Content: "", Thinking: "I need to check the weather for both cities and calculate 2+2. Let me start with the weather calls.", ToolCalls: []api.ToolCall{
{Function: api.ToolCallFunction{Name: "get_weather", Arguments: testArgs(map[string]any{"city": "Paris"})}},
{Function: api.ToolCallFunction{Name: "get_weather", Arguments: testArgs(map[string]any{"city": "London"})}},
}},
{Role: "tool", Content: "Sunny, 22°C", ToolCallID: "call1"},
{Role: "tool", Content: "Rainy, 15°C", ToolCallID: "call2"},
{Role: "assistant", Content: "", Thinking: "Now I have the weather data. Let me calculate 2+2.", ToolCalls: []api.ToolCall{
{Function: api.ToolCallFunction{Name: "calculate", Arguments: testArgs(map[string]any{"expression": "2+2"})}},
}},
{Role: "tool", Content: "4", ToolCallID: "call3"},
{Role: "assistant", Content: "Based on the weather data, Paris is sunny at 22°C and London is rainy at 15°C. Also, 2+2 equals 4.", Thinking: "Perfect! I have all the information needed to provide a complete answer."},
},
tools: []api.Tool{
{
Type: "function",
Function: api.ToolFunction{
Name: "get_weather",
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: testPropsMap(map[string]api.ToolProperty{
"city": {Type: api.PropertyType{"string"}},
}),
},
},
},
{
Type: "function",
Function: api.ToolFunction{
Name: "calculate",
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: testPropsMap(map[string]api.ToolProperty{
"expression": {Type: api.PropertyType{"string"}},
}),
},
},
},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\n" +
"# Tools\n\nYou have access to the following functions:\n\n<tools>\n" +
"<function>\n<name>get_weather</name>\n" +
"<parameters>\n" +
"<parameter>\n<name>city</name>\n<type>string</type>\n</parameter>\n" +
"</parameters>\n</function>\n" +
"<function>\n<name>calculate</name>\n" +
"<parameters>\n" +
"<parameter>\n<name>expression</name>\n<type>string</type>\n</parameter>\n" +
"</parameters>\n</function>\n</tools>\n\n" +
"If you choose to call a function ONLY reply in the following format with NO suffix:\n\n" +
"<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n" +
"<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n" +
"</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n" +
"- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n" +
"- Required parameters MUST be specified\n" +
"- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n" +
"- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n" +
"</IMPORTANT><|im_end|>\n" +
"<|im_start|>user\nWhat's the weather in Paris and London? Also, what's 2+2?<|im_end|>\n" +
"<|im_start|>assistant\n" +
"<think>\nI need to check the weather for both cities and calculate 2+2. Let me start with the weather calls.\n</think>\n" +
"<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>\n" +
"<tool_call>\n<function=get_weather>\n<parameter=city>\nLondon\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n" +
"<|im_start|>user\n<tool_response>\nSunny, 22°C\n</tool_response>\n<tool_response>\nRainy, 15°C\n</tool_response>\n<|im_end|>\n" +
"<|im_start|>assistant\n" +
"<think>\nNow I have the weather data. Let me calculate 2+2.\n</think>\n" +
"<tool_call>\n<function=calculate>\n<parameter=expression>\n2+2\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n" +
"<|im_start|>user\n<tool_response>\n4\n</tool_response>\n<|im_end|>\n" +
"<|im_start|>assistant\n" +
"<think>\nPerfect! I have all the information needed to provide a complete answer.\n</think>\n" +
"Based on the weather data, Paris is sunny at 22°C and London is rainy at 15°C. Also, 2+2 equals 4.<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
},
{
name: "empty messages list",
msgs: []api.Message{},
thinkValue: nil,
expected: "<|im_start|>system\n<|im_end|>\n<|im_start|>assistant\n<think></think>",
},
{
name: "tool result with JSON content",
msgs: []api.Message{
{Role: "user", Content: "Get user info"},
{
Role: "assistant",
ToolCalls: []api.ToolCall{
{Function: api.ToolCallFunction{Name: "get_user", Arguments: testArgs(map[string]any{"id": "123"})}},
},
},
{Role: "tool", Content: `{"name": "John", "age": 30, "active": true}`},
},
tools: []api.Tool{
{
Type: "function",
Function: api.ToolFunction{
Name: "get_user",
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: testPropsMap(map[string]api.ToolProperty{"id": {Type: api.PropertyType{"string"}}}),
},
},
},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\n" +
"# Tools\n\nYou have access to the following functions:\n\n<tools>\n" +
"<function>\n<name>get_user</name>\n<parameters>\n" +
"<parameter>\n<name>id</name>\n<type>string</type>\n</parameter>\n" +
"</parameters>\n</function>\n</tools>\n\n" +
"If you choose to call a function ONLY reply in the following format with NO suffix:\n\n" +
"<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n" +
"<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n" +
"</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n" +
"- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n" +
"- Required parameters MUST be specified\n" +
"- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n" +
"- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n" +
"</IMPORTANT><|im_end|>\n" +
"<|im_start|>user\nGet user info<|im_end|>\n" +
"<|im_start|>assistant\n<think></think>\n" +
"<tool_call>\n<function=get_user>\n<parameter=id>\n123\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n" +
"<|im_start|>user\n<tool_response>\n{\"name\": \"John\", \"age\": 30, \"active\": true}\n</tool_response>\n<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
},
{
name: "assistant message with only thinking no content",
msgs: []api.Message{
{Role: "user", Content: "Think about this"},
{Role: "assistant", Thinking: "Deep thoughts here...", Content: ""},
{Role: "user", Content: "What did you think?"},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\n<|im_end|>\n" +
"<|im_start|>user\nThink about this<|im_end|>\n" +
"<|im_start|>assistant\n<think></think><|im_end|>\n" +
"<|im_start|>user\nWhat did you think?<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
},
{
name: "tool call with complex nested argument",
msgs: []api.Message{
{Role: "user", Content: "Create data"},
{
Role: "assistant",
ToolCalls: []api.ToolCall{
{Function: api.ToolCallFunction{
Name: "create",
Arguments: testArgs(map[string]any{
"data": map[string]any{"nested": "value", "count": 42},
}),
}},
},
},
{Role: "tool", Content: "Created"},
},
tools: []api.Tool{
{
Type: "function",
Function: api.ToolFunction{
Name: "create",
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: testPropsMap(map[string]api.ToolProperty{"data": {Type: api.PropertyType{"object"}}}),
},
},
},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\n" +
"# Tools\n\nYou have access to the following functions:\n\n<tools>\n" +
"<function>\n<name>create</name>\n<parameters>\n" +
"<parameter>\n<name>data</name>\n<type>object</type>\n</parameter>\n" +
"</parameters>\n</function>\n</tools>\n\n" +
"If you choose to call a function ONLY reply in the following format with NO suffix:\n\n" +
"<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n" +
"<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n" +
"</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n" +
"- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n" +
"- Required parameters MUST be specified\n" +
"- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n" +
"- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n" +
"</IMPORTANT><|im_end|>\n" +
"<|im_start|>user\nCreate data<|im_end|>\n" +
"<|im_start|>assistant\n<think></think>\n" +
"<tool_call>\n<function=create>\n<parameter=data>\n{\"count\":42,\"nested\":\"value\"}\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n" +
"<|im_start|>user\n<tool_response>\nCreated\n</tool_response>\n<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
},
{
name: "content explaining the format itself",
msgs: []api.Message{
{Role: "user", Content: "How do I format a tool call?"},
{Role: "assistant", Content: "To call a tool, use <tool_call> tags with <function=name> inside."},
{Role: "user", Content: "Thanks!"},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\n<|im_end|>\n" +
"<|im_start|>user\nHow do I format a tool call?<|im_end|>\n" +
"<|im_start|>assistant\n<think></think>To call a tool, use <tool_call> tags with <function=name> inside.<|im_end|>\n" +
"<|im_start|>user\nThanks!<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
},
{
name: "unicode in content and tool args",
msgs: []api.Message{
{Role: "user", Content: "Translate 你好"},
{
Role: "assistant",
ToolCalls: []api.ToolCall{
{Function: api.ToolCallFunction{Name: "translate", Arguments: testArgs(map[string]any{"text": "你好"})}},
},
},
{Role: "tool", Content: "Hello"},
},
tools: []api.Tool{
{
Type: "function",
Function: api.ToolFunction{
Name: "translate",
Parameters: api.ToolFunctionParameters{
Type: "object",
Properties: testPropsMap(map[string]api.ToolProperty{
"text": {Type: api.PropertyType{"string"}},
}),
},
},
},
},
thinkValue: &api.ThinkValue{Value: true},
expected: "<|im_start|>system\n" +
"# Tools\n\nYou have access to the following functions:\n\n<tools>\n" +
"<function>\n<name>translate</name>\n<parameters>\n" +
"<parameter>\n<name>text</name>\n<type>string</type>\n</parameter>\n" +
"</parameters>\n</function>\n</tools>\n\n" +
"If you choose to call a function ONLY reply in the following format with NO suffix:\n\n" +
"<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n" +
"<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n" +
"</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n" +
"- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n" +
"- Required parameters MUST be specified\n" +
"- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n" +
"- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n" +
"</IMPORTANT><|im_end|>\n" +
"<|im_start|>user\nTranslate 你好<|im_end|>\n" +
"<|im_start|>assistant\n<think></think>\n" +
"<tool_call>\n<function=translate>\n<parameter=text>\n你好\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n" +
"<|im_start|>user\n<tool_response>\nHello\n</tool_response>\n<|im_end|>\n" +
"<|im_start|>assistant\n<think>\n",
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\n[img-0]Describe the first image.<|im_end|>\n<|im_start|>assistant\n<think></think>It shows something.<|im_end|>\n<|im_start|>user\n[img-1][img-2]Compare these.<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
renderer := &Nemotron3NanoRenderer{}
rendered, err := renderer.Render(tt.msgs, tt.tools, tt.thinkValue)
rendered, err := renderer.Render(tt.msgs, nil, nil)
if err != nil {
t.Fatal(err)
}
if diff := cmp.Diff(rendered, tt.expected); diff != "" {
t.Errorf("mismatch (-got +want):\n%s", diff)
if diff := cmp.Diff(tt.expected, rendered); diff != "" {
t.Fatalf("mismatch (-want +got):\n%s", diff)
}
})
}
model/renderers/renderer.go
+2
View File
@@ -95,6 +95,8 @@ func rendererForName(name string) Renderer {
return &LFM2Renderer{IsThinking: false, useImgTags: RenderImgTags}
case "lfm2-thinking":
return &LFM2Renderer{IsThinking: true, useImgTags: RenderImgTags}
case "laguna":
return &LagunaRenderer{}
default:
return nil
}
model/renderers/testdata/nemotron3nano_chat_template.jinja2
Vendored
+222
View File
@@ -0,0 +1,222 @@
{% macro render_extra_keys(json_dict, handled_keys) %}
{%- if json_dict is mapping %}
{%- for json_key in json_dict if json_key not in handled_keys %}
{%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
{{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}
{%- else %}
{{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
{%- endif %}
{%- endfor %}
{%- endif %}
{% endmacro %}
{%- set enable_thinking = enable_thinking if enable_thinking is defined else True %}
{%- set reasoning_budget = reasoning_budget if reasoning_budget is defined else None %}
{%- set truncate_history_thinking = truncate_history_thinking if truncate_history_thinking is defined else True %}
{%- set response_format = response_format if response_format is defined else None %}
{# Scan messages for VLM thinking toggles to override enable_thinking #}
{%- set toggle = namespace(enable=enable_thinking) %}
{%- for m in messages %}
{%- if m['role'] == 'user' or m['role'] == 'system' -%}
{%- if m['content'] is string -%}
{%- set c = m['content'] %}
{%- if '/think' in c.replace('</think>', '') -%}
{%- set toggle.enable = true -%}
{%- elif '/no_think' in c -%}
{%- set toggle.enable = false -%}
{%- endif -%}
{%- endif -%}
{%- endif -%}
{%- endfor %}
{# Prepare message iteration similar to LM template #}
{%- set ns = namespace(last_user_idx = -1) %}
{%- set loop_messages = messages %}
{%- for m in loop_messages %}
{%- if m["role"] == "user" %}
{%- set ns.last_user_idx = loop.index0 %}
{%- endif %}
{%- endfor %}
{%- if messages[0]["role"] == "system" %}
{%- set system_message = messages[0]["content"] %}
{%- set loop_messages = messages[1:] %}
{%- else %}
{%- set system_message = "" %}
{%- set loop_messages = messages %}
{%- endif %}
{%- if not tools is defined %}
{%- set tools = [] %}
{%- endif %}
{# Recompute last_user_idx relative to loop_messages after handling system #}
{%- set ns = namespace(last_user_idx = -1) %}
{%- for m in loop_messages %}
{%- if m["role"] == "user" %}
{%- set ns.last_user_idx = loop.index0 %}
{%- endif %}
{%- endfor %}
{# System preamble with LM formatting, sanitize thinking toggles #}
{%- if system_message is defined %}
{%- set sys_content = system_message | string %}
{%- set sys_content = sys_content.replace('</think>', '<_end_think>').replace('/think', '').replace('/no_think', '').replace('<_end_think>', '</think>') %}
{{- "<|im_start|>system\n" + sys_content }}
{%- else %}
{%- if tools is iterable and tools | length > 0 %}
{{- "<|im_start|>system\n" }}
{%- endif %}
{%- endif %}
{%- if tools is iterable and tools | length > 0 %}
{%- if system_message is defined and system_message | length > 0 %}
{{- "\n\n" }}
{%- endif %}
{{- "# Tools\n\nYou have access to the following functions:\n\n" }}
{{- "<tools>" }}
{%- for tool in tools %}
{%- if tool.function is defined %}
{%- set tool = tool.function %}
{%- endif %}
{{- "\n<function>\n<name>" ~ tool.name ~ "</name>" }}
{%- if tool.description is defined %}
{{- '\n<description>' ~ (tool.description | trim) ~ '</description>' }}
{%- endif %}
{{- '\n<parameters>' }}
{%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
{%- for param_name, param_fields in tool.parameters.properties|items %}
{{- '\n<parameter>' }}
{{- '\n<name>' ~ param_name ~ '</name>' }}
{%- if param_fields.type is defined %}
{{- '\n<type>' ~ (param_fields.type | string) ~ '</type>' }}
{%- endif %}
{%- if param_fields.description is defined %}
{{- '\n<description>' ~ (param_fields.description | trim) ~ '</description>' }}
{%- endif %}
{%- if param_fields.enum is defined %}
{{- '\n<enum>' ~ (param_fields.enum | tojson | safe) ~ '</enum>' }}
{%- endif %}
{%- set handled_keys = ['name', 'type', 'description', 'enum'] %}
{{- render_extra_keys(param_fields, handled_keys) }}
{{- '\n</parameter>' }}
{%- endfor %}
{%- endif %}
{% set handled_keys = ['type', 'properties', 'required'] %}
{{- render_extra_keys(tool.parameters, handled_keys) }}
{%- if tool.parameters is defined and tool.parameters.required is defined %}
{{- '\n<required>' ~ (tool.parameters.required | tojson | safe) ~ '</required>' }}
{%- endif %}
{{- '\n</parameters>' }}
{%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
{{- render_extra_keys(tool, handled_keys) }}
{{- '\n</function>' }}
{%- endfor %}
{{- "\n</tools>" }}
{{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
{%- endif %}
{%- if system_message is defined %}
{{- '<|im_end|>\n' }}
{%- else %}
{%- if tools is iterable and tools | length > 0 %}
{{- '<|im_end|>\n' }}
{%- endif %}
{%- endif %}
{# Iterate conversation #}
{%- for message in loop_messages %}
{%- if message.role == "assistant" %}
{# Use LM assistant handling #}
{%- if message.reasoning_content is defined and message.reasoning_content is string and message.reasoning_content | trim | length > 0 %}
{%- set content = "<think>\n" ~ message.reasoning_content ~ "\n</think>\n" ~ (message.content | default('', true)) %}
{%- else %}
{%- set content = message.content | default('', true) %}
{%- if content is string -%}
{%- if '<think>' not in content and '</think>' not in content -%}
{%- set content = "<think></think>" ~ content -%}
{%- endif -%}
{%- else -%}
{%- set content = content -%}
{%- endif -%}
{%- endif %}
{%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
{{- '<|im_start|>assistant\n' }}
{%- set include_content = not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
{%- if content is string and content | trim | length > 0 %}
{%- if include_content %}
{{- (content | trim) ~ '\n' -}}
{%- else %}
{%- set c = (content | string) %}
{%- if '</think>' in c %}
{%- set c = c.split('</think>')[-1] %}
{%- elif '<think>' in c %}
{%- set c = c.split('<think>')[0] %}
{%- endif %}
{%- set c = "<think></think>" ~ c | trim %}
{%- if c | length > 0 %}
{{- c ~ '\n' -}}
{%- endif %}
{%- endif %}
{%- else %}
{{- "<think></think>" -}}
{%- endif %}
{%- for tool_call in message.tool_calls %}
{%- if tool_call.function is defined %}
{%- set tool_call = tool_call.function %}
{%- endif %}
{{- '<tool_call>\n<function=' ~ tool_call.name ~ '>\n' -}}
{%- if tool_call.arguments is defined %}
{%- for args_name, args_value in tool_call.arguments|items %}
{{- '<parameter=' ~ args_name ~ '>\n' -}}
{%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
{{- args_value ~ '\n</parameter>\n' -}}
{%- endfor %}
{%- endif %}
{{- '</function>\n</tool_call>\n' -}}
{%- endfor %}
{{- '<|im_end|>\n' }}
{%- else %}
{%- if not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
{{- '<|im_start|>assistant\n' ~ (content | default('', true) | string | trim) ~ '<|im_end|>\n' }}
{%- else %}
{%- set c = (content | default('', true) | string) %}
{%- if '<think>' in c and '</think>' in c %}
{%- set c = "<think></think>" ~ c.split('</think>')[-1] %}
{%- endif %}
{%- set c = c | trim %}
{%- if c | length > 0 %}
{{- '<|im_start|>assistant\n' ~ c ~ '<|im_end|>\n' }}
{%- else %}
{{- '<|im_start|>assistant\n<|im_end|>\n' }}
{%- endif %}
{%- endif %}
{%- endif %}
{%- elif message.role == "user" or message.role == "system" %}
{{- '<|im_start|>' + message.role + '\n' }}
{%- set content = message.content | string %}
{{- content }}
{{- '<|im_end|>\n' }}
{%- elif message.role == "tool" %}
{%- if loop.previtem and loop.previtem.role != "tool" %}
{{- '<|im_start|>user\n' }}
{%- endif %}
{{- '<tool_response>\n' }}
{{- message.content }}
{{- '\n</tool_response>\n' }}
{%- if not loop.last and loop.nextitem.role != "tool" %}
{{- '<|im_end|>\n' }}
{%- elif loop.last %}
{{- '<|im_end|>\n' }}
{%- endif %}
{%- else %}
{{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' }}
{%- endif %}
{%- endfor %}
{# Generation prompt using computed thinking toggle #}
{%- if add_generation_prompt %}
{%- if toggle.enable %}
{{- '<|im_start|>assistant\n<think>\n' }}
{%- else %}
{{- '<|im_start|>assistant\n<think></think>' }}
{%- endif %}
{%- endif %}