diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 3a7ba8555..3c6bfb2a7 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -215,8 +215,18 @@ jobs: name: Install CUDA ${{ matrix.cuda-version }} run: | $ErrorActionPreference = "Stop" + $ProgressPreference = 'SilentlyContinue' if ("${{ steps.cache-install.outputs.cache-hit }}" -ne 'true') { - Invoke-WebRequest -Uri "${{ matrix.install }}" -OutFile "install.exe" + for ($attempt = 1; $attempt -le 3; $attempt++) { + try { + Invoke-WebRequest -Uri "${{ matrix.install }}" -OutFile "install.exe" + break + } catch { + if ($attempt -eq 3) { throw } + Write-Host "CUDA installer download attempt $attempt failed: $($_.Exception.Message); retrying in 15s" + Start-Sleep -Seconds 15 + } + } $subpackages = @(${{ join(matrix.cuda-components, ', ') }}) | Foreach-Object {"${_}_${{ matrix.cuda-version }}"} Start-Process -FilePath .\install.exe -ArgumentList (@("-s") + $subpackages) -NoNewWindow -Wait } diff --git a/.github/workflows/test-llamacpp-update.yaml b/.github/workflows/test-llamacpp-update.yaml index 004791c0a..d37530466 100644 --- a/.github/workflows/test-llamacpp-update.yaml +++ b/.github/workflows/test-llamacpp-update.yaml @@ -381,8 +381,18 @@ jobs: name: Install CUDA ${{ matrix.cuda-version }} run: | $ErrorActionPreference = "Stop" + $ProgressPreference = 'SilentlyContinue' if ("${{ steps.cache-install.outputs.cache-hit }}" -ne 'true') { - Invoke-WebRequest -Uri "${{ matrix.install }}" -OutFile "install.exe" + for ($attempt = 1; $attempt -le 3; $attempt++) { + try { + Invoke-WebRequest -Uri "${{ matrix.install }}" -OutFile "install.exe" + break + } catch { + if ($attempt -eq 3) { throw } + Write-Host "CUDA installer download attempt $attempt failed: $($_.Exception.Message); retrying in 15s" + Start-Sleep -Seconds 15 + } + } $subpackages = @(${{ join(matrix.cuda-components, ', ') }}) | Foreach-Object {"${_}_${{ matrix.cuda-version }}"} Start-Process -FilePath .\install.exe -ArgumentList (@("-s") + $subpackages) -NoNewWindow -Wait } diff --git a/integration/reg_groups_test.go b/integration/reg_groups_test.go index ed5a9d8ff..9b86dde4f 100644 --- a/integration/reg_groups_test.go +++ b/integration/reg_groups_test.go @@ -84,6 +84,7 @@ func TestVision(t *testing.T) { "vision-detail", "vision-multi-image", "vision-description", + "vision-ocr-document", "vision-split-batch", "vision-text", ) diff --git a/integration/reg_release_test.go b/integration/reg_release_test.go index 5f2b2e46e..de8b5ea74 100644 --- a/integration/reg_release_test.go +++ b/integration/reg_release_test.go @@ -117,7 +117,6 @@ func init() { integrationTestCase("create-safetensors", "", runCreateSafetensorsLLM), integrationTestCase("create-gguf", "", runCreateGGUF), integrationTestCase("quantization", "qwen2.5:0.5b-instruct-fp16", runQuantization), - integrationTestCase("image-generation", "", runImageGeneration), ) // Model-parametric cases @@ -128,5 +127,6 @@ func init() { registerVisionTextCases(testModels(releaseVisionTextModels)) registerToolCases(testModels(releaseToolsModels)) registerToolStressCases(testModels(releaseToolsModels)) + registerVisionOCRDocumentCases(testModels(releaseVisionModels)) registerAudioTranscriptionCases(testModels(releaseAudioModels)) } diff --git a/integration/vision_ocr_test.go b/integration/vision_ocr_test.go new file mode 100644 index 000000000..62699edd4 --- /dev/null +++ b/integration/vision_ocr_test.go @@ -0,0 +1,166 @@ +//go:build integration + +package integration + +import ( + "bytes" + "context" + "fmt" + "image" + "image/color" + "image/draw" + "image/png" + "strings" + "testing" + "time" + + "github.com/ollama/ollama/api" + "golang.org/x/image/font" + "golang.org/x/image/font/gofont/goregular" + "golang.org/x/image/font/opentype" + "golang.org/x/image/math/fixed" +) + +const visionOCRReference = "HARBOR-DELTA-8061" + +func registerVisionOCRDocumentCases(models []string) { + registerModelIntegrationCases("vision-ocr-document", models, runVisionOCRDocument) +} + +func runVisionOCRDocument(t *testing.T, model string) { + t.Helper() + + skipUnderMinVRAM(t, 8) + skipKnownIntegrationFlake(t, "vision-ocr-document", model) + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Minute) + defer cancel() + client, _, cleanup := InitServerConnection(ctx, t) + defer cleanup() + + requireCapability(ctx, t, client, model, "vision") + pullOrSkip(ctx, t, client, model) + + document, err := visionOCRDocument() + if err != nil { + t.Fatal(err) + } + req := api.ChatRequest{ + Model: model, + Messages: []api.Message{ + { + Role: "user", + Content: "Scan the entire document and read the unique FINAL AUDIT CODE field. " + + "Reply with only its value.", + Images: []api.ImageData{document}, + }, + }, + Stream: &stream, + Options: map[string]any{ + "seed": 42, + "temperature": 0.0, + }, + KeepAlive: &api.Duration{Duration: 10 * time.Second}, + } + + preloadGenerateModel(ctx, t, client, api.GenerateRequest{Model: req.Model}) + skipIfNotGPULoaded(ctx, t, client, req.Model, 80) + + // The shared response normalization collapses whitespace runs but keeps + // single spaces, and quantized models occasionally split the code with a + // stray space ("HARBO R-..."). Accept any stream and compare with all + // whitespace removed instead — every glyph still has to be right. + msg := DoChat(ctx, t, client, req, []string{""}, 240*time.Second, 30*time.Second) + if msg == nil { + return + } + despace := func(s string) string { return strings.Join(strings.Fields(s), "") } + if !strings.Contains(despace(msg.Content), despace(visionOCRReference)) { + t.Fatalf("%s: audit code %q not found in %q", model, visionOCRReference, msg.Content) + } +} + +func visionOCRDocument() ([]byte, error) { + const ( + tileWidth = 800 + tileHeight = 800 + columns = 3 + rows = 4 + ) + + parsedFont, err := opentype.Parse(goregular.TTF) + if err != nil { + return nil, fmt.Errorf("parse document font: %w", err) + } + bodyFace, err := opentype.NewFace(parsedFont, &opentype.FaceOptions{ + Size: 40, + DPI: 72, + Hinting: font.HintingFull, + }) + if err != nil { + return nil, fmt.Errorf("create document font: %w", err) + } + referenceFace, err := opentype.NewFace(parsedFont, &opentype.FaceOptions{ + Size: 48, + DPI: 72, + Hinting: font.HintingFull, + }) + if err != nil { + return nil, fmt.Errorf("create reference font: %w", err) + } + + white := color.RGBA{R: 255, G: 255, B: 255, A: 255} + ink := image.NewUniform(color.RGBA{R: 24, G: 29, B: 36, A: 255}) + rule := color.RGBA{R: 180, G: 186, B: 194, A: 255} + tile := image.NewRGBA(image.Rect(0, 0, tileWidth, tileHeight)) + draw.Draw(tile, tile.Bounds(), image.NewUniform(white), image.Point{}, draw.Src) + drawOCRRule(tile, image.Rect(40, 40, tileWidth-40, 44), rule) + drawOCRRule(tile, image.Rect(40, 130, tileWidth-40, 134), rule) + drawOCRRule(tile, image.Rect(40, 520, tileWidth-40, 524), rule) + drawOCRText(tile, bodyFace, ink, 48, 110, "FIELD SERVICE RECORD") + drawOCRText(tile, bodyFace, ink, 48, 220, "Equipment: turbine assembly") + drawOCRText(tile, bodyFace, ink, 48, 290, "Inspection: pressure and seals") + drawOCRText(tile, bodyFace, ink, 48, 360, "Status: passed") + drawOCRText(tile, bodyFace, ink, 48, 470, "Technician notes: no defects") + + page := image.NewRGBA(image.Rect(0, 0, columns*tileWidth, rows*tileHeight)) + draw.Draw(page, page.Bounds(), image.NewUniform(white), image.Point{}, draw.Src) + for row := range rows { + for column := range columns { + cell := image.NewRGBA(tile.Bounds()) + draw.Draw(cell, cell.Bounds(), tile, image.Point{}, draw.Src) + drawOCRText(cell, bodyFace, ink, 48, 190, fmt.Sprintf("ROW %d COLUMN %d", row+1, column+1)) + + label := "REFERENCE:" + reference := fmt.Sprintf("CELL-R%d-C%d-%04d", row+1, column+1, (row+1)*100+column+1) + if row == rows-1 && column == columns-1 { + label = "FINAL AUDIT CODE:" + reference = visionOCRReference + } + drawOCRText(cell, referenceFace, ink, 48, 640, label) + drawOCRText(cell, referenceFace, ink, 48, 710, reference) + + origin := image.Pt(column*tileWidth, row*tileHeight) + draw.Draw(page, image.Rectangle{Min: origin, Max: origin.Add(cell.Bounds().Size())}, cell, image.Point{}, draw.Src) + } + } + + var encoded bytes.Buffer + if err := png.Encode(&encoded, page); err != nil { + return nil, fmt.Errorf("encode OCR document: %w", err) + } + return encoded.Bytes(), nil +} + +func drawOCRText(dst draw.Image, face font.Face, source image.Image, x, y int, text string) { + drawer := font.Drawer{ + Dst: dst, + Src: source, + Face: face, + Dot: fixed.P(x, y), + } + drawer.DrawString(text) +} + +func drawOCRRule(dst draw.Image, bounds image.Rectangle, fill color.Color) { + draw.Draw(dst, bounds, image.NewUniform(fill), image.Point{}, draw.Src) +} diff --git a/llm/llama_server.go b/llm/llama_server.go index c99c29556..1b13a5b0d 100644 --- a/llm/llama_server.go +++ b/llm/llama_server.go @@ -169,6 +169,7 @@ type llamaServerRunner struct { type llamaServerLaunchConfig struct { modelPath string modelArch string + draftType string projectors []string mmprojMemory uint64 modelLayers uint64 @@ -376,7 +377,7 @@ func startLlamaServer(launch llamaServerLaunchConfig, out io.Writer) (cmd *exec. params = appendJinjaArgs(params, launch.config) params = appendMMProjArgs(params, launch) - params = appendMTPDraftArgs(params, launch.config, launch.opts) + params = appendDraftArgs(params, launch.draftType, launch.config.DraftModelPath, launch.opts) params = append(params, qwenVLServerArgs(launch.modelArch)...) @@ -796,23 +797,41 @@ func appendContextShiftArgs(params []string, opts api.Options, enabled bool) []s return params } -func appendMTPDraftArgs(params []string, config LlamaServerConfig, opts api.Options) []string { - if !config.EnableMTP && config.DraftModelPath == "" { +const ( + draftTypeMTP = "draft-mtp" + draftTypeDFlash = "draft-dflash" +) + +func appendDraftArgs(params []string, draftType, draftModelPath string, opts api.Options) []string { + if draftType == "" { return params } if opts.DraftNumPredict <= 0 { return params } - params = append(params, "--spec-type", "draft-mtp") + params = append(params, "--spec-type", draftType) params = append(params, "--spec-draft-n-max", strconv.Itoa(opts.DraftNumPredict)) - params = append(params, "--spec-draft-backend-sampling") - if config.DraftModelPath != "" { - params = append(params, "--spec-draft-model", config.DraftModelPath) + if draftType == draftTypeMTP { + params = append(params, "--spec-draft-backend-sampling") + } + if draftModelPath != "" { + params = append(params, "--spec-draft-model", draftModelPath) } return params } +func externalDraftType(path string) (string, error) { + f, err := LoadModel(path, 1) + if err != nil { + return "", fmt.Errorf("load draft model metadata: %w", err) + } + if f.KV().Architecture() == "dflash" { + return draftTypeDFlash, nil + } + return draftTypeMTP, nil +} + func hasMTPDraft(f *ggml.GGML) bool { if f.KV().Uint("nextn_predict_layers") > 0 { return true @@ -881,6 +900,17 @@ func NewLlamaServerRunner( config.EnableMTP = true } + draftType := "" + if config.EnableMTP { + draftType = draftTypeMTP + } + if config.DraftModelPath != "" { + draftType, err = externalDraftType(config.DraftModelPath) + if err != nil { + return nil, err + } + } + gpuLibs := ml.LibraryPaths(gpus) status := NewStatusWriter(os.Stderr) @@ -898,6 +928,7 @@ func NewLlamaServerRunner( launch := llamaServerLaunchConfig{ modelPath: modelPath, modelArch: arch, + draftType: draftType, projectors: slices.Clone(projectors), mmprojMemory: mmprojMemory, modelLayers: f.KV().BlockCount() + 1, diff --git a/llm/llama_server_test.go b/llm/llama_server_test.go index 744b4322e..127c11f2d 100644 --- a/llm/llama_server_test.go +++ b/llm/llama_server_test.go @@ -2469,12 +2469,13 @@ func TestAppendContextShiftArgs(t *testing.T) { } } -func TestAppendMTPDraftArgs(t *testing.T) { +func TestAppendDraftArgs(t *testing.T) { tests := []struct { - name string - config LlamaServerConfig - opts api.Options - want []string + name string + draftType string + draftPath string + opts api.Options + want []string }{ { name: "no draft model leaves speculative decoding disabled", @@ -2482,30 +2483,63 @@ func TestAppendMTPDraftArgs(t *testing.T) { want: []string{"base"}, }, { - name: "embedded draft uses configured draft depth", - config: LlamaServerConfig{EnableMTP: true}, - opts: api.Options{Runner: api.Runner{DraftNumPredict: 4}}, - want: []string{"base", "--spec-type", "draft-mtp", "--spec-draft-n-max", "4", "--spec-draft-backend-sampling"}, + name: "embedded MTP draft uses configured draft depth", + draftType: draftTypeMTP, + opts: api.Options{Runner: api.Runner{DraftNumPredict: 4}}, + want: []string{"base", "--spec-type", "draft-mtp", "--spec-draft-n-max", "4", "--spec-draft-backend-sampling"}, }, { - name: "separate draft model uses configured draft depth", - config: LlamaServerConfig{DraftModelPath: "draft.gguf"}, - opts: api.Options{Runner: api.Runner{DraftNumPredict: 8}}, - want: []string{"base", "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-draft-backend-sampling", "--spec-draft-model", "draft.gguf"}, + name: "separate MTP draft uses configured draft depth", + draftType: draftTypeMTP, + draftPath: "draft.gguf", + opts: api.Options{Runner: api.Runner{DraftNumPredict: 8}}, + want: []string{"base", "--spec-type", "draft-mtp", "--spec-draft-n-max", "8", "--spec-draft-backend-sampling", "--spec-draft-model", "draft.gguf"}, }, { - name: "zero draft depth disables speculative decoding", - config: LlamaServerConfig{EnableMTP: true, DraftModelPath: "draft.gguf"}, - opts: api.Options{Runner: api.Runner{DraftNumPredict: 0}}, - want: []string{"base"}, + name: "DFlash draft omits MTP backend sampling", + draftType: draftTypeDFlash, + draftPath: "draft.gguf", + opts: api.Options{Runner: api.Runner{DraftNumPredict: 4}}, + want: []string{"base", "--spec-type", "draft-dflash", "--spec-draft-n-max", "4", "--spec-draft-model", "draft.gguf"}, + }, + { + name: "zero draft depth disables speculative decoding", + draftType: draftTypeDFlash, + draftPath: "draft.gguf", + opts: api.Options{Runner: api.Runner{DraftNumPredict: 0}}, + want: []string{"base"}, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - got := appendMTPDraftArgs([]string{"base"}, tt.config, tt.opts) + got := appendDraftArgs([]string{"base"}, tt.draftType, tt.draftPath, tt.opts) if !slices.Equal(got, tt.want) { - t.Fatalf("appendMTPDraftArgs = %v, want %v", got, tt.want) + t.Fatalf("appendDraftArgs = %v, want %v", got, tt.want) + } + }) + } +} + +func TestExternalDraftType(t *testing.T) { + tests := []struct { + architecture string + want string + }{ + {architecture: "dflash", want: draftTypeDFlash}, + {architecture: "qwen35", want: draftTypeMTP}, + {architecture: "unknown", want: draftTypeMTP}, + } + + for _, tt := range tests { + t.Run(tt.architecture, func(t *testing.T) { + path, _ := writeTestGGML(t, ggml.KV{"general.architecture": tt.architecture}, nil) + got, err := externalDraftType(path) + if err != nil { + t.Fatal(err) + } + if got != tt.want { + t.Fatalf("externalDraftType = %q, want %q", got, tt.want) } }) } diff --git a/model/parsers/glimmer.go b/model/parsers/glimmer.go new file mode 100644 index 000000000..06d80549d --- /dev/null +++ b/model/parsers/glimmer.go @@ -0,0 +1,516 @@ +package parsers + +import ( + "fmt" + "log/slog" + "strings" + + "github.com/ollama/ollama/api" +) + +const ( + glimmerStartTag = "<|start|>" + glimmerMessageTag = "<|message|>" + glimmerEndMessageTag = "<|eom|>" + glimmerEndTurnTag = "<|eot|>" + + glimmerAssistantHeaderPrefix = glimmerStartTag + "assistant" + + glimmerATEMCallsOpen = "" + glimmerATEMCallsClose = "" + glimmerATEMInvokeOpen = `" + glimmerATEMParamOpen = `" +) + +type glimmerParserState int + +const ( + glimmerParserHeader glimmerParserState = iota + glimmerParserContent + glimmerParserThinking + glimmerParserTool +) + +type GlimmerParser struct { + state glimmerParserState + buffer strings.Builder + recipient string + tools map[string]api.Tool + callIndex int + emitThinking bool + + // contentStreamed records whether any of the current message's body has + // already been emitted as content. The fumbled-recipient tool-call + // fallback only applies while nothing has streamed, so a message either + // becomes a tool call or streams as content — never half of each. + contentStreamed bool +} + +func (p *GlimmerParser) HasToolSupport() bool { return true } +func (p *GlimmerParser) HasThinkingSupport() bool { return true } + +func (p *GlimmerParser) PreservedTokens() []string { + return []string{glimmerStartTag, glimmerMessageTag, glimmerEndMessageTag, glimmerEndTurnTag} +} + +func (p *GlimmerParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool { + p.state = glimmerParserHeader + p.buffer.Reset() + p.recipient = "" + p.tools = glimmerToolsByName(tools) + p.callIndex = 0 + p.emitThinking = thinkValue == nil || thinkValue.Bool() + p.contentStreamed = false + return tools +} + +func (p *GlimmerParser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) { + p.buffer.WriteString(s) + var contentSB, thinkingSB strings.Builder + + for { + if p.state == glimmerParserHeader { + progress, fallback := p.consumeHeader(done) + contentSB.WriteString(fallback) + if !progress { + break + } + continue + } + + progress, body, complete := p.consumeBody(done) + switch p.state { + case glimmerParserContent: + if call, ok := p.contentToolCallFallback(body, complete); ok { + calls = append(calls, call) + } else if body != "" { + contentSB.WriteString(body) + p.contentStreamed = true + } + case glimmerParserThinking: + if p.emitThinking { + thinkingSB.WriteString(body) + } + case glimmerParserTool: + if complete { + call, parseErr := p.parseToolCall(body) + if parseErr != nil { + return contentSB.String(), thinkingSB.String(), calls, parseErr + } + calls = append(calls, call) + } + } + if complete { + p.state = glimmerParserHeader + p.recipient = "" + p.contentStreamed = false + } + if !progress { + break + } + } + + return contentSB.String(), thinkingSB.String(), calls, nil +} + +func (p *GlimmerParser) consumeHeader(done bool) (progress bool, fallback string) { + acc := p.buffer.String() + if strings.HasPrefix(acc, glimmerAssistantHeaderPrefix) { + acc = acc[len(glimmerAssistantHeaderPrefix):] + p.buffer.Reset() + p.buffer.WriteString(acc) + } else if strings.HasPrefix(glimmerAssistantHeaderPrefix, acc) && strings.HasPrefix(acc, "<") && !done { + return false, "" + } + + idx := strings.Index(acc, glimmerMessageTag) + if idx < 0 { + if !done { + return false, "" + } + p.buffer.Reset() + return acc != "", acc + } + + header := strings.TrimSpace(acc[:idx]) + p.buffer.Reset() + p.buffer.WriteString(acc[idx+len(glimmerMessageTag):]) + p.recipient = strings.TrimSpace(strings.TrimPrefix(header, "to=")) + p.setStateForRecipient() + return true, "" +} + +func (p *GlimmerParser) setStateForRecipient() { + _, isTool := glimmerResolveToolName(p.tools, p.recipient) + switch { + case p.recipient == "self": + p.state = glimmerParserThinking + case isTool: + p.state = glimmerParserTool + default: + p.state = glimmerParserContent + } +} + +func (p *GlimmerParser) consumeBody(done bool) (progress bool, body string, complete bool) { + acc := p.buffer.String() + hold := p.state == glimmerParserTool || p.holdContent(acc) + controlStart := 0 + if hold { + if close := strings.Index(acc, glimmerATEMCallsClose); close >= 0 { + controlStart = close + len(glimmerATEMCallsClose) + } else { + controlStart = len(acc) + } + } + + idx, markerLen := glimmerBodyTerminator(acc[controlStart:]) + if idx >= 0 { + idx += controlStart + } + implicitIdx, waitForImplicit := p.implicitHeaderStart(acc, controlStart, done) + if implicitIdx >= 0 && !waitForImplicit && (idx < 0 || implicitIdx < idx) { + body = acc[:implicitIdx] + p.buffer.Reset() + p.buffer.WriteString(acc[implicitIdx:]) + return true, body, true + } + if idx >= 0 { + body = acc[:idx] + p.buffer.Reset() + p.buffer.WriteString(acc[idx+markerLen:]) + return true, body, true + } + if done { + p.buffer.Reset() + return acc != "" || p.state == glimmerParserTool, acc, true + } + if hold { + return false, "", false + } + if waitForImplicit { + body = acc[:implicitIdx] + p.buffer.Reset() + p.buffer.WriteString(acc[implicitIdx:]) + return body != "", body, false + } + + keep := 0 + for _, marker := range []string{glimmerEndMessageTag, glimmerEndTurnTag} { + keep = max(keep, overlap(acc, marker)) + } + if len(acc) == keep { + return false, "", false + } + body = acc[:len(acc)-keep] + p.buffer.Reset() + p.buffer.WriteString(acc[len(acc)-keep:]) + return body != "", body, false +} + +func (p *GlimmerParser) implicitHeaderStart(s string, search int, done bool) (idx int, wait bool) { + for { + i := strings.Index(s[search:], glimmerStartTag) + if i < 0 { + break + } + i += search + + switch p.assistantHeaderStatus(s[i:], done) { + case glimmerHeaderValid: + return i, false + case glimmerHeaderPending: + return i, true + } + + search = i + len(glimmerStartTag) + } + + if done { + return -1, false + } + if keep := overlap(s[search:], glimmerAssistantHeaderPrefix); keep > 0 { + return len(s) - keep, true + } + return -1, false +} + +type glimmerHeaderStatus int + +const ( + glimmerHeaderInvalid glimmerHeaderStatus = iota + glimmerHeaderPending + glimmerHeaderValid +) + +func (p *GlimmerParser) assistantHeaderStatus(s string, done bool) glimmerHeaderStatus { + if len(s) < len(glimmerAssistantHeaderPrefix) { + if !done && strings.HasPrefix(glimmerAssistantHeaderPrefix, s) { + return glimmerHeaderPending + } + return glimmerHeaderInvalid + } + if !strings.HasPrefix(s, glimmerAssistantHeaderPrefix) { + return glimmerHeaderInvalid + } + + rest := s[len(glimmerAssistantHeaderPrefix):] + if strings.HasPrefix(rest, glimmerMessageTag) { + return glimmerHeaderValid + } + if !done && strings.HasPrefix(glimmerMessageTag, rest) { + return glimmerHeaderPending + } + if rest == "" { + if done { + return glimmerHeaderInvalid + } + return glimmerHeaderPending + } + + const recipientPrefix = " to=" + if len(rest) < len(recipientPrefix) { + if !done && strings.HasPrefix(recipientPrefix, rest) { + return glimmerHeaderPending + } + return glimmerHeaderInvalid + } + if !strings.HasPrefix(rest, recipientPrefix) { + return glimmerHeaderInvalid + } + + recipientAndAfter := rest[len(recipientPrefix):] + messageIdx := strings.Index(recipientAndAfter, glimmerMessageTag) + if messageIdx >= 0 { + if p.validAssistantRecipient(recipientAndAfter[:messageIdx]) { + return glimmerHeaderValid + } + return glimmerHeaderInvalid + } + if done { + return glimmerHeaderInvalid + } + if recipientAndAfter == "" { + return glimmerHeaderPending + } + if strings.ContainsAny(recipientAndAfter, " \t\r\n") { + return glimmerHeaderInvalid + } + if lt := strings.IndexByte(recipientAndAfter, '<'); lt >= 0 { + if strings.HasPrefix(glimmerMessageTag, recipientAndAfter[lt:]) { + return glimmerHeaderPending + } + return glimmerHeaderInvalid + } + return glimmerHeaderPending +} + +func (p *GlimmerParser) validAssistantRecipient(recipient string) bool { + return recipient == "self" || recipient == "user" || p.tools[recipient].Function.Name != "" +} + +func glimmerBodyTerminator(s string) (int, int) { + idx, markerLen := -1, 0 + for _, marker := range []string{glimmerEndMessageTag, glimmerEndTurnTag} { + if i := strings.Index(s, marker); i >= 0 && (idx < 0 || i < idx) { + idx = i + markerLen = len(marker) + } + } + return idx, markerLen +} + +// holdContent reports whether the content accumulated so far may be a fumbled +// tool call: a message that opens with the ATEM wrapper despite a non-tool +// recipient (the model sometimes omits the recipient, addresses the namespace +// or the user, or names the wrapper element itself). Such content is withheld +// from streaming until the message completes, so contentToolCallFallback can +// decide whether it is a tool call or ordinary content. Anything else streams +// immediately. +func (p *GlimmerParser) holdContent(acc string) bool { + if p.state != glimmerParserContent || p.contentStreamed || len(p.tools) == 0 { + return false + } + trimmed := strings.TrimLeft(acc, " \t\r\n") + return strings.HasPrefix(glimmerATEMCallsOpen, trimmed) || strings.HasPrefix(trimmed, glimmerATEMCallsOpen) +} + +// contentToolCallFallback recovers a tool call from a content-position message +// whose complete body is solely a well-formed ATEM block invoking a declared +// tool. The recipient header is authoritative when it names a tool; this +// fallback exists for the fumbled headers holdContent describes, where +// streaming the raw XML at the client is never the right answer. The function +// name is read from the invoke element, not the header. +func (p *GlimmerParser) contentToolCallFallback(body string, complete bool) (api.ToolCall, bool) { + if !complete || p.contentStreamed || len(p.tools) == 0 { + return api.ToolCall{}, false + } + trimmed := strings.TrimSpace(body) + if !strings.HasPrefix(trimmed, glimmerATEMCallsOpen) || !strings.HasSuffix(trimmed, glimmerATEMCallsClose) { + return api.ToolCall{}, false + } + resolved, ok := glimmerResolveToolName(p.tools, glimmerATEMInvokeName(trimmed)) + if !ok { + return api.ToolCall{}, false + } + tool := p.tools[resolved] + if _, args, err := parseGlimmerATEM(trimmed, tool); err == nil { + return p.newToolCall(resolved, args), true + } + return api.ToolCall{}, false +} + +// glimmerATEMInvokeName extracts the invoke element's function name from a body +// already known to carry the ATEM wrapper, or "" when the invoke is malformed. +func glimmerATEMInvokeName(body string) string { + inner := strings.TrimSpace(body[len(glimmerATEMCallsOpen) : len(body)-len(glimmerATEMCallsClose)]) + if !strings.HasPrefix(inner, glimmerATEMInvokeOpen) { + return "" + } + name, _, ok := strings.Cut(inner[len(glimmerATEMInvokeOpen):], `">`) + if !ok { + return "" + } + return glimmerTrimStrayMessageTag(name) +} + +// glimmerTrimStrayMessageTag drops <|message|> boundary tokens the model +// occasionally fumbles into the invoke name (e.g. `name="read<|message|>">`), +// echoing the header form `to=read<|message|>`. Function names are +// identifiers, so the tag is never legitimate there; parameter values are +// left untouched (control-token text is preserved in values). +func glimmerTrimStrayMessageTag(name string) string { + if !strings.Contains(name, glimmerMessageTag) { + return name + } + slog.Warn("glimmer parser recovered stray message boundary token in ATEM invoke name", "name", name) + return strings.ReplaceAll(name, glimmerMessageTag, "") +} + +func (p *GlimmerParser) parseToolCall(body string) (api.ToolCall, error) { + if p.recipient == "" { + return api.ToolCall{}, fmt.Errorf("empty Glimmer function name") + } + + recipient, ok := glimmerResolveToolName(p.tools, p.recipient) + if !ok { + return api.ToolCall{}, fmt.Errorf("undeclared Glimmer function %q", p.recipient) + } + tool := p.tools[recipient] + + name, args, err := parseGlimmerATEM(body, tool) + if err != nil { + return api.ToolCall{}, fmt.Errorf("parse Glimmer call to %s: %w", p.recipient, err) + } + if resolved, ok := glimmerResolveToolName(p.tools, name); !ok || resolved != recipient { + return api.ToolCall{}, fmt.Errorf("Glimmer recipient %q does not match ATEM invoke %q", p.recipient, name) + } + return p.newToolCall(recipient, args), nil +} + +// glimmerResolveToolName resolves a recipient or ATEM invoke name to a declared +// tool name. Exact matches win. The chat template derives a namespace from +// the first dot-component of each declared tool name and advertises +// recipients ".*" to the model, so an undotted tool `read` is +// legitimately addressable as `read.read`; resolve `ns.fn` to the declared +// `fn` when ns is fn's own derived namespace. +func glimmerResolveToolName(tools map[string]api.Tool, name string) (string, bool) { + if tools[name].Function.Name != "" { + return name, true + } + if ns, fn, ok := strings.Cut(name, "."); ok { + if tools[fn].Function.Name != "" { + fnNS, _, _ := strings.Cut(fn, ".") + if ns == fnNS { + return fn, true + } + } + } + return name, false +} + +func glimmerToolsByName(tools []api.Tool) map[string]api.Tool { + if len(tools) == 0 { + return nil + } + + byName := make(map[string]api.Tool, len(tools)) + for _, tool := range tools { + name := strings.TrimSpace(tool.Function.Name) + if name == "" { + continue + } + byName[name] = tool + } + return byName +} + +func parseGlimmerATEM(body string, tool api.Tool) (string, api.ToolCallFunctionArguments, error) { + body = strings.TrimSpace(body) + if !strings.HasPrefix(body, glimmerATEMCallsOpen) || !strings.HasSuffix(body, glimmerATEMCallsClose) { + return "", api.ToolCallFunctionArguments{}, fmt.Errorf("missing ATEM function_calls wrapper") + } + + invoke := strings.TrimSpace(body[len(glimmerATEMCallsOpen) : len(body)-len(glimmerATEMCallsClose)]) + if !strings.HasPrefix(invoke, glimmerATEMInvokeOpen) || !strings.HasSuffix(invoke, glimmerATEMInvokeClose) { + return "", api.ToolCallFunctionArguments{}, fmt.Errorf("missing ATEM invoke wrapper") + } + invoke = invoke[len(glimmerATEMInvokeOpen):] + nameEnd := strings.Index(invoke, `">`) + if nameEnd < 0 { + return "", api.ToolCallFunctionArguments{}, fmt.Errorf("malformed ATEM invoke name") + } + name := glimmerTrimStrayMessageTag(invoke[:nameEnd]) + params := invoke[nameEnd+2 : len(invoke)-len(glimmerATEMInvokeClose)] + params = strings.TrimPrefix(params, "\n") + params = strings.TrimSuffix(params, "\n") + + args := api.NewToolCallFunctionArguments() + for params != "" { + if !strings.HasPrefix(params, glimmerATEMParamOpen) { + return "", api.ToolCallFunctionArguments{}, fmt.Errorf("malformed ATEM parameter") + } + params = params[len(glimmerATEMParamOpen):] + paramNameEnd := strings.Index(params, `">`) + if paramNameEnd < 0 { + return "", api.ToolCallFunctionArguments{}, fmt.Errorf("malformed ATEM parameter name") + } + paramName := params[:paramNameEnd] + params = params[paramNameEnd+2:] + valueEnd := strings.Index(params, glimmerATEMParamClose) + if valueEnd < 0 { + return "", api.ToolCallFunctionArguments{}, fmt.Errorf("unterminated ATEM parameter %q", paramName) + } + value := params[:valueEnd] + params = params[valueEnd+len(glimmerATEMParamClose):] + params = strings.TrimPrefix(params, "\n") + + var paramType api.PropertyType + if tool.Function.Parameters.Properties != nil { + if property, ok := tool.Function.Parameters.Properties.Get(paramName); ok { + if len(property.AnyOf) > 0 { + for _, option := range property.AnyOf { + paramType = append(paramType, option.Type...) + } + } else { + paramType = property.Type + } + } + } + args.Set(paramName, parseTypedToolValue(value, paramType)) + } + + return name, args, nil +} + +func (p *GlimmerParser) newToolCall(name string, args api.ToolCallFunctionArguments) api.ToolCall { + call := api.ToolCall{Function: api.ToolCallFunction{ + Name: name, + Arguments: args, + Index: p.callIndex, + }} + p.callIndex++ + return call +} diff --git a/model/parsers/glimmer_test.go b/model/parsers/glimmer_test.go new file mode 100644 index 000000000..40bcb3205 --- /dev/null +++ b/model/parsers/glimmer_test.go @@ -0,0 +1,605 @@ +package parsers + +import ( + "reflect" + "strings" + "testing" + + "github.com/ollama/ollama/api" +) + +func glimmerTestTool(name string, properties map[string]api.ToolProperty) api.Tool { + return api.Tool{Type: "function", Function: api.ToolFunction{ + Name: name, + Parameters: api.ToolFunctionParameters{ + Type: "object", + Properties: testPropsMap(properties), + }, + }} +} + +func glimmerTestATEM(name, parameters string) string { + return ` + +` + parameters + ` +` +} + +func TestGlimmerParserFinalAnswer(t *testing.T) { + p := &GlimmerParser{} + p.Init(nil, nil, nil) + content, thinking, calls, err := p.Add(` to=user<|message|>Hello`, true) + if err != nil { + t.Fatal(err) + } + if content != "Hello" || thinking != "" || len(calls) != 0 { + t.Fatalf("got content=%q thinking=%q calls=%v", content, thinking, calls) + } +} + +func TestGlimmerParserSelfRecipientIsThinking(t *testing.T) { + p := &GlimmerParser{} + p.Init(nil, nil, nil) + content, thinking, calls, err := p.Add(` to=self<|message|>Check the facts.`, true) + if err != nil { + t.Fatal(err) + } + if content != "" || thinking != "Check the facts." || len(calls) != 0 { + t.Fatalf("got content=%q thinking=%q calls=%v", content, thinking, calls) + } +} + +func TestGlimmerParserSuppressesThinking(t *testing.T) { + p := &GlimmerParser{} + p.Init(nil, nil, &api.ThinkValue{Value: false}) + content, thinking, calls, err := p.Add( + ` to=self<|message|>Check the facts.<|eom|><|start|>assistant to=user<|message|>Answer<|eot|>`, + true, + ) + if err != nil { + t.Fatal(err) + } + if content != "Answer" || thinking != "" || len(calls) != 0 { + t.Fatalf("got content=%q thinking=%q calls=%v", content, thinking, calls) + } +} + +func TestGlimmerParserStreamingATEMAtEveryBoundary(t *testing.T) { + tool := glimmerTestTool("get_weather", map[string]api.ToolProperty{ + "city": {Type: api.PropertyType{"string"}}, + }) + input := ` to=get_weather<|message|>` + glimmerTestATEM( + "get_weather", + `SF +`, + ) + `<|eot|>` + + for split := 0; split <= len(input); split++ { + p := &GlimmerParser{} + p.Init([]api.Tool{tool}, nil, nil) + + var content, thinking string + var calls []api.ToolCall + for i, chunk := range []string{input[:split], input[split:]} { + gotContent, gotThinking, gotCalls, err := p.Add(chunk, i == 1) + if err != nil { + t.Fatalf("split %d: %v", split, err) + } + content += gotContent + thinking += gotThinking + calls = append(calls, gotCalls...) + } + + if content != "" || thinking != "" || len(calls) != 1 { + t.Fatalf("split %d: content=%q thinking=%q calls=%v", split, content, thinking, calls) + } + if calls[0].Function.Name != "get_weather" || calls[0].Function.Index != 0 { + t.Fatalf("split %d: unexpected call: %#v", split, calls[0]) + } + if city, ok := calls[0].Function.Arguments.Get("city"); !ok || city != "SF" { + t.Fatalf("split %d: city = %#v, %v; want SF", split, city, ok) + } + } +} + +func TestGlimmerParserATEMValueTypesAndWhitespace(t *testing.T) { + tool := glimmerTestTool("tools.run", map[string]api.ToolProperty{ + "text": {Type: api.PropertyType{"string"}}, + "count": {Type: api.PropertyType{"integer"}}, + "enabled": {Type: api.PropertyType{"boolean"}}, + "items": {Type: api.PropertyType{"array"}}, + "config": {Type: api.PropertyType{"object"}}, + "choice": {AnyOf: []api.ToolProperty{ + {Type: api.PropertyType{"null"}}, + {Type: api.PropertyType{"string"}}, + }}, + }) + input := ` to=tools.run<|message|>` + glimmerTestATEM( + "tools.run", + ` keep spaces +3 +true +["one", "two"] +{"mode":"fast"} +null +`, + ) + + p := &GlimmerParser{} + p.Init([]api.Tool{tool}, nil, nil) + content, thinking, calls, err := p.Add(input, true) + if err != nil { + t.Fatal(err) + } + if content != "" || thinking != "" || len(calls) != 1 { + t.Fatalf("got content=%q thinking=%q calls=%v", content, thinking, calls) + } + + want := map[string]any{ + "text": " keep spaces ", + "count": 3, + "enabled": true, + "items": []any{"one", "two"}, + "config": map[string]any{"mode": "fast"}, + "choice": nil, + } + for name, expected := range want { + got, ok := calls[0].Function.Arguments.Get(name) + if !ok || !reflect.DeepEqual(got, expected) { + t.Errorf("%s = %#v, %v; want %#v", name, got, ok, expected) + } + } +} + +func TestGlimmerParserATEMValueMayContainControlTokenText(t *testing.T) { + tool := glimmerTestTool("echo", map[string]api.ToolProperty{ + "text": {Type: api.PropertyType{"string"}}, + }) + value := `literal <|eot|> and <|start|>assistant to=user<|message|> text` + input := ` to=echo<|message|>` + glimmerTestATEM( + "echo", + ``+value+` +`, + ) + `<|eot|>` + + p := &GlimmerParser{} + p.Init([]api.Tool{tool}, nil, nil) + content, thinking, calls, err := p.Add(input, true) + if err != nil { + t.Fatal(err) + } + if content != "" || thinking != "" || len(calls) != 1 { + t.Fatalf("got content=%q thinking=%q calls=%v", content, thinking, calls) + } + if got, ok := calls[0].Function.Arguments.Get("text"); !ok || got != value { + t.Fatalf("text = %#v, %v; want %q", got, ok, value) + } +} + +func TestGlimmerParserMultipleMessagesAndIndices(t *testing.T) { + first := glimmerTestTool("first", map[string]api.ToolProperty{ + "x": {Type: api.PropertyType{"integer"}}, + }) + second := glimmerTestTool("second", nil) + input := " to=self<|message|>Think<|eom|>" + + "<|start|>assistant to=first<|message|>" + + glimmerTestATEM("first", `1 +`) + "<|eom|>" + + "<|start|>assistant to=second<|message|>" + + glimmerTestATEM("second", "") + "<|eot|>" + + for split := 0; split <= len(input); split++ { + p := &GlimmerParser{} + p.Init([]api.Tool{first, second}, nil, nil) + + var content, thinking string + var calls []api.ToolCall + for i, chunk := range []string{input[:split], input[split:]} { + gotContent, gotThinking, gotCalls, err := p.Add(chunk, i == 1) + if err != nil { + t.Fatalf("split %d: %v", split, err) + } + content += gotContent + thinking += gotThinking + calls = append(calls, gotCalls...) + } + if content != "" || thinking != "Think" || len(calls) != 2 { + t.Fatalf("split %d: content=%q thinking=%q calls=%v", split, content, thinking, calls) + } + if calls[0].Function.Name != "first" || calls[0].Function.Index != 0 { + t.Fatalf("split %d: first call = %#v", split, calls[0]) + } + if calls[1].Function.Name != "second" || calls[1].Function.Index != 1 { + t.Fatalf("split %d: second call = %#v", split, calls[1]) + } + } +} + +// TestGlimmerParserStrayMessageTagInInvokeName covers the observed stress +// failure where the model fumbles a <|message|> boundary token into the +// invoke name (`name="read<|message|>">`), echoing the header form +// `to=read<|message|>`. The call must parse as the intended tool. +func TestGlimmerParserStrayMessageTagInInvokeName(t *testing.T) { + tool := glimmerTestTool("read", map[string]api.ToolProperty{ + "path": {Type: api.PropertyType{"string"}}, + }) + input := ` to=read<|message|>` + glimmerTestATEM( + "read<|message|>", + `go.mod +`, + ) + `<|eot|>` + + p := &GlimmerParser{} + p.Init([]api.Tool{tool}, nil, nil) + content, thinking, calls, err := p.Add(input, true) + if err != nil { + t.Fatal(err) + } + if content != "" || thinking != "" || len(calls) != 1 { + t.Fatalf("got content=%q thinking=%q calls=%v", content, thinking, calls) + } + if calls[0].Function.Name != "read" { + t.Fatalf("call name = %q, want %q", calls[0].Function.Name, "read") + } + if got, ok := calls[0].Function.Arguments.Get("path"); !ok || got != "go.mod" { + t.Fatalf("path = %#v, %v; want %q", got, ok, "go.mod") + } +} + +// TestGlimmerParserNamespacedSelfReference covers the observed stress failure +// where the model addresses an undotted tool through its own derived +// namespace (`read.read` for a tool declared as `read`) — the chat template +// advertises recipients ".*" with ns = first dot-component of each +// declared name, so this form is in-protocol and must resolve. +func TestGlimmerParserNamespacedSelfReference(t *testing.T) { + tool := glimmerTestTool("read", map[string]api.ToolProperty{ + "path": {Type: api.PropertyType{"string"}}, + }) + params := `go.mod +` + for _, tc := range []struct{ recipient, invoke string }{ + {"read", "read.read"}, + {"read.read", "read.read"}, + {"read.read", "read"}, + } { + p := &GlimmerParser{} + p.Init([]api.Tool{tool}, nil, nil) + input := ` to=` + tc.recipient + `<|message|>` + glimmerTestATEM(tc.invoke, params) + `<|eot|>` + content, thinking, calls, err := p.Add(input, true) + if err != nil { + t.Fatalf("recipient=%q invoke=%q: %v", tc.recipient, tc.invoke, err) + } + if content != "" || thinking != "" || len(calls) != 1 { + t.Fatalf("recipient=%q invoke=%q: content=%q thinking=%q calls=%v", tc.recipient, tc.invoke, content, thinking, calls) + } + if calls[0].Function.Name != "read" { + t.Fatalf("recipient=%q invoke=%q: call name = %q, want %q", tc.recipient, tc.invoke, calls[0].Function.Name, "read") + } + } +} + +// TestGlimmerParserNamespaceResolutionStaysStrict pins the limits of namespace +// resolution: undeclared namespaces do not resolve, and dotted declared names +// match exactly. +func TestGlimmerParserNamespaceResolutionStaysStrict(t *testing.T) { + read := glimmerTestTool("read", nil) + repoRead := glimmerTestTool("repo.read", nil) + tools := map[string]api.Tool{"read": read, "repo.read": repoRead} + + for _, tc := range []struct { + name, want string + ok bool + }{ + {"read", "read", true}, + {"read.read", "read", true}, + {"repo.read", "repo.read", true}, + {"functions.read", "functions.read", false}, + {"write.write", "write.write", false}, + // double-prefix through the declared name's own namespace still + // resolves: ns "repo" + declared "repo.read" + {"repo.repo.read", "repo.read", true}, + } { + got, ok := glimmerResolveToolName(tools, tc.name) + if got != tc.want || ok != tc.ok { + t.Errorf("glimmerResolveToolName(%q) = %q, %v; want %q, %v", tc.name, got, ok, tc.want, tc.ok) + } + } +} + +func TestGlimmerParserRejectsMismatchedATEMInvoke(t *testing.T) { + tool := glimmerTestTool("get_weather", nil) + p := &GlimmerParser{} + p.Init([]api.Tool{tool}, nil, nil) + _, _, _, err := p.Add( + ` to=get_weather<|message|>`+glimmerTestATEM("read_file", ""), + true, + ) + if err == nil || !strings.Contains(err.Error(), "does not match") { + t.Fatalf("error = %v, want recipient mismatch", err) + } +} + +func TestGlimmerParserRejectsMalformedATEM(t *testing.T) { + tool := glimmerTestTool("get_weather", nil) + p := &GlimmerParser{} + p.Init([]api.Tool{tool}, nil, nil) + _, _, _, err := p.Add(` to=get_weather<|message|>{"city":"SF"}`, true) + if err == nil || !strings.Contains(err.Error(), "function_calls wrapper") { + t.Fatalf("error = %v, want missing wrapper", err) + } +} + +// TestGlimmerParserFumbledRecipientToolCallFallback covers the observed failure +// mode where the model emits a well-formed ATEM block for a declared tool but +// fumbles the recipient header — omitting it, addressing the namespace or the +// user, or naming the wrapper element. The block must come back as the tool +// call named by the invoke element, never as raw XML content. +func TestGlimmerParserFumbledRecipientToolCallFallback(t *testing.T) { + atem := glimmerTestATEM("muse.bash", `pwd && ls -la +`) + cases := map[string]struct { + input string + thinking string + content string + }{ + "namespace recipient": {input: ` to=muse<|message|>` + atem + `<|eom|>`}, + "missing recipient": {input: `<|message|>` + atem + `<|eom|>`}, + "user recipient": {input: ` to=user<|message|>` + atem + `<|eot|>`}, + "wrapper as recipient": { + input: ` to=atem:function_calls<|message|>` + atem + `<|eom|>`, + }, + "after thinking": { + input: ` to=self<|message|>let me look<|eom|><|start|>assistant<|message|>` + atem + `<|eom|>`, + thinking: "let me look", + }, + "implicit boundary": { + input: `<|message|>` + atem + `<|start|>assistant to=user<|message|>Done<|eot|>`, + content: "Done", + }, + } + tools := []api.Tool{glimmerTestTool("muse.bash", map[string]api.ToolProperty{ + "command": {Type: api.PropertyType{"string"}}, + })} + + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + p := &GlimmerParser{} + p.Init(tools, nil, nil) + content, thinking, calls, err := p.Add(tc.input, true) + if err != nil { + t.Fatal(err) + } + if content != tc.content || thinking != tc.thinking || len(calls) != 1 { + t.Fatalf("got content=%q thinking=%q calls=%v", content, thinking, calls) + } + if calls[0].Function.Name != "muse.bash" { + t.Fatalf("call name = %q, want muse.bash", calls[0].Function.Name) + } + if command, ok := calls[0].Function.Arguments.Get("command"); !ok || command != "pwd && ls -la" { + t.Fatalf("command = %#v, %v; want pwd && ls -la", command, ok) + } + }) + } +} + +// TestGlimmerParserFumbledRecipientStreaming re-runs the fallback at every chunk +// boundary: the held block must never leak partial XML into the content +// stream regardless of where the model output is split. +func TestGlimmerParserFumbledRecipientStreaming(t *testing.T) { + tools := []api.Tool{glimmerTestTool("muse.bash", map[string]api.ToolProperty{ + "command": {Type: api.PropertyType{"string"}}, + })} + input := `<|message|>` + glimmerTestATEM("muse.bash", `pwd +`) + `<|eom|>` + + for split := 0; split <= len(input); split++ { + p := &GlimmerParser{} + p.Init(tools, nil, nil) + + var content, thinking string + var calls []api.ToolCall + for i, chunk := range []string{input[:split], input[split:]} { + gotContent, gotThinking, gotCalls, err := p.Add(chunk, i == 1) + if err != nil { + t.Fatalf("split %d: %v", split, err) + } + content += gotContent + thinking += gotThinking + calls = append(calls, gotCalls...) + } + + if content != "" || thinking != "" || len(calls) != 1 { + t.Fatalf("split %d: content=%q thinking=%q calls=%v", split, content, thinking, calls) + } + if calls[0].Function.Name != "muse.bash" { + t.Fatalf("split %d: call name = %q", split, calls[0].Function.Name) + } + } +} + +// The fallback must not fire for content that merely resembles a tool call: +// blocks for undeclared tools, blocks embedded in surrounding prose, and +// sessions with no tools at all stay ordinary content. +func TestGlimmerParserContentThatResemblesToolCallStaysContent(t *testing.T) { + atem := glimmerTestATEM("muse.bash", `pwd +`) + declared := []api.Tool{glimmerTestTool("muse.bash", map[string]api.ToolProperty{ + "command": {Type: api.PropertyType{"string"}}, + })} + cases := map[string]struct { + tools []api.Tool + input string + content string + }{ + "undeclared tool": { + tools: declared, + input: ` to=user<|message|>` + glimmerTestATEM("other.run", "") + `<|eot|>`, + content: glimmerTestATEM("other.run", ""), + }, + "leading prose": { + tools: declared, + input: ` to=user<|message|>Running: +` + atem + `<|eot|>`, + content: "Running:\n" + atem, + }, + "trailing prose": { + tools: declared, + input: ` to=user<|message|>` + atem + ` +Done.<|eot|>`, + content: atem + "\nDone.", + }, + "no tools declared": { + tools: nil, + input: ` to=user<|message|>` + atem + `<|eot|>`, + content: atem, + }, + } + + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + p := &GlimmerParser{} + p.Init(tc.tools, nil, nil) + content, thinking, calls, err := p.Add(tc.input, true) + if err != nil { + t.Fatal(err) + } + if content != tc.content || thinking != "" || len(calls) != 0 { + t.Fatalf("got content=%q thinking=%q calls=%v", content, thinking, calls) + } + }) + } +} + +func TestGlimmerParserUndeclaredRecipientRemainsContent(t *testing.T) { + p := &GlimmerParser{} + p.Init(nil, nil, nil) + body := glimmerTestATEM("not_a_tool", "") + content, thinking, calls, err := p.Add(` to=not_a_tool<|message|>`+body, true) + if err != nil { + t.Fatal(err) + } + if content != body || thinking != "" || len(calls) != 0 { + t.Fatalf("got content=%q thinking=%q calls=%v", content, thinking, calls) + } +} + +func TestGlimmerParserImplicitMessageBoundary(t *testing.T) { + tool := glimmerTestTool("get_weather", map[string]api.ToolProperty{ + "city": {Type: api.PropertyType{"string"}}, + }) + input := " to=self<|message|>Think" + + "<|start|>assistant to=get_weather<|message|>" + + glimmerTestATEM("get_weather", `SF +`) + + "<|start|>assistant to=user<|message|>Done<|eot|>" + + p := &GlimmerParser{} + p.Init([]api.Tool{tool}, nil, nil) + content, thinking, calls, err := p.Add(input, true) + if err != nil { + t.Fatal(err) + } + if content != "Done" || thinking != "Think" || len(calls) != 1 { + t.Fatalf("got content=%q thinking=%q calls=%v", content, thinking, calls) + } +} + +func TestGlimmerParserWithholdsSplitControlToken(t *testing.T) { + p := &GlimmerParser{} + p.Init(nil, nil, nil) + content, _, _, err := p.Add(" to=user<|message|>Hello<|eo", false) + if err != nil { + t.Fatal(err) + } + if content != "Hello" { + t.Fatalf("first content = %q, want Hello", content) + } + content, _, _, err = p.Add("t|>", true) + if err != nil { + t.Fatal(err) + } + if content != "" { + t.Fatalf("control token leaked as content: %q", content) + } +} + +func TestGlimmerParserLiteralStartTokenInMessages(t *testing.T) { + tests := []struct { + name string + recipient string + content string + thinking string + }{ + {name: "content", recipient: "user", content: "Use `<|start|>` here."}, + {name: "thinking", recipient: "self", thinking: "Use `<|start|>` here."}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + p := &GlimmerParser{} + p.Init(nil, nil, nil) + chunks := []string{" to=" + tt.recipient + "<|message|>Use `<", "|", "start", "|>` here."} + var content, thinking string + for i, chunk := range chunks { + gotContent, gotThinking, calls, err := p.Add(chunk, i == len(chunks)-1) + if err != nil { + t.Fatal(err) + } + if len(calls) != 0 { + t.Fatalf("chunk %d got calls=%v", i, calls) + } + content += gotContent + thinking += gotThinking + } + if content != tt.content || thinking != tt.thinking { + t.Fatalf("got content=%q thinking=%q", content, thinking) + } + }) + } +} + +func FuzzGlimmerParser(f *testing.F) { + tool := glimmerTestTool("echo", map[string]api.ToolProperty{ + "text": {Type: api.PropertyType{"string"}}, + }) + for _, seed := range []string{ + ` to=user<|message|>Hello<|eot|>`, + ` to=self<|message|>Think<|eom|><|start|>assistant to=user<|message|>Done<|eot|>`, + ` to=echo<|message|>` + glimmerTestATEM("echo", `value`), + ` to=echo<|message|>`, + `literal <|start|> and <|message|> text`, + } { + f.Add(seed, uint(0)) + f.Add(seed, uint(len(seed)/2)) + } + + f.Fuzz(func(t *testing.T, input string, split uint) { + p := &GlimmerParser{} + p.Init([]api.Tool{tool}, nil, nil) + + at := 0 + if len(input) > 0 { + at = int(split % uint(len(input)+1)) + } + + var calls []api.ToolCall + for i, chunk := range []string{input[:at], input[at:]} { + _, _, got, err := p.Add(chunk, i == 1) + if err != nil { + return + } + calls = append(calls, got...) + } + + lastIndex := -1 + for _, call := range calls { + if call.Function.Name != "echo" { + t.Fatalf("undeclared tool call emitted: %#v", call) + } + if call.Function.Index <= lastIndex { + t.Fatalf("tool call indices are not increasing: %#v", calls) + } + lastIndex = call.Function.Index + } + }) +} diff --git a/model/parsers/parsers.go b/model/parsers/parsers.go index 6791d15bd..cf146e4a5 100644 --- a/model/parsers/parsers.go +++ b/model/parsers/parsers.go @@ -98,6 +98,8 @@ func ParserForName(name string) Parser { return &LagunaV8Parser{} case "cohere": return &CohereParser{} + case "glimmer": + return &GlimmerParser{} default: return nil } diff --git a/model/parsers/qwen3coder.go b/model/parsers/qwen3coder.go index fe941c8b0..1a8e3cf12 100644 --- a/model/parsers/qwen3coder.go +++ b/model/parsers/qwen3coder.go @@ -294,7 +294,10 @@ func parseValue(raw string, paramType api.PropertyType) any { // they exist). This follows the reference implementation raw = strings.TrimPrefix(raw, "\n") raw = strings.TrimSuffix(raw, "\n") + return parseTypedToolValue(raw, paramType) +} +func parseTypedToolValue(raw string, paramType api.PropertyType) any { // Check for null first (case-insensitive) - this takes precedence over any type if strings.ToLower(raw) == "null" { return nil diff --git a/model/renderers/glimmer.go b/model/renderers/glimmer.go new file mode 100644 index 000000000..b5579d65d --- /dev/null +++ b/model/renderers/glimmer.go @@ -0,0 +1,347 @@ +package renderers + +import ( + "bytes" + "encoding/json" + "fmt" + "reflect" + "strings" + "time" + + "github.com/ollama/ollama/api" +) + +const ( + glimmerBOS = "<|begin_of_text|>" + glimmerStart = "<|start|>" + glimmerMessage = "<|message|>" + glimmerEndMessage = "<|eom|>" + glimmerEndTurn = "<|eot|>" + glimmerDefaultSystem = "You are a helpful AI assistant." + glimmerKnowledgeCutoff = "2026-01-04" +) + +const glimmerToolDefinitionsPrefix = `In this environment you have access to a set of tools you can use to answer the user's question. + +You can invoke a function by writing a "" block like the following: + + +$PARAMETER_VALUE +... + + + +String and scalar parameters should be specified as is, while lists and objects should use JSON format. Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with regular expressions. +Here are the functions available in JSONSchema format: +// Tool metadata +` + +const glimmerToolDefinitionsSuffix = ` + +Here's an example of how to call a function in the tool set: +(If the tool namespace is not specified, invoke the function directly as ` + "`example_function_name`" + ` rather than ` + "`example_tool_name.example_function_name`" + `) + +to=example_tool_name.example_function_name + + + +value_1 +This is the value for the second parameter +that can span +"multiple" lines + + +` + +type GlimmerRenderer struct { + useImgTags bool + currentDate string +} + +func (r *GlimmerRenderer) LeadingBOS() string { + return glimmerBOS +} + +func glimmerReasoningStrength(think *api.ThinkValue) string { + if think == nil { + return "high" + } + if !think.Bool() { + return "none" + } + if think.IsString() { + return think.String() + } + return "high" +} + +func glimmerNamespace(name string) string { + if namespace, _, ok := strings.Cut(name, "."); ok { + return namespace + } + return name +} + +func glimmerNamespaces(tools []api.Tool) []string { + seen := make(map[string]bool) + var namespaces []string + for _, tool := range tools { + namespace := glimmerNamespace(tool.Function.Name) + if !seen[namespace] { + seen[namespace] = true + namespaces = append(namespaces, namespace) + } + } + return namespaces +} + +// glimmerJSON reproduces the Transformers 5.x tojson filter. +func glimmerJSON(v any) (string, error) { + var buf bytes.Buffer + encoder := json.NewEncoder(&buf) + encoder.SetEscapeHTML(false) + if err := encoder.Encode(v); err != nil { + return "", err + } + return string(addJSONSpaces(bytes.TrimSuffix(buf.Bytes(), []byte("\n")))), nil +} + +func writeGlimmerToolDefinitions(sb *strings.Builder, tools []api.Tool) error { + sb.WriteString(glimmerToolDefinitionsPrefix) + for _, namespace := range glimmerNamespaces(tools) { + name, err := glimmerJSON(namespace) + if err != nil { + return err + } + description, err := glimmerJSON("") + if err != nil { + return err + } + sb.WriteString(`{"name": `) + sb.WriteString(name) + sb.WriteString(`, "description": `) + sb.WriteString(description) + sb.WriteString("}\n") + } + + sb.WriteString("// Function schemas") + for _, tool := range tools { + name, err := glimmerJSON(tool.Function.Name) + if err != nil { + return err + } + description, err := glimmerJSON(tool.Function.Description) + if err != nil { + return err + } + parameters, err := glimmerJSON(tool.Function.Parameters) + if err != nil { + return err + } + sb.WriteString("\n{\"name\": ") + sb.WriteString(name) + sb.WriteString(`, "description": `) + sb.WriteString(description) + sb.WriteString(`, "parameters": `) + sb.WriteString(parameters) + sb.WriteByte('}') + } + sb.WriteString(glimmerToolDefinitionsSuffix) + return nil +} + +func writeGlimmerSystemMeta(sb *strings.Builder, tools []api.Tool) { + sb.WriteString(`# Valid recipients: "self"`) + for _, namespace := range glimmerNamespaces(tools) { + sb.WriteString(`, "`) + sb.WriteString(namespace) + sb.WriteString(`.*"`) + } + sb.WriteString(`, "user".`) +} + +func writeGlimmerSystem(sb *strings.Builder, content string, tools []api.Tool, strength, currentDate string, defaultSystem bool) error { + writeGlimmerMessageStart(sb, "system") + sb.WriteString(content) + if defaultSystem { + sb.WriteString("\nKnowledge cutoff: ") + sb.WriteString(glimmerKnowledgeCutoff) + sb.WriteByte('.') + sb.WriteString("\nCurrent date: ") + sb.WriteString(currentDate) + sb.WriteByte('.') + } + sb.WriteString("\n\nReasoning strength: ") + sb.WriteString(strength) + sb.WriteByte('.') + if len(tools) > 0 { + sb.WriteString("\n\n") + if err := writeGlimmerToolDefinitions(sb, tools); err != nil { + return err + } + } + sb.WriteString("\n\n") + writeGlimmerSystemMeta(sb, tools) + sb.WriteString(glimmerEndTurn) + return nil +} + +func writeGlimmerMessageStart(sb *strings.Builder, header string) { + sb.WriteString(glimmerStart) + sb.WriteString(header) + sb.WriteString(glimmerMessage) +} + +func writeGlimmerMessage(sb *strings.Builder, header, content, end string) { + writeGlimmerMessageStart(sb, header) + sb.WriteString(content) + sb.WriteString(end) +} + +func glimmerToolResultName(message api.Message, messages []api.Message) string { + if message.ToolName != "" { + return message.ToolName + } + + name := message.ToolCallID + if message.ToolCallID == "" { + return name + } + for _, candidate := range messages { + for _, call := range candidate.ToolCalls { + if call.ID == message.ToolCallID { + name = call.Function.Name + } + } + } + return name +} + +func glimmerCompositeValue(v any) bool { + if v == nil { + return false + } + switch reflect.TypeOf(v).Kind() { + case reflect.Array, reflect.Map, reflect.Slice: + return true + default: + return false + } +} + +func writeGlimmerATEM(sb *strings.Builder, call api.ToolCall) error { + sb.WriteString("\n\n") + for name, value := range call.Function.Arguments.All() { + sb.WriteString("") + switch { + case value == nil: + sb.WriteString("null") + case glimmerCompositeValue(value): + encoded, err := glimmerJSON(value) + if err != nil { + return err + } + sb.WriteString(encoded) + default: + fmt.Fprint(sb, value) + } + sb.WriteString("\n") + } + sb.WriteString("\n") + return nil +} + +func (r *GlimmerRenderer) renderContent(message api.Message, imageOffset int) (string, int) { + if r.useImgTags { + return renderContentWithImageTags(message.Content, len(message.Images), imageOffset) + } + + var sb strings.Builder + for range message.Images { + sb.WriteString("<|patch|>") + } + sb.WriteString(message.Content) + return sb.String(), imageOffset + len(message.Images) +} + +func (r *GlimmerRenderer) Render(messages []api.Message, tools []api.Tool, think *api.ThinkValue) (string, error) { + var sb strings.Builder + sb.WriteString(glimmerBOS) + + hasSystem := false + for _, message := range messages { + if message.Role == "system" { + hasSystem = true + break + } + } + + strength := glimmerReasoningStrength(think) + if !hasSystem { + currentDate := r.currentDate + if currentDate == "" { + currentDate = time.Now().Format(time.DateOnly) + } + if err := writeGlimmerSystem(&sb, glimmerDefaultSystem, tools, strength, currentDate, true); err != nil { + return "", err + } + } + + imageOffset := 0 + for i, message := range messages { + content, nextImageOffset := r.renderContent(message, imageOffset) + imageOffset = nextImageOffset + + endToken := glimmerEndTurn + if i+1 < len(messages) && messages[i+1].Role == message.Role { + endToken = glimmerEndMessage + } + + switch message.Role { + case "system": + if err := writeGlimmerSystem(&sb, content, tools, strength, "", false); err != nil { + return "", err + } + case "user": + writeGlimmerMessage(&sb, "user", content, glimmerEndTurn) + case "tool": + name := glimmerToolResultName(message, messages) + writeGlimmerMessageStart(&sb, "tool "+name) + sb.WriteString(`\n") + sb.WriteString(content) + sb.WriteString("\n") + sb.WriteString(glimmerEndTurn) + case "assistant": + if message.Thinking != "" { + writeGlimmerMessage(&sb, "assistant to=self", message.Thinking, glimmerEndMessage) + } + // A final thinking-only message is the server's unfinished self + // segment, represented by recipient=self and end_turn=false in Jinja. + if len(message.ToolCalls) > 0 { + for j, call := range message.ToolCalls { + writeGlimmerMessageStart(&sb, "assistant to="+call.Function.Name) + if err := writeGlimmerATEM(&sb, call); err != nil { + return "", err + } + if j+1 == len(message.ToolCalls) { + sb.WriteString(endToken) + } else { + sb.WriteString(glimmerEndMessage) + } + } + } else if content != "" || message.Thinking == "" { + writeGlimmerMessage(&sb, "assistant to=user", content, glimmerEndTurn) + } + } + } + + sb.WriteString(glimmerStart) + sb.WriteString("assistant") + return sb.String(), nil +} diff --git a/model/renderers/glimmer_reference_test.go b/model/renderers/glimmer_reference_test.go new file mode 100644 index 000000000..a6843acad --- /dev/null +++ b/model/renderers/glimmer_reference_test.go @@ -0,0 +1,743 @@ +package renderers + +import ( + "encoding/json" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + + "github.com/ollama/ollama/api" + "github.com/ollama/ollama/model/parsers" +) + +// glimmerChatTemplate is copied byte-for-byte from hf/chat_template.jinja at +// publisher revision 7f1ae4102757b5303d0896ec86f6a8b77b217822. +// SHA-256: 114f55ebdc1804c1af371197b9fdf2d6bb925966c9dfe46b73782a71bc07965e. +const glimmerChatTemplate = "testdata/glimmer_chat_template.jinja" + +const ( + glimmerRefBOS = "<|begin_of_text|>" + glimmerRefStart = "<|start|>" + glimmerRefMessage = "<|message|>" + glimmerRefEOM = "<|eom|>" + glimmerRefEOT = "<|eot|>" + glimmerRefCurrentDate = "2026-07-29" +) + +const glimmerRefDefaultSystem = `You are a helpful AI assistant. +Knowledge cutoff: 2026-01-04. +Current date: 2026-07-29. + +Reasoning strength: high. + +# Valid recipients: "self", "user".` + +const glimmerRefWeatherToolDefinitions = `In this environment you have access to a set of tools you can use to answer the user's question. + +You can invoke a function by writing a "" block like the following: + + +$PARAMETER_VALUE +... + + + +String and scalar parameters should be specified as is, while lists and objects should use JSON format. Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with regular expressions. +Here are the functions available in JSONSchema format: +// Tool metadata +{"name": "get_weather", "description": ""} +// Function schemas +{"name": "get_weather", "description": "Get the weather", "parameters": {"type": "object", "required": ["city"], "properties": {"city": {"type": "string"}, "units": {"type": "string"}}}} + +Here's an example of how to call a function in the tool set: +(If the tool namespace is not specified, invoke the function directly as ` + "`example_function_name`" + ` rather than ` + "`example_tool_name.example_function_name`" + `) + +to=example_tool_name.example_function_name + + + +value_1 +This is the value for the second parameter +that can span +"multiple" lines + + +` + +type glimmerJinjaMessage struct { + Role string `json:"role"` + Content any `json:"content"` + Name string `json:"name,omitempty"` + ToolCallID string `json:"tool_call_id,omitempty"` + ReasoningContent string `json:"reasoning_content,omitempty"` + ToolCalls []api.ToolCall `json:"tool_calls,omitempty"` + Recipient string `json:"recipient,omitempty"` + EndTurn *bool `json:"end_turn,omitempty"` +} + +type glimmerJinjaContentPart struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` +} + +func glimmerRefPrompt(parts ...string) string { + return glimmerRefBOS + strings.Join(parts, "") + glimmerRefStart + "assistant" +} + +func glimmerRefMsg(header, content, end string) string { + return glimmerRefStart + header + glimmerRefMessage + content + end +} + +func glimmerReferenceWeatherTool() api.Tool { + return api.Tool{ + Type: "function", + Function: api.ToolFunction{ + Name: "get_weather", + Description: "Get the weather", + Parameters: api.ToolFunctionParameters{ + Type: "object", + Required: []string{"city"}, + Properties: testPropsOrdered([]orderedProp{ + {Key: "city", Value: api.ToolProperty{Type: api.PropertyType{"string"}}}, + {Key: "units", Value: api.ToolProperty{Type: api.PropertyType{"string"}}}, + }), + }, + }, + } +} + +func TestGlimmerRendererMatchesJinja2Reference(t *testing.T) { + verifyJinja2 := os.Getenv("VERIFY_JINJA2") != "" + if verifyJinja2 { + requireGlimmerJinja2(t) + t.Log("VERIFY_JINJA2=1: verifying expected values against Glimmer Jinja2 template") + } + + weatherTool := glimmerReferenceWeatherTool() + weatherArgs := api.NewToolCallFunctionArguments() + weatherArgs.Set("city", "SF") + weatherArgs.Set("units", "fahrenheit") + + tests := []struct { + name string + messages []api.Message + tools []api.Tool + think *api.ThinkValue + expected string + }{ + { + name: "default system", + messages: []api.Message{{Role: "user", Content: "Hello"}}, + expected: glimmerRefPrompt( + glimmerRefMsg("system", glimmerRefDefaultSystem, glimmerRefEOT), + glimmerRefMsg("user", "Hello", glimmerRefEOT), + ), + }, + { + name: "explicit system image thinking and answer", + messages: []api.Message{ + {Role: "system", Content: "Be concise."}, + {Role: "user", Content: "Read this.", Images: []api.ImageData{{1}}}, + {Role: "assistant", Thinking: "I should inspect it.", Content: "It says hello."}, + }, + expected: glimmerRefPrompt( + glimmerRefMsg("system", `Be concise. + +Reasoning strength: high. + +# Valid recipients: "self", "user".`, glimmerRefEOT), + glimmerRefMsg("user", "<|patch|>Read this.", glimmerRefEOT), + glimmerRefMsg("assistant to=self", "I should inspect it.", glimmerRefEOM), + glimmerRefMsg("assistant to=user", "It says hello.", glimmerRefEOT), + ), + }, + { + name: "tool call and result", + messages: []api.Message{ + {Role: "user", Content: "Weather?"}, + {Role: "assistant", Thinking: "Need current data.", ToolCalls: []api.ToolCall{{ + ID: "call_weather", + Function: api.ToolCallFunction{Name: "get_weather", Arguments: weatherArgs}, + }}}, + {Role: "tool", ToolCallID: "call_weather", Content: `{"temp":65}`}, + {Role: "assistant", Content: "It is 65F."}, + }, + tools: []api.Tool{weatherTool}, + expected: glimmerRefPrompt( + glimmerRefMsg("system", `You are a helpful AI assistant. +Knowledge cutoff: 2026-01-04. +Current date: 2026-07-29. + +Reasoning strength: high. + +`+glimmerRefWeatherToolDefinitions+` + +# Valid recipients: "self", "get_weather.*", "user".`, glimmerRefEOT), + glimmerRefMsg("user", "Weather?", glimmerRefEOT), + glimmerRefMsg("assistant to=self", "Need current data.", glimmerRefEOM), + glimmerRefMsg("assistant to=get_weather", ` + +SF +fahrenheit + +`, glimmerRefEOT), + glimmerRefMsg("tool get_weather", ` +{"temp":65} +`, glimmerRefEOT), + glimmerRefMsg("assistant to=user", "It is 65F.", glimmerRefEOT), + ), + }, + { + name: "adjacent assistant tool calls use message boundary", + messages: []api.Message{ + {Role: "user", Content: "Run twice."}, + {Role: "assistant", ToolCalls: []api.ToolCall{{ + Function: api.ToolCallFunction{Name: "get_weather", Arguments: weatherArgs}, + }}}, + {Role: "assistant", ToolCalls: []api.ToolCall{{ + Function: api.ToolCallFunction{Name: "get_weather", Arguments: weatherArgs}, + }}}, + {Role: "user", Content: "Continue."}, + }, + expected: glimmerRefPrompt( + glimmerRefMsg("system", glimmerRefDefaultSystem, glimmerRefEOT), + glimmerRefMsg("user", "Run twice.", glimmerRefEOT), + glimmerRefMsg("assistant to=get_weather", ` + +SF +fahrenheit + +`, glimmerRefEOM), + glimmerRefMsg("assistant to=get_weather", ` + +SF +fahrenheit + +`, glimmerRefEOT), + glimmerRefMsg("user", "Continue.", glimmerRefEOT), + ), + }, + { + name: "explicit reasoning level", + messages: []api.Message{{Role: "user", Content: "Solve."}}, + think: &api.ThinkValue{Value: "low"}, + expected: glimmerRefPrompt( + glimmerRefMsg("system", strings.Replace(glimmerRefDefaultSystem, "high", "low", 1), glimmerRefEOT), + glimmerRefMsg("user", "Solve.", glimmerRefEOT), + ), + }, + { + name: "thinking disabled", + messages: []api.Message{{Role: "user", Content: "Answer directly."}}, + think: &api.ThinkValue{Value: false}, + expected: glimmerRefPrompt( + glimmerRefMsg("system", strings.Replace(glimmerRefDefaultSystem, "high", "none", 1), glimmerRefEOT), + glimmerRefMsg("user", "Answer directly.", glimmerRefEOT), + ), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := (&GlimmerRenderer{currentDate: glimmerRefCurrentDate}).Render(tt.messages, tt.tools, tt.think) + if err != nil { + t.Fatal(err) + } + if got != tt.expected { + t.Fatalf("renderer output mismatch:\ngot: %q\nwant: %q", got, tt.expected) + } + + if verifyJinja2 { + jinja2Output := renderGlimmerWithJinja2(t, tt.messages, tt.tools, tt.think) + if jinja2Output != tt.expected { + fmt.Fprintf(os.Stderr, "\nJINJA2 OUTPUT for %s:\n%q\n\n", tt.name, jinja2Output) + t.Fatalf("hardcoded expected value doesn't match Jinja2 template:\ngot: %q\nwant: %q", tt.expected, jinja2Output) + } + } + }) + } +} + +func TestGlimmerRendererMatchesJinja2ExpandedParity(t *testing.T) { + if os.Getenv("VERIFY_JINJA2") == "" { + t.Skip("set VERIFY_JINJA2=1 to run expanded Jinja2 parity checks") + } + requireGlimmerJinja2(t) + + weatherTool := glimmerReferenceWeatherTool() + readTool := api.Tool{ + Type: "function", + Function: api.ToolFunction{ + Name: "filesystem.read", + Description: `Read a file's contents`, + Parameters: api.ToolFunctionParameters{ + Type: "object", + Required: []string{"path"}, + Properties: testPropsOrdered([]orderedProp{ + {Key: "path", Value: api.ToolProperty{Type: api.PropertyType{"string"}}}, + {Key: "lines", Value: api.ToolProperty{Type: api.PropertyType{"array"}, Items: map[string]any{"type": "integer"}}}, + }), + }, + }, + } + + weatherArgs := api.NewToolCallFunctionArguments() + weatherArgs.Set("city", "SF") + readArgs := api.NewToolCallFunctionArguments() + readArgs.Set("path", "README.md") + readArgs.Set("lines", []any{1, 5}) + allValues := api.NewToolCallFunctionArguments() + allValues.Set("string", " keep spaces ") + allValues.Set("integer", 3) + allValues.Set("number", 3.5) + allValues.Set("boolean", true) + allValues.Set("null", nil) + allValues.Set("array", []any{"héllo", 2}) + allValues.Set("object", map[string]any{"b": 2, "a": ""}) + + tests := []struct { + name string + messages []api.Message + tools []api.Tool + think *api.ThinkValue + }{ + { + name: "minimal user", + messages: []api.Message{{Role: "user", Content: "Hello"}}, + }, + { + name: "explicit system multi turn", + messages: []api.Message{ + {Role: "system", Content: "Be concise."}, + {Role: "user", Content: "One"}, + {Role: "assistant", Content: "Two"}, + {Role: "user", Content: "Three"}, + }, + }, + { + name: "late and repeated system messages", + messages: []api.Message{ + {Role: "user", Content: "One"}, + {Role: "system", Content: "First policy."}, + {Role: "system", Content: "Second policy."}, + {Role: "user", Content: "Two"}, + }, + tools: []api.Tool{weatherTool}, + }, + { + name: "multiple images", + messages: []api.Message{{ + Role: "user", + Content: "Compare them.", + Images: []api.ImageData{{1}, {2}}, + }}, + }, + { + name: "unsupported role omitted", + messages: []api.Message{ + {Role: "developer", ToolName: "policy", Content: "Follow policy."}, + {Role: "user", Content: "Hi"}, + }, + }, + { + name: "assistant thinking and final", + messages: []api.Message{ + {Role: "user", Content: "Solve"}, + {Role: "assistant", Thinking: "Private reasoning.", Content: "Answer."}, + }, + }, + { + name: "multiple tool calls and out of order results", + messages: []api.Message{ + {Role: "user", Content: "Weather and file?"}, + {Role: "assistant", Thinking: "Need both.", ToolCalls: []api.ToolCall{ + { + ID: "call_weather", + Function: api.ToolCallFunction{Name: "get_weather", Arguments: weatherArgs}, + }, + { + ID: "call_read", + Function: api.ToolCallFunction{Name: "filesystem.read", Arguments: readArgs}, + }, + }}, + {Role: "tool", ToolCallID: "call_read", Content: "README"}, + {Role: "tool", ToolCallID: "call_weather", Content: `{"temp":65}`}, + {Role: "assistant", Content: "Done."}, + }, + tools: []api.Tool{weatherTool, readTool}, + }, + { + name: "explicit tool name", + messages: []api.Message{ + {Role: "tool", ToolName: "filesystem.read", Content: "README"}, + }, + tools: []api.Tool{readTool}, + }, + { + name: "unresolved tool call id", + messages: []api.Message{ + {Role: "tool", ToolCallID: "missing_call", Content: "missing"}, + }, + }, + { + name: "empty tool result name", + messages: []api.Message{ + {Role: "tool", Content: "orphan"}, + }, + }, + { + name: "all ATEM value types", + messages: []api.Message{ + {Role: "user", Content: "Run."}, + {Role: "assistant", ToolCalls: []api.ToolCall{{ + Function: api.ToolCallFunction{Name: "filesystem.read", Arguments: allValues}, + }}}, + }, + tools: []api.Tool{readTool}, + }, + { + name: "thinking true", + messages: []api.Message{{Role: "user", Content: "Solve."}}, + think: &api.ThinkValue{Value: true}, + }, + { + name: "thinking false", + messages: []api.Message{{Role: "user", Content: "Answer."}}, + think: &api.ThinkValue{Value: false}, + }, + { + name: "reasoning max", + messages: []api.Message{{Role: "user", Content: "Solve."}}, + think: &api.ThinkValue{Value: "max"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := (&GlimmerRenderer{currentDate: glimmerRefCurrentDate}).Render(tt.messages, tt.tools, tt.think) + if err != nil { + t.Fatal(err) + } + want := renderGlimmerWithJinja2(t, tt.messages, tt.tools, tt.think) + if got != want { + t.Fatalf("renderer output doesn't match Jinja2 template:\ngot: %q\nwant: %q", got, want) + } + }) + } +} + +func TestGlimmerRendererKnownJinja2DifferenceThinkingContinuation(t *testing.T) { + messages := []api.Message{ + {Role: "user", Content: "Solve"}, + {Role: "assistant", Thinking: "Private reasoning."}, + } + expected := glimmerRefPrompt( + glimmerRefMsg("system", glimmerRefDefaultSystem, glimmerRefEOT), + glimmerRefMsg("user", "Solve", glimmerRefEOT), + glimmerRefMsg("assistant to=self", "Private reasoning.", glimmerRefEOM), + ) + + got, err := (&GlimmerRenderer{currentDate: glimmerRefCurrentDate}).Render(messages, nil, nil) + if err != nil { + t.Fatal(err) + } + if got != expected { + t.Fatalf("renderer continuation mismatch:\ngot: %q\nwant: %q", got, expected) + } + + if os.Getenv("VERIFY_JINJA2") != "" { + requireGlimmerJinja2(t) + + // api.Message cannot directly represent the reference template's + // unfinished assistant message. Mechanical conversion therefore renders + // this history as completed reasoning followed by an empty user answer. + completeExpected := glimmerRefPrompt( + glimmerRefMsg("system", glimmerRefDefaultSystem, glimmerRefEOT), + glimmerRefMsg("user", "Solve", glimmerRefEOT), + glimmerRefMsg("assistant to=self", "Private reasoning.", glimmerRefEOM), + glimmerRefMsg("assistant to=user", "", glimmerRefEOT), + ) + completeOutput := renderGlimmerWithJinja2(t, messages, nil, nil) + if completeOutput != completeExpected { + t.Fatalf("mechanical Jinja2 conversion mismatch:\ngot: %q\nwant: %q", completeOutput, completeExpected) + } + if completeOutput == got { + t.Fatal("known unfinished-message representation difference disappeared") + } + + // The renderer's partial output must still match the publisher template + // when given the reference representation of an unfinished self message. + endTurn := false + partialOutput := renderGlimmerJinjaMessages(t, []glimmerJinjaMessage{ + {Role: "user", Content: "Solve"}, + { + Role: "assistant", + Content: "Private reasoning.", + Recipient: "self", + EndTurn: &endTurn, + }, + }, nil, nil) + if partialOutput != expected { + t.Fatalf("hardcoded continuation doesn't match Jinja2 template:\ngot: %q\nwant: %q", expected, partialOutput) + } + } +} + +func TestGlimmerParserRendererToolLoopRoundTrip(t *testing.T) { + verifyJinja2 := os.Getenv("VERIFY_JINJA2") != "" + if verifyJinja2 { + requireGlimmerJinja2(t) + } + + issueTool := api.Tool{ + Type: "function", + Function: api.ToolFunction{ + Name: "issue", + Description: "Read an issue and its recent comments.", + Parameters: api.ToolFunctionParameters{ + Type: "object", + Required: []string{"number"}, + Properties: testPropsOrdered([]orderedProp{{ + Key: "number", + Value: api.ToolProperty{Type: api.PropertyType{"integer"}}, + }}), + }, + }, + } + readTool := api.Tool{ + Type: "function", + Function: api.ToolFunction{ + Name: "read", + Description: "Read a repository file.", + Parameters: api.ToolFunctionParameters{ + Type: "object", + Required: []string{"path"}, + Properties: testPropsOrdered([]orderedProp{{ + Key: "path", + Value: api.ToolProperty{Type: api.PropertyType{"string"}}, + }}), + }, + }, + } + tools := []api.Tool{issueTool, readTool} + messages := []api.Message{{ + Role: "user", + Content: "Investigate issue 1736 and inspect server/download.go.", + }} + + parse := func(raw string) api.Message { + t.Helper() + parser := parsers.ParserForName("glimmer") + parser.Init(tools, nil, nil) + content, thinking, calls, err := parser.Add(raw, true) + if err != nil { + t.Fatalf("parse Glimmer assistant turn: %v", err) + } + return api.Message{ + Role: "assistant", + Content: content, + Thinking: thinking, + ToolCalls: calls, + } + } + verifyPrompt := func(stage string) { + t.Helper() + got, err := (&GlimmerRenderer{currentDate: glimmerRefCurrentDate}).Render(messages, tools, nil) + if err != nil { + t.Fatalf("%s render: %v", stage, err) + } + if !strings.HasSuffix(got, glimmerRefStart+"assistant") { + t.Fatalf("%s prompt does not end with the assistant generation prefix", stage) + } + if verifyJinja2 { + want := renderGlimmerWithJinja2(t, messages, tools, nil) + if got != want { + t.Fatalf("%s prompt doesn't match Jinja2:\ngot: %q\nwant: %q", stage, got, want) + } + } + } + + verifyPrompt("initial") + + issueCall := parse(` to=self<|message|>I need the issue details.<|eom|>` + + `<|start|>assistant to=issue<|message|> + +1736 + +<|eot|>`) + if issueCall.Thinking != "I need the issue details." || len(issueCall.ToolCalls) != 1 { + t.Fatalf("issue turn = %#v", issueCall) + } + if number, ok := issueCall.ToolCalls[0].Function.Arguments.Get("number"); !ok || number != 1736 { + t.Fatalf("issue number = %#v, %v; want 1736", number, ok) + } + issueCall.ToolCalls[0].ID = "call_issue" + messages = append(messages, issueCall, api.Message{ + Role: "tool", + ToolCallID: "call_issue", + Content: "Issue 1736 links the proposed first-byte fix.", + }) + verifyPrompt("after issue") + + readCall := parse(` to=self<|message|>Now inspect the implementation.<|eom|>` + + `<|start|>assistant to=read<|message|> + +server/download.go + +<|eot|>`) + if readCall.Thinking != "Now inspect the implementation." || len(readCall.ToolCalls) != 1 { + t.Fatalf("read turn = %#v", readCall) + } + if path, ok := readCall.ToolCalls[0].Function.Arguments.Get("path"); !ok || path != "server/download.go" { + t.Fatalf("read path = %#v, %v; want server/download.go", path, ok) + } + readCall.ToolCalls[0].ID = "call_read" + messages = append(messages, readCall, api.Message{ + Role: "tool", + ToolCallID: "call_read", + Content: "The inactivity clock starts only after the first body byte.", + }) + verifyPrompt("after read") + + answer := parse(` to=self<|message|>I have enough evidence.<|eom|>` + + `<|start|>assistant to=user<|message|>The proposed fix is plausible.<|eot|>`) + if answer.Thinking != "I have enough evidence." || answer.Content != "The proposed fix is plausible." || len(answer.ToolCalls) != 0 { + t.Fatalf("answer turn = %#v", answer) + } + messages = append(messages, answer) + verifyPrompt("after answer") +} + +func renderGlimmerWithJinja2(t *testing.T, messages []api.Message, tools []api.Tool, think *api.ThinkValue) string { + t.Helper() + + var strength any + if think != nil { + switch { + case !think.Bool(): + strength = "none" + case think.IsString(): + strength = think.String() + default: + strength = "high" + } + } + return renderGlimmerJinjaMessages(t, glimmerMessagesForJinja(messages), tools, strength) +} + +func renderGlimmerJinjaMessages(t *testing.T, messages []glimmerJinjaMessage, tools []api.Tool, strength any) string { + t.Helper() + + templatePath, err := filepath.Abs(glimmerChatTemplate) + if err != nil { + t.Fatalf("failed to get template path: %v", err) + } + + messagesJSON, err := json.Marshal(messages) + if err != nil { + t.Fatalf("failed to marshal messages: %v", err) + } + toolsJSON, err := json.Marshal(tools) + if err != nil { + t.Fatalf("failed to marshal tools: %v", err) + } + + strengthJSON, err := json.Marshal(strength) + if err != nil { + t.Fatalf("failed to marshal reasoning strength: %v", err) + } + + script := fmt.Sprintf(` +import json +from transformers.utils.chat_template_utils import _compile_jinja_template +tmpl = _compile_jinja_template(open(%q).read()) +kwargs = { + "messages": json.loads(%q), + "tools": json.loads(%q), + "bos_token": %q, + "add_generation_prompt": True, + "current_date": %q, +} +reasoning_strength = json.loads(%q) +if reasoning_strength is not None: + kwargs["reasoning_strength"] = reasoning_strength +print(tmpl.render(**kwargs), end="") +`, templatePath, string(messagesJSON), string(toolsJSON), glimmerRefBOS, glimmerRefCurrentDate, string(strengthJSON)) + + cmd := glimmerPythonCommand(t, "-c", script) + var stdout, stderr strings.Builder + cmd.Stdout = &stdout + cmd.Stderr = &stderr + if err := cmd.Run(); err != nil { + t.Fatalf("Jinja2 render failed: %v\nstderr: %s", err, stderr.String()) + } + return stdout.String() +} + +func requireGlimmerJinja2(t *testing.T) { + t.Helper() + if err := glimmerPythonCommand(t, "-c", "import transformers").Run(); err != nil { + t.Fatal("VERIFY_JINJA2=1 requires .venv/bin/python with transformers 5.x or uv with downloadable transformers") + } +} + +func glimmerPythonCommand(t *testing.T, args ...string) *exec.Cmd { + t.Helper() + if python, ok := findGlimmerVenvPython(); ok { + return exec.CommandContext(t.Context(), python, args...) + } + + uvArgs := append([]string{"run", "--with", "transformers>=5,<6", "python"}, args...) + return exec.CommandContext(t.Context(), "uv", uvArgs...) +} + +func findGlimmerVenvPython() (string, bool) { + dir, err := os.Getwd() + if err != nil { + return "", false + } + for { + python := filepath.Join(dir, ".venv", "bin", "python") + if _, err := os.Stat(python); err == nil { + return python, true + } + parent := filepath.Dir(dir) + if parent == dir { + return "", false + } + dir = parent + } +} + +func glimmerMessagesForJinja(messages []api.Message) []glimmerJinjaMessage { + result := make([]glimmerJinjaMessage, 0, len(messages)) + for _, message := range messages { + result = append(result, glimmerJinjaMessage{ + Role: message.Role, + Content: glimmerJinjaContent(message.Content, len(message.Images)), + Name: message.ToolName, + ToolCallID: message.ToolCallID, + ReasoningContent: message.Thinking, + ToolCalls: message.ToolCalls, + }) + } + return result +} + +func glimmerJinjaContent(content string, imageCount int) any { + if imageCount == 0 { + return content + } + + parts := make([]glimmerJinjaContentPart, 0, imageCount+1) + for range imageCount { + parts = append(parts, glimmerJinjaContentPart{Type: "image"}) + } + if content != "" { + parts = append(parts, glimmerJinjaContentPart{Type: "text", Text: content}) + } + return parts +} diff --git a/model/renderers/glimmer_test.go b/model/renderers/glimmer_test.go new file mode 100644 index 000000000..c9c42a05e --- /dev/null +++ b/model/renderers/glimmer_test.go @@ -0,0 +1,143 @@ +package renderers + +import ( + "strings" + "testing" + + "github.com/ollama/ollama/api" +) + +func TestGlimmerRenderImages(t *testing.T) { + tests := []struct { + name string + renderTags bool + messages []api.Message + wantContent string + }{ + { + name: "reference patch token", + messages: []api.Message{{ + Role: "user", + Content: "Describe.", + Images: []api.ImageData{{1}}, + }}, + wantContent: "<|start|>user<|message|><|patch|>Describe.<|eot|>", + }, + { + name: "runner tags across turns", + renderTags: true, + messages: []api.Message{ + {Role: "user", Content: "First.", Images: []api.ImageData{{1}}}, + {Role: "assistant", Content: "Done."}, + {Role: "user", Content: "Compare.", Images: []api.ImageData{{2}, {3}}}, + }, + wantContent: "<|start|>user<|message|>[img-0] First.<|eot|>" + + "<|start|>assistant to=user<|message|>Done.<|eot|>" + + "<|start|>user<|message|>[img-1][img-2] Compare.<|eot|>", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := (&GlimmerRenderer{useImgTags: tt.renderTags}).Render(tt.messages, nil, nil) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(got, tt.wantContent) { + t.Fatalf("rendered prompt missing image content:\ngot: %q\nwant: %q", got, tt.wantContent) + } + }) + } +} + +func TestGlimmerRenderATEMValues(t *testing.T) { + arguments := api.NewToolCallFunctionArguments() + arguments.Set("text", "keep spaces") + arguments.Set("count", 3) + arguments.Set("enabled", true) + arguments.Set("fallback", nil) + arguments.Set("items", []any{"one", "two"}) + arguments.Set("config", map[string]any{"mode": "fast"}) + + got, err := (&GlimmerRenderer{}).Render([]api.Message{ + {Role: "user", Content: "Run it."}, + {Role: "assistant", ToolCalls: []api.ToolCall{{ + Function: api.ToolCallFunction{Name: "tools.run", Arguments: arguments}, + }}}, + }, nil, nil) + if err != nil { + t.Fatal(err) + } + + want := ` + +keep spaces +3 +true +null +["one", "two"] +{"mode": "fast"} + +` + if !strings.Contains(got, want) { + t.Fatalf("rendered prompt missing ATEM call:\ngot: %q\nwant: %q", got, want) + } +} + +func TestGlimmerRenderReasoningStrength(t *testing.T) { + tests := []struct { + name string + think *api.ThinkValue + want string + }{ + {name: "default", want: "Reasoning strength: high."}, + {name: "enabled", think: &api.ThinkValue{Value: true}, want: "Reasoning strength: high."}, + {name: "disabled", think: &api.ThinkValue{Value: false}, want: "Reasoning strength: none."}, + {name: "level", think: &api.ThinkValue{Value: "medium"}, want: "Reasoning strength: medium."}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := (&GlimmerRenderer{}).Render([]api.Message{{Role: "user", Content: "Hello"}}, nil, tt.think) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(got, tt.want) { + t.Fatalf("rendered prompt missing %q:\n%s", tt.want, got) + } + }) + } +} + +func TestGlimmerRenderToolResultName(t *testing.T) { + arguments := api.NewToolCallFunctionArguments() + got, err := (&GlimmerRenderer{}).Render([]api.Message{ + {Role: "assistant", ToolCalls: []api.ToolCall{{ + ID: "call_1", + Function: api.ToolCallFunction{Name: "get_weather", Arguments: arguments}, + }}}, + {Role: "tool", ToolCallID: "call_1", Content: `{"temp":65}`}, + {Role: "tool", ToolCallID: "missing", Content: "not found"}, + }, nil, nil) + if err != nil { + t.Fatal(err) + } + + for _, want := range []string{ + `<|start|>tool get_weather<|message|>` + "\n" + `{"temp":65}` + "\n<|eot|>", + `<|start|>tool missing<|message|>` + "\nnot found\n<|eot|>", + } { + if !strings.Contains(got, want) { + t.Fatalf("rendered prompt missing tool result %q:\n%s", want, got) + } + } +} + +func TestGlimmerRendererRegistered(t *testing.T) { + if rendererForName("glimmer") == nil { + t.Fatal("glimmer renderer is not registered") + } + if got := LeadingBOSForRenderer("glimmer"); got != glimmerBOS { + t.Fatalf("LeadingBOSForRenderer(glimmer) = %q, want %q", got, glimmerBOS) + } +} diff --git a/model/renderers/json.go b/model/renderers/json.go index 76d46a90b..8dc71673c 100644 --- a/model/renderers/json.go +++ b/model/renderers/json.go @@ -10,7 +10,10 @@ func marshalWithSpaces(v any) ([]byte, error) { if err != nil { return nil, err } + return addJSONSpaces(b), nil +} +func addJSONSpaces(b []byte) []byte { out := make([]byte, 0, len(b)+len(b)/8) inStr, esc := false, false for _, c := range b { @@ -41,5 +44,5 @@ func marshalWithSpaces(v any) ([]byte, error) { out = append(out, c) } } - return out, nil + return out } diff --git a/model/renderers/renderer.go b/model/renderers/renderer.go index db4d9528d..3d0a8674d 100644 --- a/model/renderers/renderer.go +++ b/model/renderers/renderer.go @@ -113,6 +113,8 @@ func rendererForName(name string) Renderer { return &LagunaV8Renderer{} case "cohere": return &CohereRenderer{} + case "glimmer": + return &GlimmerRenderer{useImgTags: RenderImgTags} default: return nil } diff --git a/model/renderers/testdata/glimmer_chat_template.jinja b/model/renderers/testdata/glimmer_chat_template.jinja new file mode 100644 index 000000000..8a8673897 --- /dev/null +++ b/model/renderers/testdata/glimmer_chat_template.jinja @@ -0,0 +1 @@ +{%- macro render_content(content) -%}{%- if content is string -%}{{- content -}}{%- elif content is not none -%}{%- for part in content -%}{%- if part['type'] == 'image' -%}{{- '<|patch|>' -}}{%- elif part['type'] == 'video' -%}{{- '<|video|>' -}}{%- elif part['type'] == 'text' -%}{{- part['text'] -}}{%- endif -%}{%- endfor -%}{%- endif -%}{%- endmacro -%}{%- macro render_atem(tc) -%}{%- set args = tc.function.arguments -%}{%- if args is not mapping -%}{{- raise_exception('Onyx ATEM chat template requires tool_call.function.arguments to be a dict (mapping); a JSON string cannot be parsed in the HF jinja sandbox.') -}}{%- endif -%}{{- '\n\n' -}}{%- for k, v in args.items() -%}{{- '' -}}{%- if v is boolean -%}{%- if v -%}true{%- else -%}false{%- endif -%}{%- elif v is none -%}null{%- elif v is mapping or (v is iterable and v is not string) -%}{{- v | tojson -}}{%- else -%}{{- v -}}{%- endif -%}{{- '\n' -}}{%- endfor -%}{{- '\n' -}}{%- endmacro -%}{%- macro render_tool_defs(tools) -%}{{- 'In this environment you have access to a set of tools you can use to answer the user\'s question.\n\n' -}}{{- 'You can invoke a function by writing a "" block like the following:\n' -}}{{- '\n\n$PARAMETER_VALUE\n...\n\n\n\n' -}}{{- 'String and scalar parameters should be specified as is, while lists and objects should use JSON format. Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with regular expressions.\n' -}}{{- 'Here are the functions available in JSONSchema format:\n' -}}{{- '// Tool metadata\n' -}}{%- set nsns = namespace(seen=[]) -%}{%- for tool in tools -%}{%- set fn = tool.function if tool.function is defined else tool -%}{%- set tns = fn.name.split('.')[0] -%}{%- if tns not in nsns.seen -%}{%- set nsns.seen = nsns.seen + [tns] -%}{%- endif -%}{%- endfor -%}{%- set nd = tool_namespace_descriptions if tool_namespace_descriptions is defined else {} -%}{%- for tns in nsns.seen -%}{{- '{"name": ' + (tns | tojson) + ', "description": ' + ((nd[tns] if tns in nd else '') | tojson) + '}\n' -}}{%- endfor -%}{{- '// Function schemas' -}}{%- for tool in tools -%}{%- set fn = tool.function if tool.function is defined else tool -%}{{- '\n{"name": ' + (fn.name | tojson) + ', "description": ' + (fn.description | tojson) + ', "parameters": ' + (fn.parameters | tojson) + '}' -}}{%- endfor -%}{{- '\n\nHere\'s an example of how to call a function in the tool set:\n' -}}{{- '(If the tool namespace is not specified, invoke the function directly as `example_function_name` rather than `example_tool_name.example_function_name`)\n\n' -}}{{- 'to=example_tool_name.example_function_name\n\n' -}}{{- '\n\n' -}}{{- 'value_1\n' -}}{{- 'This is the value for the second parameter\nthat can span\n"multiple" lines\n\n' -}}{{- '\n' -}}{%- endmacro -%}{%- macro render_reasoning() -%}{%- set rs = reasoning_strength if reasoning_strength is defined and reasoning_strength else 'high' -%}{{- 'Reasoning strength: ' + rs + '.' -}}{%- endmacro -%}{%- macro render_system_meta(tools) -%}{%- set rns = namespace(recipients=['"self"'], nslist=[]) -%}{%- if tools -%}{%- for tool in tools -%}{%- set fn = tool.function if tool.function is defined else tool -%}{%- set tns = fn.name.split('.')[0] -%}{%- if tns not in rns.nslist -%}{%- set rns.nslist = rns.nslist + [tns] -%}{%- endif -%}{%- endfor -%}{%- for tns in rns.nslist -%}{%- set rns.recipients = rns.recipients + ['"' + tns + '.*"'] -%}{%- endfor -%}{%- endif -%}{%- set rns.recipients = rns.recipients + ['"user"'] -%}{{- '# Valid recipients: ' + rns.recipients | join(', ') + '.' -}}{%- endmacro -%}{{- bos_token -}}{%- set ns = namespace(has_system=false) -%}{%- for m in messages -%}{%- if m['role'] == 'system' -%}{%- set ns.has_system = true -%}{%- endif -%}{%- endfor -%}{%- if not ns.has_system -%}{{- '<|start|>system<|message|>You are a helpful AI assistant.' -}}{%- set kc = knowledge_cutoff if knowledge_cutoff is defined and knowledge_cutoff else '2026-01-04' -%}{{- '\nKnowledge cutoff: ' + kc + '.' -}}{%- if current_date is defined and current_date -%}{{- '\nCurrent date: ' + current_date + '.' -}}{%- elif strftime_now is defined -%}{{- '\nCurrent date: ' + strftime_now('%Y-%m-%d') + '.' -}}{%- endif -%}{{- '\n\n' -}}{{- render_reasoning() -}}{%- if tools -%}{{- '\n\n' -}}{{- render_tool_defs(tools) -}}{%- endif -%}{{- '\n\n' -}}{{- render_system_meta(tools) -}}{{- '<|eot|>' -}}{%- endif -%}{%- for message in messages -%}{%- set role = message['role'] -%}{%- set end_token = '<|eom|>' if (not loop.last and messages[loop.index0 + 1]['role'] == role) else '<|eot|>' -%}{%- if role == 'system' -%}{{- '<|start|>system<|message|>' -}}{{- render_content(message['content']) -}}{{- '\n\n' -}}{{- render_reasoning() -}}{%- if tools -%}{{- '\n\n' -}}{{- render_tool_defs(tools) -}}{%- endif -%}{{- '\n\n' -}}{{- render_system_meta(tools) -}}{{- '<|eot|>' -}}{%- elif role == 'user' -%}{{- '<|start|>user<|message|>' -}}{{- render_content(message['content']) -}}{{- '<|eot|>' -}}{%- elif role == 'tool' -%}{%- set tname = message.get('name') -%}{%- if not tname -%}{%- set tcid = message.get('tool_call_id') -%}{%- set rns = namespace(name=tcid if tcid else '') -%}{%- for m in messages -%}{%- if m.get('tool_calls') -%}{%- for tc in m['tool_calls'] -%}{%- if tcid is not none and tc.id is defined and tc.id == tcid -%}{%- set rns.name = tc.function.name -%}{%- endif -%}{%- endfor -%}{%- endif -%}{%- endfor -%}{%- set tname = rns.name -%}{%- endif -%}{{- '<|start|>tool ' + tname + '<|message|>\n' -}}{{- render_content(message['content']) -}}{{- '\n<|eot|>' -}}{%- elif role == 'assistant' -%}{%- if message.get('reasoning_content') -%}{{- '<|start|>assistant to=self<|message|>' + message['reasoning_content'] + '<|eom|>' -}}{%- endif -%}{%- if message.get('tool_calls') -%}{%- for tc in message['tool_calls'] -%}{{- '<|start|>assistant to=' + tc.function.name + '<|message|>' -}}{{- render_atem(tc) -}}{%- if loop.last -%}{{- end_token -}}{%- else -%}{{- '<|eom|>' -}}{%- endif -%}{%- endfor -%}{%- else -%}{%- set recipient = message.get('recipient') or 'user' -%}{%- set end_turn = message.get('end_turn') -%}{%- if end_turn is none -%}{%- set end_turn = not (recipient and recipient != 'user') -%}{%- endif -%}{{- '<|start|>assistant' -}}{%- if recipient -%}{{- ' to=' + recipient -}}{%- endif -%}{{- '<|message|>' -}}{{- render_content(message['content']) -}}{{- ('<|eot|>' if end_turn else '<|eom|>') -}}{%- endif -%}{%- endif -%}{%- endfor -%}{%- if add_generation_prompt -%}{{- '<|start|>assistant' -}}{%- endif -%} \ No newline at end of file diff --git a/server/glimmer_images_test.go b/server/glimmer_images_test.go new file mode 100644 index 000000000..30dcdcf43 --- /dev/null +++ b/server/glimmer_images_test.go @@ -0,0 +1,29 @@ +package server + +import ( + "slices" + "testing" + + "github.com/ollama/ollama/types/model" +) + +func TestGlimmerSafetensorsCapabilities(t *testing.T) { + m := Model{ + Config: model.ConfigV2{ + ModelFormat: "safetensors", + Renderer: "glimmer", + Parser: "glimmer", + Capabilities: []string{"completion", "vision", "audio"}, + }, + } + + want := []model.Capability{ + model.CapabilityCompletion, + model.CapabilityVision, + model.CapabilityTools, + model.CapabilityThinking, + } + if got := m.Capabilities(); !slices.Equal(got, want) { + t.Fatalf("capabilities = %v, want %v", got, want) + } +} diff --git a/server/images.go b/server/images.go index b6ce6a777..e8493b14a 100644 --- a/server/images.go +++ b/server/images.go @@ -468,6 +468,9 @@ func suppressAudioCapability(m *Model, arch string) bool { if isGemma4Renderer(m.Config.Renderer) && m.Config.ModelFormat == "safetensors" { return true } + if m.Config.ModelFormat == "safetensors" && m.Config.Renderer == "glimmer" { + return true + } if arch == "nemotron_h_omni" || m.Config.ModelFamily == "nemotron_h_omni" || diff --git a/server/model_list_cache.go b/server/model_list_cache.go index b110659d4..6775382aa 100644 --- a/server/model_list_cache.go +++ b/server/model_list_cache.go @@ -397,6 +397,14 @@ func buildModelListSummary(name model.Name, mf *manifest.Manifest) (modelListSum }) } + // Mirrors suppressAudioCapability in images.go so /api/tags and /api/show + // agree for safetensors models whose MLX runner serves vision but not audio. + if cfg.ModelFormat == "safetensors" && cfg.Renderer == "glimmer" { + summary.Capabilities = slices.DeleteFunc(summary.Capabilities, func(c model.Capability) bool { + return c == model.CapabilityAudio + }) + } + return summary, nil } diff --git a/x/create/client/create.go b/x/create/client/create.go index 8f6014ce0..c391fdeb0 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -516,13 +516,14 @@ func detectCapabilities(modelDir string) modelCapabilities { ModelType string `json:"model_type"` VisionConfig *map[string]any `json:"vision_config"` AudioConfig *map[string]any `json:"audio_config"` + HasVision bool `json:"has_vision"` } if data, err := os.ReadFile(filepath.Join(modelDir, "config.json")); err == nil { _ = json.Unmarshal(data, &cfg) } return modelCapabilities{ - vision: cfg.VisionConfig != nil, + vision: cfg.VisionConfig != nil || cfg.HasVision, audio: cfg.AudioConfig != nil, thinking: chatTemplateHasThinkingSupport(readChatTemplate(modelDir)) || alwaysSupportsThinking(cfg.Architectures, cfg.ModelType), @@ -617,6 +618,9 @@ func getParserName(modelDir string) string { // Check architectures for known parsers for _, arch := range cfg.Architectures { archLower := strings.ToLower(arch) + if strings.HasPrefix(arch, "MuseGlimmer") { + return "glimmer" + } if strings.Contains(archLower, "laguna") { return lagunaRendererParserName(modelDir) } @@ -643,6 +647,9 @@ func getParserName(modelDir string) string { // Also check model_type if cfg.ModelType != "" { typeLower := strings.ToLower(cfg.ModelType) + if typeLower == "muse_glimmer" { + return "glimmer" + } if strings.Contains(typeLower, "laguna") { return lagunaRendererParserName(modelDir) } @@ -689,6 +696,9 @@ func getRendererName(modelDir string) string { // Check architectures for known renderers for _, arch := range cfg.Architectures { archLower := strings.ToLower(arch) + if strings.HasPrefix(arch, "MuseGlimmer") { + return "glimmer" + } if strings.Contains(archLower, "laguna") { return lagunaRendererParserName(modelDir) } @@ -715,6 +725,9 @@ func getRendererName(modelDir string) string { // Also check model_type if cfg.ModelType != "" { typeLower := strings.ToLower(cfg.ModelType) + if typeLower == "muse_glimmer" { + return "glimmer" + } if strings.Contains(typeLower, "laguna") { return lagunaRendererParserName(modelDir) } diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index e7e5d8b82..50f54d316 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -597,6 +597,11 @@ func TestDetectCapabilities(t *testing.T) { configJSON: `{"architectures": ["Gemma4ForConditionalGeneration"], "vision_config": {}}`, want: modelCapabilities{vision: true}, }, + { + name: "flat vision flag", + configJSON: `{"architectures": ["MuseGlimmerForConditionalGeneration"], "model_type": "muse_glimmer", "has_vision": true}`, + want: modelCapabilities{vision: true}, + }, { name: "audio config", configJSON: `{"architectures": ["Qwen3OmniForConditionalGeneration"], "audio_config": {}}`, @@ -660,6 +665,11 @@ func TestInferSafetensorsCapabilitiesFromParser(t *testing.T) { parserName: "functiongemma", want: []string{"completion", "tools"}, }, + { + name: "glimmer tools and thinking", + parserName: "glimmer", + want: []string{"completion", "tools", "thinking"}, + }, } for _, tt := range tests { @@ -676,6 +686,23 @@ func TestInferSafetensorsCapabilitiesFromParser(t *testing.T) { } } +func TestInferSafetensorsCapabilitiesGlimmerPreservesVisionMetadata(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(`{ + "architectures": ["MuseGlimmerForConditionalGeneration"], + "model_type": "muse_glimmer", + "has_vision": true + }`), 0o644); err != nil { + t.Fatal(err) + } + + got := inferSafetensorsCapabilities(dir, "glimmer") + want := []string{"completion", "vision", "tools", "thinking"} + if !slices.Equal(got, want) { + t.Fatalf("inferSafetensorsCapabilities() = %#v, want %#v", got, want) + } +} + func TestInferSafetensorsCapabilitiesLaguna(t *testing.T) { dir := t.TempDir() if err := os.WriteFile(filepath.Join(dir, "config.json"), []byte(`{"architectures": ["LagunaForCausalLM"], "model_type": "laguna"}`), 0o644); err != nil { @@ -734,6 +761,11 @@ func TestGetParserName(t *testing.T) { configJSON: `{"architectures": ["LagunaForCausalLM"], "model_type": "laguna"}`, want: "laguna", }, + { + name: "glimmer model", + configJSON: `{"architectures": ["MuseGlimmerForConditionalGeneration"], "model_type": "muse_glimmer"}`, + want: "glimmer", + }, { name: "no config", configJSON: `{}`, @@ -789,6 +821,11 @@ func TestGetRendererName(t *testing.T) { configJSON: `{"architectures": ["LagunaForCausalLM"], "model_type": "laguna"}`, want: "laguna", }, + { + name: "glimmer model", + configJSON: `{"architectures": ["MuseGlimmerForConditionalGeneration"], "model_type": "muse_glimmer"}`, + want: "glimmer", + }, } for _, tt := range tests { diff --git a/x/create/create.go b/x/create/create.go index bfa687fa9..8ba6f3753 100644 --- a/x/create/create.go +++ b/x/create/create.go @@ -485,6 +485,7 @@ var tensorImportTransformRegistry = map[string]tensorImportTransformFactory{ "gemma4_unified": newGemma4ImportTransform, "gemma4_unified_text": newGemma4ImportTransform, "LagunaForCausalLM": newLagunaImportTransform, + "MuseGlimmerForConditionalGeneration": newGlimmerImportTransform, "Cohere2MoeForCausalLM": newCohere2MoeImportTransform, "Gemma4AssistantForCausalLM": newGemma4ImportTransform, "Gemma4UnifiedAssistantForCausalLM": newGemma4ImportTransform, diff --git a/x/create/glimmer.go b/x/create/glimmer.go new file mode 100644 index 000000000..8a4f2fe15 --- /dev/null +++ b/x/create/glimmer.go @@ -0,0 +1,80 @@ +package create + +import ( + "encoding/json" + "fmt" + "strings" +) + +type glimmerImportTransform struct { + numLayers int +} + +type glimmerConfig struct { + NumHiddenLayers int `json:"num_hidden_layers"` + TextConfig struct { + NumHiddenLayers int `json:"num_hidden_layers"` + } `json:"text_config"` +} + +func newGlimmerImportTransform(rawConfig json.RawMessage) (quantizePolicy, error) { + var cfg glimmerConfig + if err := json.Unmarshal(rawConfig, &cfg); err != nil { + return nil, fmt.Errorf("glimmer: parse config.json: %w", err) + } + numLayers := cfg.NumHiddenLayers + if numLayers == 0 { + numLayers = cfg.TextConfig.NumHiddenLayers + } + return glimmerImportTransform{numLayers: numLayers}, nil +} + +func (t glimmerImportTransform) quantizationType(name string, shape []int32, quantize string) string { + // Preserve vision tensors at source precision so image text and fine detail + // are not degraded by the language model's quantization policy. + if isGlimmerVisionTensor(name) { + return "" + } + + base := normalizeQuantType(quantize) + if isEmbedTokensWeight(name) { + if e := promoteEmbedding(shape, base); e != "" { + return e + } + if isAligned(shape, base) { + return base + } + return "" + } + + if isGlimmerSensitiveProjection(name) && eightBit(base) != base { + return sensitiveType(t.promoteSensitive(name), shape, base) + } + + return GetTensorQuantization(name, shape, quantize) +} + +func isGlimmerVisionTensor(name string) bool { + return isVision(name) +} + +func isGlimmerSensitiveProjection(name string) bool { + return strings.Contains(name, ".self_attn.q_proj") || + strings.Contains(name, ".self_attn.o_proj") || + strings.Contains(name, ".self_attn.k_proj") || + strings.Contains(name, ".self_attn.v_proj") || + strings.Contains(name, ".self_attn.gate_proj") || + strings.Contains(name, ".self_attn.output_gate_proj") || + strings.Contains(name, ".mlp.down_proj") +} + +func (t glimmerImportTransform) promoteSensitive(name string) bool { + if strings.Contains(name, ".self_attn.q_proj") || + strings.Contains(name, ".self_attn.o_proj") || + strings.Contains(name, ".self_attn.k_proj") || + strings.Contains(name, ".self_attn.v_proj") { + return true + } + layer := layerIndex(name) + return t.numLayers > 0 && layer >= 0 && useMoreBits(layer, t.numLayers) +} diff --git a/x/create/glimmer_test.go b/x/create/glimmer_test.go new file mode 100644 index 000000000..cf2010deb --- /dev/null +++ b/x/create/glimmer_test.go @@ -0,0 +1,125 @@ +package create + +import ( + "encoding/json" + "testing" +) + +func TestGlimmerImportTransformPreservesMultimodalTensors(t *testing.T) { + transform := glimmerImportTransform{numLayers: 52} + shape := []int32{256, 256} + + for _, name := range []string{ + "model.vision_tower.layers.0.attn.q_proj.weight", + "model.vision_adapter.fc1.weight", + "model.vision_projection.weight", + } { + if got := transform.quantizationType(name, shape, "int4"); got != "" { + t.Errorf("quantizationType(%q) = %q, want source precision", name, got) + } + } + + if got := transform.quantizationType("model.language_model.layers.10.self_attn.q_proj.weight", shape, "int4"); got == "" { + t.Fatal("text decoder projection unexpectedly kept at source precision") + } +} + +func TestGlimmerPlanKeepsMultimodalTensors(t *testing.T) { + inv := newInventory(sourceModelConfig{Architectures: []string{"MuseGlimmerForConditionalGeneration"}}, map[string]string{ + "model.language_model.layers.0.self_attn.q_proj.weight": "BF16", + "model.vision_tower.layers.0.attn.q_proj.weight": "BF16", + "model.vision_adapter.fc1.weight": "BF16", + "model.vision_projection.weight": "BF16", + }) + inv.RawConfig = json.RawMessage(`{"text_config":{"num_hidden_layers":52}}`) + + policy, err := newTensorImportTransform(inv) + if err != nil { + t.Fatal(err) + } + specs, err := Plan(inv, Classification{Kind: SourceFloat, Quantize: "nvfp4"}, policy) + if err != nil { + t.Fatal(err) + } + + for _, name := range []string{ + "model.vision_tower.layers.0.attn.q_proj.weight", + "model.vision_adapter.fc1.weight", + "model.vision_projection.weight", + } { + spec, ok := specByName(specs, name) + if !ok { + t.Fatalf("missing multimodal tensor %q in plan", name) + } + if len(spec.Tensors) != 1 || spec.Tensors[0].Quantize != "" { + t.Fatalf("planned tensor %q = %+v, want source precision", name, spec.Tensors) + } + } + + text, ok := specByName(specs, "model.language_model.layers.0.self_attn.q_proj.weight") + if !ok { + t.Fatal("missing text decoder tensor in plan") + } + if len(text.Tensors) != 1 || text.Tensors[0].Quantize == "" { + t.Fatalf("text decoder tensor = %+v, want quantized", text.Tensors) + } +} + +func TestGlimmerQuantizationType(t *testing.T) { + transform := glimmerImportTransform{numLayers: 52} + large := []int32{6656, 6656} + lmHead := []int32{202048, 6656} + ffnDown := []int32{6656, 19968} + + tests := []struct { + name string + tensor string + shape []int32 + quantize string + want string + }{ + {"embed_tokens nvfp4 promotes", "model.language_model.embed_tokens.weight", lmHead, "nvfp4", "mxfp8"}, + {"embed_tokens mxfp8 stays", "model.language_model.embed_tokens.weight", lmHead, "mxfp8", "mxfp8"}, + {"lm_head nvfp4 promotes", "lm_head.weight", lmHead, "nvfp4", "mxfp8"}, + {"lm_head mxfp8 stays", "lm_head.weight", lmHead, "mxfp8", "mxfp8"}, + {"lm_head int4 quantizes", "lm_head.weight", lmHead, "int4", "int8"}, + + {"q_proj nvfp4 promotes", "model.language_model.layers.8.self_attn.q_proj.weight", large, "nvfp4", "mxfp8"}, + {"q_proj nvfp4 always promotes", "model.language_model.layers.6.self_attn.q_proj.weight", large, "nvfp4", "mxfp8"}, + {"o_proj nvfp4 promotes", "model.language_model.layers.8.self_attn.o_proj.weight", large, "nvfp4", "mxfp8"}, + {"o_proj nvfp4 always promotes", "model.language_model.layers.6.self_attn.o_proj.weight", large, "nvfp4", "mxfp8"}, + {"k_proj nvfp4 promotes", "model.language_model.layers.8.self_attn.k_proj.weight", []int32{256, 6656}, "nvfp4", "mxfp8"}, + {"v_proj nvfp4 promotes", "model.language_model.layers.8.self_attn.v_proj.weight", []int32{256, 6656}, "nvfp4", "mxfp8"}, + + {"down_proj nvfp4 first layer promotes", "model.language_model.layers.0.mlp.down_proj.weight", ffnDown, "nvfp4", "mxfp8"}, + {"down_proj nvfp4 non-promoted layer", "model.language_model.layers.6.mlp.down_proj.weight", ffnDown, "nvfp4", "nvfp4"}, + {"down_proj nvfp4 periodic layer promotes", "model.language_model.layers.8.mlp.down_proj.weight", ffnDown, "nvfp4", "mxfp8"}, + {"down_proj mxfp8 stays", "model.language_model.layers.6.mlp.down_proj.weight", ffnDown, "mxfp8", "mxfp8"}, + {"output_gate nvfp4 promoted layer", "model.language_model.layers.0.self_attn.gate_proj.weight", large, "nvfp4", "mxfp8"}, + {"output_gate nvfp4 non-promoted layer", "model.language_model.layers.6.self_attn.gate_proj.weight", large, "nvfp4", "nvfp4"}, + + {"vision projection preserved", "model.vision_projection.weight", large, "nvfp4", ""}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := transform.quantizationType(tt.tensor, tt.shape, tt.quantize) + if got != tt.want { + t.Fatalf("quantizationType(%q, %v, %q) = %q, want %q", tt.tensor, tt.shape, tt.quantize, got, tt.want) + } + }) + } +} + +func TestGlimmerImportTransformRegistered(t *testing.T) { + transform, err := newTensorImportTransform(Inventory{ + Config: sourceModelConfig{Architectures: []string{"MuseGlimmerForConditionalGeneration"}}, + RawConfig: json.RawMessage(`{"text_config":{"num_hidden_layers":52}}`), + }) + if err != nil { + t.Fatal(err) + } + if _, ok := transform.(glimmerImportTransform); !ok { + t.Fatalf("newTensorImportTransform() = %T, want glimmerImportTransform", transform) + } +} diff --git a/x/internal/mlxtest/mlxtest.go b/x/internal/mlxtest/mlxtest.go new file mode 100644 index 000000000..ef006abb2 --- /dev/null +++ b/x/internal/mlxtest/mlxtest.go @@ -0,0 +1,41 @@ +// Package mlxtest provides shared scaffolding for tests that exercise MLX +// through the cgo wrapper in x/mlxrunner/mlx. +package mlxtest + +import ( + "runtime" + "testing" + + "github.com/ollama/ollama/x/mlxrunner/mlx" +) + +// SkipIfUnavailable skips the test when the MLX dynamic library cannot be +// loaded (e.g. no MLX backend built for this platform). +func SkipIfUnavailable(t *testing.T) { + t.Helper() + if err := mlx.CheckInit(); err != nil { + t.Skipf("MLX not available: %v", err) + } +} + +// Setup prepares a test that calls into MLX natively: it skips when MLX is +// unavailable and pins the test goroutine to its OS thread for the duration +// of the test. +// +// The thread pin is load-bearing, not defensive: MLX's default stream cache +// is thread-local, and anything that migrates the goroutine mid-test (the +// race detector's scheduler in particular) otherwise panics with +// "There is no Stream(gpu, 0) in current thread". +// +// Setup deliberately does not switch devices or sweep caches: switching the +// default device re-creates the process-wide default stream, and sweeping the +// allocator cache between tests changes allocator reuse — both perturbed +// tests that share lazy arrays with subtests running on other threads. +func Setup(t *testing.T) { + t.Helper() + + SkipIfUnavailable(t) + + runtime.LockOSThread() + t.Cleanup(runtime.UnlockOSThread) +} diff --git a/x/mlxrunner/imports.go b/x/mlxrunner/imports.go index 7ea16e29f..8761a391c 100644 --- a/x/mlxrunner/imports.go +++ b/x/mlxrunner/imports.go @@ -5,6 +5,7 @@ import ( _ "github.com/ollama/ollama/x/models/dflash" _ "github.com/ollama/ollama/x/models/gemma3" _ "github.com/ollama/ollama/x/models/gemma4" + _ "github.com/ollama/ollama/x/models/glimmer" _ "github.com/ollama/ollama/x/models/glm4_moe_lite" _ "github.com/ollama/ollama/x/models/laguna" _ "github.com/ollama/ollama/x/models/llama" diff --git a/x/mlxrunner/mlx/act.go b/x/mlxrunner/mlx/act.go index 4fa21e270..deb1d5a75 100644 --- a/x/mlxrunner/mlx/act.go +++ b/x/mlxrunner/mlx/act.go @@ -25,6 +25,17 @@ var GELUApprox = Compile1( Shapeless(), ) +func gelu(x *Array) *Array { + dt := x.DType() + half := FromValue[float32](0.5).AsType(dt) + one := FromValue[float32](1).AsType(dt) + invSqrt2 := FromValue(float32(1 / math.Sqrt2)).AsType(dt) + return half.Multiply(x).Multiply(one.Add(erf(x.Multiply(invSqrt2)))) +} + +// GELU returns the exact erf formulation used by torch.nn.functional.gelu. +var GELU = Compile1("GELU", gelu, Shapeless()) + // SiLU returns a * sigmoid(a) as a fused kernel. var SiLU = Compile1( "SiLU", diff --git a/x/mlxrunner/mlx/act_test.go b/x/mlxrunner/mlx/act_test.go new file mode 100644 index 000000000..90d4fc4ba --- /dev/null +++ b/x/mlxrunner/mlx/act_test.go @@ -0,0 +1,99 @@ +package mlx + +import ( + "math" + "testing" + + "github.com/ollama/ollama/x/internal/mlxthread" +) + +func TestGELUCompiledMatchesEager(t *testing.T) { + values := []float32{-6, -2, -0.5, 0, 0.5, 2, 6} + tests := []struct { + name string + dtype DType + tolerance float32 + }{ + {name: "float32", dtype: DTypeFloat32, tolerance: 1e-6}, + {name: "bfloat16", dtype: DTypeBFloat16, tolerance: 1e-2}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + withMLXThread(t, func() { + EnableCompile() + input := FromValues(values, len(values)).AsType(tt.dtype) + Pin(input) + + want := gelu(input) + got := GELU(input) + wantF32 := want.AsType(DTypeFloat32) + gotF32 := got.AsType(DTypeFloat32) + Eval(wantF32, gotF32) + + wantValues := wantF32.Floats() + gotValues := gotF32.Floats() + for i := range wantValues { + if delta := float32(math.Abs(float64(gotValues[i] - wantValues[i]))); delta > tt.tolerance { + t.Fatalf("%s GELU[%d] = %v, want %v (delta %v)", tt.name, i, gotValues[i], wantValues[i], delta) + } + } + Unpin(input) + }) + }) + } +} + +func BenchmarkGELUEager(b *testing.B) { + benchmarkGELU(b, gelu) +} + +func BenchmarkGELUCompiled(b *testing.B) { + benchmarkGELU(b, GELU) +} + +func benchmarkGELU(b *testing.B, fn func(*Array) *Array) { + thread, err := mlxthread.Start("mlx-gelu-benchmark", func() error { + if err := CheckInit(); err != nil { + return err + } + if GPUIsAvailable() { + SetDefaultDeviceGPU() + } + EnableCompile() + return nil + }) + if err != nil { + b.Skipf("MLX not available: %v", err) + } + defer func() { + if err := thread.Stop(b.Context(), func() { + Sweep() + ClearCache() + resetDefaultStreamCache() + }); err != nil { + b.Fatal(err) + } + }() + + if err := thread.Do(b.Context(), func() error { + input := AddScalar(Zeros(DTypeBFloat16, 1, 4096, 8192), 1) + Eval(input) + Pin(input) + defer Unpin(input) + + warmup := fn(input) + Eval(warmup) + Sweep() + + b.ResetTimer() + for range b.N { + output := fn(input) + Eval(output) + Sweep() + } + return nil + }); err != nil { + b.Fatal(err) + } +} diff --git a/x/mlxrunner/mlx/ops_extra.go b/x/mlxrunner/mlx/ops_extra.go index b37322023..585a55dc7 100644 --- a/x/mlxrunner/mlx/ops_extra.go +++ b/x/mlxrunner/mlx/ops_extra.go @@ -491,6 +491,12 @@ func Cos(a *Array) *Array { return out } +func erf(a *Array) *Array { + out := New("ERF") + C.mlx_erf(&out.ctx, a.ctx, DefaultStream().ctx) + return out +} + func Clip(a, aMin, aMax *Array) *Array { out := New("CLIP") C.mlx_clip(&out.ctx, a.ctx, aMin.ctx, aMax.ctx, DefaultStream().ctx) diff --git a/x/mlxrunner/pipeline.go b/x/mlxrunner/pipeline.go index 4e6df5b6b..0cf9e3d14 100644 --- a/x/mlxrunner/pipeline.go +++ b/x/mlxrunner/pipeline.go @@ -181,13 +181,17 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *specu Media: manifest, Layout: media.rowLayout(), }, caches) - spec.committed(chunkIDs, auxHidden, position, manifest) - // Unpin finished items before the sweep: the consuming forward's - // graph retains their data until evaluation, so the buffers die with - // this chunk's eval instead of surviving into the next chunk. - media.release(position + n) + // Report to the drafter only after the chunk's eval: a draft flush + // evaluates, and an eval before the sweep cannot free any buffer the + // chunk's live handles retain — on media chunks, the whole vision tower. + mlx.Pin(chunkIDs, auxHidden) mlx.Sweep() materializeCaches() + spec.committed(chunkIDs, auxHidden, position, manifest) + mlx.Unpin(chunkIDs, auxHidden) + // Released after committed so the drafter can capture rows its + // deferred flush still embeds. + media.release(position + n) processed += n position += n slog.Info("Prompt processing progress", "processed", processed, "total", total) diff --git a/x/models/dflash/dflash.go b/x/models/dflash/dflash.go index 3283a9e7f..0b8f727d0 100644 --- a/x/models/dflash/dflash.go +++ b/x/models/dflash/dflash.go @@ -23,6 +23,9 @@ func init() { base.RegisterDraft("DFlashLagunaForCausalLM", func(root *model.Root, target base.Model) (base.DraftModel, error) { return newModel(root, target, true) }) + base.RegisterDraft("MuseGlimmerAssistantModel", func(root *model.Root, target base.Model) (base.DraftModel, error) { + return newModel(root, target, false) + }) } var _ base.BlockDraft = (*Model)(nil) @@ -44,6 +47,13 @@ type Config struct { VocabSize int32 TargetLayerIDs []int + // RopeInterleaved selects the draft's rotary pairing convention: + // true pairs adjacent dims (torch view_as_complex over pairs, the glimmer + // publisher convention); false pairs split halves (HF rotate_half, the + // laguna convention). Defaults to false for backwards compatibility with + // laguna drafts. + RopeInterleaved bool + // Causal, when set, overrides every layer's attention direction; // otherwise only sliding layers run causal. Causal *bool @@ -148,6 +158,7 @@ func parseConfig(data []byte) (*Config, error) { VocabSize int32 `json:"vocab_size"` LayerTypes []string `json:"layer_types"` SlidingWindow int32 `json:"sliding_window"` + RopeInterleaved *bool `json:"rope_interleaved"` } if err := json.Unmarshal(data, &raw); err != nil { return nil, fmt.Errorf("parse dflash config: %w", err) @@ -168,6 +179,9 @@ func parseConfig(data []byte) (*Config, error) { TargetLayerIDs: raw.DFlashConfig.TargetLayerIDs, Causal: raw.DFlashConfig.Causal, } + if raw.RopeInterleaved != nil { + cfg.RopeInterleaved = *raw.RopeInterleaved + } if cfg.RopeTheta == 0 { cfg.RopeTheta = raw.RopeParameters.RopeTheta } @@ -609,7 +623,7 @@ func (a *Attention) contextKV(hctx *mlx.Array, positions *mlx.Array, cfg *Config k = a.KNorm.Forward(k, cfg.RMSNormEps) k = mlx.Transpose(k, 0, 2, 1, 3) v = mlx.Transpose(v, 0, 2, 1, 3) - k = mlx.RoPEWithBase(k, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions) + k = mlx.RoPEWithBase(k, int(cfg.HeadDim), cfg.RopeInterleaved, cfg.RopeTheta, 1.0, positions) return k, v } @@ -631,8 +645,8 @@ func (a *Attention) Forward(x, ctxK, ctxV *mlx.Array, c cache.Cache, bb *batch.B k = mlx.Transpose(k, 0, 2, 1, 3) v = mlx.Transpose(v, 0, 2, 1, 3) - q = mlx.RoPEWithBase(q, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions) - k = mlx.RoPEWithBase(k, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions) + q = mlx.RoPEWithBase(q, int(cfg.HeadDim), cfg.RopeInterleaved, cfg.RopeTheta, 1.0, positions) + k = mlx.RoPEWithBase(k, int(cfg.HeadDim), cfg.RopeInterleaved, cfg.RopeTheta, 1.0, positions) if ctxK != nil { k = ctxK.Concatenate(2, k) diff --git a/x/models/glimmer/glimmer.go b/x/models/glimmer/glimmer.go new file mode 100644 index 000000000..5098b8811 --- /dev/null +++ b/x/models/glimmer/glimmer.go @@ -0,0 +1,685 @@ +// Package glimmer provides the Meta Glimmer text model implementation for MLX. +package glimmer + +import ( + "encoding/json" + "fmt" + "math" + "slices" + + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/cache" + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/mlxrunner/model" + "github.com/ollama/ollama/x/mlxrunner/model/base" + "github.com/ollama/ollama/x/models/nn" + "github.com/ollama/ollama/x/tokenizer" +) + +func init() { + base.Register("MuseGlimmerForConditionalGeneration", newModel) +} + +var _ base.Model = (*Model)(nil) + +// Config holds the decoder and vision fields from an Glimmer configuration. +type Config struct { + HiddenSize int32 `json:"hidden_size"` + NumHiddenLayers int32 `json:"num_hidden_layers"` + IntermediateSize int32 `json:"intermediate_size"` + NumAttentionHeads int32 `json:"num_attention_heads"` + NumKeyValueHeads int32 `json:"num_key_value_heads"` + HeadDim int32 `json:"head_dim"` + VocabSize int32 `json:"vocab_size"` + RMSNormEps float32 `json:"rms_norm_eps"` + PostNormEps float32 `json:"post_norm_eps"` + RopeTheta float32 `json:"rope_theta"` + MaxPositionEmbeddings int32 `json:"max_position_embeddings"` + SlidingWindow int32 `json:"sliding_window"` + LayerTypes []string `json:"layer_types"` + NoRopeLayers []int32 `json:"no_rope_layers"` + NormalizeTokenEmbeddings bool `json:"normalize_tok_embeddings"` + UseQKNorm bool `json:"use_qk_norm"` + QKScaleFactor float32 `json:"qk_scale_factor"` + UseAttentionOutputGate bool `json:"use_attn_output_gate"` + OutputMultiplier float32 `json:"output_multiplier"` + OutputSoftCapTemp float32 `json:"output_soft_cap_temp"` + TieWordEmbeddings bool `json:"tie_word_embeddings"` + + HasVision bool `json:"has_vision"` + PatchTokenID int32 `json:"patch_token_id"` + VisionAdapterDim int32 `json:"vision_adapter_dim"` + VisionDownsampleFactor int32 `json:"vision_downsample_factor"` + VisionHeads int32 `json:"vision_heads"` + VisionLatentDim int32 `json:"vision_latent_dim"` + VisionLayers int32 `json:"vision_layers"` + VisionMLPRatio float32 `json:"vision_mlp_ratio"` + VisionOutputDim int32 `json:"vision_output_dim"` + VisionPatchSize int32 `json:"vision_patch_size"` + VisionPatchTemporal int32 `json:"vision_patch_temporal"` + VisionPosEmbeddingGridH int32 `json:"vision_pos_emb_grid_h"` + VisionPosEmbeddingGridW int32 `json:"vision_pos_emb_grid_w"` + VisionSparseAttentionFactor int32 `json:"vision_sparse_attention_factor"` + VisionLayerTypes []string `json:"-"` + + // Quantization parameters are populated from the created model metadata. + QuantGroupSize int `json:"-"` + QuantBits int `json:"-"` + QuantMode string `json:"-"` + TensorQuant map[string]*model.TensorQuantInfo `json:"-"` + + // Computed fields. + AttentionScale float32 `json:"-"` + QueryScale float32 `json:"-"` + + NormalizeVisionEmbeddings bool `json:"-"` + VisionRoPEInterleaved bool `json:"-"` + + ImageStartTokenID int32 `json:"-"` + ImageEndTokenID int32 `json:"-"` +} + +type glimmerHFTextConfig struct { + HiddenSize int32 `json:"hidden_size"` + NumHiddenLayers int32 `json:"num_hidden_layers"` + IntermediateSize int32 `json:"intermediate_size"` + NumAttentionHeads int32 `json:"num_attention_heads"` + NumKeyValueHeads int32 `json:"num_key_value_heads"` + HeadDim int32 `json:"head_dim"` + VocabSize int32 `json:"vocab_size"` + RMSNormEps float32 `json:"rms_norm_eps"` + PostNormEps float32 `json:"post_norm_eps"` + MaxPositionEmbeddings int32 `json:"max_position_embeddings"` + SlidingWindow int32 `json:"sliding_window"` + LayerTypes []string `json:"layer_types"` + LayerRoPETheta []float32 `json:"layer_rope_theta"` + QKScaleFactor float32 `json:"qk_scale_factor"` + OutputMultiplier float32 `json:"output_multiplier"` + FinalLogitSoftcapping float32 `json:"final_logit_softcapping"` + TieWordEmbeddings bool `json:"tie_word_embeddings"` + RopeParameters struct { + RopeTheta float32 `json:"rope_theta"` + } `json:"rope_parameters"` +} + +type glimmerHFVisionConfig struct { + HiddenSize int32 `json:"hidden_size"` + IntermediateSize int32 `json:"intermediate_size"` + NumAttentionHeads int32 `json:"num_attention_heads"` + NumHiddenLayers int32 `json:"num_hidden_layers"` + PatchSize int32 `json:"patch_size"` + PatchTemporal int32 `json:"patch_temporal"` + MergeSize int32 `json:"merge_size"` + PosEmbeddingHeight int32 `json:"pos_emb_height"` + PosEmbeddingWidth int32 `json:"pos_emb_width"` + LayerTypes []string `json:"layer_types"` +} + +type glimmerHFConfig struct { + TextConfig *glimmerHFTextConfig `json:"text_config"` + VisionConfig *glimmerHFVisionConfig `json:"vision_config"` + ImageTokenID int32 `json:"image_token_id"` + OutHiddenSize int32 `json:"out_hidden_size"` + ProjectorHiddenSize int32 `json:"projector_hidden_size"` +} + +func configFromHF(hf glimmerHFConfig) Config { + t := hf.TextConfig + cfg := Config{ + HiddenSize: t.HiddenSize, + NumHiddenLayers: t.NumHiddenLayers, + IntermediateSize: t.IntermediateSize, + NumAttentionHeads: t.NumAttentionHeads, + NumKeyValueHeads: t.NumKeyValueHeads, + HeadDim: t.HeadDim, + VocabSize: t.VocabSize, + RMSNormEps: t.RMSNormEps, + PostNormEps: t.PostNormEps, + RopeTheta: t.RopeParameters.RopeTheta, + MaxPositionEmbeddings: t.MaxPositionEmbeddings, + SlidingWindow: t.SlidingWindow, + LayerTypes: t.LayerTypes, + UseQKNorm: true, + QKScaleFactor: t.QKScaleFactor, + UseAttentionOutputGate: true, + OutputMultiplier: t.OutputMultiplier, + OutputSoftCapTemp: t.FinalLogitSoftcapping, + TieWordEmbeddings: t.TieWordEmbeddings, + NormalizeVisionEmbeddings: true, + // The official HF config omits normalize_tok_embeddings because the HF + // checkpoint bakes that norm into the embed table (verified against the + // publisher checkpoint: hf == rms_norm(publisher) to bf16 noise), so the + // HF path must NOT re-normalize. + } + for _, theta := range t.LayerRoPETheta { + if theta == 0 { + cfg.NoRopeLayers = append(cfg.NoRopeLayers, 0) + } else { + cfg.NoRopeLayers = append(cfg.NoRopeLayers, 1) + } + } + + if v := hf.VisionConfig; v != nil { + cfg.HasVision = true + cfg.PatchTokenID = hf.ImageTokenID + cfg.VisionAdapterDim = hf.ProjectorHiddenSize + cfg.VisionDownsampleFactor = v.MergeSize + cfg.VisionHeads = v.NumAttentionHeads + cfg.VisionLatentDim = v.HiddenSize + cfg.VisionLayers = v.NumHiddenLayers + cfg.VisionMLPRatio = float32(v.IntermediateSize) / float32(v.HiddenSize) + cfg.VisionOutputDim = hf.OutHiddenSize + cfg.VisionPatchSize = v.PatchSize + cfg.VisionPatchTemporal = v.PatchTemporal + cfg.VisionPosEmbeddingGridH = v.PosEmbeddingHeight + cfg.VisionPosEmbeddingGridW = v.PosEmbeddingWidth + cfg.VisionLayerTypes = v.LayerTypes + for i, layerType := range v.LayerTypes { + if layerType == "full_attention" { + cfg.VisionSparseAttentionFactor = int32(i + 1) + break + } + } + } + + return cfg +} + +// Model is the Glimmer text decoder with an optional vision encoder. +type Model struct { + EmbedTokens nn.EmbeddingLayer + Layers []*Layer + Norm *nn.RMSNorm + LMHead nn.LinearLayer + + // auxHiddenLayers are the tapped layers for a DFlash draft; empty means + // the final hidden. + auxHiddenLayers []int + + VisionEncoder *VisionEncoder + VisionAdapter *VisionAdapter + VisionProjection nn.LinearLayer + + tok *tokenizer.Tokenizer + + *Config +} + +// Layer is one Glimmer decoder block. +type Layer struct { + Attention *Attention + MLP *MLP + + InputNorm *nn.RMSNorm + PostAttentionNorm *nn.RMSNorm + PostAttnNorm *nn.RMSNorm + PostFFNNorm *nn.RMSNorm + + IsSliding bool + UseRope bool +} + +// Attention implements grouped-query attention with scaleless Q/K norm and an +// output gate applied before the output projection. +type Attention struct { + QProj nn.LinearLayer + KProj nn.LinearLayer + VProj nn.LinearLayer + OProj nn.LinearLayer + OutputGateProj nn.LinearLayer +} + +// MLP is the Glimmer SwiGLU feed-forward block. +type MLP struct { + GateProj nn.LinearLayer + UpProj nn.LinearLayer + DownProj nn.LinearLayer +} + +func parseConfig(data []byte) (Config, error) { + var cfg Config + if err := json.Unmarshal(data, &cfg); err != nil { + return Config{}, fmt.Errorf("parse config: %w", err) + } + + var hf glimmerHFConfig + if err := json.Unmarshal(data, &hf); err != nil { + return Config{}, fmt.Errorf("parse config: %w", err) + } + officialHF := hf.TextConfig != nil + if officialHF { + cfg = configFromHF(hf) + } else { + cfg.NormalizeVisionEmbeddings = cfg.NormalizeTokenEmbeddings + cfg.VisionRoPEInterleaved = true + if cfg.HasVision { + cfg.VisionLayerTypes = make([]string, cfg.VisionLayers) + for i := range cfg.VisionLayerTypes { + cfg.VisionLayerTypes[i] = "window_attention" + if i == len(cfg.VisionLayerTypes)-1 || (cfg.VisionSparseAttentionFactor > 0 && (i+1)%int(cfg.VisionSparseAttentionFactor) == 0) { + cfg.VisionLayerTypes[i] = "full_attention" + } + } + } + } + + if cfg.HiddenSize <= 0 { + return Config{}, fmt.Errorf("invalid hidden_size: %d", cfg.HiddenSize) + } + if cfg.NumHiddenLayers <= 0 { + return Config{}, fmt.Errorf("invalid num_hidden_layers: %d", cfg.NumHiddenLayers) + } + if cfg.IntermediateSize <= 0 { + return Config{}, fmt.Errorf("invalid intermediate_size: %d", cfg.IntermediateSize) + } + if cfg.NumAttentionHeads <= 0 { + return Config{}, fmt.Errorf("invalid num_attention_heads: %d", cfg.NumAttentionHeads) + } + if cfg.NumKeyValueHeads <= 0 { + return Config{}, fmt.Errorf("invalid num_key_value_heads: %d", cfg.NumKeyValueHeads) + } + if cfg.HeadDim <= 0 { + return Config{}, fmt.Errorf("invalid head_dim: %d", cfg.HeadDim) + } + if cfg.VocabSize <= 0 { + return Config{}, fmt.Errorf("invalid vocab_size: %d", cfg.VocabSize) + } + if cfg.MaxPositionEmbeddings <= 0 { + return Config{}, fmt.Errorf("invalid max_position_embeddings: %d", cfg.MaxPositionEmbeddings) + } + if cfg.NumAttentionHeads%cfg.NumKeyValueHeads != 0 { + return Config{}, fmt.Errorf("num_attention_heads (%d) must be divisible by num_key_value_heads (%d)", cfg.NumAttentionHeads, cfg.NumKeyValueHeads) + } + if len(cfg.LayerTypes) != int(cfg.NumHiddenLayers) { + return Config{}, fmt.Errorf("layer_types has %d entries, want %d", len(cfg.LayerTypes), cfg.NumHiddenLayers) + } + if len(cfg.NoRopeLayers) != int(cfg.NumHiddenLayers) { + return Config{}, fmt.Errorf("no_rope_layers has %d entries, want %d", len(cfg.NoRopeLayers), cfg.NumHiddenLayers) + } + hasSlidingLayer := false + for i, layerType := range cfg.LayerTypes { + if layerType != "sliding_attention" && layerType != "full_attention" { + return Config{}, fmt.Errorf("layer_types[%d] has unsupported value %q", i, layerType) + } + hasSlidingLayer = hasSlidingLayer || layerType == "sliding_attention" + if cfg.NoRopeLayers[i] != 0 && cfg.NoRopeLayers[i] != 1 { + return Config{}, fmt.Errorf("no_rope_layers[%d] has unsupported value %d", i, cfg.NoRopeLayers[i]) + } + } + if hasSlidingLayer && cfg.SlidingWindow <= 0 { + return Config{}, fmt.Errorf("invalid sliding_window: %d", cfg.SlidingWindow) + } + if cfg.RMSNormEps == 0 { + cfg.RMSNormEps = 1e-5 + } + if cfg.PostNormEps == 0 { + cfg.PostNormEps = 1e-8 + } + if cfg.RopeTheta == 0 { + cfg.RopeTheta = 500000 + } + if cfg.OutputMultiplier == 0 { + cfg.OutputMultiplier = 1 + } + if cfg.QKScaleFactor == 0 { + cfg.QKScaleFactor = 1 + } + if cfg.HasVision { + if cfg.PatchTokenID <= 0 { + return Config{}, fmt.Errorf("invalid patch_token_id: %d", cfg.PatchTokenID) + } + if cfg.VisionAdapterDim <= 0 || cfg.VisionLatentDim <= 0 || cfg.VisionOutputDim <= 0 { + return Config{}, fmt.Errorf("invalid vision dimensions") + } + if cfg.VisionHeads <= 0 || cfg.VisionLayers <= 0 { + return Config{}, fmt.Errorf("invalid vision architecture") + } + if cfg.VisionLatentDim%cfg.VisionHeads != 0 { + return Config{}, fmt.Errorf("vision_latent_dim (%d) must be divisible by vision_heads (%d)", cfg.VisionLatentDim, cfg.VisionHeads) + } + if cfg.VisionPatchSize <= 0 || cfg.VisionPatchTemporal <= 0 || cfg.VisionDownsampleFactor <= 0 { + return Config{}, fmt.Errorf("invalid vision patch configuration") + } + if cfg.VisionPosEmbeddingGridH <= 0 || cfg.VisionPosEmbeddingGridW <= 0 { + return Config{}, fmt.Errorf("invalid vision position embedding grid") + } + if cfg.VisionSparseAttentionFactor <= 0 { + cfg.VisionSparseAttentionFactor = 1 + } + if len(cfg.VisionLayerTypes) != int(cfg.VisionLayers) { + return Config{}, fmt.Errorf("vision layer_types has %d entries, want %d", len(cfg.VisionLayerTypes), cfg.VisionLayers) + } + for i, layerType := range cfg.VisionLayerTypes { + if layerType != "window_attention" && layerType != "full_attention" { + return Config{}, fmt.Errorf("vision layer_types[%d] has unsupported value %q", i, layerType) + } + } + } + + cfg.AttentionScale = float32(1 / math.Sqrt(float64(cfg.HeadDim))) + cfg.QueryScale = cfg.QKScaleFactor + if !officialHF { + cfg.QueryScale = float32(float64(cfg.QKScaleFactor) / math.Sqrt(float64(cfg.HeadDim))) + } + return cfg, nil +} + +func newModel(root *model.Root) (base.Model, error) { + configData, err := root.Manifest.ReadConfig("config.json") + if err != nil { + return nil, fmt.Errorf("load config: %w", err) + } + + cfg, err := parseConfig(configData) + if err != nil { + return nil, err + } + + if qt := root.QuantType(); qt != "" { + cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode = model.QuantizationParams(qt) + if gs := root.GroupSize(); gs > 0 { + cfg.QuantGroupSize = gs + } + } else { + cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode = model.QuantizationParams("") + } + cfg.TensorQuant = root.AllTensorQuant() + + tokData, err := root.Manifest.ReadConfig("tokenizer.json") + if err != nil { + return nil, fmt.Errorf("load tokenizer config: %w", err) + } + tokConfig := &tokenizer.TokenizerConfig{ConfigJSON: configData} + if data, err := root.Manifest.ReadConfig("generation_config.json"); err == nil { + tokConfig.GenerationConfigJSON = data + } + if data, err := root.Manifest.ReadConfig("tokenizer_config.json"); err == nil { + tokConfig.TokenizerConfigJSON = data + } + + tok, err := tokenizer.LoadFromBytesWithConfig(tokData, tokConfig) + if err != nil { + return nil, fmt.Errorf("parse tokenizer: %w", err) + } + for _, token := range []string{"<|start|>", "<|message|>", "<|eom|>", "<|eot|>", "<|end_of_text|>"} { + if _, ok := tok.GetSpecialToken(token); !ok { + return nil, fmt.Errorf("missing special token %q", token) + } + } + if err := validateTokenizer(tok); err != nil { + return nil, err + } + if cfg.HasVision { + var ok bool + if _, ok = tok.GetSpecialToken("<|image|>"); !ok { + return nil, fmt.Errorf("missing image sentinel token") + } + if cfg.ImageStartTokenID, ok = tok.GetSpecialToken("<|image_start|>"); !ok { + return nil, fmt.Errorf("missing image start token") + } + if cfg.ImageEndTokenID, ok = tok.GetSpecialToken("<|image_end|>"); !ok { + return nil, fmt.Errorf("missing image end token") + } + } + + m := &Model{ + Layers: make([]*Layer, cfg.NumHiddenLayers), + Config: &cfg, + tok: tok, + } + for i := range cfg.NumHiddenLayers { + m.Layers[i] = &Layer{ + IsSliding: cfg.LayerTypes[i] == "sliding_attention", + // Despite its historical name, 1 means RoPE and 0 means NoPE. + UseRope: cfg.NoRopeLayers[i] == 1, + } + } + return m, nil +} + +func validateTokenizer(tok *tokenizer.Tokenizer) error { + eom, _ := tok.GetSpecialToken("<|eom|>") + if tok.IsEOS(eom) { + return fmt.Errorf("invalid generation config: <|eom|> must be a non-terminal message boundary") + } + + for _, token := range []string{"<|eot|>", "<|end_of_text|>"} { + id, _ := tok.GetSpecialToken(token) + if !tok.IsEOS(id) { + return fmt.Errorf("invalid generation config: %s must be an EOS token", token) + } + } + return nil +} + +func shiftedRMSNorm(weight *mlx.Array, eps float32) *nn.RMSNorm { + if weight == nil { + return nil + } + return nn.NewRMSNorm(mlx.AddScalar(weight, 1), eps) +} + +func firstLinear(linears model.LinearFactory, paths ...string) nn.LinearLayer { + for _, path := range paths { + if linear := linears.Make(path); linear != nil { + return linear + } + } + return nil +} + +func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { + linears := model.NewLinearFactory(tensors, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + officialHF := tensors["model.language_model.embed_tokens.weight"] != nil + textPrefix := "model" + if officialHF { + textPrefix = "model.language_model" + } + + m.EmbedTokens = model.MakeEmbeddingLayer(tensors, textPrefix+".embed_tokens", m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + if m.EmbedTokens == nil { + return fmt.Errorf("missing embedding weight: %s.embed_tokens.weight", textPrefix) + } + + normWeight := tensors[textPrefix+".norm.weight"] + if normWeight == nil { + return fmt.Errorf("missing final norm weight: %s.norm.weight", textPrefix) + } + m.Norm = nn.NewRMSNorm(normWeight, m.RMSNormEps) + + if m.TieWordEmbeddings { + m.LMHead = m.EmbedTokens.AsLinear() + } else { + m.LMHead = linears.Make("lm_head") + } + if m.LMHead == nil { + return fmt.Errorf("missing language model head: lm_head.weight") + } + + for i := range m.NumHiddenLayers { + prefix := fmt.Sprintf("%s.layers.%d", textPrefix, i) + layer := m.Layers[i] + layer.Attention = &Attention{ + QProj: linears.Make(prefix + ".self_attn.q_proj"), + KProj: linears.Make(prefix + ".self_attn.k_proj"), + VProj: linears.Make(prefix + ".self_attn.v_proj"), + OProj: linears.Make(prefix + ".self_attn.o_proj"), + OutputGateProj: firstLinear(linears, prefix+".self_attn.gate_proj", prefix+".self_attn.output_gate_proj"), + } + layer.MLP = &MLP{ + GateProj: linears.Make(prefix + ".mlp.gate_proj"), + UpProj: linears.Make(prefix + ".mlp.up_proj"), + DownProj: linears.Make(prefix + ".mlp.down_proj"), + } + layer.InputNorm = shiftedRMSNorm(tensors[prefix+".input_layernorm.weight"], m.RMSNormEps) + if officialHF { + layer.PostAttentionNorm = shiftedRMSNorm(tensors[prefix+".pre_feedforward_layernorm.weight"], m.RMSNormEps) + layer.PostAttnNorm = shiftedRMSNorm(tensors[prefix+".post_attention_layernorm.weight"], m.PostNormEps) + layer.PostFFNNorm = shiftedRMSNorm(tensors[prefix+".post_feedforward_layernorm.weight"], m.PostNormEps) + } else { + layer.PostAttentionNorm = shiftedRMSNorm(tensors[prefix+".post_attention_layernorm.weight"], m.RMSNormEps) + layer.PostAttnNorm = shiftedRMSNorm(tensors[prefix+".post_attn_norm.weight"], m.PostNormEps) + layer.PostFFNNorm = shiftedRMSNorm(tensors[prefix+".post_ffn_norm.weight"], m.PostNormEps) + } + + if layer.InputNorm == nil || layer.PostAttentionNorm == nil || layer.PostAttnNorm == nil || layer.PostFFNNorm == nil { + return fmt.Errorf("layer %d: missing normalization weight", i) + } + if layer.Attention.QProj == nil || layer.Attention.KProj == nil || layer.Attention.VProj == nil || layer.Attention.OProj == nil { + return fmt.Errorf("layer %d: missing attention projection", i) + } + if m.UseAttentionOutputGate && layer.Attention.OutputGateProj == nil { + return fmt.Errorf("layer %d: missing attention output gate projection", i) + } + if layer.MLP.GateProj == nil || layer.MLP.UpProj == nil || layer.MLP.DownProj == nil { + return fmt.Errorf("layer %d: missing MLP projection", i) + } + } + + if m.HasVision { + if err := m.loadVisionWeights(tensors, linears); err != nil { + return err + } + } + + return nil +} + +func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden *mlx.Array) { + h := m.embed(b.InputIDs) + if len(b.Media) > 0 { + h = m.scatterMedia(h, b) + } + return m.forwardEmbeddings(h, b, caches) +} + +func (m *Model) embed(inputIDs *mlx.Array) *mlx.Array { + h := m.EmbedTokens.Forward(inputIDs) + if m.NormalizeTokenEmbeddings { + h = mlx.RMSNormFn(h, nil, m.RMSNormEps) + } + return h +} + +func (m *Model) forwardEmbeddings(h *mlx.Array, b *batch.Batch, caches []cache.Cache) (hidden, auxHidden *mlx.Array) { + dims := b.InputIDs.Dims() + B, L := int32(dims[0]), int32(dims[1]) + positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets)) + + var features []*mlx.Array + for i, layer := range m.Layers { + var c cache.Cache + if caches != nil && i < len(caches) { + c = caches[i] + } + h = layer.Forward(h, b, c, positions, B, L, m.Config) + if slices.Contains(m.auxHiddenLayers, i) { + features = append(features, h) + } + } + + out := m.Norm.Forward(h, m.RMSNormEps) + if features != nil { + return out, mlx.Concatenate(features, -1) + } + return out, out +} + +// SetAuxHiddenLayers taps the listed layers' outputs, which Forward then +// returns as the draft-conditioning state in place of the final hidden. +func (m *Model) SetAuxHiddenLayers(layers []int) { + m.auxHiddenLayers = layers +} + +// TokenEmbeddings is the raw table lookup, for a draft that embeds with the +// target's table. Note the model's optional embedding normalization does not +// apply here; the draft mirrors what the target was trained with. +func (m *Model) TokenEmbeddings(ids *mlx.Array) *mlx.Array { + return m.EmbedTokens.Forward(ids) +} + +// RawLogits is the raw head projection, skipping the float32 cast, output +// multiplier and softcap the target's own Unembed applies. +func (m *Model) RawLogits(hidden *mlx.Array) *mlx.Array { + return m.LMHead.Forward(hidden) +} + +func (m *Model) Unembed(x *mlx.Array) *mlx.Array { + logits := m.LMHead.Forward(x).AsType(mlx.DTypeFloat32) + if m.OutputMultiplier != 1 { + logits = mlx.MulScalar(logits, m.OutputMultiplier) + } + if m.OutputSoftCapTemp > 0 { + cap := mlx.FromValue(m.OutputSoftCapTemp).AsType(logits.DType()) + logits = mlx.LogitSoftcap(logits, cap) + } + return logits +} + +func (m *Model) NumLayers() int { return len(m.Layers) } +func (m *Model) MaxContextLength() int { return int(m.MaxPositionEmbeddings) } +func (m *Model) Tokenizer() *tokenizer.Tokenizer { return m.tok } + +// NewCaches uses bounded rotating caches for sliding layers and full KV caches +// for the periodic NoPE/global-attention layers. +func (m *Model) NewCaches() []cache.Cache { + caches := make([]cache.Cache, len(m.Layers)) + for i, layer := range m.Layers { + if layer.IsSliding { + caches[i] = cache.NewRotatingKVCache(int(m.SlidingWindow)) + } else { + caches[i] = cache.NewKVCache() + } + } + return caches +} + +func (l *Layer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { + attnInput := l.InputNorm.Forward(x, cfg.RMSNormEps) + attnOut := l.Attention.Forward(attnInput, b, c, positions, B, L, l.UseRope, cfg) + h := mlx.Add(x, l.PostAttnNorm.Forward(attnOut, cfg.PostNormEps)) + + mlpInput := l.PostAttentionNorm.Forward(h, cfg.RMSNormEps) + mlpOut := l.MLP.Forward(mlpInput) + return mlx.Add(h, l.PostFFNNorm.Forward(mlpOut, cfg.PostNormEps)) +} + +func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, useRope bool, cfg *Config) *mlx.Array { + q := mlx.Reshape(a.QProj.Forward(x), B, L, cfg.NumAttentionHeads, cfg.HeadDim) + k := mlx.Reshape(a.KProj.Forward(x), B, L, cfg.NumKeyValueHeads, cfg.HeadDim) + v := mlx.Reshape(a.VProj.Forward(x), B, L, cfg.NumKeyValueHeads, cfg.HeadDim) + + if cfg.UseQKNorm { + q = mlx.MulScalar(mlx.RMSNormFn(q, nil, cfg.RMSNormEps), cfg.QueryScale) + k = mlx.RMSNormFn(k, nil, cfg.RMSNormEps) + } + + q = mlx.Transpose(q, 0, 2, 1, 3) + k = mlx.Transpose(k, 0, 2, 1, 3) + v = mlx.Transpose(v, 0, 2, 1, 3) + if useRope { + q = mlx.RoPEWithBase(q, int(cfg.HeadDim), true, cfg.RopeTheta, 1, positions) + k = mlx.RoPEWithBase(k, int(cfg.HeadDim), true, cfg.RopeTheta, 1, positions) + } + + var kv nn.SDPAOption + if c != nil { + kv = nn.WithKVHistory(c.(cache.Attention).Update(b, k, v)) + } else { + kv = nn.WithKV(k, v, b.SeqQueryLens) + } + out := nn.ScaledDotProductAttention(b, q, cfg.AttentionScale, kv, nn.WithMask(nn.CausalMask())) + out = mlx.Transpose(out, 0, 2, 1, 3) + + if cfg.UseAttentionOutputGate { + gate := mlx.Reshape(a.OutputGateProj.Forward(x), B, L, cfg.NumAttentionHeads, cfg.HeadDim) + out = mlx.Mul(out, mlx.Sigmoid(gate)) + } + + out = mlx.Reshape(out, B, L, cfg.NumAttentionHeads*cfg.HeadDim) + return a.OProj.Forward(out) +} + +func (m *MLP) Forward(x *mlx.Array) *mlx.Array { + return m.DownProj.Forward(mlx.SwiGLU(m.GateProj.Forward(x), m.UpProj.Forward(x))) +} diff --git a/x/models/glimmer/glimmer_test.go b/x/models/glimmer/glimmer_test.go new file mode 100644 index 000000000..a5676d67c --- /dev/null +++ b/x/models/glimmer/glimmer_test.go @@ -0,0 +1,336 @@ +package glimmer + +import ( + "encoding/json" + "fmt" + "math" + "slices" + "strings" + "testing" + + "github.com/ollama/ollama/x/internal/mlxtest" + "github.com/ollama/ollama/x/mlxrunner/cache" + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/models/nn" + "github.com/ollama/ollama/x/tokenizer" +) + +func testConfig() Config { + return Config{ + HiddenSize: 6656, + NumHiddenLayers: 4, + IntermediateSize: 19968, + NumAttentionHeads: 32, + NumKeyValueHeads: 2, + HeadDim: 128, + VocabSize: 202048, + RMSNormEps: 1e-5, + PostNormEps: 1e-8, + RopeTheta: 500000, + MaxPositionEmbeddings: 16384, + SlidingWindow: 2048, + LayerTypes: []string{"sliding_attention", "sliding_attention", "sliding_attention", "full_attention"}, + NoRopeLayers: []int32{1, 1, 1, 0}, + NormalizeTokenEmbeddings: true, + UseQKNorm: true, + QKScaleFactor: 43.7840518911, + UseAttentionOutputGate: true, + OutputMultiplier: 0.19611613, + OutputSoftCapTemp: 20, + } +} + +func marshalConfig(t *testing.T, cfg Config) []byte { + t.Helper() + data, err := json.Marshal(cfg) + if err != nil { + t.Fatal(err) + } + return data +} + +func TestParseConfig(t *testing.T) { + cfg, err := parseConfig(marshalConfig(t, testConfig())) + if err != nil { + t.Fatal(err) + } + + if got, want := cfg.AttentionScale, float32(1/math.Sqrt(128)); math.Abs(float64(got-want)) > 1e-7 { + t.Errorf("AttentionScale = %g, want %g", got, want) + } + if got, want := cfg.QueryScale, float32(float64(cfg.QKScaleFactor)/math.Sqrt(128)); math.Abs(float64(got-want)) > 1e-7 { + t.Errorf("QueryScale = %g, want %g", got, want) + } +} + +func TestParseOfficialHFConfig(t *testing.T) { + data := []byte(`{ + "text_config": { + "hidden_size": 6656, + "num_hidden_layers": 4, + "intermediate_size": 19968, + "num_attention_heads": 32, + "num_key_value_heads": 2, + "head_dim": 128, + "vocab_size": 202048, + "rms_norm_eps": 0.00001, + "post_norm_eps": 1e-8, + "max_position_embeddings": 131072, + "sliding_window": 2048, + "layer_types": ["sliding_attention", "sliding_attention", "sliding_attention", "full_attention"], + "layer_rope_theta": [500000, 500000, 500000, 0], + "qk_scale_factor": 3.87, + "output_multiplier": 0.19611613513818404, + "final_logit_softcapping": 20, + "rope_parameters": {"rope_theta": 500000} + }, + "vision_config": { + "hidden_size": 1536, + "intermediate_size": 8960, + "num_attention_heads": 16, + "num_hidden_layers": 4, + "patch_size": 14, + "patch_temporal": 2, + "merge_size": 2, + "pos_emb_height": 32, + "pos_emb_width": 32, + "layer_types": ["window_attention", "window_attention", "window_attention", "full_attention"] + }, + "image_token_id": 200092, + "out_hidden_size": 6144, + "projector_hidden_size": 4096 + }`) + + cfg, err := parseConfig(data) + if err != nil { + t.Fatal(err) + } + if cfg.QueryScale != 3.87 { + t.Errorf("QueryScale = %g, want 3.87", cfg.QueryScale) + } + if cfg.MaxPositionEmbeddings != 131072 { + t.Errorf("MaxPositionEmbeddings = %d, want 131072", cfg.MaxPositionEmbeddings) + } + if !slices.Equal(cfg.NoRopeLayers, []int32{1, 1, 1, 0}) { + t.Errorf("NoRopeLayers = %v, want [1 1 1 0]", cfg.NoRopeLayers) + } + if cfg.NormalizeTokenEmbeddings { + t.Error("NormalizeTokenEmbeddings = true, want false for pre-normalized HF embeddings") + } + if !cfg.NormalizeVisionEmbeddings { + t.Error("NormalizeVisionEmbeddings = false, want true") + } + if cfg.VisionRoPEInterleaved { + t.Error("VisionRoPEInterleaved = true, want half-rotation layout") + } + if cfg.VisionSparseAttentionFactor != 4 || cfg.VisionOutputDim != 6144 || cfg.VisionAdapterDim != 4096 { + t.Errorf("vision config = factor %d, output %d, adapter %d", cfg.VisionSparseAttentionFactor, cfg.VisionOutputDim, cfg.VisionAdapterDim) + } +} + +func TestValidateTokenizerEOS(t *testing.T) { + const tokenizerJSON = `{ + "model": { + "type": "BPE", + "vocab": {"x": 0}, + "merges": [] + }, + "added_tokens": [ + {"id": 1, "content": "<|start|>", "special": true}, + {"id": 2, "content": "<|message|>", "special": true}, + {"id": 3, "content": "<|eom|>", "special": true}, + {"id": 4, "content": "<|eot|>", "special": true}, + {"id": 5, "content": "<|end_of_text|>", "special": true} + ] + }` + + tests := []struct { + name string + eos string + want string + }{ + {name: "official boundaries", eos: `[5, 4]`}, + {name: "message boundary is terminal", eos: `[5, 3, 4]`, want: "<|eom|> must be a non-terminal"}, + {name: "end of turn is not terminal", eos: `[5]`, want: "<|eot|> must be an EOS token"}, + {name: "end of text is not terminal", eos: `[4]`, want: "<|end_of_text|> must be an EOS token"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + tok, err := tokenizer.LoadFromBytesWithConfig([]byte(tokenizerJSON), &tokenizer.TokenizerConfig{ + GenerationConfigJSON: []byte(`{"eos_token_id":` + tt.eos + `}`), + }) + if err != nil { + t.Fatal(err) + } + + err = validateTokenizer(tok) + if tt.want == "" { + if err != nil { + t.Fatal(err) + } + return + } + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("validateTokenizer() error = %v, want containing %q", err, tt.want) + } + }) + } +} + +func TestParseConfigRejectsInvalidSchedules(t *testing.T) { + tests := []struct { + name string + mutate func(*Config) + want string + }{ + {"layer count", func(cfg *Config) { cfg.LayerTypes = cfg.LayerTypes[:3] }, "layer_types has 3 entries"}, + {"layer type", func(cfg *Config) { cfg.LayerTypes[0] = "linear_attention" }, `unsupported value "linear_attention"`}, + {"RoPE flag", func(cfg *Config) { cfg.NoRopeLayers[0] = 2 }, "no_rope_layers[0] has unsupported value 2"}, + {"sliding window", func(cfg *Config) { cfg.SlidingWindow = 0 }, "invalid sliding_window: 0"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cfg := testConfig() + tt.mutate(&cfg) + _, err := parseConfig(marshalConfig(t, cfg)) + if err == nil || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("parseConfig() error = %v, want containing %q", err, tt.want) + } + }) + } +} + +func TestNewCachesMatchesAttentionSchedule(t *testing.T) { + cfg := testConfig() + m := Model{ + Config: &cfg, + Layers: []*Layer{ + {IsSliding: true}, + {IsSliding: true}, + {IsSliding: true}, + {IsSliding: false}, + }, + } + + caches := m.NewCaches() + for i := range 3 { + if _, ok := caches[i].(*cache.RotatingKVCache); !ok { + t.Errorf("cache %d = %T, want *cache.RotatingKVCache", i, caches[i]) + } + } + if _, ok := caches[3].(*cache.KVCache); !ok { + t.Errorf("cache 3 = %T, want *cache.KVCache", caches[3]) + } +} + +func TestUnembedReturnsFloat32Logits(t *testing.T) { + mlxtest.Setup(t) + + input := mlx.FromValues([]float32{1}, 1, 1, 1).AsType(mlx.DTypeBFloat16) + weight := mlx.FromValues([]float32{1, 2}, 2, 1).AsType(mlx.DTypeBFloat16) + m := Model{ + LMHead: nn.NewLinear(weight, nil), + Config: &Config{ + OutputMultiplier: 0.19611613, + OutputSoftCapTemp: 20, + }, + } + + if got := m.Unembed(input).DType(); got != mlx.DTypeFloat32 { + t.Fatalf("Unembed() dtype = %v, want %v", got, mlx.DTypeFloat32) + } +} + +func TestComputeImageSizeMatchesReference(t *testing.T) { + tests := []struct { + width, height int + targetWidth, targetHeight int + tokens int + }{ + {1, 1, 28, 28, 1}, + {100, 100, 112, 112, 16}, + {640, 480, 644, 476, 391}, + {480, 640, 476, 644, 391}, + {1920, 1080, 1932, 1092, 2691}, + {1080, 1920, 1092, 1932, 2691}, + {4000, 4000, 1792, 1792, 4096}, + {8192, 512, 7168, 448, 4096}, + {512, 8192, 448, 7168, 4096}, + {1234, 987, 1260, 1008, 1620}, + } + + for _, tt := range tests { + t.Run(fmt.Sprintf("%dx%d", tt.width, tt.height), func(t *testing.T) { + width, height, tokens := computeImageSize(tt.width, tt.height, 28, maxImageTokens) + if width != tt.targetWidth || height != tt.targetHeight || tokens != tt.tokens { + t.Fatalf("computeImageSize(%d, %d) = (%d, %d, %d), want (%d, %d, %d)", + tt.width, tt.height, width, height, tokens, + tt.targetWidth, tt.targetHeight, tt.tokens) + } + }) + } +} + +func TestComputeImageSizeDeterministicTieBreak(t *testing.T) { + width, height, tokens := computeImageSize(128, 128, 28, maxImageTokens) + if width != 140 || height != 140 || tokens != 25 { + t.Fatalf("computeImageSize() = %dx%d (%d tokens), want 140x140 (25 tokens)", width, height, tokens) + } +} + +func TestSparseVisionPermutation(t *testing.T) { + permutation, lengths := sparseVisionPermutation(3, 5, 2, 3) + wantPermutation := []int32{0, 1, 2, 5, 6, 7, 3, 4, 8, 9, 10, 11, 12, 13, 14} + wantLengths := []int{6, 4, 3, 2} + if !slices.Equal(permutation, wantPermutation) { + t.Fatalf("permutation = %v, want %v", permutation, wantPermutation) + } + if !slices.Equal(lengths, wantLengths) { + t.Fatalf("lengths = %v, want %v", lengths, wantLengths) + } +} + +func TestApplyVisionRoPELayouts(t *testing.T) { + tests := []struct { + name string + interleaved bool + want []float32 + }{ + { + name: "official half rotation", + want: []float32{ + 1*0.5 - 3*0.75, + 2*0.25 - 4*0.125, + 3*0.5 + 1*0.75, + 4*0.25 + 2*0.125, + }, + }, + { + name: "legacy interleaved rotation", + interleaved: true, + want: []float32{ + 1*0.5 - 2*0.75, + 1*0.75 + 2*0.5, + 3*0.25 - 4*0.125, + 3*0.125 + 4*0.25, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mlxtest.Setup(t) + + x := mlx.FromValues([]float32{1, 2, 3, 4}, 1, 1, 1, 4) + cos := mlx.FromValues([]float32{0.5, 0.25}, 1, 2) + sin := mlx.FromValues([]float32{0.75, 0.125}, 1, 2) + got := applyVisionRoPE(x, cos, sin, tt.interleaved) + mlx.Eval(got) + if !slices.Equal(got.Floats(), tt.want) { + t.Fatalf("applyVisionRoPE() = %v, want %v", got.Floats(), tt.want) + } + }) + } +} diff --git a/x/models/glimmer/media.go b/x/models/glimmer/media.go new file mode 100644 index 000000000..97ea6f2d4 --- /dev/null +++ b/x/models/glimmer/media.go @@ -0,0 +1,199 @@ +package glimmer + +import ( + "bytes" + "fmt" + "image" + _ "image/gif" + _ "image/jpeg" + _ "image/png" + "math" + + "golang.org/x/image/draw" + + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/mlxrunner/model/base" +) + +// The reference processor uses the full 4096-token budget for still images. +// Reducing it here discards source detail and disproportionately hurts OCR. +const maxImageTokens = 4096 + +var _ base.MediaModel = (*Model)(nil) + +// preparedImage is glimmer's model-private media state: the patch grid the +// encoder consumes and the soft-token count of its pixel-shuffled output. +type preparedImage struct { + gridH int + gridW int + outputTokens int +} + +// PrepareMedia implements base.MediaModel: splice each image segment's +// placeholder expansion — image_start, the patch-token run, image_end — into +// the stream, decoding, resizing, and patchifying the image on the CPU. +func (m *Model) PrepareMedia(segments []base.Segment) (*base.PreparedRequest, error) { + prepared := &base.PreparedRequest{} + for s, seg := range segments { + if seg.Data == nil { + prepared.Tokens = append(prepared.Tokens, seg.Tokens...) + continue + } + if !m.HasVision || m.VisionEncoder == nil { + return nil, fmt.Errorf("this model does not support %s input", seg.Kind) + } + if seg.Kind != "image" { + return nil, fmt.Errorf("glimmer does not support %s input", seg.Kind) + } + + patches, geom, err := m.preprocessImage(seg.Data) + if err != nil { + return nil, fmt.Errorf("preprocess image: %w", err) + } + + start := len(prepared.Tokens) + prepared.Tokens = append(prepared.Tokens, m.ImageStartTokenID) + for range geom.outputTokens { + prepared.Tokens = append(prepared.Tokens, m.PatchTokenID) + } + prepared.Tokens = append(prepared.Tokens, m.ImageEndTokenID) + + n := geom.gridH * geom.gridW + prepared.Items = append(prepared.Items, base.PreparedItem{ + Range: [2]int{start, len(prepared.Tokens)}, + Source: s, + MediaData: patches, + Dims: []int{1, n, len(patches) / n}, + Opaque: geom, + // Patch rows attend causally in the text stack (there is no + // relaxed image-span mask), so chunked prefill may split the run. + Causal: true, + }) + } + return prepared, nil +} + +// EncodeMedia implements base.MediaModel: run the vision tower over one +// prepared image, returning the lazy [outputTokens, hidden] features. +func (m *Model) EncodeMedia(item *base.PreparedItem, data *mlx.Array) *mlx.Array { + geom := item.Opaque.(preparedImage) + return m.encodeVision(data, geom.gridH, geom.gridW) +} + +// softRun returns a media item's feature-bearing token range: the expansion +// is image_start + patch*N + image_end, so the run starts one past the splice. +func softRun(item batch.MediaItem) (start, end int) { + geom := item.Opaque.(preparedImage) + return item.Pos + 1, item.Pos + 1 + geom.outputTokens +} + +// scatterMedia overwrites the patch-token rows this chunk covers with the +// item's projected features. +func (m *Model) scatterMedia(h *mlx.Array, b *batch.Batch) *mlx.Array { + for _, item := range b.Media { + if item.Features == nil { + continue + } + start, end := softRun(item) + off := int(b.SeqOffsets[item.Seq]) + qLo := max(start, off) + qHi := min(end, off+int(b.SeqQueryLens[item.Seq])) + if qHi <= qLo { + continue + } + + feat := item.Features.Slice(mlx.Slice(qLo-start, qHi-start), mlx.Slice()) + feat = mlx.Reshape(feat.AsType(h.DType()), 1, int32(qHi-qLo), m.HiddenSize) + h = h.SliceUpdate(feat, mlx.Slice(item.Seq, item.Seq+1), mlx.Slice(qLo-off, qHi-off), mlx.Slice()) + } + return h +} + +func computeImageSize(width, height, patchStride, maxTokens int) (targetWidth, targetHeight, tokens int) { + if width <= 0 || height <= 0 || patchStride <= 0 || maxTokens <= 0 { + return 0, 0, 0 + } + + patchRows := float64(height) / float64(patchStride) + patchCols := float64(width) / float64(patchStride) + ratio := patchCols / patchRows + if patchRows*patchCols > float64(maxTokens) { + patchRows = math.Sqrt(float64(maxTokens) / ratio) + patchCols = patchRows * ratio + } + + rows := []int{int(math.Floor(patchRows)), int(math.Ceil(patchRows))} + cols := []int{int(math.Floor(patchCols)), int(math.Ceil(patchCols))} + bestRows, bestCols := 0, 0 + bestDelta := math.Inf(1) + seen := make(map[[2]int]bool, 4) + for _, r := range rows { + for _, c := range cols { + pair := [2]int{r, c} + if seen[pair] || r < 1 || c < 1 || r*c > maxTokens { + continue + } + seen[pair] = true + delta := math.Abs(float64(r)/float64(c) - float64(height)/float64(width)) + // The reference deduplicates candidates with a Python set, leaving + // exact ties implementation-dependent. Prefer the larger grid to + // preserve source detail deterministically. + if delta < bestDelta || (delta == bestDelta && r*c > bestRows*bestCols) { + bestRows, bestCols, bestDelta = r, c, delta + } + } + } + if bestRows == 0 { + bestRows = max(1, int(math.Round(patchRows))) + bestCols = max(1, int(math.Round(patchCols))) + } + + return bestCols * patchStride, bestRows * patchStride, bestRows * bestCols +} + +// preprocessImage decodes and prepares one image on the CPU: budgeted +// resize, [-1,1] rescale, and patchify to the tower's layout — one row per +// patch, (temporal, RGB, pixel row, pixel column) within it. +func (m *Model) preprocessImage(data []byte) ([]float32, preparedImage, error) { + src, _, err := image.Decode(bytes.NewReader(data)) + if err != nil { + return nil, preparedImage{}, fmt.Errorf("decode: %w", err) + } + + bounds := src.Bounds() + patchStride := int(m.VisionPatchSize * m.VisionDownsampleFactor) + targetW, targetH, outputTokens := computeImageSize(bounds.Dx(), bounds.Dy(), patchStride, maxImageTokens) + if targetW == 0 || targetH == 0 || outputTokens == 0 { + return nil, preparedImage{}, fmt.Errorf("invalid image dimensions %dx%d", bounds.Dx(), bounds.Dy()) + } + + resized := image.NewNRGBA(image.Rect(0, 0, targetW, targetH)) + draw.CatmullRom.Scale(resized, resized.Bounds(), src, bounds, draw.Src, nil) + + patchSize := int(m.VisionPatchSize) + temporal := int(m.VisionPatchTemporal) + gridH, gridW := targetH/patchSize, targetW/patchSize + patchDim := temporal * 3 * patchSize * patchSize + patches := make([]float32, gridH*gridW*patchDim) + + at := 0 + for patchY := range gridH { + for patchX := range gridW { + for range temporal { + for channel := range 3 { + for y := range patchSize { + row := (patchY*patchSize+y)*resized.Stride + patchX*patchSize*4 + for x := range patchSize { + value := resized.Pix[row+x*4+channel] + patches[at] = 2*float32(value)/255 - 1 + at++ + } + } + } + } + } + } + + return patches, preparedImage{gridH: gridH, gridW: gridW, outputTokens: outputTokens}, nil +} diff --git a/x/models/glimmer/media_test.go b/x/models/glimmer/media_test.go new file mode 100644 index 000000000..380fcf0de --- /dev/null +++ b/x/models/glimmer/media_test.go @@ -0,0 +1,98 @@ +package glimmer + +import ( + "bytes" + "image" + "image/png" + "testing" + + "github.com/ollama/ollama/x/mlxrunner/model/base" +) + +func testVisionModel() *Model { + return &Model{ + VisionEncoder: &VisionEncoder{}, + Config: &Config{ + HasVision: true, + PatchTokenID: 7, + ImageStartTokenID: 8, + ImageEndTokenID: 9, + VisionPatchSize: 14, + VisionPatchTemporal: 1, + VisionDownsampleFactor: 2, + }, + } +} + +func testPNG(t *testing.T, w, h int) []byte { + t.Helper() + var buf bytes.Buffer + if err := png.Encode(&buf, image.NewNRGBA(image.Rect(0, 0, w, h))); err != nil { + t.Fatal(err) + } + return buf.Bytes() +} + +func TestPrepareMediaSplicesExpansion(t *testing.T) { + m := testVisionModel() + prepared, err := m.PrepareMedia([]base.Segment{ + {Tokens: []int32{1, 2}}, + {Kind: "image", Data: testPNG(t, 56, 56)}, + {Tokens: []int32{3}}, + }) + if err != nil { + t.Fatal(err) + } + + // 56x56 at patch stride 28 (14*2) is a 2x2 downsampled grid: 4 output + // tokens over a 4x4 patch grid. + wantTokens := []int32{1, 2, 8, 7, 7, 7, 7, 9, 3} + if len(prepared.Tokens) != len(wantTokens) { + t.Fatalf("tokens = %v, want %v", prepared.Tokens, wantTokens) + } + for i, tok := range wantTokens { + if prepared.Tokens[i] != tok { + t.Fatalf("tokens = %v, want %v", prepared.Tokens, wantTokens) + } + } + + if len(prepared.Items) != 1 { + t.Fatalf("items = %d, want 1", len(prepared.Items)) + } + item := prepared.Items[0] + if item.Range != [2]int{2, 8} { + t.Fatalf("range = %v, want [2 8]", item.Range) + } + if item.Source != 1 { + t.Fatalf("source = %d, want 1", item.Source) + } + if !item.Causal { + t.Fatal("expected causal expansion") + } + + geom := item.Opaque.(preparedImage) + if geom.gridH != 4 || geom.gridW != 4 || geom.outputTokens != 4 { + t.Fatalf("geometry = %+v, want 4x4 grid with 4 output tokens", geom) + } + want := 1 + for _, d := range item.Dims { + want *= d + } + if len(item.MediaData) != want { + t.Fatalf("media data length %d does not match dims %v", len(item.MediaData), item.Dims) + } +} + +func TestPrepareMediaRejectsUnsupportedKind(t *testing.T) { + m := testVisionModel() + _, err := m.PrepareMedia([]base.Segment{{Kind: "audio", Data: []byte{1}}}) + if err == nil { + t.Fatal("expected error for audio input") + } + + text := &Model{Config: &Config{}} + _, err = text.PrepareMedia([]base.Segment{{Kind: "image", Data: []byte{1}}}) + if err == nil { + t.Fatal("expected error for text-only model") + } +} diff --git a/x/models/glimmer/vision.go b/x/models/glimmer/vision.go new file mode 100644 index 000000000..737e8e6f5 --- /dev/null +++ b/x/models/glimmer/vision.go @@ -0,0 +1,411 @@ +package glimmer + +import ( + "fmt" + "math" + + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/mlxrunner/model" + "github.com/ollama/ollama/x/models/nn" +) + +type VisionEncoder struct { + Conv1Linear nn.LinearLayer + PositionalEmbeddingVLM *mlx.Array + LNPre *nn.LayerNorm + Layers []*VisionBlock + LNPost *nn.LayerNorm +} + +type VisionBlock struct { + LN1 *nn.LayerNorm + Attention *VisionAttention + LN2 *nn.LayerNorm + MLP *VisionMLP +} + +type VisionAttention struct { + QProj nn.LinearLayer + KProj nn.LinearLayer + VProj nn.LinearLayer + OProj nn.LinearLayer + InterleavedRoPE bool +} + +type VisionMLP struct { + FC nn.LinearLayer + Proj nn.LinearLayer +} + +type VisionAdapter struct { + FC nn.LinearLayer + Proj nn.LinearLayer +} + +func (m *Model) loadVisionWeights(tensors map[string]*mlx.Array, linears model.LinearFactory) error { + officialHF := tensors["model.vision_tower.patch_embedder.position_embedding_table.weight"] != nil + encoderPrefix := "model.vision_encoder." + convPath := encoderPrefix + "conv1_linear" + positionPath := encoderPrefix + "positional_embedding_vlm" + if officialHF { + encoderPrefix = "model.vision_tower." + convPath = encoderPrefix + "patch_embedder.patch_embedding" + positionPath = encoderPrefix + "patch_embedder.position_embedding_table.weight" + } + encoder := &VisionEncoder{ + Conv1Linear: linears.Make(convPath), + PositionalEmbeddingVLM: tensors[positionPath], + LNPre: &nn.LayerNorm{ + Weight: tensors[encoderPrefix+"ln_pre.weight"], + Bias: tensors[encoderPrefix+"ln_pre.bias"], + Eps: 1e-5, + }, + Layers: make([]*VisionBlock, m.VisionLayers), + LNPost: &nn.LayerNorm{ + Weight: tensors[encoderPrefix+"ln_post.weight"], + Bias: tensors[encoderPrefix+"ln_post.bias"], + Eps: 1e-5, + }, + } + if encoder.Conv1Linear == nil || encoder.PositionalEmbeddingVLM == nil || + encoder.LNPre.Weight == nil || encoder.LNPre.Bias == nil || + encoder.LNPost.Weight == nil || encoder.LNPost.Bias == nil { + return fmt.Errorf("missing vision encoder input or normalization weights") + } + + for i := range m.VisionLayers { + prefix := fmt.Sprintf("%stransformer.%d.", encoderPrefix, i) + ln1, ln2 := "ln_1", "ln_2" + attnQ, attnK, attnV, attnO := "attn.q_proj", "attn.k_proj", "attn.v_proj", "attn.o_proj" + mlpFC, mlpProj := "mlp.c_fc", "mlp.c_proj" + if officialHF { + prefix = fmt.Sprintf("%slayers.%d.", encoderPrefix, i) + ln1, ln2 = "norm1", "norm2" + attnQ, attnK, attnV, attnO = "attn.q_proj", "attn.k_proj", "attn.v_proj", "attn.proj" + mlpFC, mlpProj = "mlp.fc1", "mlp.fc2" + } + block := &VisionBlock{ + LN1: &nn.LayerNorm{ + Weight: tensors[prefix+ln1+".weight"], + Bias: tensors[prefix+ln1+".bias"], + Eps: 1e-5, + }, + Attention: &VisionAttention{ + QProj: linears.Make(prefix + attnQ), + KProj: linears.Make(prefix + attnK), + VProj: linears.Make(prefix + attnV), + OProj: linears.Make(prefix + attnO), + InterleavedRoPE: m.VisionRoPEInterleaved, + }, + LN2: &nn.LayerNorm{ + Weight: tensors[prefix+ln2+".weight"], + Bias: tensors[prefix+ln2+".bias"], + Eps: 1e-5, + }, + MLP: &VisionMLP{ + FC: linears.Make(prefix + mlpFC), + Proj: linears.Make(prefix + mlpProj), + }, + } + if block.LN1.Weight == nil || block.LN1.Bias == nil || + block.LN2.Weight == nil || block.LN2.Bias == nil || + block.Attention.QProj == nil || block.Attention.KProj == nil || + block.Attention.VProj == nil || block.Attention.OProj == nil || + block.MLP.FC == nil || block.MLP.Proj == nil { + return fmt.Errorf("vision layer %d: missing weights", i) + } + encoder.Layers[i] = block + } + + adapterFC, adapterProj := "model.vision_adapter.c_fc", "model.vision_adapter.c_proj" + if officialHF { + adapterFC, adapterProj = "model.vision_adapter.fc1", "model.vision_adapter.fc2" + } + adapter := &VisionAdapter{FC: linears.Make(adapterFC), Proj: linears.Make(adapterProj)} + projection := linears.Make("model.vision_projection") + if adapter.FC == nil || adapter.Proj == nil || projection == nil { + return fmt.Errorf("missing vision adapter or projection weights") + } + + m.VisionEncoder = encoder + m.VisionAdapter = adapter + m.VisionProjection = projection + return nil +} + +// encodeVision runs the vision tower over one uploaded image, returning the +// lazy [outputTokens, hidden] features. Graph construction only, per the +// base.MediaModel contract — the consuming forward's evaluation pulls the +// encoder, and its intermediates are transient within that single eval. +func (m *Model) encodeVision(patches *mlx.Array, gridH, gridW int) *mlx.Array { + encoder := m.VisionEncoder + inputCount := gridH * gridW + permutation, sparseLens := sparseVisionPermutation( + gridH, + gridW, + int(m.VisionPosEmbeddingGridH), + int(m.VisionPosEmbeddingGridW), + ) + positional := interpolateVisionPositions( + encoder.PositionalEmbeddingVLM, + gridH, + gridW, + int(m.VisionPosEmbeddingGridH), + int(m.VisionPosEmbeddingGridW), + ) + cos, sin := visionRoPE(gridH, gridW, int(m.VisionLatentDim/m.VisionHeads)) + + x := encoder.Conv1Linear.Forward(patches.AsType(encoder.PositionalEmbeddingVLM.DType())) + x = mlx.Add(x, positional.ExpandDims(0)) + x = encoder.LNPre.Forward(x) + + var sparseMask *mlx.Array + if m.VisionSparseAttentionFactor > 1 { + indices := mlx.FromValues(permutation, len(permutation)) + x = mlx.Take(x, indices, 1) + cos = mlx.Take(cos, indices, 0) + sin = mlx.Take(sin, indices, 0) + sparseMask = visionBlockDiagonalMask(sparseLens, x.DType()) + } + + for i, layer := range encoder.Layers { + // Full-attention layers attend across the whole single image: no mask. + mask := sparseMask + if m.VisionLayerTypes[i] == "full_attention" { + mask = nil + } + x = layer.Forward(x, cos, sin, mask, int(m.VisionHeads)) + } + + if m.VisionSparseAttentionFactor > 1 { + restore := make([]int32, inputCount) + for reordered, original := range permutation { + restore[original] = int32(reordered) + } + x = mlx.Take(x, mlx.FromValues(restore, len(restore)), 1) + } + + x = encoder.LNPost.Forward(x) + features := pixelShuffleVision(x.Squeeze(0), gridH, gridW, int(m.VisionDownsampleFactor)) + features = m.VisionAdapter.Forward(features) + features = m.VisionProjection.Forward(features) + if m.NormalizeVisionEmbeddings { + features = mlx.RMSNormFn(features, nil, m.RMSNormEps) + } + return features +} + +func (b *VisionBlock) Forward(x, cos, sin, mask *mlx.Array, heads int) *mlx.Array { + h := b.LN1.Forward(x) + h = b.Attention.Forward(h, cos, sin, mask, heads) + x = mlx.Add(x, h) + h = b.LN2.Forward(x) + return mlx.Add(x, b.MLP.Forward(h)) +} + +func (a *VisionAttention) Forward(x, cos, sin, mask *mlx.Array, heads int) *mlx.Array { + shape := x.Dims() + batchSize, sequence, width := int32(shape[0]), int32(shape[1]), int32(shape[2]) + headDim := width / int32(heads) + + q := mlx.Reshape(a.QProj.Forward(x), batchSize, sequence, int32(heads), headDim) + k := mlx.Reshape(a.KProj.Forward(x), batchSize, sequence, int32(heads), headDim) + v := mlx.Reshape(a.VProj.Forward(x), batchSize, sequence, int32(heads), headDim) + q = applyVisionRoPE(q, cos, sin, a.InterleavedRoPE) + k = applyVisionRoPE(k, cos, sin, a.InterleavedRoPE) + q = mlx.Transpose(q, 0, 2, 1, 3) + k = mlx.Transpose(k, 0, 2, 1, 3) + v = mlx.Transpose(v, 0, 2, 1, 3) + + out := mlx.FastScaledDotProductAttention(q, k, v, float32(1/math.Sqrt(float64(headDim))), "", mask) + out = mlx.Transpose(out, 0, 2, 1, 3) + return a.OProj.Forward(mlx.Reshape(out, batchSize, sequence, width)) +} + +func (m *VisionMLP) Forward(x *mlx.Array) *mlx.Array { + return m.Proj.Forward(mlx.GELU(m.FC.Forward(x))) +} + +func (a *VisionAdapter) Forward(x *mlx.Array) *mlx.Array { + return mlx.GELU(a.Proj.Forward(mlx.GELU(a.FC.Forward(x)))) +} + +func sparseVisionPermutation(gridH, gridW, windowH, windowW int) ([]int32, []int) { + if windowH <= 0 || windowW <= 0 { + permutation := make([]int32, gridH*gridW) + for i := range permutation { + permutation[i] = int32(i) + } + return permutation, []int{len(permutation)} + } + + var permutation []int32 + var lengths []int + for top := 0; top < gridH; top += windowH { + for left := 0; left < gridW; left += windowW { + before := len(permutation) + for y := top; y < min(top+windowH, gridH); y++ { + for x := left; x < min(left+windowW, gridW); x++ { + permutation = append(permutation, int32(y*gridW+x)) + } + } + lengths = append(lengths, len(permutation)-before) + } + } + return permutation, lengths +} + +func visionBlockDiagonalMask(lengths []int, dtype mlx.DType) *mlx.Array { + if len(lengths) <= 1 { + return nil + } + total := 0 + for _, length := range lengths { + total += length + } + if total == 0 { + return nil + } + // Expand a per-block matrix by each patch's block index, so the + // quadratic tensor is built on the device. + blocks := make([]float32, len(lengths)*len(lengths)) + for i := range blocks { + blocks[i] = float32(math.Inf(-1)) + } + ids := make([]int32, total) + offset := 0 + for b, length := range lengths { + blocks[b*len(lengths)+b] = 0 + for i := offset; i < offset+length; i++ { + ids[i] = int32(b) + } + offset += length + } + index := mlx.FromValues(ids, total) + m := mlx.Take(mlx.FromValues(blocks, len(lengths), len(lengths)), index, 0) + m = mlx.Take(m, index, 1) + return mlx.Reshape(m, 1, 1, int32(total), int32(total)).AsType(dtype) +} + +func interpolateVisionPositions(table *mlx.Array, gridH, gridW, tableH, tableW int) *mlx.Array { + count := gridH * gridW + indices := [4][]int32{ + make([]int32, count), + make([]int32, count), + make([]int32, count), + make([]int32, count), + } + weights := [4][]float32{ + make([]float32, count), + make([]float32, count), + make([]float32, count), + make([]float32, count), + } + + at := 0 + // This order and coordinate assignment intentionally match the reference's + // meshgrid(ys, xs, indexing="xy") call. + for x := range gridW { + for y := range gridH { + sourceX := (float64(y)+0.5)*float64(tableW)/float64(gridH) - 0.5 + sourceY := (float64(x)+0.5)*float64(tableH)/float64(gridW) - 0.5 + x0, y0 := int(math.Floor(sourceX)), int(math.Floor(sourceY)) + dx, dy := float32(sourceX-float64(x0)), float32(sourceY-float64(y0)) + neighbors := [4]struct { + x, y int + weight float32 + }{ + {x0, y0, (1 - dx) * (1 - dy)}, + {x0 + 1, y0, dx * (1 - dy)}, + {x0, y0 + 1, (1 - dx) * dy}, + {x0 + 1, y0 + 1, dx * dy}, + } + for i, neighbor := range neighbors { + if neighbor.x >= 0 && neighbor.x < tableW && neighbor.y >= 0 && neighbor.y < tableH { + indices[i][at] = int32(neighbor.y*tableW + neighbor.x) + weights[i][at] = neighbor.weight + } + } + at++ + } + } + + var result *mlx.Array + for i := range indices { + selected := mlx.Take(table, mlx.FromValues(indices[i], len(indices[i])), 0) + weight := mlx.FromValues(weights[i], len(weights[i]), 1).AsType(table.DType()) + term := mlx.Mul(selected, weight) + if result == nil { + result = term + } else { + result = mlx.Add(result, term) + } + } + return result +} + +func visionRoPE(gridH, gridW, headDim int) (*mlx.Array, *mlx.Array) { + half := headDim / 2 + quarter := half / 2 + cosValues := make([]float32, gridH*gridW*half) + sinValues := make([]float32, len(cosValues)) + for y := range gridH { + for x := range gridW { + base := (y*gridW + x) * half + for i := range quarter { + frequency := math.Pow(10000, -2*float64(i)/float64(half)) + widthAngle := float64(x+1) * frequency + heightAngle := float64(y+1) * frequency + cosValues[base+i] = float32(math.Cos(widthAngle)) + sinValues[base+i] = float32(math.Sin(widthAngle)) + cosValues[base+quarter+i] = float32(math.Cos(heightAngle)) + sinValues[base+quarter+i] = float32(math.Sin(heightAngle)) + } + } + } + return mlx.FromValues(cosValues, gridH*gridW, half), mlx.FromValues(sinValues, gridH*gridW, half) +} + +func applyVisionRoPE(x, cos, sin *mlx.Array, interleaved bool) *mlx.Array { + shape := x.Dims() + batchSize, sequence, heads, headDim := int32(shape[0]), int32(shape[1]), int32(shape[2]), int32(shape[3]) + dtype := x.DType() + if !interleaved { + x = x.AsType(mlx.DTypeFloat32) + cos = mlx.Concatenate([]*mlx.Array{cos, cos}, 1).AsType(mlx.DTypeFloat32).ExpandDims(0).ExpandDims(2) + sin = mlx.Concatenate([]*mlx.Array{sin, sin}, 1).AsType(mlx.DTypeFloat32).ExpandDims(0).ExpandDims(2) + half := int(headDim / 2) + first := x.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(0, half)) + second := x.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(half, int(headDim))) + rotated := mlx.Concatenate([]*mlx.Array{mlx.Neg(second), first}, 3) + return mlx.Add(mlx.Mul(x, cos), mlx.Mul(rotated, sin)).AsType(dtype) + } + + pairs := mlx.Reshape(x.AsType(mlx.DTypeFloat32), batchSize, sequence, heads, headDim/2, 2) + real := pairs.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(0)).Squeeze(4) + imag := pairs.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(), mlx.Slice(1)).Squeeze(4) + cos = cos.AsType(mlx.DTypeFloat32).ExpandDims(0).ExpandDims(2) + sin = sin.AsType(mlx.DTypeFloat32).ExpandDims(0).ExpandDims(2) + outReal := mlx.Sub(mlx.Mul(real, cos), mlx.Mul(imag, sin)) + outImag := mlx.Add(mlx.Mul(real, sin), mlx.Mul(imag, cos)) + return mlx.Reshape(mlx.Stack([]*mlx.Array{outReal, outImag}, 4), batchSize, sequence, heads, headDim).AsType(dtype) +} + +func pixelShuffleVision(x *mlx.Array, gridH, gridW, factor int) *mlx.Array { + outputH, outputW := gridH/factor, gridW/factor + permutation := make([]int32, 0, gridH*gridW) + for y := range outputH { + for x := range outputW { + for innerY := range factor { + for innerX := range factor { + permutation = append(permutation, int32((y*factor+innerY)*gridW+x*factor+innerX)) + } + } + } + } + x = mlx.Take(x, mlx.FromValues(permutation, len(permutation)), 0) + width := int32(x.Dim(1)) + x = mlx.Reshape(x, int32(outputH*outputW), int32(factor*factor), width) + x = mlx.Transpose(x, 0, 2, 1) + return mlx.Reshape(x, int32(outputH*outputW), width*int32(factor*factor)) +} diff --git a/x/tokenizer/tokenizer.go b/x/tokenizer/tokenizer.go index 465b5f9a2..b4c8a1c0e 100644 --- a/x/tokenizer/tokenizer.go +++ b/x/tokenizer/tokenizer.go @@ -35,11 +35,12 @@ type Vocabulary struct { // Tokenizer handles BPE and SentencePiece tokenization type Tokenizer struct { - vocab *Vocabulary - pretokenizer *regexp.Regexp - specialTokens map[string]int32 // Special tokens for direct lookup - sortedSpecialTokens []string // Special tokens sorted by length, longest first - typ TokenizerType // Algorithm type + vocab *Vocabulary + pretokenizer *regexp.Regexp + pretokenizerSpaceBeforePunctuation bool + specialTokens map[string]int32 // Special tokens for direct lookup + sortedSpecialTokens []string // Special tokens sorted by length, longest first + typ TokenizerType // Algorithm type } // Precomputed GPT-2 byte-level encoding table diff --git a/x/tokenizer/tokenizer_correctness_test.go b/x/tokenizer/tokenizer_correctness_test.go index 91adc167d..842b01103 100644 --- a/x/tokenizer/tokenizer_correctness_test.go +++ b/x/tokenizer/tokenizer_correctness_test.go @@ -111,6 +111,38 @@ func TestSplitBySpecialTokensFallbackWithoutCache(t *testing.T) { } } +func TestAdjustWhitespaceBoundary(t *testing.T) { + tests := []struct { + name string + part string + boundary int + spaceBeforePunct bool + want int + }{ + {name: "letter", part: " word", boundary: 2, want: 1}, + {name: "punctuation without optional prefix", part: " }", boundary: 2, want: 2}, + {name: "punctuation with optional prefix", part: " }", boundary: 2, spaceBeforePunct: true, want: 1}, + {name: "tab before letter", part: "\tword", boundary: 1, want: 0}, + {name: "tab before mark", part: "\t\u0301", boundary: 1, spaceBeforePunct: true, want: 1}, + {name: "tab before punctuation", part: "\t}", boundary: 1, spaceBeforePunct: true, want: 1}, + {name: "number", part: " 1", boundary: 2, want: 2}, + {name: "newline", part: " \nword", boundary: 2, want: 2}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + curr := tokenMatch{end: tt.boundary} + next := tokenMatch{start: tt.boundary, end: len(tt.part)} + + adjustWhitespaceBoundary(tt.part, &curr, &next, tt.spaceBeforePunct) + + if curr.end != tt.want || next.start != tt.want { + t.Fatalf("boundary = (%d, %d), want (%d, %d)", curr.end, next.start, tt.want, tt.want) + } + }) + } +} + func TestEncodeDeterministicAcrossGOMAXPROCS(t *testing.T) { tok := benchmarkLoadMiniLlama(t) diff --git a/x/tokenizer/tokenizer_encode.go b/x/tokenizer/tokenizer_encode.go index 73e7a2f1f..5ee658083 100644 --- a/x/tokenizer/tokenizer_encode.go +++ b/x/tokenizer/tokenizer_encode.go @@ -89,7 +89,7 @@ func (t *Tokenizer) splitBySpecialTokens(s string) []string { return result } -func adjustWhitespaceBoundary(part string, curr, next *tokenMatch) { +func adjustWhitespaceBoundary(part string, curr, next *tokenMatch, spaceBeforePunct bool) { m := part[curr.start:curr.end] nextText := part[next.start:next.end] @@ -98,7 +98,8 @@ func adjustWhitespaceBoundary(part string, curr, next *tokenMatch) { } firstRune, _ := utf8.DecodeRuneInString(nextText) - if !unicode.IsLetter(firstRune) { + shiftASCIIOnly := !unicode.IsLetter(firstRune) + if shiftASCIIOnly && (!spaceBeforePunct || unicode.IsNumber(firstRune) || unicode.IsSpace(firstRune)) { return } @@ -106,6 +107,9 @@ func adjustWhitespaceBoundary(part string, curr, next *tokenMatch) { for j := curr.end; j > curr.start; { r, size := utf8.DecodeLastRuneInString(part[curr.start:j]) if unicode.IsSpace(r) { + if shiftASCIIOnly && r != ' ' { + return + } lastSpaceStart = j - size break } @@ -153,7 +157,7 @@ func (t *Tokenizer) forEachPartChunk(part string, fn func(encodeChunk)) { next := tokenMatch{start: offset + loc[0], end: offset + loc[1]} offset += loc[1] - adjustWhitespaceBoundary(part, &curr, &next) + adjustWhitespaceBoundary(part, &curr, &next, t.pretokenizerSpaceBeforePunctuation) if curr.end > curr.start { fn(encodeChunk{text: part[curr.start:curr.end], isSpecial: false}) diff --git a/x/tokenizer/tokenizer_load.go b/x/tokenizer/tokenizer_load.go index 03767ee22..c0b6bc25e 100644 --- a/x/tokenizer/tokenizer_load.go +++ b/x/tokenizer/tokenizer_load.go @@ -152,6 +152,7 @@ func loadFromTokenizerJSON(data []byte) (*Tokenizer, error) { return nil, fmt.Errorf("failed to compile pretokenizer regex %q: %w", pattern, err) } t.pretokenizer = re + t.pretokenizerSpaceBeforePunctuation = strings.Contains(pattern, ` ?[^\s\p{L}\p{N}]`) } cacheSortedSpecialTokens(t) diff --git a/x/tokenizer/tokenizer_load_test.go b/x/tokenizer/tokenizer_load_test.go index 54e5023c1..c0d52e1a4 100644 --- a/x/tokenizer/tokenizer_load_test.go +++ b/x/tokenizer/tokenizer_load_test.go @@ -1,6 +1,7 @@ package tokenizer import ( + "encoding/json" "strings" "testing" ) @@ -50,3 +51,56 @@ func TestExtractPretokenizerSkipsUnsupportedSequenceSplit(t *testing.T) { t.Fatalf("selected unsupported newline splitter: %q", pattern) } } + +func TestLoadPretokenizerOptionalPunctuationSpace(t *testing.T) { + tests := []struct { + name string + pattern string + want []string + }{ + { + name: "o200k optional space", + pattern: ` ?[^\s\p{L}\p{N}]+[\r\n/]*|\s+(?!\S)|\s+`, + want: []string{" ", " }\n"}, + }, + { + name: "punctuation without optional space", + pattern: `[^\s\p{L}\p{N}]+[\r\n/]*|\s+(?!\S)|\s+`, + want: []string{" ", "}\n"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, err := json.Marshal(map[string]any{ + "model": map[string]any{ + "type": "BPE", + "vocab": map[string]int{"}": 0}, + "merges": []string{}, + }, + "pre_tokenizer": map[string]any{ + "type": "Split", + "pattern": map[string]string{ + "Regex": tt.pattern, + }, + }, + }) + if err != nil { + t.Fatal(err) + } + + tok, err := LoadFromBytes(data) + if err != nil { + t.Fatal(err) + } + + var got []string + tok.forEachPartChunk(" }\n", func(chunk encodeChunk) { + got = append(got, chunk.text) + }) + if strings.Join(got, "\x00") != strings.Join(tt.want, "\x00") { + t.Fatalf("chunks = %q, want %q", got, tt.want) + } + }) + } +}