diff --git a/model/parsers/cohere.go b/model/parsers/cohere.go new file mode 100644 index 00000000..54f7216a --- /dev/null +++ b/model/parsers/cohere.go @@ -0,0 +1,315 @@ +package parsers + +import ( + "encoding/json" + "log/slog" + "strings" + "unicode" + + "github.com/ollama/ollama/api" +) + +// CohereParser parses output from Cohere North / Command A 2026 models +// (e.g. North-Mini-Code-1.0). The generation prompt ends with +// <|START_THINKING|> (reasoning on) or <|START_THINKING|><|END_THINKING|> +// (reasoning off), so output begins inside the thinking block when reasoning +// is enabled. After thinking, the model emits either +// <|START_TEXT|>content<|END_TEXT|> or an <|START_ACTION|>[...]<|END_ACTION|> +// tool call array, then <|END_OF_TURN_TOKEN|>. +type CohereParser struct { + state cohereParserState + buffer strings.Builder + callIndex int +} + +type cohereParserState int + +const ( + cohereCollectingThinking cohereParserState = iota + cohereAwaitingBlock + cohereCollectingContent + cohereCollectingAction +) + +const ( + cohereEndThinking = "<|END_THINKING|>" + cohereStartText = "<|START_TEXT|>" + cohereEndText = "<|END_TEXT|>" + cohereStartAction = "<|START_ACTION|>" + cohereEndAction = "<|END_ACTION|>" + cohereEndOfTurn = "<|END_OF_TURN_TOKEN|>" + + // Legacy response markers from the older Command A chat template, which + // also ships in these models' tokenizer_config and shows up in sampled + // output; treated as aliases for START_TEXT / END_TEXT. + cohereStartResponse = "<|START_RESPONSE|>" + cohereEndResponse = "<|END_RESPONSE|>" +) + +func (p *CohereParser) HasToolSupport() bool { + return true +} + +func (p *CohereParser) HasThinkingSupport() bool { + return true +} + +func (p *CohereParser) PreservedTokens() []string { + return []string{ + "<|START_THINKING|>", cohereEndThinking, + cohereStartText, cohereEndText, + cohereStartAction, cohereEndAction, + cohereStartResponse, cohereEndResponse, + } +} + +func (p *CohereParser) Init(tools []api.Tool, lastMessage *api.Message, thinkValue *api.ThinkValue) []api.Tool { + p.buffer.Reset() + p.callIndex = 0 + + // The template enables reasoning by default; nil means default. + thinkingEnabled := thinkValue == nil || thinkValue.Bool() + + assistantPrefill := lastMessage != nil && lastMessage.Role == "assistant" && lastMessage.Content != "" + switch { + case assistantPrefill: + // The prompt left an open <|START_TEXT|> for continuation. + p.state = cohereCollectingContent + case thinkingEnabled: + p.state = cohereCollectingThinking + default: + p.state = cohereAwaitingBlock + } + return tools +} + +func (p *CohereParser) Add(s string, done bool) (content string, thinking string, calls []api.ToolCall, err error) { + p.buffer.WriteString(s) + + var contentSb, thinkingSb strings.Builder + for { + c, t, tc, more := p.eat(done) + contentSb.WriteString(c) + thinkingSb.WriteString(t) + calls = append(calls, tc...) + if !more { + break + } + } + + for i := range calls { + calls[i].Function.Index = p.callIndex + p.callIndex++ + } + + return contentSb.String(), thinkingSb.String(), calls, nil +} + +// eat consumes what it can from the buffer for the current state. It returns +// more=true when a state transition happened and the remaining buffer should +// be reprocessed. +func (p *CohereParser) eat(done bool) (content string, thinking string, calls []api.ToolCall, more bool) { + buf := p.buffer.String() + if buf == "" { + return "", "", nil, false + } + + switch p.state { + case cohereCollectingThinking: + if idx := strings.Index(buf, cohereEndThinking); idx != -1 { + thinking = strings.TrimRightFunc(buf[:idx], unicode.IsSpace) + p.resetBuffer(buf[idx+len(cohereEndThinking):]) + p.state = cohereAwaitingBlock + return "", thinking, nil, true + } + // Emit all but a possible partial tag / trailing whitespace. + keep := overlap(buf, cohereEndThinking) + emitEnd := len(buf) - keep + emitEnd -= trailingWhitespaceLen(buf[:emitEnd]) + if emitEnd > 0 { + thinking = buf[:emitEnd] + p.resetBuffer(buf[emitEnd:]) + } + return "", thinking, nil, false + + case cohereAwaitingBlock: + // Between blocks: look for the next text or action opener, skipping + // whitespace and the end-of-turn token. + for _, open := range []string{cohereStartText, cohereStartResponse} { + if idx := strings.Index(buf, open); idx != -1 { + p.resetBuffer(buf[idx+len(open):]) + p.state = cohereCollectingContent + return "", "", nil, true + } + } + if idx := strings.Index(buf, cohereStartAction); idx != -1 { + p.resetBuffer(buf[idx+len(cohereStartAction):]) + p.state = cohereCollectingAction + return "", "", nil, true + } + trimmed := strings.TrimLeftFunc(buf, unicode.IsSpace) + if strings.HasPrefix(trimmed, cohereEndOfTurn) { + p.resetBuffer(trimmed[len(cohereEndOfTurn):]) + return "", "", nil, true + } + if trimmed == "" { + if done { + p.buffer.Reset() + } + return "", "", nil, false + } + // Wait only while the buffer could still grow into one of the + // expected tags. Anything else — bare content or an unrecognized + // tag — streams out as content rather than buffering forever. + if !done && maybePartialTag(trimmed) { + return "", "", nil, false + } + p.buffer.Reset() + p.state = cohereCollectingContent + p.buffer.WriteString(trimmed) + return "", "", nil, true + + case cohereCollectingContent: + // END_OF_TURN also closes content for models that skip END_TEXT. + for _, close := range []string{cohereEndText, cohereEndResponse, cohereEndOfTurn} { + if idx := strings.Index(buf, close); idx != -1 { + content = buf[:idx] + p.resetBuffer(buf[idx+len(close):]) + p.state = cohereAwaitingBlock + return content, "", nil, true + } + } + keep := max(overlap(buf, cohereEndText), overlap(buf, cohereEndResponse), overlap(buf, cohereEndOfTurn)) + emitEnd := len(buf) - keep + if emitEnd > 0 { + content = buf[:emitEnd] + p.resetBuffer(buf[emitEnd:]) + } + if done && p.buffer.Len() > 0 { + content += p.buffer.String() + p.buffer.Reset() + } + return content, "", nil, false + + case cohereCollectingAction: + if idx := strings.Index(buf, cohereEndAction); idx != -1 { + payload := buf[:idx] + p.resetBuffer(buf[idx+len(cohereEndAction):]) + p.state = cohereAwaitingBlock + calls = parseCohereActions(payload) + return "", "", calls, true + } + if done { + // Best effort on truncated output. + calls = parseCohereActions(buf) + p.buffer.Reset() + return "", "", calls, false + } + return "", "", nil, false + } + + return "", "", nil, false +} + +func (p *CohereParser) resetBuffer(s string) { + p.buffer.Reset() + p.buffer.WriteString(s) +} + +// maybePartialTag reports whether s is a proper prefix of one of the tags the +// awaiting-block state recognizes — the only case worth waiting on for more +// output before treating the buffer as content. +func maybePartialTag(s string) bool { + for _, tag := range []string{cohereStartText, cohereStartResponse, cohereStartAction, cohereEndOfTurn} { + if len(s) < len(tag) && strings.HasPrefix(tag, s) { + return true + } + } + return false +} + +type cohereToolCall struct { + ToolCallID string `json:"tool_call_id"` + ToolName string `json:"tool_name"` + Parameters api.ToolCallFunctionArguments `json:"parameters"` +} + +func (c cohereToolCall) toolCall() api.ToolCall { + return api.ToolCall{ + ID: c.ToolCallID, + Function: api.ToolCallFunction{ + Name: c.ToolName, + Arguments: c.Parameters, + }, + } +} + +// parseCohereActions parses the JSON array inside an action block: +// [{"tool_call_id": "0", "tool_name": ..., "parameters": {...}}, ...] +// +// Sampled output occasionally malforms the JSON (a missing comma between +// calls, an unquoted value). When the array as a whole fails to parse, fall +// back to scanning its balanced top-level objects and parse each +// independently, so one bad call doesn't drop its siblings. +func parseCohereActions(payload string) []api.ToolCall { + payload = strings.TrimSpace(payload) + if payload == "" { + return nil + } + + var parsed []cohereToolCall + if err := json.Unmarshal([]byte(payload), &parsed); err == nil { + calls := make([]api.ToolCall, 0, len(parsed)) + for _, c := range parsed { + calls = append(calls, c.toolCall()) + } + return calls + } + + var calls []api.ToolCall + for _, obj := range scanJSONObjects(payload) { + var c cohereToolCall + if err := json.Unmarshal([]byte(obj), &c); err != nil { + if len(obj) > 200 { + obj = obj[:200] + "…" + } + slog.Warn("cohere action parsing failed", "error", err, "action", obj) + continue + } + calls = append(calls, c.toolCall()) + } + return calls +} + +// scanJSONObjects returns the balanced top-level {...} chunks of s, tracking +// strings and escapes so braces inside values don't split objects. +func scanJSONObjects(s string) []string { + var objects []string + depth, start := 0, -1 + inString, escaped := false, false + for i := range len(s) { + switch c := s[i]; { + case escaped: + escaped = false + case c == '\\' && inString: + escaped = true + case c == '"': + inString = !inString + case inString: + case c == '{': + if depth == 0 { + start = i + } + depth++ + case c == '}': + if depth > 0 { + depth-- + if depth == 0 && start >= 0 { + objects = append(objects, s[start:i+1]) + start = -1 + } + } + } + } + return objects +} diff --git a/model/parsers/cohere_test.go b/model/parsers/cohere_test.go new file mode 100644 index 00000000..31d0bd74 --- /dev/null +++ b/model/parsers/cohere_test.go @@ -0,0 +1,268 @@ +package parsers + +import ( + "testing" + + "github.com/ollama/ollama/api" +) + +// Model output begins inside <|START_THINKING|> when reasoning is on (the +// generation prompt ends with that tag). + +func cohereAddAll(t *testing.T, p *CohereParser, chunks []string) (content, thinking string, calls []api.ToolCall) { + t.Helper() + for i, c := range chunks { + done := i == len(chunks)-1 + ct, th, tc, err := p.Add(c, done) + if err != nil { + t.Fatal(err) + } + content += ct + thinking += th + calls = append(calls, tc...) + } + return content, thinking, calls +} + +func TestCohereParseThinkingThenText(t *testing.T) { + p := &CohereParser{} + p.Init(nil, nil, nil) + + content, thinking, calls := cohereAddAll(t, p, []string{ + "Let me think", " about this.<|END_THINKING|>", + "<|START_TEXT|>Hello", " world!<|END_TEXT|>", + }) + if thinking != "Let me think about this." { + t.Errorf("thinking = %q", thinking) + } + if content != "Hello world!" { + t.Errorf("content = %q", content) + } + if len(calls) != 0 { + t.Errorf("unexpected calls: %v", calls) + } +} + +func TestCohereParseSplitTags(t *testing.T) { + p := &CohereParser{} + p.Init(nil, nil, nil) + + // Tags split across chunk boundaries must not leak into output. + content, thinking, _ := cohereAddAll(t, p, []string{ + "think<|END_TH", "INKING|><|STAR", "T_TEXT|>ans", "wer<|END_", "TEXT|>", + }) + if thinking != "think" { + t.Errorf("thinking = %q", thinking) + } + if content != "answer" { + t.Errorf("content = %q", content) + } +} + +func TestCohereParseToolCall(t *testing.T) { + p := &CohereParser{} + p.Init(nil, nil, nil) + + content, thinking, calls := cohereAddAll(t, p, []string{ + "plan<|END_THINKING|><|START_ACTION|>[\n", + ` {"tool_call_id": "0", "tool_name": "get_weather", "parameters": {"city": "Paris"}},`, + ` {"tool_call_id": "1", "tool_name": "get_time", "parameters": {}}`, + "\n]<|END_ACTION|>", + }) + if thinking != "plan" { + t.Errorf("thinking = %q", thinking) + } + if content != "" { + t.Errorf("content = %q", content) + } + if len(calls) != 2 { + t.Fatalf("calls = %d, want 2", len(calls)) + } + if calls[0].Function.Name != "get_weather" || calls[1].Function.Name != "get_time" { + t.Errorf("call names = %q, %q", calls[0].Function.Name, calls[1].Function.Name) + } + if v, ok := calls[0].Function.Arguments.Get("city"); !ok || v != "Paris" { + t.Errorf("call 0 city = %v %v", v, ok) + } + if calls[0].Function.Index != 0 || calls[1].Function.Index != 1 { + t.Errorf("call indices = %d, %d", calls[0].Function.Index, calls[1].Function.Index) + } +} + +func TestCohereParseReasoningOff(t *testing.T) { + p := &CohereParser{} + think := &api.ThinkValue{Value: false} + p.Init(nil, nil, think) + + content, thinking, _ := cohereAddAll(t, p, []string{ + "<|START_TEXT|>direct answer<|END_TEXT|>", + }) + if thinking != "" { + t.Errorf("thinking = %q", thinking) + } + if content != "direct answer" { + t.Errorf("content = %q", content) + } +} + +func TestCohereParseBareContent(t *testing.T) { + // Models occasionally skip the START_TEXT wrapper; treat raw text after + // thinking as content. + p := &CohereParser{} + p.Init(nil, nil, nil) + + content, thinking, _ := cohereAddAll(t, p, []string{ + "thought<|END_THINKING|>", "Just plain text", " output", + }) + if thinking != "thought" { + t.Errorf("thinking = %q", thinking) + } + if content != "Just plain text output" { + t.Errorf("content = %q", content) + } +} + +func TestCohereParseEndOfTurnWithoutEndText(t *testing.T) { + p := &CohereParser{} + p.Init(nil, nil, nil) + + content, _, _ := cohereAddAll(t, p, []string{ + "t<|END_THINKING|><|START_TEXT|>answer<|END_OF_TURN_TOKEN|>", + }) + if content != "answer" { + t.Errorf("content = %q", content) + } +} + +func TestCohereParsePrefillContinuation(t *testing.T) { + p := &CohereParser{} + last := &api.Message{Role: "assistant", Content: "partial"} + p.Init(nil, last, nil) + + content, thinking, _ := cohereAddAll(t, p, []string{" continued<|END_TEXT|>"}) + if thinking != "" { + t.Errorf("thinking = %q", thinking) + } + if content != " continued" { + t.Errorf("content = %q", content) + } +} + +func TestCohereParserRegistered(t *testing.T) { + p := ParserForName("cohere") + if p == nil { + t.Fatal("cohere parser not registered") + } + if !p.HasToolSupport() || !p.HasThinkingSupport() { + t.Error("cohere parser should support tools and thinking") + } +} + +func TestCohereParseMalformedActions(t *testing.T) { + // One malformed call (unquoted value) must not drop its well-formed + // sibling, and a missing comma between calls must not drop either. + p := &CohereParser{} + p.Init(nil, nil, nil) + _, _, calls := cohereAddAll(t, p, []string{ + "plan<|END_THINKING|><|START_ACTION|>[\n", + ` {"tool_call_id": "0", "tool_name": "set_alarm", "parameters": {"time": 15:30}},`, + ` {"tool_call_id": "1", "tool_name": "get_weather", "parameters": {"city": "Oslo"}}`, + "\n]<|END_ACTION|>", + }) + if len(calls) != 1 || calls[0].Function.Name != "get_weather" { + t.Fatalf("calls = %v, want the well-formed get_weather call", calls) + } + + p = &CohereParser{} + p.Init(nil, nil, nil) + _, _, calls = cohereAddAll(t, p, []string{ + `<|END_THINKING|><|START_ACTION|>[`, + `{"tool_call_id": "0", "tool_name": "a", "parameters": {}}`, + `{"tool_call_id": "1", "tool_name": "b", "parameters": {"x": "{not json}"}}`, + `]<|END_ACTION|>`, + }) + if len(calls) != 2 || calls[0].Function.Name != "a" || calls[1].Function.Name != "b" { + t.Fatalf("calls = %v, want both calls despite missing comma", calls) + } + if v, ok := calls[1].Function.Arguments.Get("x"); !ok || v != "{not json}" { + t.Fatalf("braces inside string values must not split objects, got %v", v) + } + + // Unparseable garbage yields no calls and no panic. + p = &CohereParser{} + p.Init(nil, nil, nil) + _, _, calls = cohereAddAll(t, p, []string{"x<|END_THINKING|><|START_ACTION|>[!!!]<|END_ACTION|>"}) + if len(calls) != 0 { + t.Fatalf("calls = %v, want none", calls) + } +} + +func TestCohereParseLegacyResponseMarkers(t *testing.T) { + // Models trained on the older Command A template sometimes emit + // <|START_RESPONSE|>/<|END_RESPONSE|> instead of START_TEXT/END_TEXT. + p := &CohereParser{} + p.Init(nil, nil, nil) + content, thinking, _ := cohereAddAll(t, p, []string{ + "plan<|END_THINKING|>", "<|START_RESPONSE|>Hello", " there<|END_RESPONSE|>", + }) + if thinking != "plan" { + t.Errorf("thinking = %q", thinking) + } + if content != "Hello there" { + t.Errorf("content = %q", content) + } +} + +func TestCohereParseStreamsBeforeDone(t *testing.T) { + // Regression test: output after END_THINKING that opens with an + // unrecognized tag must stream as it arrives, not buffer until the + // generation finishes (it previously buffered forever, presenting as a + // hung response). + for _, tc := range []struct { + name string + chunk string + }{ + {"legacy response marker", "<|START_RESPONSE|>The answer is 42."}, + {"unrecognized tag", "<|TOOL_PLAN|>I should call the tool."}, + {"bare content", "Just plain text."}, + } { + t.Run(tc.name, func(t *testing.T) { + p := &CohereParser{} + p.Init(nil, nil, nil) + if _, _, _, err := p.Add("x<|END_THINKING|>", false); err != nil { + t.Fatal(err) + } + content, _, _, err := p.Add(tc.chunk, false) // done=false: still streaming + if err != nil { + t.Fatal(err) + } + if content == "" { + t.Fatalf("content did not stream before done for %q", tc.chunk) + } + }) + } +} + +func TestCohereParseEndOfTurnBetweenBlocks(t *testing.T) { + // A literal end-of-turn marker between blocks is consumed, not shown. + p := &CohereParser{} + p.Init(nil, nil, nil) + content, thinking, _ := cohereAddAll(t, p, []string{ + "t<|END_THINKING|><|START_TEXT|>hi<|END_TEXT|><|END_OF_TURN_TOKEN|>", + }) + if thinking != "t" || content != "hi" { + t.Errorf("thinking = %q, content = %q", thinking, content) + } +} + +func TestCohereParseBareContentBeforeEndOfTurn(t *testing.T) { + // Bare content followed by an end-of-turn marker keeps the content. + p := &CohereParser{} + p.Init(nil, nil, nil) + content, _, _ := cohereAddAll(t, p, []string{ + "t<|END_THINKING|>plain answer<|END_OF_TURN_TOKEN|>", + }) + if content != "plain answer" { + t.Errorf("content = %q, want %q", content, "plain answer") + } +} diff --git a/model/parsers/parsers.go b/model/parsers/parsers.go index c7fea831..ae12bada 100644 --- a/model/parsers/parsers.go +++ b/model/parsers/parsers.go @@ -92,6 +92,8 @@ func ParserForName(name string) Parser { return &LFM2Parser{hasThinkingSupport: true} case "laguna": return &LagunaParser{} + case "cohere": + return &CohereParser{} default: return nil } diff --git a/model/renderers/cohere.go b/model/renderers/cohere.go new file mode 100644 index 00000000..770322ad --- /dev/null +++ b/model/renderers/cohere.go @@ -0,0 +1,233 @@ +package renderers + +import ( + "encoding/json" + "strconv" + "strings" + + "github.com/ollama/ollama/api" +) + +// CohereRenderer renders the Cohere North / Command A 2026 chat template +// (Cohere2 MoE models such as CohereLabs/North-Mini-Code-1.0): a platform +// system turn with an Available Tools section, <|START_TEXT|>-wrapped message +// bodies, <|START_THINKING|> reasoning, <|START_ACTION|> tool calls, and +// <|START_TOOL_RESULT|> tool results. +// +// The chat template's model-specific platform instructions (identity, default +// policies) belong in the model's Modelfile SYSTEM prompt, which arrives here +// as the first system message. +type CohereRenderer struct{} + +func (r *CohereRenderer) LeadingBOS() string { + return "" +} + +// cohereToolJSON renders one tool entry exactly as the template's tojson +// filter does: {"name": ..., "description": ..., "parameters": {...}, +// "responses": null} with ", " / ": " separators. +func cohereToolJSON(tool api.Tool) (string, error) { + params, err := marshalWithSpaces(tool.Function.Parameters) + if err != nil { + return "", err + } + name, err := json.Marshal(tool.Function.Name) + if err != nil { + return "", err + } + desc, err := json.Marshal(tool.Function.Description) + if err != nil { + return "", err + } + var sb strings.Builder + sb.WriteString(`{"name": `) + sb.Write(name) + sb.WriteString(`, "description": `) + sb.Write(desc) + sb.WriteString(`, "parameters": `) + sb.WriteString(string(params)) + sb.WriteString(`, "responses": null}`) + return sb.String(), nil +} + +// writeToolsSection writes the "# Available Tools" block, reproducing the +// template's whitespace for the empty and populated cases. +func writeToolsSection(sb *strings.Builder, tools []api.Tool) error { + sb.WriteString("# Available Tools\n```json\n[\n") + if len(tools) == 0 { + sb.WriteString("\n\n") + } else { + for i, tool := range tools { + entry, err := cohereToolJSON(tool) + if err != nil { + return err + } + sb.WriteString("\n ") + sb.WriteString(entry) + if i < len(tools)-1 { + sb.WriteString(",") + } + sb.WriteString("\n\n") + } + } + sb.WriteString("\n]\n```") + return nil +} + +// cohereToolResult is one entry of a <|START_TOOL_RESULT|> array. +func writeToolResult(sb *strings.Builder, callID string, content string) error { + wrapped, err := marshalWithSpaces(map[string]string{"content": content}) + if err != nil { + return err + } + sb.WriteString("\n {\n \"tool_call_id\": \"") + sb.WriteString(callID) + sb.WriteString("\",\n \"results\": {\n\n \n \"0\": ") + sb.WriteString(string(wrapped)) + sb.WriteString("\n\n },\n \"is_error\": null\n }") + return nil +} + +func (r *CohereRenderer) Render(messages []api.Message, tools []api.Tool, think *api.ThinkValue) (string, error) { + var sb strings.Builder + + // The template defaults reasoning to true; an explicit think=false + // disables it. + reasoning := think == nil || think.Bool() + + // The first system message — the request's system prompt, or the model's + // Modelfile SYSTEM when the request has none — fills the template's + // platform instruction slot. + var system string + rest := messages + if len(messages) > 0 && strings.EqualFold(messages[0].Role, "system") { + system = messages[0].Content + rest = messages[1:] + } + + // Platform system turn: system prompt plus the Available Tools section + // (rendered even when no tools are defined). + sb.WriteString("<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TEXT|>") + if system != "" { + sb.WriteString(system) + sb.WriteString("\n\n\n\n") + } + if err := writeToolsSection(&sb, tools); err != nil { + return "", err + } + sb.WriteString("<|END_TEXT|><|END_OF_TURN_TOKEN|>") + + // Tool call ids regenerate as sequential indices across the whole + // conversation (regen_tool_call_ids default); results reference the + // index of their originating call. + callIndex := 0 + callIDToIndex := map[string]string{} + nextCallID := func(id string) string { + idx := strconv.Itoa(callIndex) + callIndex++ + if id != "" { + if _, seen := callIDToIndex[id]; !seen { + callIDToIndex[id] = idx + } + } + return idx + } + resolveResultID := func(m api.Message) string { + if idx, ok := callIDToIndex[m.ToolCallID]; ok { + return idx + } + // Fall back to call order when ids are absent. + return m.ToolCallID + } + + prefill := false + for i := 0; i < len(rest); i++ { + message := rest[i] + switch strings.ToLower(message.Role) { + case "system": + sb.WriteString("<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TEXT|>") + sb.WriteString(message.Content) + sb.WriteString("<|END_TEXT|><|END_OF_TURN_TOKEN|>") + case "user": + sb.WriteString("<|START_OF_TURN_TOKEN|><|USER_TOKEN|><|START_TEXT|>") + sb.WriteString(message.Content) + sb.WriteString("<|END_TEXT|><|END_OF_TURN_TOKEN|>") + case "assistant", "chatbot": + sb.WriteString("<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>") + if len(message.ToolCalls) > 0 { + // Whitespace before THINKING/ACTION matches the template's + // (untrimmed jinja blocks). + sb.WriteString("\n \n ") + if message.Thinking != "" { + sb.WriteString("<|START_THINKING|>") + sb.WriteString(message.Thinking) + sb.WriteString("<|END_THINKING|>") + } + sb.WriteString("<|START_ACTION|>[") + for j, tc := range message.ToolCalls { + args, err := marshalWithSpaces(tc.Function.Arguments) + if err != nil { + return "", err + } + sb.WriteString("\n\n {\"tool_call_id\": \"") + sb.WriteString(nextCallID(toolCallID(tc))) + sb.WriteString("\", \"tool_name\": \"") + sb.WriteString(tc.Function.Name) + sb.WriteString("\", \"parameters\": ") + sb.WriteString(string(args)) + sb.WriteString("}") + if j < len(message.ToolCalls)-1 { + sb.WriteString(",") + } + } + sb.WriteString("\n\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>") + } else { + if message.Thinking != "" { + sb.WriteString("<|START_THINKING|>") + sb.WriteString(message.Thinking) + sb.WriteString("<|END_THINKING|>") + } + sb.WriteString("<|START_TEXT|>") + sb.WriteString(message.Content) + if i == len(rest)-1 { + // Assistant prefill: leave the text open for + // continuation. + prefill = true + } else { + sb.WriteString("<|END_TEXT|><|END_OF_TURN_TOKEN|>") + } + } + case "tool": + // Consecutive tool messages merge into one TOOL_RESULT array. + sb.WriteString("<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TOOL_RESULT|>[") + if err := writeToolResult(&sb, resolveResultID(message), message.Content); err != nil { + return "", err + } + for i+1 < len(rest) && strings.EqualFold(rest[i+1].Role, "tool") { + i++ + sb.WriteString(",") + if err := writeToolResult(&sb, resolveResultID(rest[i]), rest[i].Content); err != nil { + return "", err + } + } + sb.WriteString("\n\n]<|END_TOOL_RESULT|><|END_OF_TURN_TOKEN|>") + } + } + + if !prefill { + sb.WriteString("<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>") + if reasoning { + sb.WriteString("<|START_THINKING|>") + } else { + sb.WriteString("<|START_THINKING|><|END_THINKING|>") + } + } + + return sb.String(), nil +} + +// toolCallID returns the tool call's id when the client supplied one. The +// api.ToolCall ID field may be empty for calls synthesized by ollama. +func toolCallID(tc api.ToolCall) string { + return tc.ID +} diff --git a/model/renderers/cohere_test.go b/model/renderers/cohere_test.go new file mode 100644 index 00000000..f29b6ac4 --- /dev/null +++ b/model/renderers/cohere_test.go @@ -0,0 +1,189 @@ +package renderers + +import ( + "testing" + + "github.com/ollama/ollama/api" +) + +// Ground truth in these tests comes from rendering North-Mini-Code-1.0's +// chat_template.jinja with HF jinja semantics (add_generation_prompt=true), +// minus the leading "" from {{ bos_token }} (the tokenizer adds BOS as a +// token at encode time). + +const cohereSystemTurnNoTools = "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TEXT|>" + + "# Available Tools\n```json\n[\n\n\n\n]\n```" + + "<|END_TEXT|><|END_OF_TURN_TOKEN|>" + +func TestCohereRenderUserOnly(t *testing.T) { + r := &CohereRenderer{} + got, err := r.Render([]api.Message{{Role: "user", Content: "USERMSG"}}, nil, nil) + if err != nil { + t.Fatal(err) + } + + want := cohereSystemTurnNoTools + + "<|START_OF_TURN_TOKEN|><|USER_TOKEN|><|START_TEXT|>USERMSG<|END_TEXT|><|END_OF_TURN_TOKEN|>" + + "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_THINKING|>" + if got != want { + t.Errorf("render mismatch:\ngot: %q\nwant: %q", got, want) + } +} + +func TestCohereRenderSystemHistoryAndThinking(t *testing.T) { + r := &CohereRenderer{} + got, err := r.Render([]api.Message{ + {Role: "system", Content: "DEVPREAMBLE"}, + {Role: "user", Content: "Q1"}, + {Role: "assistant", Content: "A1", Thinking: "THINK1"}, + {Role: "user", Content: "Q2"}, + }, nil, nil) + if err != nil { + t.Fatal(err) + } + + want := "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TEXT|>" + + "DEVPREAMBLE\n\n\n\n# Available Tools\n```json\n[\n\n\n\n]\n```" + + "<|END_TEXT|><|END_OF_TURN_TOKEN|>" + + "<|START_OF_TURN_TOKEN|><|USER_TOKEN|><|START_TEXT|>Q1<|END_TEXT|><|END_OF_TURN_TOKEN|>" + + "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_THINKING|>THINK1<|END_THINKING|><|START_TEXT|>A1<|END_TEXT|><|END_OF_TURN_TOKEN|>" + + "<|START_OF_TURN_TOKEN|><|USER_TOKEN|><|START_TEXT|>Q2<|END_TEXT|><|END_OF_TURN_TOKEN|>" + + "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_THINKING|>" + if got != want { + t.Errorf("render mismatch:\ngot: %q\nwant: %q", got, want) + } +} + +func TestCohereRenderReasoningOff(t *testing.T) { + r := &CohereRenderer{} + think := &api.ThinkValue{Value: false} + got, err := r.Render([]api.Message{{Role: "user", Content: "Q"}}, nil, think) + if err != nil { + t.Fatal(err) + } + wantSuffix := "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_THINKING|><|END_THINKING|>" + if got[len(got)-len(wantSuffix):] != wantSuffix { + t.Errorf("reasoning-off generation prompt mismatch, got tail %q", got[len(got)-len(wantSuffix):]) + } +} + +func TestCohereRenderToolFlow(t *testing.T) { + r := &CohereRenderer{} + tools := []api.Tool{{ + Type: "function", + Function: api.ToolFunction{ + Name: "get_weather", + Description: "Get weather", + Parameters: api.ToolFunctionParameters{ + Type: "object", + Required: []string{"city"}, + Properties: testPropsOrdered([]orderedProp{ + {Key: "city", Value: api.ToolProperty{Type: api.PropertyType{"string"}}}, + }), + }, + }, + }} + args := api.ToolCallFunctionArguments{} + args.Set("city", "Paris") + got, err := r.Render([]api.Message{ + {Role: "user", Content: "weather in Paris?"}, + {Role: "assistant", Thinking: "I should call the tool", ToolCalls: []api.ToolCall{ + {ID: "call_x", Function: api.ToolCallFunction{Name: "get_weather", Arguments: args}}, + }}, + {Role: "tool", ToolCallID: "call_x", Content: "15C sunny"}, + {Role: "user", Content: "thanks"}, + }, tools, nil) + if err != nil { + t.Fatal(err) + } + + // Tools section (from jinja): one entry, surrounded by the template's + // whitespace. Key order inside "parameters" follows Go's struct order + // (jinja preserves whatever order the client sent; both are valid JSON + // schema). + wantTools := "# Available Tools\n```json\n[\n\n {\"name\": \"get_weather\", \"description\": \"Get weather\", \"parameters\": {\"type\": \"object\", \"required\": [\"city\"], \"properties\": {\"city\": {\"type\": \"string\"}}}, \"responses\": null}\n\n\n]\n```" + if !contains(t, got, wantTools, "tools section") { + return + } + + // Assistant tool call turn. + wantAction := "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>\n \n <|START_THINKING|>I should call the tool<|END_THINKING|><|START_ACTION|>[\n\n {\"tool_call_id\": \"0\", \"tool_name\": \"get_weather\", \"parameters\": {\"city\": \"Paris\"}}\n\n]<|END_ACTION|><|END_OF_TURN_TOKEN|>" + if !contains(t, got, wantAction, "action turn") { + return + } + + // Tool result turn. + wantResult := "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|><|START_TOOL_RESULT|>[\n {\n \"tool_call_id\": \"0\",\n \"results\": {\n\n \n \"0\": {\"content\": \"15C sunny\"}\n\n },\n \"is_error\": null\n }\n\n]<|END_TOOL_RESULT|><|END_OF_TURN_TOKEN|>" + contains(t, got, wantResult, "tool result turn") +} + +func TestCohereRenderMultiToolCallsAndResults(t *testing.T) { + r := &CohereRenderer{} + emptyArgs := api.ToolCallFunctionArguments{} + xArgs := api.ToolCallFunctionArguments{} + xArgs.Set("x", 1) + got, err := r.Render([]api.Message{ + {Role: "user", Content: "go"}, + {Role: "assistant", ToolCalls: []api.ToolCall{ + {ID: "a", Function: api.ToolCallFunction{Name: "t1", Arguments: emptyArgs}}, + {ID: "b", Function: api.ToolCallFunction{Name: "t2", Arguments: xArgs}}, + }}, + {Role: "tool", ToolCallID: "a", Content: "r1"}, + {Role: "tool", ToolCallID: "b", Content: "r2"}, + }, nil, nil) + if err != nil { + t.Fatal(err) + } + + wantAction := "<|START_ACTION|>[\n\n {\"tool_call_id\": \"0\", \"tool_name\": \"t1\", \"parameters\": {}},\n\n {\"tool_call_id\": \"1\", \"tool_name\": \"t2\", \"parameters\": {\"x\": 1}}\n\n]<|END_ACTION|>" + if !contains(t, got, wantAction, "multi action") { + return + } + + // Consecutive tool messages merge into one result block with sequential + // regenerated ids. + wantResults := "<|START_TOOL_RESULT|>[\n {\n \"tool_call_id\": \"0\",\n \"results\": {\n\n \n \"0\": {\"content\": \"r1\"}\n\n },\n \"is_error\": null\n },\n {\n \"tool_call_id\": \"1\",\n \"results\": {\n\n \n \"0\": {\"content\": \"r2\"}\n\n },\n \"is_error\": null\n }\n\n]<|END_TOOL_RESULT|>" + contains(t, got, wantResults, "merged tool results") +} + +func TestCohereRenderAssistantPrefill(t *testing.T) { + r := &CohereRenderer{} + got, err := r.Render([]api.Message{ + {Role: "user", Content: "Q"}, + {Role: "assistant", Content: "partial"}, + }, nil, nil) + if err != nil { + t.Fatal(err) + } + wantSuffix := "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|><|START_TEXT|>partial" + if got[len(got)-len(wantSuffix):] != wantSuffix { + t.Errorf("prefill should leave the text open, got tail %q", got[max(0, len(got)-90):]) + } +} + +func TestCohereRendererRegistered(t *testing.T) { + if rendererForName("cohere") == nil { + t.Fatal("cohere renderer not registered") + } + if got := LeadingBOSForRenderer("cohere"); got != "" { + t.Errorf("LeadingBOS = %q, want ", got) + } +} + +func contains(t *testing.T, haystack, needle, what string) bool { + t.Helper() + if idx := indexOf(haystack, needle); idx == -1 { + t.Errorf("missing %s:\nwant substring: %q\nin: %q", what, needle, haystack) + return false + } + return true +} + +func indexOf(s, sub string) int { + for i := 0; i+len(sub) <= len(s); i++ { + if s[i:i+len(sub)] == sub { + return i + } + } + return -1 +} diff --git a/model/renderers/renderer.go b/model/renderers/renderer.go index d11f583a..12c8a827 100644 --- a/model/renderers/renderer.go +++ b/model/renderers/renderer.go @@ -107,6 +107,8 @@ func rendererForName(name string) Renderer { return &LFM2Renderer{IsThinking: true, useImgTags: RenderImgTags} case "laguna": return &LagunaRenderer{} + case "cohere": + return &CohereRenderer{} default: return nil } diff --git a/x/create/client/create.go b/x/create/client/create.go index 1a5968c5..01a6b4d1 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -745,6 +745,9 @@ func getParserName(modelDir string) string { if strings.Contains(archLower, "laguna") { return "laguna" } + if strings.Contains(archLower, "cohere2moe") || strings.Contains(archLower, "cohere2_moe") { + return "cohere" + } if strings.Contains(archLower, "glm4") || strings.Contains(archLower, "glm-4") { return "glm-4.7" } @@ -765,6 +768,9 @@ func getParserName(modelDir string) string { if strings.Contains(typeLower, "laguna") { return "laguna" } + if strings.Contains(typeLower, "cohere2_moe") { + return "cohere" + } if strings.Contains(typeLower, "glm4") || strings.Contains(typeLower, "glm-4") { return "glm-4.7" } @@ -805,6 +811,9 @@ func getRendererName(modelDir string) string { if strings.Contains(archLower, "laguna") { return "laguna" } + if strings.Contains(archLower, "cohere2moe") || strings.Contains(archLower, "cohere2_moe") { + return "cohere" + } if strings.Contains(archLower, "gemma4") { return "gemma4" } @@ -825,6 +834,9 @@ func getRendererName(modelDir string) string { if strings.Contains(typeLower, "laguna") { return "laguna" } + if strings.Contains(typeLower, "cohere2_moe") { + return "cohere" + } if strings.Contains(typeLower, "gemma4") { return "gemma4" } diff --git a/x/create/cohere2moe.go b/x/create/cohere2moe.go new file mode 100644 index 00000000..c111eed0 --- /dev/null +++ b/x/create/cohere2moe.go @@ -0,0 +1,106 @@ +package create + +import ( + "encoding/json" + "os" + "path/filepath" + "regexp" + "strconv" + "strings" + + "github.com/ollama/ollama/x/safetensors" +) + +// cohere2MoeImportTransform adjusts quantization for Cohere2 MoE imports +// (Command A family / North models). +type cohere2MoeImportTransform struct { + numLayers int +} + +func newCohere2MoeImportTransform(modelDir string, _ sourceModelConfig) (tensorImportTransform, error) { + data, err := os.ReadFile(filepath.Join(modelDir, "config.json")) + if err != nil { + return cohere2MoeImportTransform{}, nil //nolint:nilerr // fallback to no heuristic + } + var cfg struct { + NumHiddenLayers int `json:"num_hidden_layers"` + } + if err := json.Unmarshal(data, &cfg); err != nil { + return cohere2MoeImportTransform{}, nil //nolint:nilerr // fallback to no heuristic + } + return cohere2MoeImportTransform{numLayers: cfg.NumHiddenLayers}, nil +} + +func (cohere2MoeImportTransform) skipTensor(string) bool { return false } + +func (cohere2MoeImportTransform) transformTensor(td *safetensors.TensorData) ([]*safetensors.TensorData, error) { + return []*safetensors.TensorData{td}, nil +} + +var cohere2MoeLayerIndexRe = regexp.MustCompile(`\.layers\.(\d+)\.`) + +func (t cohere2MoeImportTransform) quantizationType(name string, shape []int32, quantize string) string { + quantNorm := normalizeQuantType(quantize) + + // The embedding serves double duty: lookup (via QuantizedEmbedding) and the + // tied lm_head projection (via AsLinear). With a 262k vocab the bf16 + // embedding dominates decode bandwidth through the lm_head matmul, so + // quantize it to the 8-bit variant of the requested mode. + if strings.HasSuffix(name, "embed_tokens.weight") && len(shape) == 2 { + switch quantNorm { + case "int4", "int8": + if isAligned(shape, "int8") { + return "int8" + } + case "mxfp4", "nvfp4", "mxfp8": + if isAligned(shape, "mxfp8") { + return "mxfp8" + } + } + return "" + } + + // The MoE router picks the top-k expert set; quantization noise there can + // flip expert selection and compound downstream. It is tiny, so keep it in + // source precision. (GetTensorQuantization already skips "mlp.gate.weight"; + // kept explicit here so renames in the default policy cannot regress this.) + if strings.HasSuffix(name, ".mlp.gate.weight") { + return "" + } + + // Sensitive tensors (v_proj, k_proj, down_proj) get higher precision only + // at quantization-sensitive layer positions (gemma4's useMoreBits + // heuristic) instead of the default policy's blanket promotion. The + // blanket int8 down_proj costs ~25% of decode bandwidth on a top-8 MoE; + // the layer-position heuristic keeps the early/late layers (and every + // third in between) at 8 bits where residual-stream error matters most. + promote := "" + switch quantNorm { + case "int4": + promote = "int8" + case "mxfp4", "nvfp4": + promote = "mxfp8" + } + isSensitive := strings.Contains(name, ".v_proj") || strings.Contains(name, ".k_proj") || strings.Contains(name, "down_proj") + if promote != "" && isSensitive && t.numLayers > 0 { + layerIdx := -1 + if m := cohere2MoeLayerIndexRe.FindStringSubmatch(name); m != nil { + if idx, err := strconv.Atoi(m[1]); err == nil { + layerIdx = idx + } + } + if layerIdx >= 0 { + if useMoreBits(layerIdx, t.numLayers) && isAligned(shape, promote) { + return promote + } + if !isAligned(shape, quantNorm) { + return "" + } + // Bypass GetTensorQuantization's blanket promotion — the + // layer-position heuristic is authoritative here. + return quantNorm + } + } + + return GetTensorQuantization(name, shape, quantize) +} diff --git a/x/create/create.go b/x/create/create.go index 6d9f944c..acf2d244 100644 --- a/x/create/create.go +++ b/x/create/create.go @@ -821,6 +821,7 @@ var tensorImportTransformRegistry = map[string]tensorImportTransformFactory{ "gemma4_unified": newGemma4ImportTransform, "gemma4_unified_text": newGemma4ImportTransform, "LagunaForCausalLM": newLagunaImportTransform, + "Cohere2MoeForCausalLM": newCohere2MoeImportTransform, "Gemma4AssistantForCausalLM": newGemma4ImportTransform, "Gemma4UnifiedAssistantForCausalLM": newGemma4ImportTransform, "gemma4_unified_assistant": newGemma4ImportTransform, diff --git a/x/mlxrunner/imports.go b/x/mlxrunner/imports.go index 316691d9..83451010 100644 --- a/x/mlxrunner/imports.go +++ b/x/mlxrunner/imports.go @@ -1,6 +1,7 @@ package mlxrunner import ( + _ "github.com/ollama/ollama/x/models/cohere2_moe" _ "github.com/ollama/ollama/x/models/gemma3" _ "github.com/ollama/ollama/x/models/gemma4" _ "github.com/ollama/ollama/x/models/glm4_moe_lite" diff --git a/x/models/cohere2_moe/cohere2_moe.go b/x/models/cohere2_moe/cohere2_moe.go new file mode 100644 index 00000000..5b14b13b --- /dev/null +++ b/x/models/cohere2_moe/cohere2_moe.go @@ -0,0 +1,770 @@ +// Package cohere2_moe provides the Cohere2 MoE (Command A family, North) text +// model implementation for MLX. +// +// Architecture notes (matches transformers' Cohere2MoeForCausalLM): +// - Parallel residual blocks: a single input layernorm feeds both attention +// and the MLP, and their outputs are summed onto the residual. +// - Interleaved sliding-window and full attention layers. Sliding layers use +// interleaved ("traditional") RoPE; full-attention layers use no positional +// encoding (NoPE), except prefix dense layers when +// prefix_dense_sliding_window_pattern == 1, which force RoPE. +// - The first first_k_dense_replace layers use a dense SwiGLU MLP with +// prefix_dense_intermediate_size; the rest are sparse MoE layers routed by +// a linear gate with sigmoid or softmax selection over the top-k logits. +// - Logits are scaled by logit_scale. Embeddings are tied by default. +package cohere2_moe + +import ( + "encoding/json" + "fmt" + "math" + + "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("Cohere2MoeForCausalLM", NewModel) +} + +// Config holds the Cohere2 MoE configuration (HuggingFace config.json). +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"` + MaxPositionEmbeddings int32 `json:"max_position_embeddings"` + LayerNormEps float32 `json:"layer_norm_eps"` + RMSNormEps *float32 `json:"rms_norm_eps"` + RopeTheta float32 `json:"rope_theta"` + LogitScale float32 `json:"logit_scale"` + AttentionBias bool `json:"attention_bias"` + TieWordEmbeddings *bool `json:"tie_word_embeddings"` + + SlidingWindow int32 `json:"sliding_window"` + SlidingWindowPattern int32 `json:"sliding_window_pattern"` + PrefixDenseSlidingWindowPattern int32 `json:"prefix_dense_sliding_window_pattern"` + LayerTypes []string `json:"layer_types"` + MLPLayerTypes []string `json:"mlp_layer_types"` + FirstKDenseReplace int32 `json:"first_k_dense_replace"` + PrefixDenseIntermediateSize int32 `json:"prefix_dense_intermediate_size"` + + NumExperts int32 `json:"num_experts"` + NumExpertsPerTok int32 `json:"num_experts_per_tok"` + NumSharedExperts int32 `json:"num_shared_experts"` + SharedExpertCombinationStrategy string `json:"shared_expert_combination_strategy"` + ExpertSelectionFn string `json:"expert_selection_fn"` + NormTopKProb bool `json:"norm_topk_prob"` + + // Quantization metadata (set at load, not from config.json). + QuantGroupSize int `json:"-"` + QuantBits int `json:"-"` + QuantMode string `json:"-"` + TensorQuant map[string]*model.TensorQuantInfo `json:"-"` + + // Computed fields. + Scale float32 `json:"-"` +} + +// normLayer abstracts the per-config choice between RMSNorm (rms_norm_eps set) +// and Cohere-style bias-free LayerNorm. +type normLayer interface { + Forward(x *mlx.Array) *mlx.Array +} + +type rmsNorm struct { + Weight *mlx.Array + Eps float32 +} + +func (n *rmsNorm) Forward(x *mlx.Array) *mlx.Array { return mlx.RMSNormFn(x, n.Weight, n.Eps) } + +type layerNorm struct { + Weight *mlx.Array + Eps float32 +} + +func (n *layerNorm) Forward(x *mlx.Array) *mlx.Array { + return mlx.LayerNormFn(x, n.Weight, nil, n.Eps) +} + +// Model is the Cohere2 MoE model. +type Model struct { + EmbedTokens nn.EmbeddingLayer + Layers []*Layer + Norm normLayer + LMHead nn.LinearLayer + + tok *tokenizer.Tokenizer + *Config +} + +// Layer is a parallel-residual transformer block. +type Layer struct { + InputNorm normLayer + Attention *Attention + MLP MLPBlock + + IsSliding bool + UseRope bool +} + +// Attention implements Cohere2 attention (no q/k norm). +type Attention struct { + QProj nn.LinearLayer + KProj nn.LinearLayer + VProj nn.LinearLayer + OProj nn.LinearLayer +} + +// MLPBlock is the feed-forward interface for dense and MoE blocks. +type MLPBlock interface { + Forward(x *mlx.Array, cfg *Config) *mlx.Array +} + +// DenseMLP is a SwiGLU feed-forward block. +type DenseMLP struct { + GateProj nn.LinearLayer + UpProj nn.LinearLayer + DownProj nn.LinearLayer +} + +// SparseMoE routes each token to the top-k of NumExperts expert MLPs. +type SparseMoE struct { + Router nn.LinearLayer + SwitchMLP *SwitchMLP + SharedExpert *DenseMLP +} + +// SwitchMLP executes the selected expert MLPs with stacked expert weights. +type SwitchMLP struct { + GateWeight *mlx.Array + UpWeight *mlx.Array + DownWeight *mlx.Array + + GateWeightQ, GateScales, GateBiases *mlx.Array + UpWeightQ, UpScales, UpBiases *mlx.Array + DownWeightQ, DownScales, DownBiases *mlx.Array + + GateBits, UpBits, DownBits int + GateGroupSize, UpGroupSize, DownGroupSize int + GateMode, UpMode, DownMode string + + UseQuantized bool +} + +type stackedExpertWeights struct { + Weight *mlx.Array + Scales *mlx.Array + Biases *mlx.Array + Bits int + GroupSize int + Mode string +} + +func parseConfig(configData []byte) (Config, error) { + var raw map[string]json.RawMessage + if err := json.Unmarshal(configData, &raw); err != nil { + return Config{}, fmt.Errorf("parse config envelope: %w", err) + } + + var cfg Config + if err := json.Unmarshal(configData, &cfg); err != nil { + return Config{}, fmt.Errorf("parse config: %w", err) + } + + 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.NumAttentionHeads <= 0 { + return Config{}, fmt.Errorf("invalid num_attention_heads: %d", cfg.NumAttentionHeads) + } + if cfg.NumKeyValueHeads <= 0 { + cfg.NumKeyValueHeads = cfg.NumAttentionHeads + } + if cfg.HeadDim <= 0 { + if cfg.HiddenSize%cfg.NumAttentionHeads != 0 { + return Config{}, fmt.Errorf("hidden_size (%d) must be divisible by num_attention_heads (%d)", cfg.HiddenSize, cfg.NumAttentionHeads) + } + cfg.HeadDim = cfg.HiddenSize / cfg.NumAttentionHeads + } + + // Defaults follow transformers' Cohere2MoeConfig. + if cfg.LayerNormEps == 0 { + cfg.LayerNormEps = 1e-5 + } + if cfg.RopeTheta == 0 { + cfg.RopeTheta = 10000 + } + if cfg.LogitScale == 0 { + cfg.LogitScale = 0.0625 + } + if _, ok := raw["sliding_window"]; !ok { + cfg.SlidingWindow = 4096 + } + if cfg.SlidingWindowPattern <= 0 { + cfg.SlidingWindowPattern = 4 + } + if cfg.PrefixDenseSlidingWindowPattern <= 0 { + cfg.PrefixDenseSlidingWindowPattern = 1 + } + if cfg.MaxPositionEmbeddings <= 0 { + cfg.MaxPositionEmbeddings = 8192 + } + + if cfg.NumExperts <= 0 { + cfg.NumExperts = 8 + } + if cfg.NumExpertsPerTok <= 0 { + cfg.NumExpertsPerTok = 2 + } + if cfg.NumExpertsPerTok > cfg.NumExperts { + return Config{}, fmt.Errorf("num_experts_per_tok (%d) exceeds num_experts (%d)", cfg.NumExpertsPerTok, cfg.NumExperts) + } + if cfg.ExpertSelectionFn == "" { + cfg.ExpertSelectionFn = "softmax" + } + if cfg.ExpertSelectionFn != "softmax" && cfg.ExpertSelectionFn != "sigmoid" { + return Config{}, fmt.Errorf("unsupported expert_selection_fn: %q", cfg.ExpertSelectionFn) + } + if cfg.SharedExpertCombinationStrategy == "" { + cfg.SharedExpertCombinationStrategy = "average" + } + if cfg.SharedExpertCombinationStrategy != "average" && cfg.SharedExpertCombinationStrategy != "sum" { + return Config{}, fmt.Errorf("unsupported shared_expert_combination_strategy: %q", cfg.SharedExpertCombinationStrategy) + } + if _, ok := raw["norm_topk_prob"]; !ok { + cfg.NormTopKProb = true + } + if cfg.PrefixDenseIntermediateSize <= 0 { + cfg.PrefixDenseIntermediateSize = cfg.IntermediateSize + } + + // Derive per-layer attention types when absent: the first + // first_k_dense_replace layers follow prefix_dense_sliding_window_pattern, + // the rest follow sliding_window_pattern (full attention every Nth layer). + if len(cfg.LayerTypes) == 0 { + cfg.LayerTypes = make([]string, cfg.NumHiddenLayers) + for i := range cfg.NumHiddenLayers { + if i < cfg.FirstKDenseReplace { + cfg.LayerTypes[i] = patternLayerType(i, cfg.PrefixDenseSlidingWindowPattern) + } else { + cfg.LayerTypes[i] = patternLayerType(i-cfg.FirstKDenseReplace, cfg.SlidingWindowPattern) + } + } + } + if len(cfg.LayerTypes) != int(cfg.NumHiddenLayers) { + return Config{}, fmt.Errorf("layer_types has %d entries, want %d", len(cfg.LayerTypes), cfg.NumHiddenLayers) + } + + // Derive per-layer MLP types when absent: the first first_k_dense_replace + // layers are dense, the rest sparse. + if len(cfg.MLPLayerTypes) == 0 { + cfg.MLPLayerTypes = make([]string, cfg.NumHiddenLayers) + for i := range cfg.NumHiddenLayers { + if i < cfg.FirstKDenseReplace { + cfg.MLPLayerTypes[i] = "dense" + } else { + cfg.MLPLayerTypes[i] = "sparse" + } + } + } + if len(cfg.MLPLayerTypes) != int(cfg.NumHiddenLayers) { + return Config{}, fmt.Errorf("mlp_layer_types has %d entries, want %d", len(cfg.MLPLayerTypes), cfg.NumHiddenLayers) + } + + cfg.Scale = float32(1.0 / math.Sqrt(float64(cfg.HeadDim))) + return cfg, nil +} + +func patternLayerType(i, pattern int32) string { + if pattern > 0 && (i+1)%pattern == 0 { + return "full_attention" + } + return "sliding_attention" +} + +func (cfg *Config) layerIsSliding(i int32) bool { + return cfg.LayerTypes[i] == "sliding_attention" +} + +func (cfg *Config) layerIsDense(i int32) bool { + return cfg.MLPLayerTypes[i] == "dense" +} + +// layerUsesRope reports whether layer i applies rotary embeddings: all sliding +// layers do, and prefix dense layers force RoPE even with full attention when +// prefix_dense_sliding_window_pattern == 1 (matching Cohere2MoeAttention's +// force_rope). Other full-attention layers use no positional encoding. +func (cfg *Config) layerUsesRope(i int32) bool { + if cfg.layerIsSliding(i) { + return true + } + return cfg.layerIsDense(i) && cfg.PrefixDenseSlidingWindowPattern == 1 +} + +func (cfg *Config) newNorm(weight *mlx.Array) normLayer { + if cfg.RMSNormEps != nil { + return &rmsNorm{Weight: weight, Eps: *cfg.RMSNormEps} + } + return &layerNorm{Weight: weight, Eps: cfg.LayerNormEps} +} + +func (cfg *Config) tieEmbeddings() bool { + return cfg.TieWordEmbeddings == nil || *cfg.TieWordEmbeddings +} + +// NewModel creates a Cohere2 MoE model from a manifest root. +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 genConfigData, err := root.Manifest.ReadConfig("generation_config.json"); err == nil { + tokConfig.GenerationConfigJSON = genConfigData + } + if tokConfigData, err := root.Manifest.ReadConfig("tokenizer_config.json"); err == nil { + tokConfig.TokenizerConfigJSON = tokConfigData + } + + tok, err := tokenizer.LoadFromBytesWithConfig(tokData, tokConfig) + if err != nil { + return nil, fmt.Errorf("parse tokenizer: %w", err) + } + + m := &Model{ + Layers: make([]*Layer, cfg.NumHiddenLayers), + Config: &cfg, + tok: tok, + } + + for i := range cfg.NumHiddenLayers { + m.Layers[i] = &Layer{ + IsSliding: cfg.layerIsSliding(i), + UseRope: cfg.layerUsesRope(i), + } + } + + return m, nil +} + +func supportsGatherQMM(mode string, bits int) bool { + switch mode { + case "affine": + return bits == 4 || bits == 8 + case "mxfp8": + return bits == 8 + case "nvfp4", "mxfp4": + return bits == 4 + default: + return false + } +} + +// transposeExpertWeightForGatherMM converts stacked [E, out, in] expert +// weights to the [E, in, out] layout GatherMM consumes, materialized once at +// load so the forward path avoids per-call transposes. +func transposeExpertWeightForGatherMM(w *mlx.Array) *mlx.Array { + if w == nil || !w.Valid() || w.NumDims() != 3 { + return w + } + t := mlx.Transpose(w, 0, 2, 1) + cloned := t.Clone() + mlx.Eval(cloned) + return cloned +} + +// loadStackedProjection returns expert weights already stacked as a single 3D +// tensor (layers.N.mlp.switch_mlp..weight) — the layout `ollama create` +// writes when it packs per-expert tensors at import. +func loadStackedProjection(tensors map[string]*mlx.Array, cfg *Config, useQuantized bool, base string) *stackedExpertWeights { + key := base + ".weight" + w := tensors[key] + if w == nil { + return nil + } + + scales := tensors[key+"_scale"] + if scales == nil { + return &stackedExpertWeights{Weight: w} + } + + qbiases := tensors[key+"_qbias"] + groupSize, bits, mode := model.ResolveLinearQuantParams( + cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant, + key, w, scales, + ) + if useQuantized && supportsGatherQMM(mode, bits) { + return &stackedExpertWeights{ + Weight: w, + Scales: scales, + Biases: qbiases, + Bits: bits, + GroupSize: groupSize, + Mode: mode, + } + } + + return &stackedExpertWeights{ + Weight: mlx.Dequantize(w, scales, qbiases, groupSize, bits, mode), + Bits: bits, + GroupSize: groupSize, + Mode: mode, + } +} + +// LoadWeights assigns tensors to model fields. +func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { + cfg := m.Config + + linears := model.NewLinearFactory(tensors, cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant) + + embedTokens := model.MakeEmbeddingLayer(tensors, "model.embed_tokens", cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode, cfg.TensorQuant) + if embedTokens == nil { + return fmt.Errorf("missing embedding weight: model.embed_tokens.weight") + } + m.EmbedTokens = embedTokens + + normWeight := tensors["model.norm.weight"] + if normWeight == nil { + return fmt.Errorf("missing final norm weight: model.norm.weight") + } + m.Norm = cfg.newNorm(normWeight) + + if cfg.tieEmbeddings() { + m.LMHead = m.EmbedTokens.AsLinear() + } else if lmHead := linears.Make("lm_head"); lmHead != nil { + m.LMHead = lmHead + } else { + m.LMHead = m.EmbedTokens.AsLinear() + } + + useQuantizedExperts := supportsGatherQMM(cfg.QuantMode, cfg.QuantBits) + if !useQuantizedExperts && cfg.TensorQuant != nil { + for _, tq := range cfg.TensorQuant { + if tq == nil { + continue + } + _, bits, mode := model.QuantizationParams(tq.QuantType) + if supportsGatherQMM(mode, bits) { + useQuantizedExperts = true + break + } + } + } + + for i := range cfg.NumHiddenLayers { + layerPrefix := fmt.Sprintf("model.layers.%d", i) + layer := &Layer{ + IsSliding: cfg.layerIsSliding(i), + UseRope: cfg.layerUsesRope(i), + } + + normWeight := tensors[layerPrefix+".input_layernorm.weight"] + if normWeight == nil { + return fmt.Errorf("layer %d: missing input_layernorm", i) + } + layer.InputNorm = cfg.newNorm(normWeight) + + attn := &Attention{ + QProj: linears.Make(layerPrefix + ".self_attn.q_proj"), + KProj: linears.Make(layerPrefix + ".self_attn.k_proj"), + VProj: linears.Make(layerPrefix + ".self_attn.v_proj"), + OProj: linears.Make(layerPrefix + ".self_attn.o_proj"), + } + if attn.QProj == nil || attn.KProj == nil || attn.VProj == nil || attn.OProj == nil { + return fmt.Errorf("layer %d: missing attention projections", i) + } + layer.Attention = attn + + if cfg.layerIsDense(i) { + mlp := &DenseMLP{ + GateProj: linears.Make(layerPrefix + ".mlp.gate_proj"), + UpProj: linears.Make(layerPrefix + ".mlp.up_proj"), + DownProj: linears.Make(layerPrefix + ".mlp.down_proj"), + } + if mlp.GateProj == nil || mlp.UpProj == nil || mlp.DownProj == nil { + return fmt.Errorf("layer %d: missing dense mlp projections", i) + } + layer.MLP = mlp + } else { + moe := &SparseMoE{} + moe.Router = linears.Make(layerPrefix + ".mlp.gate") + if moe.Router == nil { + return fmt.Errorf("layer %d: missing moe router gate", i) + } + + gateW := loadStackedProjection(tensors, cfg, useQuantizedExperts, layerPrefix+".mlp.switch_mlp.gate_proj") + upW := loadStackedProjection(tensors, cfg, useQuantizedExperts, layerPrefix+".mlp.switch_mlp.up_proj") + downW := loadStackedProjection(tensors, cfg, useQuantizedExperts, layerPrefix+".mlp.switch_mlp.down_proj") + if gateW == nil || upW == nil || downW == nil { + return fmt.Errorf("layer %d: missing stacked switch_mlp expert weights (import the model with `ollama create`)", i) + } + + switchMLP := &SwitchMLP{} + if gateW.Scales != nil && upW.Scales != nil && downW.Scales != nil { + switchMLP.UseQuantized = true + switchMLP.GateWeightQ = gateW.Weight + switchMLP.GateScales = gateW.Scales + switchMLP.GateBiases = gateW.Biases + switchMLP.GateBits = gateW.Bits + switchMLP.GateGroupSize = gateW.GroupSize + switchMLP.GateMode = gateW.Mode + switchMLP.UpWeightQ = upW.Weight + switchMLP.UpScales = upW.Scales + switchMLP.UpBiases = upW.Biases + switchMLP.UpBits = upW.Bits + switchMLP.UpGroupSize = upW.GroupSize + switchMLP.UpMode = upW.Mode + switchMLP.DownWeightQ = downW.Weight + switchMLP.DownScales = downW.Scales + switchMLP.DownBiases = downW.Biases + switchMLP.DownBits = downW.Bits + switchMLP.DownGroupSize = downW.GroupSize + switchMLP.DownMode = downW.Mode + } else { + switchMLP.GateWeight = transposeExpertWeightForGatherMM(gateW.Weight) + switchMLP.UpWeight = transposeExpertWeightForGatherMM(upW.Weight) + switchMLP.DownWeight = transposeExpertWeightForGatherMM(downW.Weight) + } + moe.SwitchMLP = switchMLP + + if cfg.NumSharedExperts > 0 { + shared := &DenseMLP{ + GateProj: linears.Make(layerPrefix + ".mlp.shared_experts.gate_proj"), + UpProj: linears.Make(layerPrefix + ".mlp.shared_experts.up_proj"), + DownProj: linears.Make(layerPrefix + ".mlp.shared_experts.down_proj"), + } + if shared.GateProj == nil { + shared.GateProj = linears.Make(layerPrefix + ".mlp.shared_expert.gate_proj") + shared.UpProj = linears.Make(layerPrefix + ".mlp.shared_expert.up_proj") + shared.DownProj = linears.Make(layerPrefix + ".mlp.shared_expert.down_proj") + } + if shared.GateProj == nil || shared.UpProj == nil || shared.DownProj == nil { + return fmt.Errorf("layer %d: missing shared expert projections", i) + } + moe.SharedExpert = shared + } + + layer.MLP = moe + } + + m.Layers[i] = layer + } + + return nil +} + +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 := a.QProj.Forward(x) + k := a.KProj.Forward(x) + v := a.VProj.Forward(x) + + q = mlx.Transpose(mlx.Reshape(q, B, L, cfg.NumAttentionHeads, cfg.HeadDim), 0, 2, 1, 3) + k = mlx.Transpose(mlx.Reshape(k, B, L, cfg.NumKeyValueHeads, cfg.HeadDim), 0, 2, 1, 3) + v = mlx.Transpose(mlx.Reshape(v, B, L, cfg.NumKeyValueHeads, cfg.HeadDim), 0, 2, 1, 3) + + // Cohere uses interleaved pairs (traditional RoPE). Full-attention layers + // outside the forced-RoPE prefix use no positional encoding. + if useRope { + q = mlx.RoPEWithBase(q, int(cfg.HeadDim), true, cfg.RopeTheta, 1.0, positions) + k = mlx.RoPEWithBase(k, int(cfg.HeadDim), true, cfg.RopeTheta, 1.0, positions) + } + + var kv nn.SDPAOption + if c != nil { + history := c.(cache.Attention).Update(b, k, v) + kv = nn.WithKVHistory(history) + } else { + kv = nn.WithKV(k, v, b.SeqQueryLens) + } + out := nn.ScaledDotProductAttention(b, q, cfg.Scale, kv, nn.WithMask(nn.CausalMask())) + out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*cfg.HeadDim) + return a.OProj.Forward(out) +} + +func (m *DenseMLP) Forward(x *mlx.Array, _ *Config) *mlx.Array { + return m.DownProj.Forward(mlx.SwiGLU(m.GateProj.Forward(x), m.UpProj.Forward(x))) +} + +// route selects the top-k experts. Selection happens on the raw router logits +// and the activation (sigmoid or softmax) is applied to just the selected +// entries, matching Cohere2MoeTopKRouter (both activations are monotonic, so +// selection order is unchanged). +func (moe *SparseMoE) route(x *mlx.Array, cfg *Config) (inds, scores *mlx.Array) { + logits := moe.Router.Forward(x) + + inds = mlx.Argpartition(mlx.Neg(logits), int(cfg.NumExpertsPerTok)-1, -1) + dims := inds.Dims() + inds = mlx.SliceStartStop(inds, []int32{0, 0, 0}, []int32{int32(dims[0]), int32(dims[1]), cfg.NumExpertsPerTok}) + + selected := mlx.TakeAlongAxis(logits, inds, -1) + if cfg.ExpertSelectionFn == "sigmoid" { + scores = mlx.Sigmoid(selected) + if cfg.NormTopKProb && cfg.NumExpertsPerTok > 1 { + scores = mlx.Div(scores, mlx.Sum(scores, -1, true)) + } + } else { + scores = mlx.SoftmaxAxis(selected, -1, true) + } + return inds, scores +} + +func (moe *SparseMoE) Forward(x *mlx.Array, cfg *Config) *mlx.Array { + dims := x.Dims() + B, L := int32(dims[0]), int32(dims[1]) + + inds, scores := moe.route(x, cfg) + + expertOut := moe.SwitchMLP.Forward(x, inds, cfg) + y := mlx.Sum(mlx.Mul(expertOut, mlx.ExpandDims(scores, -1)), 2, false) + + if moe.SharedExpert != nil { + y = mlx.Add(y, moe.SharedExpert.Forward(x, cfg)) + if cfg.SharedExpertCombinationStrategy == "average" { + y = mlx.MulScalar(y, 0.5) + } + } + + return mlx.Reshape(y, B, L, cfg.HiddenSize) +} + +func (s *SwitchMLP) Forward(x *mlx.Array, indices *mlx.Array, cfg *Config) *mlx.Array { + dims := x.Dims() + B, L := int32(dims[0]), int32(dims[1]) + topK := cfg.NumExpertsPerTok + + xFlat := mlx.Reshape(x, B*L, 1, 1, cfg.HiddenSize) + idxFlat := mlx.Reshape(indices, B*L, topK) + + // Sorting tokens by expert improves gather matmul locality for prefill + // batches; the cost outweighs the benefit for small decode batches. + doSort := B*L >= 64 + var invOrder *mlx.Array + n := B * L * topK + + if doSort { + idxAll := mlx.Flatten(idxFlat) + order := mlx.Argsort(idxAll, 0) + invOrder = mlx.Argsort(order, 0) + xFlat = mlx.ExpandDims(mlx.Take(mlx.Squeeze(xFlat, 1), mlx.FloorDivideScalar(order, topK), 0), 1) + idxFlat = mlx.Reshape(mlx.Take(idxAll, order, 0), n, 1) + } + + var gate, up, hidden, down *mlx.Array + if s.UseQuantized { + gate = mlx.GatherQMM(xFlat, s.GateWeightQ, s.GateScales, s.GateBiases, + nil, idxFlat, true, s.GateGroupSize, s.GateBits, s.GateMode, doSort) + up = mlx.GatherQMM(xFlat, s.UpWeightQ, s.UpScales, s.UpBiases, + nil, idxFlat, true, s.UpGroupSize, s.UpBits, s.UpMode, doSort) + hidden = mlx.SwiGLU(gate, up) + down = mlx.GatherQMM(hidden, s.DownWeightQ, s.DownScales, s.DownBiases, + nil, idxFlat, true, s.DownGroupSize, s.DownBits, s.DownMode, doSort) + } else { + gate = mlx.GatherMM(xFlat, s.GateWeight, nil, idxFlat, doSort) + up = mlx.GatherMM(xFlat, s.UpWeight, nil, idxFlat, doSort) + hidden = mlx.SwiGLU(gate, up) + down = mlx.GatherMM(hidden, s.DownWeight, nil, idxFlat, doSort) + } + + if doSort { + down = mlx.Reshape(mlx.Take(mlx.Squeeze(mlx.Squeeze(down, 2), 1), invOrder, 0), B*L, topK, cfg.HiddenSize) + } else { + down = mlx.Squeeze(down, 2) + } + + return mlx.Reshape(down, B, L, topK, cfg.HiddenSize) +} + +// Forward runs a parallel-residual block: one shared layernorm feeds both +// attention and the MLP, and the residual adds both outputs. +func (l *Layer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { + normed := l.InputNorm.Forward(x) + attnOut := l.Attention.Forward(normed, b, c, positions, B, L, l.UseRope, cfg) + mlpOut := l.MLP.Forward(normed, cfg) + return mlx.Add(x, mlx.Add(attnOut, mlpOut)) +} + +func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) *mlx.Array { + dims := b.InputIDs.Dims() + B, L := int32(dims[0]), int32(dims[1]) + positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets)) + + h := m.EmbedTokens.Forward(b.InputIDs) + 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) + } + + return m.Norm.Forward(h) +} + +func (m *Model) Unembed(x *mlx.Array) *mlx.Array { + logits := m.LMHead.Forward(x) + if m.LogitScale != 1.0 { + logits = mlx.MulScalar(logits, m.LogitScale) + } + 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 creates per-layer caches: rotating (bounded) caches for sliding +// window layers and standard KV caches for full attention layers. +func (m *Model) NewCaches() []cache.Cache { + caches := make([]cache.Cache, len(m.Layers)) + for i, layer := range m.Layers { + if m.SlidingWindow > 0 && layer.IsSliding { + caches[i] = cache.NewRotatingKVCache(int(m.SlidingWindow)) + } else { + caches[i] = cache.NewKVCache() + } + } + return caches +}